Felipe97/llama-cpp-compiled
01.1k
1from __future__ import annotations2 3import re4from pathlib import Path5from typing import Any, Iterable, TYPE_CHECKING6 7import torch8 9if TYPE_CHECKING:10 from torch import Tensor11 12from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger13 14# Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one15# 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 tokens20# they are appended to the embedding table as extra tokens, to be looked up like any other row21# - bos_emb lives in latent space, so input_linear is folded into it here22# - the backbone has no lm_head, the embedding table is reused as output for the unused logits23#24# pipeline stage mapping:25# mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder26# flow_lm.transformer --> mapped to normal libllama text model (autoregressive)27# flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE28# mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV29 30# indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder31_ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E73132_ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E73133_DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E73134_DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E73135 36_N_SEANET_STAGES = 337_SAMPLE_RATE = 2400038 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 cache63 "max_position_embeddings": 4096,64 # not in the checkpoint, but every released variant uses head_dim 6465 "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 here81class PocketTTSModel(TextModel):82 model_arch = gguf.MODEL_ARCH.POCKETTTS83 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 do94 # unigram segmentation, so use the UGM tokenizer instead95 from sentencepiece import sentencepiece_model_pb2 as model96 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 pieces104 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.CONTROL108 scores[len(tokens) - len(extra) + i] = -1000.0109 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 mmproj125 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 return129 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 return134 135 if name.startswith("flow_lm.transformer.layers."):136 assert bid is not None137 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 return146 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 return151 152 return153 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_voice157 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 names162 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_linear169 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 here179class PocketTTSMmprojModel(MmprojModel):180 has_audio_encoder = True181 has_vision_encoder = False182 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 None200 201 # voice-prompt encoder: mimi encoder + speaker_proj202 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 models204 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 decoder215 # the SEANet and flow net hparams are constant across the family, clip.cpp holds them216 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_dims229 # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path230 if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"):231 return gguf.GGMLQuantizationType.F16232 return False233 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 class236 T = gguf.MODEL_TENSOR237 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 table240 if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."):241 return # backbone242 243 if name == "flow_lm.speaker_proj_weight":244 yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch)245 return246 if name == "flow_lm.input_linear.weight":247 yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch)248 return249 if name == "flow_lm.emb_mean":250 yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch)251 return252 if name == "flow_lm.emb_std":253 yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch)254 return255 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 return259 260 if name.startswith("flow_lm.flow_net."):261 yield from self._flow_net_tensor(name, data_torch)262 return263 264 if name == "mimi.downsample.conv.conv.weight":265 yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch)266 return267 if name == "mimi.upsample.convtr.convtr.weight":268 yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch)269 return270 if name == "mimi.quantizer.output_proj.weight":271 yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1))272 return273 274 if "_transformer.transformer.layers." in name:275 yield from self._mimi_tfm_tensor(name, data_torch)276 return277 278 if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."):279 yield from self._seanet_tensor(name, data_torch)280 return281 282 return283 284 def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]:285 T = gguf.MODEL_TENSOR286 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 return299 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 return313 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 return327 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 return339 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 return344 tensor = entry[1 if is_decoder else 0]345 suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix346 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_TENSOR350 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 else358 (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 return365 if idx == 3 * _N_SEANET_STAGES + 2:366 yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch)367 return368 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 return375 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 res2379 yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch)380 return381