char-fork/llama.cpp

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

Xuan-Son Nguyenconvert: add @ModelBase.example (#27208)fa88ae936

char/aliceai-t5-moe
brancheschar/aliceai-t5-moeall refs →
18.5 KiB380 linesraw
1from __future__ import annotations
2
3import re
4from pathlib import Path
5from typing import Any, Iterable, TYPE_CHECKING
6
7import torch
8
9if TYPE_CHECKING:
10    from torch import Tensor
11
12from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger
13
14# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one
15# continuous 32-d latent per frame. There is no codebook in this model.
16# The checkpoint ships no config.json, hparams come from _load_hparams() below.
17#
18# Tricks being used to support this model via existing llama.cpp code paths:
19# - bos_before_voice and bos_emb are learned input vectors, not tokens
20#   they are appended to the embedding table as extra tokens, to be looked up like any other row
21# - bos_emb lives in latent space, so input_linear is folded into it here
22# - the backbone has no lm_head, the embedding table is reused as output for the unused logits
23#
24# pipeline stage mapping:
25#   mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder
26#   flow_lm.transformer         --> mapped to normal libllama text model (autoregressive)
27#   flow_lm.flow_net + out_eos  --> MTMD_GEN_PROCESS_TYPE_GEN_CODE
28#   mimi decoder                --> MTMD_GEN_PROCESS_TYPE_GEN_WAV
29
30# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder
31_ENC_RES_IDX   = lambda i: 1 + 3 * i  # noqa: E731
32_ENC_SCALE_IDX = lambda i: 3 + 3 * i  # noqa: E731
33_DEC_SCALE_IDX = lambda i: 2 + 3 * i  # noqa: E731
34_DEC_RES_IDX   = lambda i: 3 + 3 * i  # noqa: E731
35
36_N_SEANET_STAGES = 3
37_SAMPLE_RATE = 24000
38
39
40def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]:
41    part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
42    if len(part_names) != 1:
43        return {}
44    with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
45        return {name: tuple(part[name].shape) for name in part.keys()}
46
47
48@ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model))
49def _load_hparams(dir_model: Path) -> dict[str, Any]:
50    logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
51    shapes = _tensor_shapes(dir_model)
52    n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"]
53    n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name))
54    n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name))
55    n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0]
56    return {
57        "architectures": ["PocketTTSModel"],
58        "model_type": "pockettts",
59        "num_hidden_layers": n_layer,
60        "hidden_size": n_embd,
61        "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0],
62        # the transformer is fully causal with no context limit, this only bounds the KV cache
63        "max_position_embeddings": 4096,
64        # not in the checkpoint, but every released variant uses head_dim 64
65        "num_attention_heads": n_embd // 64,
66        # extra rows for the learned input vectors, see _embd_table()
67        "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1),
68        "rope_theta": 10000.0,
69        "layer_norm_eps": 1e-5,
70        "audio_config": {
71            "num_hidden_layers": n_layer_a,
72            "hidden_size": n_embd_a,
73            "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0],
74            "num_attention_heads": n_embd_a // 64,
75        },
76    }
77
78
79@ModelBase.register("PocketTTSModel")
80# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here
81class PocketTTSModel(TextModel):
82    model_arch = gguf.MODEL_ARCH.POCKETTTS
83
84    _LAYER_TENSOR_MAP = {
85        "norm1":               gguf.MODEL_TENSOR.ATTN_NORM,
86        "norm2":               gguf.MODEL_TENSOR.FFN_NORM,
87        "self_attn.out_proj":  gguf.MODEL_TENSOR.ATTN_OUT,
88        "linear1":             gguf.MODEL_TENSOR.FFN_UP,
89        "linear2":             gguf.MODEL_TENSOR.FFN_DOWN,
90    }
91
92    def set_vocab(self):
93        # this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do
94        # unigram segmentation, so use the UGM tokenizer instead
95        from sentencepiece import sentencepiece_model_pb2 as model
96
97        proto = model.ModelProto()  # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute]
98        proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read())
99        assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer"
100
101        tokens, scores, toktypes = self._create_vocab_sentencepiece()
102
103        # the last rows of the embedding table are not sentencepiece pieces
104        extra = self._extra_tokens()
105        for i, name in enumerate(extra):
106            tokens[len(tokens) - len(extra) + i] = name.encode("utf-8")
107            toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL
108            scores[len(tokens) - len(extra) + i] = -1000.0
109
110        self.gguf_writer.add_tokenizer_model("t5")
111        self.gguf_writer.add_tokenizer_pre("default")
112        self.gguf_writer.add_token_list(tokens)
113        self.gguf_writer.add_token_scores(scores)
114        self.gguf_writer.add_token_types(toktypes)
115        self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix)
116        self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces)
117        if proto.normalizer_spec.precompiled_charsmap:
118            self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap)
119        self.gguf_writer.add_add_bos_token(False)
120        self.gguf_writer.add_add_eos_token(False)
121
122    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
123        if not name.startswith("flow_lm."):
124            return  # mimi and the flow net go to the mmproj
125
126        if name == "flow_lm.conditioner.embed.weight":
127            yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch))
128            return
129
130        if name.startswith("flow_lm.out_norm."):
131            suffix = "." + name.rsplit(".", 1)[1]
132            yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch)
133            return
134
135        if name.startswith("flow_lm.transformer.layers."):
136            assert bid is not None
137            key_with_suffix = name.split(f"layers.{bid}.", 1)[1]
138            key, suffix = key_with_suffix.rsplit(".", 1)
139
140            if key == "self_attn.in_proj":
141                q, k, v = data_torch.chunk(3, dim=0)
142                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q)
143                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k)
144                yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v)
145                return
146
147            tensor = self._LAYER_TENSOR_MAP.get(key)
148            if tensor is not None:
149                yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch)
150                return
151
152        return
153
154    def _extra_tokens(self) -> list[str]:
155        # the conditioner's padding row, then the learned vectors appended by _embd_table().
156        # bos_before_voice only exists when the pack sets insert_bos_before_voice
157        names = ["<|pad|>"]
158        if "flow_lm.bos_before_voice" in self.model_tensors:
159            names.append("<|bos_before_voice|>")
160        names.append("<|audio_bos|>")
161        return names
162
163    def _embd_table(self, embed: Tensor) -> Tensor:
164        rows = [embed]
165        if "flow_lm.bos_before_voice" in self.model_tensors:
166            rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype))
167
168        # bos_emb is a latent, it only enters the backbone through input_linear
169        bos_emb = self.model_tensors["flow_lm.bos_emb"]()
170        input_linear = self.model_tensors["flow_lm.input_linear.weight"]()
171        audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1)
172        rows.append(audio_bos.to(embed.dtype))
173
174        return torch.cat(rows, dim=0)
175
176
177@ModelBase.register("PocketTTSModel")
178# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here
179class PocketTTSMmprojModel(MmprojModel):
180    has_audio_encoder = True
181    has_vision_encoder = False
182
183    _MIMI_TFM_MAP = {
184        "norm1":              (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM,  gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM),
185        "norm2":              (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM),
186        "self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT,      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT),
187        "linear1":            (gguf.MODEL_TENSOR.A_ENC_FFN_UP,      gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP),
188        "linear2":            (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN,    gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN),
189        "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE),
190        "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE),
191    }
192    _MIMI_TFM_QKV = (
193        (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V),
194        (gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V),
195    )
196
197    def set_gguf_parameters(self):
198        self.gguf_writer.add_file_type(self.ftype)
199        assert self.hparams_audio is not None
200
201        # voice-prompt encoder: mimi encoder + speaker_proj
202        self.gguf_writer.add_clip_has_audio_encoder(True)
203        # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models
204        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC)
205        self.gguf_writer.add_audio_projection_dim(self.n_embd_text)
206        self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"])
207        self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"])
208        self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
209        self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"])
210        self.gguf_writer.add_audio_attention_layernorm_eps(1e-5)
211        # mimi convolves the waveform directly, it is passed around as a 1-row "mel"
212        self.gguf_writer.add_audio_num_mel_bins(1)
213
214        # generation: flow-matching decoder + mimi decoder
215        # the SEANet and flow net hparams are constant across the family, clip.cpp holds them
216        self.gguf_writer.add_clip_has_gen_audio_encoder(True)
217        self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN)
218        self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text)
219        self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"])
220        self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"])
221        self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"])
222        self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"])
223        self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5)
224
225        self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name)
226
227    def tensor_force_quant(self, name, new_name, bid, n_dims):
228        del name, bid, n_dims
229        # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path
230        if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):
231            return gguf.GGMLQuantizationType.F16
232        return False
233
234    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
235        del bid  # the block index of the mimi transformers is parsed here, not by the base class
236        T = gguf.MODEL_TENSOR
237
238        if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"):
239            return  # folded into the backbone embedding table
240        if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):
241            return  # backbone
242
243        if name == "flow_lm.speaker_proj_weight":
244            yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch)
245            return
246        if name == "flow_lm.input_linear.weight":
247            yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch)
248            return
249        if name == "flow_lm.emb_mean":
250            yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch)
251            return
252        if name == "flow_lm.emb_std":
253            yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch)
254            return
255        if name.startswith("flow_lm.out_eos."):
256            suffix = "." + name.rsplit(".", 1)[1]
257            yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch)
258            return
259
260        if name.startswith("flow_lm.flow_net."):
261            yield from self._flow_net_tensor(name, data_torch)
262            return
263
264        if name == "mimi.downsample.conv.conv.weight":
265            yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch)
266            return
267        if name == "mimi.upsample.convtr.convtr.weight":
268            yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch)
269            return
270        if name == "mimi.quantizer.output_proj.weight":
271            yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1))
272            return
273
274        if "_transformer.transformer.layers." in name:
275            yield from self._mimi_tfm_tensor(name, data_torch)
276            return
277
278        if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."):
279            yield from self._seanet_tensor(name, data_torch)
280            return
281
282        return
283
284    def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
285        T = gguf.MODEL_TENSOR
286        key = name.split("flow_lm.flow_net.", 1)[1]
287        suffix = "." + key.rsplit(".", 1)[1]
288
289        simple = {
290            "input_proj":                 T.A_GEN_FLOW_INPUT_PROJ,
291            "cond_embed":                 T.A_GEN_FLOW_COND_EMBD,
292            "final_layer.linear":         T.A_GEN_FLOW_FINAL_PROJ,
293            "final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA,
294        }
295        tensor = simple.get(key.rsplit(".", 1)[0])
296        if tensor is not None:
297            yield (self.format_tensor_name(tensor, suffix=suffix), data_torch)
298            return
299
300        if key.startswith("time_embed."):
301            bid = int(key.split(".")[1])
302            rest = key.split(f"time_embed.{bid}.", 1)[1]
303            time_map = {
304                "freqs":       (T.A_GEN_FLOW_TIME_FREQS, ""),
305                "mlp.0":       (T.A_GEN_FLOW_TIME_UP,    suffix),
306                "mlp.2":       (T.A_GEN_FLOW_TIME_DOWN,  suffix),
307                "mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM,  ""),
308            }
309            entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0])
310            if entry is not None:
311                yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch)
312            return
313
314        if key.startswith("res_blocks."):
315            bid = int(key.split(".")[1])
316            rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0]
317            blk_map = {
318                "in_ln":                T.A_GEN_FLOW_BLK_NORM,
319                "mlp.0":                T.A_GEN_FLOW_BLK_UP,
320                "mlp.2":                T.A_GEN_FLOW_BLK_DOWN,
321                "adaLN_modulation.1":   T.A_GEN_FLOW_BLK_ADA,
322            }
323            tensor = blk_map.get(rest)
324            if tensor is not None:
325                yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
326            return
327
328    def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
329        is_decoder = name.startswith("mimi.decoder_transformer.")
330        bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0])
331        key_with_suffix = name.split(f".layers.{bid}.", 1)[1]
332
333        if key_with_suffix == "self_attn.in_proj.weight":
334            q, k, v = data_torch.chunk(3, dim=0)
335            names = self._MIMI_TFM_QKV[1 if is_decoder else 0]
336            for tensor, part in zip(names, (q, k, v)):
337                yield (self.format_tensor_name(tensor, bid), part)
338            return
339
340        key, suffix = key_with_suffix.rsplit(".", 1)
341        entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix)
342        if entry is None:
343            return
344        tensor = entry[1 if is_decoder else 0]
345        suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix
346        yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch)
347
348    def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:
349        T = gguf.MODEL_TENSOR
350        is_decoder = name.startswith("mimi.decoder.")
351        idx = int(name.split(".model.", 1)[1].split(".")[0])
352        suffix = "." + name.rsplit(".", 1)[1]
353
354        conv_in, conv_out, res1, res2, scale = (
355            (T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1,
356             T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV)
357            if is_decoder else
358            (T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1,
359             T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV)
360        )
361
362        if idx == 0:
363            yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch)
364            return
365        if idx == 3 * _N_SEANET_STAGES + 2:
366            yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch)
367            return
368
369        for stage in range(_N_SEANET_STAGES):
370            res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage)
371            scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage)
372            if idx == scale_idx:
373                yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch)
374                return
375            if idx == res_idx:
376                # block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU)
377                inner = int(name.split(".block.", 1)[1].split(".")[0])
378                tensor = res1 if inner == 1 else res2
379                yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch)
380                return