char-fork/llama.cpp

git clone https://git.t4t.associates/char-fork/llama.cpp

Pascalmtmd: Fix Qwen3-tts-0.6b (#28231)0f3a71be1

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
25.0 KiB477 linesraw
1from __future__ import annotations
2
3import json
4from pathlib import Path
5from typing import Any, Callable, Iterable, TYPE_CHECKING
6
7import torch
8import torch.nn.functional as F
9
10if TYPE_CHECKING:
11    from torch import Tensor
12
13from .base import ModelBase, MmprojModel, TextModel, gguf
14
15# Tricks being used to support this model via existing llama.cpp code paths:
16# - Text projection MLP is folded into the embedding table
17# - codec_embedding is concat to the text embedding table, vocab is extended
18#   example: codec_bos_id(2149) --> "<|codec_bos|>"
19#            codec_eos_token_id(2150) --> "<|codec_eos_token|>"
20#            codec_language_id.chinese(2055) --> "<|codec_language_chinese|>"
21#            other rows --> "<|codec_0|>", "<|codec_1|>", ..., "<|codec_1023|>"
22# - output tensor codec_head is smaller than vocab, so logits will be padded at inference time
23# - suppress_tokens is used to limit the backbone to only sample either semantic or EOS (stop) token
24
25# pipeline stage mapping:
26#   speaker reference encoder --> mapped to normal mtmd audio encoder
27#   backbone --> mapped to normal libllama text model (autoregressive)
28#   code_predictor --> MTMD_GEN_PROCESS_TYPE_GEN_CODE
29#   code2wav --> MTMD_GEN_PROCESS_TYPE_GEN_WAV
30
31# torch activation functions used by Qwen3TTSTalkerResizeMLP (config's hidden_act)
32_ACT2FN = {
33    "silu": F.silu,
34    "gelu": F.gelu,
35    "relu": F.relu,
36}
37
38
39@ModelBase.register("Qwen3TTSForConditionalGeneration")
40@ModelBase.example("Qwen/Qwen3-TTS-12Hz-1.7B-Base")
41class Qwen3TTSTalkerModel(TextModel):
42    model_arch = gguf.MODEL_ARCH.QWEN3TTS
43
44    _TEXT_PROJ_KEYS = (
45        "model.text_embedding.weight",
46        "text_projection.linear_fc1.weight",
47        "text_projection.linear_fc1.bias",
48        "text_projection.linear_fc2.weight",
49        "text_projection.linear_fc2.bias",
50    )
51
52    _text_proj_buffer: dict[str, Tensor]
53    _folded_text_embed: Tensor | None
54    _codec_embed: Tensor | None
55
56    def __init__(self, dir_model: Path, *args, **kwargs):
57        hparams = kwargs.pop("hparams", None)
58        if hparams is None:
59            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
60        raw_talker_config = dict(hparams["talker_config"])
61        self._talker_config = raw_talker_config
62        self.n_codec_vocab = raw_talker_config["vocab_size"]
63        talker_config = dict(raw_talker_config)
64        talker_config["vocab_size"] = talker_config["text_vocab_size"]
65        hparams["text_config"] = talker_config
66        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
67        self._text_proj_buffer = {}
68        self._folded_text_embed = None
69        self._codec_embed = None
70
71    def _codec_token_names(self) -> list[str]:
72        # start every row with a generic name, then override the ones with a
73        # known meaning (bos/eos/language/etc, derived from the *_id fields
74        # of talker_config) with a more descriptive one
75        names = [f"<|codec_{i}|>" for i in range(self.n_codec_vocab)]
76        for key, val in self._talker_config.items():
77            if not key.endswith("_id"):
78                continue
79            prefix = key[:-len("_id")]
80            if isinstance(val, int):
81                names[val] = f"<|{prefix}|>"
82            elif isinstance(val, dict):
83                for subkey, subval in val.items():
84                    names[subval] = f"<|{prefix}_{subkey}|>"
85        return names
86
87    def set_vocab(self):
88        codec_tokens = self._codec_token_names()
89        codec_toktypes = [gguf.TokenType.CONTROL] * len(codec_tokens)
90
91        try:
92            tokens, scores, toktypes = self._create_vocab_sentencepiece()
93            self.gguf_writer.add_tokenizer_model("llama")
94            self.gguf_writer.add_tokenizer_pre("default")
95            tokens += [t.encode("utf-8") for t in codec_tokens]
96            scores += [0.0] * len(codec_tokens)
97            toktypes += codec_toktypes
98            self.gguf_writer.add_token_list(tokens)
99            self.gguf_writer.add_token_scores(scores)
100            self.gguf_writer.add_token_types(toktypes)
101            special_vocab = gguf.SpecialVocab(self.dir_model, n_vocab=len(tokens))
102            special_vocab.add_to_gguf(self.gguf_writer)
103            return
104        except FileNotFoundError:
105            pass
106
107        tokens, toktypes, tokpre = self.get_vocab_base()
108        tokens += codec_tokens
109        toktypes += codec_toktypes
110        self.gguf_writer.add_tokenizer_model("gpt2")
111        self.gguf_writer.add_tokenizer_pre(tokpre)
112        self.gguf_writer.add_token_list(tokens)
113        self.gguf_writer.add_token_types(toktypes)
114
115        special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True)
116        special_vocab.add_to_gguf(self.gguf_writer)
117
118        # make sure that the model has no chat template, so chat will be disabled
119        self.gguf_writer.add_chat_template(None)
120
121    def set_gguf_parameters(self):
122        super().set_gguf_parameters()
123
124        # note: final vocab layout is [text_vocab | codec_vocab], with text_vocab is actually padded with -inf in cgraph
125        # for codec_vocab, only first 2048 rows can be sampled for semantic code
126        # plus codec_eos_token_id that used for signaling end of generation
127        # ref: https://github.com/QwenLM/Qwen3-TTS/blob/022e286b98fbec7e1e916cb940cdf532cd9f488e/qwen_tts/core/models/modeling_qwen3_tts.py#L2059-L2063
128
129        vocab_size = self.hparams["vocab_size"] + self.n_codec_vocab
130        codec_eos_token_id = self.hparams["vocab_size"] + self._talker_config["codec_eos_token_id"]
131        self.gguf_writer.add_suppress_tokens([
132            i for i in range(vocab_size - 1024, vocab_size)
133            if i != codec_eos_token_id
134        ])
135        self.gguf_writer.add_eos_token_id(codec_eos_token_id)
136        self.gguf_writer.add_add_eos_token(False)
137
138    @classmethod
139    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
140        name, gen = item
141
142        if not name.startswith("talker.") or name.startswith("talker.code_predictor."):
143            return None
144
145        name = name[len("talker."):]
146        return super().filter_tensors((name, gen))
147
148    def _maybe_emit_token_embd(self) -> Iterable[tuple[str, Tensor]]:
149        if self._folded_text_embed is None or self._codec_embed is None:
150            return
151        combined = torch.cat([self._folded_text_embed, self._codec_embed], dim=0)
152        yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), combined)
153
154    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
155        # codec_embedding rows are appended after the text vocab, extending the embedding table
156        if name == "model.codec_embedding.weight":
157            self._codec_embed = data_torch
158            yield from self._maybe_emit_token_embd()
159            return
160
161        # codec_head is the output head for the (smaller) codec vocab; logits get padded to
162        # the extended vocab size at inference time
163        if name == "codec_head.weight":
164            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT), data_torch)
165            return
166
167        if name in self._TEXT_PROJ_KEYS:
168            self._text_proj_buffer[name] = data_torch
169            if len(self._text_proj_buffer) < len(self._TEXT_PROJ_KEYS):
170                return
171
172            # fold MLP into the embedding table at conversion time, MLP won't be used at inference time anyway
173            act_fn = _ACT2FN[self.hparams["hidden_act"]]
174            embed = self._text_proj_buffer["model.text_embedding.weight"]
175            hidden = act_fn(F.linear(embed,
176                                     self._text_proj_buffer["text_projection.linear_fc1.weight"],
177                                     self._text_proj_buffer["text_projection.linear_fc1.bias"]))
178            folded = F.linear(hidden,
179                              self._text_proj_buffer["text_projection.linear_fc2.weight"],
180                              self._text_proj_buffer["text_projection.linear_fc2.bias"])
181            self._folded_text_embed = folded
182            yield from self._maybe_emit_token_embd()
183            return
184
185        yield from super().modify_tensors(data_torch, name, bid)
186
187
188@ModelBase.register("Qwen3TTSForConditionalGeneration")
189@ModelBase.example("Qwen/Qwen3-TTS-12Hz-1.7B-Base")
190class Qwen3TTSSpeakerEncoderModel(MmprojModel):
191    has_vision_encoder = False
192    has_audio_encoder = True
193
194    # talker.code_predictor.model.layers.{bid}.<key> -> A_GEN_CODE_*
195    # bypass tensor_mapping.py for now to make it simple
196    _CODE_LAYER_TENSOR_MAP = {
197        "input_layernorm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM,
198        "self_attn.q_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q,
199        "self_attn.q_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM,
200        "self_attn.k_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K,
201        "self_attn.k_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM,
202        "self_attn.v_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_V,
203        "self_attn.o_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT,
204        "post_attention_layernorm": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM,
205        "mlp.gate_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE,
206        "mlp.up_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP,
207        "mlp.down_proj": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN,
208    }
209
210    # note: codebook pages will be stacked to 3D
211    _CODE_GEN_N_CODEBOOKS = 15
212    _code_embed_buffer: dict[int, Tensor] = {}
213    _code_head_buffer: dict[int, Tensor] = {}
214    _wav_config_cache: dict[str, Any] | None = None
215
216    def __init__(self, dir_model: Path, *args, **kwargs):
217        hparams = kwargs.pop("hparams", None)
218        if hparams is None:
219            hparams = ModelBase.load_hparams(dir_model, is_mistral_format=False)
220        hparams["text_config"] = {"hidden_size": hparams["talker_config"]["hidden_size"]}
221        # ECAPA-TDNN has a fixed 4-stage backbone, but MmprojModel.__init__ needs a n_block_keys
222        hparams["speaker_encoder_config"]["n_layers"] = 4
223        super().__init__(dir_model, *args, hparams=hparams, **kwargs)
224        self._wav_config_cache = None
225
226    def get_audio_config(self) -> dict[str, Any] | None:
227        return self.global_config.get("speaker_encoder_config")
228
229    def set_gguf_parameters(self):
230        self.gguf_writer.add_file_type(self.ftype)
231        self.gguf_writer.add_clip_has_audio_encoder(True)
232        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.QWEN3TTS_SPKENC)
233
234        # handle speaker encoder config
235        self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
236        # mel_spectrogram() front-end: sr=24000, n_fft=1024, hop=256, n_mels=128, fmin=0, fmax=12000 (=sr/2, the clip.cpp default)
237        self.gguf_writer.add_audio_num_mel_bins(128)
238        # 3 SE-Res2Net stages; the stem conv, mfa, asp and fc are not counted here
239        self.gguf_writer.add_audio_block_count(3)
240        # ECAPA-TDNN has no attention/FFN, these are dummy to allow clip.cpp to load it
241        self.gguf_writer.add_audio_embedding_length(1536)
242        self.gguf_writer.add_audio_head_count(1)
243        self.gguf_writer.add_audio_feed_forward_length(1536)
244        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
245
246        # handle code predictor config
247        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
248        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.QWEN3TTS_GEN)
249        code_predictor_config = self.global_config["talker_config"]["code_predictor_config"]
250        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
251        self.gguf_writer.add_gen_audio_embedding_length(code_predictor_config["hidden_size"])
252        self.gguf_writer.add_gen_audio_feed_forward_length(code_predictor_config["intermediate_size"])
253        self.gguf_writer.add_gen_audio_block_count(code_predictor_config["num_hidden_layers"])
254        self.gguf_writer.add_gen_audio_head_count(code_predictor_config["num_attention_heads"])
255        self.gguf_writer.add_gen_audio_head_count_kv(code_predictor_config["num_key_value_heads"])
256        self.gguf_writer.add_gen_audio_attention_layernorm_eps(code_predictor_config["rms_norm_eps"])
257        # note: code2wav hparams are hardcoded on the mtmd/clip.cpp side for now, not written here
258
259    def _wav_decoder_config(self) -> dict[str, Any] | None:
260        # code2wav has its own config.json, inside the speech_tokenizer dir
261        if self._wav_config_cache is None:
262            path = self.dir_model / "speech_tokenizer" / "config.json"
263            with open(path, "r", encoding="utf-8") as f:
264                cfg = json.load(f)
265            self._wav_config_cache = cfg["decoder_config"]
266        return self._wav_config_cache
267
268    def tensor_force_quant(self, name, new_name, bid, n_dims):
269        # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
270        if new_name.endswith(".weight") and (
271            new_name in ("a.gen.wav.pre_conv.weight", "a.gen.wav.dac.entry.weight", "a.gen.wav.dac.post_conv.weight")
272            or (".up.blk." in new_name and new_name.endswith(".dwconv.weight"))
273            or (".dac.blk." in new_name and (new_name.endswith(".conv1.weight") or new_name.endswith(".conv2.weight")))
274        ):
275            return gguf.GGMLQuantizationType.F16
276        # ConvTranspose1d kernels: only F16/F32 are implemented, no BF16
277        if new_name.endswith(".conv.weight") and (".up.blk." in new_name or ".dac.blk." in new_name):
278            return gguf.GGMLQuantizationType.F32
279        # the code predictor FFN intermediate peaks around 1.5e5, above the F16 range, and mul_mat
280        # casts its input to the weight type
281        if new_name.startswith("a.gen.code.blk.") and new_name.endswith(".ffn_down.weight"):
282            return gguf.GGMLQuantizationType.F32
283        return super().tensor_force_quant(name, new_name, bid, n_dims)
284
285    @classmethod
286    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
287        name, gen = item
288
289        if not (
290            name.startswith("speaker_encoder.")
291            or name.startswith("talker.code_predictor.")
292            or name == "talker.model.codec_embedding.weight"
293        ):
294            return None
295
296        return super().filter_tensors((name, gen))
297
298    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
299        # code2wav tensors are already named by generate_extra_tensors(), pass them through
300        if name.startswith("a.gen.wav."):
301            yield (name, data_torch)
302            return
303
304        # codebook-0 embedding, fed back to the talker backbone (codebooks 1-15 live in code_predictor)
305        if name == "talker.model.codec_embedding.weight":
306            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUT_EMBD), data_torch)
307            return
308
309        if name == "talker.code_predictor.model.norm.weight":
310            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM), data_torch)
311            return
312
313        if name.startswith("talker.code_predictor.small_to_mtp_projection."):
314            suffix = "." + name.rsplit(".", 1)[1]
315            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_PROJ_IN, suffix=suffix), data_torch)
316            return
317
318        if name.startswith("talker.code_predictor.model.codec_embedding."):
319            idx = int(name.split("codec_embedding.")[1].split(".")[0])
320            self._code_embed_buffer[idx] = data_torch
321            if len(self._code_embed_buffer) < self._CODE_GEN_N_CODEBOOKS:
322                return
323            stacked = torch.stack([self._code_embed_buffer.pop(i) for i in range(self._CODE_GEN_N_CODEBOOKS)], dim=0)
324            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_EMBD), stacked)
325            return
326
327        if name.startswith("talker.code_predictor.lm_head."):
328            idx = int(name.split("lm_head.")[1].split(".")[0])
329            self._code_head_buffer[idx] = data_torch
330            if len(self._code_head_buffer) < self._CODE_GEN_N_CODEBOOKS:
331                return
332            stacked = torch.stack([self._code_head_buffer.pop(i) for i in range(self._CODE_GEN_N_CODEBOOKS)], dim=0)
333            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_GEN_CODE_HEAD), stacked)
334            return
335
336        if name.startswith("talker.code_predictor.model.layers."):
337            rest = name.split("model.layers.")[1]        # "{bid}.<key>.weight"
338            _, key_with_suffix = rest.split(".", 1)       # "<key>.weight"
339            key = key_with_suffix.rsplit(".", 1)[0]        # "<key>"
340            tensor = self._CODE_LAYER_TENSOR_MAP.get(key)
341            if tensor is not None:
342                yield (self.format_tensor_name(tensor, bid), data_torch)
343                return
344
345        if "res2net_block.blocks." in name:
346            assert bid is not None  # the outer stage index, picked up from the tensor name automatically
347            xid = int(name.split("res2net_block.blocks.")[1].split(".")[0])
348            suffix = "." + name.rsplit(".", 1)[1]
349            new_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_ENC_CONV_RES2].format(bid=bid, xid=xid) + suffix
350            yield (new_name, data_torch)
351            return
352
353        yield from super().modify_tensors(data_torch, name, bid)
354
355    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
356        yield from self._generate_code2wav_tensors()
357
358    def _generate_code2wav_tensors(self) -> Iterable[tuple[str, Tensor]]:
359        # code2wav weights live in speech_tokenizer/model.safetensors, not the main safetensors
360        from safetensors.torch import load_file
361
362        wav_config = self._wav_decoder_config()
363        state_dict = load_file(self.dir_model / "speech_tokenizer" / "model.safetensors")
364
365        def get(name: str) -> Tensor:
366            return state_dict[name]
367
368        def snake_fold(alpha: Tensor, beta: Tensor) -> tuple[Tensor, Tensor]:
369            # fold SnakeBeta's exp()/reciprocal here, so the graph is only mul/sin/sqr/mul/add
370            return torch.exp(alpha), 1.0 / (torch.exp(beta) + 1e-9)
371
372        def rvq_codebook(prefix: str, n_layers: int) -> Tensor:
373            # checkpoint has EMA accumulators, so codebook[i] = embedding_sum[i] / cluster_usage[i]
374            books = []
375            for i in range(n_layers):
376                embedding_sum = get(f"{prefix}.vq.layers.{i}._codebook.embedding_sum")
377                cluster_usage = get(f"{prefix}.vq.layers.{i}._codebook.cluster_usage")
378                books.append(embedding_sum / cluster_usage.clamp_min(1e-5).unsqueeze(-1))
379            return torch.stack(books, dim=0) if n_layers > 1 else books[0]
380
381        T = gguf.MODEL_TENSOR
382
383        # --- quantizer: RVQ codebook decode ---
384        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_IN), get("decoder.quantizer.rvq_first.input_proj.weight").squeeze(-1))
385        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_OUT), get("decoder.quantizer.rvq_first.output_proj.weight").squeeze(-1))
386        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), rvq_codebook("decoder.quantizer.rvq_first", 1))
387        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_IN), get("decoder.quantizer.rvq_rest.input_proj.weight").squeeze(-1))
388        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_OUT), get("decoder.quantizer.rvq_rest.output_proj.weight").squeeze(-1))
389        yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_REST_CB), rvq_codebook("decoder.quantizer.rvq_rest", self._CODE_GEN_N_CODEBOOKS))
390
391        # --- pre_conv ---
392        yield (self.format_tensor_name(T.A_GEN_WAV_PRE_CONV, suffix=".weight"), get("decoder.pre_conv.conv.weight"))
393        yield (self.format_tensor_name(T.A_GEN_WAV_PRE_CONV, suffix=".bias"), get("decoder.pre_conv.conv.bias"))
394
395        # --- pre_transformer ---
396        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=".weight"), get("decoder.pre_transformer.input_proj.weight"))
397        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=".bias"), get("decoder.pre_transformer.input_proj.bias"))
398        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=".weight"), get("decoder.pre_transformer.output_proj.weight"))
399        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=".bias"), get("decoder.pre_transformer.output_proj.bias"))
400        yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), get("decoder.pre_transformer.norm.weight"))
401
402        tfm_layer_map = {
403            "input_layernorm.weight":         T.A_GEN_WAV_TFM_ATTN_NORM,
404            "self_attn.q_proj.weight":        T.A_GEN_WAV_TFM_ATTN_Q,
405            "self_attn.k_proj.weight":        T.A_GEN_WAV_TFM_ATTN_K,
406            "self_attn.v_proj.weight":        T.A_GEN_WAV_TFM_ATTN_V,
407            "self_attn.o_proj.weight":        T.A_GEN_WAV_TFM_ATTN_OUT,
408            "self_attn_layer_scale.scale":    T.A_GEN_WAV_TFM_ATTN_SCALE,
409            "post_attention_layernorm.weight": T.A_GEN_WAV_TFM_FFN_NORM,
410            "mlp.gate_proj.weight":           T.A_GEN_WAV_TFM_FFN_GATE,
411            "mlp.up_proj.weight":             T.A_GEN_WAV_TFM_FFN_UP,
412            "mlp.down_proj.weight":           T.A_GEN_WAV_TFM_FFN_DOWN,
413            "mlp_layer_scale.scale":          T.A_GEN_WAV_TFM_FFN_SCALE,
414        }
415        assert wav_config is not None
416        for bid in range(wav_config["num_hidden_layers"]):
417            for key, tensor_id in tfm_layer_map.items():
418                yield (self.format_tensor_name(tensor_id, bid), get(f"decoder.pre_transformer.layers.{bid}.{key}"))
419
420        # --- upsample: 2x (causal ConvTranspose1d + ConvNeXt block) ---
421        up_map = {
422            "0.conv.weight":     (T.A_GEN_WAV_UP_CONV, ".weight"),
423            "0.conv.bias":       (T.A_GEN_WAV_UP_CONV, ".bias"),
424            "1.dwconv.conv.weight": (T.A_GEN_WAV_UP_DWCONV, ".weight"),
425            "1.dwconv.conv.bias":   (T.A_GEN_WAV_UP_DWCONV, ".bias"),
426            "1.norm.weight":     (T.A_GEN_WAV_UP_NORM, ".weight"),
427            "1.norm.bias":       (T.A_GEN_WAV_UP_NORM, ".bias"),
428            "1.pwconv1.weight":  (T.A_GEN_WAV_UP_PW1, ".weight"),
429            "1.pwconv1.bias":    (T.A_GEN_WAV_UP_PW1, ".bias"),
430            "1.pwconv2.weight":  (T.A_GEN_WAV_UP_PW2, ".weight"),
431            "1.pwconv2.bias":    (T.A_GEN_WAV_UP_PW2, ".bias"),
432            "1.gamma":           (T.A_GEN_WAV_UP_GAMMA, ""),
433        }
434        for bid in range(len(wav_config["upsampling_ratios"])):
435            for key, (tensor_id, suffix) in up_map.items():
436                yield (self.format_tensor_name(tensor_id, bid, suffix=suffix), get(f"decoder.upsample.{bid}.{key}"))
437
438        # --- DAC decoder ---
439        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".weight"), get("decoder.decoder.0.conv.weight"))
440        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_ENTRY, suffix=".bias"), get("decoder.decoder.0.conv.bias"))
441
442        n_dac_blocks = len(wav_config["upsample_rates"])
443        for bid in range(n_dac_blocks):
444            py = bid + 1  # decoder.decoder.0 is the entry conv, blocks start at 1
445
446            a, b = snake_fold(get(f"decoder.decoder.{py}.block.0.alpha"), get(f"decoder.decoder.{py}.block.0.beta"))
447            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".alpha"), a)
448            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_SNAKE, bid, suffix=".beta"), b)
449            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".weight"), get(f"decoder.decoder.{py}.block.1.conv.weight"))
450            yield (self.format_tensor_name(T.A_GEN_WAV_DAC_UP_CONV, bid, suffix=".bias"), get(f"decoder.decoder.{py}.block.1.conv.bias"))
451
452            for xid in range(3):
453                ridx = xid + 2  # block.2/3/4 are the 3 residual units
454
455                a1, b1 = snake_fold(get(f"decoder.decoder.{py}.block.{ridx}.act1.alpha"), get(f"decoder.decoder.{py}.block.{ridx}.act1.beta"))
456                name1 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_ACT1].format(bid=bid, xid=xid)
457                yield (name1 + ".alpha", a1)
458                yield (name1 + ".beta", b1)
459
460                name_conv1 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_CONV1].format(bid=bid, xid=xid)
461                yield (name_conv1 + ".weight", get(f"decoder.decoder.{py}.block.{ridx}.conv1.conv.weight"))
462                yield (name_conv1 + ".bias", get(f"decoder.decoder.{py}.block.{ridx}.conv1.conv.bias"))
463
464                a2, b2 = snake_fold(get(f"decoder.decoder.{py}.block.{ridx}.act2.alpha"), get(f"decoder.decoder.{py}.block.{ridx}.act2.beta"))
465                name2 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_ACT2].format(bid=bid, xid=xid)
466                yield (name2 + ".alpha", a2)
467                yield (name2 + ".beta", b2)
468
469                name_conv2 = gguf.TENSOR_NAMES[T.A_GEN_WAV_DAC_RES_CONV2].format(bid=bid, xid=xid)
470                yield (name_conv2 + ".weight", get(f"decoder.decoder.{py}.block.{ridx}.conv2.conv.weight"))
471                yield (name_conv2 + ".bias", get(f"decoder.decoder.{py}.block.{ridx}.conv2.conv.bias"))
472
473        a5, b5 = snake_fold(get("decoder.decoder.5.alpha"), get("decoder.decoder.5.beta"))
474        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".alpha"), a5)
475        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_SNAKE, suffix=".beta"), b5)
476        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".weight"), get("decoder.decoder.6.conv.weight"))
477        yield (self.format_tensor_name(T.A_GEN_WAV_DAC_POST_CONV, suffix=".bias"), get("decoder.decoder.6.conv.bias"))