from __future__ import annotations import re from pathlib import Path from typing import Any, Iterable, TYPE_CHECKING import torch if TYPE_CHECKING: from torch import Tensor from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf, logger # Pocket TTS is a CALM: the backbone conditions a flow-matching decoder that generates one # continuous 32-d latent per frame. There is no codebook in this model. # The checkpoint ships no config.json, hparams come from _load_hparams() below. # # Tricks being used to support this model via existing llama.cpp code paths: # - bos_before_voice and bos_emb are learned input vectors, not tokens # they are appended to the embedding table as extra tokens, to be looked up like any other row # - bos_emb lives in latent space, so input_linear is folded into it here # - the backbone has no lm_head, the embedding table is reused as output for the unused logits # # pipeline stage mapping: # mimi encoder + speaker_proj --> mapped to normal mtmd audio encoder # flow_lm.transformer --> mapped to normal libllama text model (autoregressive) # flow_lm.flow_net + out_eos --> MTMD_GEN_PROCESS_TYPE_GEN_CODE # mimi decoder --> MTMD_GEN_PROCESS_TYPE_GEN_WAV # indices into mimi.encoder.model / mimi.decoder.model for stage i, see SEANetEncoder/SEANetDecoder _ENC_RES_IDX = lambda i: 1 + 3 * i # noqa: E731 _ENC_SCALE_IDX = lambda i: 3 + 3 * i # noqa: E731 _DEC_SCALE_IDX = lambda i: 2 + 3 * i # noqa: E731 _DEC_RES_IDX = lambda i: 3 + 3 * i # noqa: E731 _N_SEANET_STAGES = 3 _SAMPLE_RATE = 24000 def _tensor_shapes(dir_model: Path) -> dict[str, tuple[int, ...]]: part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors") if len(part_names) != 1: return {} with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part: return {name: tuple(part[name].shape) for name in part.keys()} @ModelBase.register_hparams_loader(lambda dir_model: "flow_lm.bos_emb" in _tensor_shapes(dir_model)) def _load_hparams(dir_model: Path) -> dict[str, Any]: logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes") shapes = _tensor_shapes(dir_model) n_vocab, n_embd = shapes["flow_lm.conditioner.embed.weight"] n_layer = sum(1 for name in shapes if re.fullmatch(r"flow_lm\.transformer\.layers\.\d+\.norm1\.weight", name)) n_layer_a = sum(1 for name in shapes if re.fullmatch(r"mimi\.encoder_transformer\.transformer\.layers\.\d+\.norm1\.weight", name)) n_embd_a = shapes["mimi.encoder_transformer.transformer.layers.0.norm1.weight"][0] return { "architectures": ["PocketTTSModel"], "model_type": "pockettts", "num_hidden_layers": n_layer, "hidden_size": n_embd, "intermediate_size": shapes["flow_lm.transformer.layers.0.linear1.weight"][0], # the transformer is fully causal with no context limit, this only bounds the KV cache "max_position_embeddings": 4096, # not in the checkpoint, but every released variant uses head_dim 64 "num_attention_heads": n_embd // 64, # extra rows for the learned input vectors, see _embd_table() "vocab_size": n_vocab + (2 if "flow_lm.bos_before_voice" in shapes else 1), "rope_theta": 10000.0, "layer_norm_eps": 1e-5, "audio_config": { "num_hidden_layers": n_layer_a, "hidden_size": n_embd_a, "intermediate_size": shapes["mimi.encoder_transformer.transformer.layers.0.linear1.weight"][0], "num_attention_heads": n_embd_a // 64, }, } @ModelBase.register("PocketTTSModel") class PocketTTSModel(TextModel): model_arch = gguf.MODEL_ARCH.POCKETTTS _LAYER_TENSOR_MAP = { "norm1": gguf.MODEL_TENSOR.ATTN_NORM, "norm2": gguf.MODEL_TENSOR.FFN_NORM, "self_attn.out_proj": gguf.MODEL_TENSOR.ATTN_OUT, "linear1": gguf.MODEL_TENSOR.FFN_UP, "linear2": gguf.MODEL_TENSOR.FFN_DOWN, } def set_vocab(self): # this is a unigram sentencepiece model, llama.cpp's SPM tokenizer cannot do # unigram segmentation, so use the UGM tokenizer instead from sentencepiece import sentencepiece_model_pb2 as model proto = model.ModelProto() # pyright: ignore[reportAttributeAccessIssue] # ty: ignore[unresolved-attribute] proto.ParseFromString(open(self.dir_model / "tokenizer.model", "rb").read()) assert proto.trainer_spec.model_type == 1, "expected a unigram tokenizer" tokens, scores, toktypes = self._create_vocab_sentencepiece() # the last rows of the embedding table are not sentencepiece pieces extra = self._extra_tokens() for i, name in enumerate(extra): tokens[len(tokens) - len(extra) + i] = name.encode("utf-8") toktypes[len(tokens) - len(extra) + i] = SentencePieceTokenTypes.CONTROL scores[len(tokens) - len(extra) + i] = -1000.0 self.gguf_writer.add_tokenizer_model("t5") self.gguf_writer.add_tokenizer_pre("default") self.gguf_writer.add_token_list(tokens) self.gguf_writer.add_token_scores(scores) self.gguf_writer.add_token_types(toktypes) self.gguf_writer.add_add_space_prefix(proto.normalizer_spec.add_dummy_prefix) self.gguf_writer.add_remove_extra_whitespaces(proto.normalizer_spec.remove_extra_whitespaces) if proto.normalizer_spec.precompiled_charsmap: self.gguf_writer.add_precompiled_charsmap(proto.normalizer_spec.precompiled_charsmap) self.gguf_writer.add_add_bos_token(False) self.gguf_writer.add_add_eos_token(False) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if not name.startswith("flow_lm."): return # mimi and the flow net go to the mmproj if name == "flow_lm.conditioner.embed.weight": yield (self.format_tensor_name(gguf.MODEL_TENSOR.TOKEN_EMBD), self._embd_table(data_torch)) return if name.startswith("flow_lm.out_norm."): suffix = "." + name.rsplit(".", 1)[1] yield (self.format_tensor_name(gguf.MODEL_TENSOR.OUTPUT_NORM, suffix=suffix), data_torch) return if name.startswith("flow_lm.transformer.layers."): assert bid is not None key_with_suffix = name.split(f"layers.{bid}.", 1)[1] key, suffix = key_with_suffix.rsplit(".", 1) if key == "self_attn.in_proj": q, k, v = data_torch.chunk(3, dim=0) yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_Q, bid), q) yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K, bid), k) yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V, bid), v) return tensor = self._LAYER_TENSOR_MAP.get(key) if tensor is not None: yield (self.format_tensor_name(tensor, bid, suffix="." + suffix), data_torch) return return def _extra_tokens(self) -> list[str]: # the conditioner's padding row, then the learned vectors appended by _embd_table(). # bos_before_voice only exists when the pack sets insert_bos_before_voice names = ["<|pad|>"] if "flow_lm.bos_before_voice" in self.model_tensors: names.append("<|bos_before_voice|>") names.append("<|audio_bos|>") return names def _embd_table(self, embed: Tensor) -> Tensor: rows = [embed] if "flow_lm.bos_before_voice" in self.model_tensors: rows.append(self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1).to(embed.dtype)) # bos_emb is a latent, it only enters the backbone through input_linear bos_emb = self.model_tensors["flow_lm.bos_emb"]() input_linear = self.model_tensors["flow_lm.input_linear.weight"]() audio_bos = torch.nn.functional.linear(bos_emb.float(), input_linear.float()).reshape(1, -1) rows.append(audio_bos.to(embed.dtype)) return torch.cat(rows, dim=0) @ModelBase.register("PocketTTSModel") class PocketTTSMmprojModel(MmprojModel): has_audio_encoder = True has_vision_encoder = False _MIMI_TFM_MAP = { "norm1": (gguf.MODEL_TENSOR.A_ENC_INPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM), "norm2": (gguf.MODEL_TENSOR.A_ENC_OUTPUT_NORM, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM), "self_attn.out_proj": (gguf.MODEL_TENSOR.A_ENC_OUTPUT, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT), "linear1": (gguf.MODEL_TENSOR.A_ENC_FFN_UP, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP), "linear2": (gguf.MODEL_TENSOR.A_ENC_FFN_DOWN, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN), "layer_scale_1.scale": (gguf.MODEL_TENSOR.A_ENC_ATTN_SCALE, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_SCALE), "layer_scale_2.scale": (gguf.MODEL_TENSOR.A_ENC_FFN_SCALE_LS, gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE), } _MIMI_TFM_QKV = ( (gguf.MODEL_TENSOR.A_ENC_ATTN_Q, gguf.MODEL_TENSOR.A_ENC_ATTN_K, gguf.MODEL_TENSOR.A_ENC_ATTN_V), (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), ) def set_gguf_parameters(self): self.gguf_writer.add_file_type(self.ftype) assert self.hparams_audio is not None # voice-prompt encoder: mimi encoder + speaker_proj self.gguf_writer.add_clip_has_audio_encoder(True) # note: the 24kHz sample rate is hardcoded on the clip.cpp side, like the other audio models self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_SPKENC) self.gguf_writer.add_audio_projection_dim(self.n_embd_text) self.gguf_writer.add_audio_block_count(self.hparams_audio["num_hidden_layers"]) self.gguf_writer.add_audio_embedding_length(self.hparams_audio["hidden_size"]) self.gguf_writer.add_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) self.gguf_writer.add_audio_head_count(self.hparams_audio["num_attention_heads"]) self.gguf_writer.add_audio_attention_layernorm_eps(1e-5) # mimi convolves the waveform directly, it is passed around as a 1-row "mel" self.gguf_writer.add_audio_num_mel_bins(1) # generation: flow-matching decoder + mimi decoder # the SEANet and flow net hparams are constant across the family, clip.cpp holds them self.gguf_writer.add_clip_has_gen_audio_encoder(True) self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.POCKETTTS_GEN) self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) self.gguf_writer.add_gen_audio_embedding_length(self.hparams_audio["hidden_size"]) self.gguf_writer.add_gen_audio_feed_forward_length(self.hparams_audio["intermediate_size"]) self.gguf_writer.add_gen_audio_block_count(self.hparams_audio["num_hidden_layers"]) self.gguf_writer.add_gen_audio_head_count(self.hparams_audio["num_attention_heads"]) self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) self.gguf_writer.add_gen_audio_model_variant(self.dir_model.name) def tensor_force_quant(self, name, new_name, bid, n_dims): del name, bid, n_dims # conv1d/conv1d_dw kernels must be F16, ggml_conv_1d(_dw) has no BF16 path if ".seanet." in new_name or new_name in ("a.downsample.conv.weight", "a.gen.wav.upsample.weight"): return gguf.GGMLQuantizationType.F16 return False def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: del bid # the block index of the mimi transformers is parsed here, not by the base class T = gguf.MODEL_TENSOR if name in ("flow_lm.bos_emb", "flow_lm.bos_before_voice", "flow_lm.conditioner.embed.weight"): return # folded into the backbone embedding table if name.startswith("flow_lm.transformer.") or name.startswith("flow_lm.out_norm."): return # backbone if name == "flow_lm.speaker_proj_weight": yield (self.format_tensor_name(T.A_ENC_SPEAKER_PROJ), data_torch) return if name == "flow_lm.input_linear.weight": yield (self.format_tensor_name(T.A_GEN_INPUT_LINEAR), data_torch) return if name == "flow_lm.emb_mean": yield (self.format_tensor_name(T.A_GEN_EMB_MEAN, suffix=""), data_torch) return if name == "flow_lm.emb_std": yield (self.format_tensor_name(T.A_GEN_EMB_STD, suffix=""), data_torch) return if name.startswith("flow_lm.out_eos."): suffix = "." + name.rsplit(".", 1)[1] yield (self.format_tensor_name(T.A_GEN_OUT_EOS, suffix=suffix), data_torch) return if name.startswith("flow_lm.flow_net."): yield from self._flow_net_tensor(name, data_torch) return if name == "mimi.downsample.conv.conv.weight": yield (self.format_tensor_name(T.A_ENC_DOWNSAMPLE_CONV), data_torch) return if name == "mimi.upsample.convtr.convtr.weight": yield (self.format_tensor_name(T.A_GEN_WAV_UPSAMPLE), data_torch) return if name == "mimi.quantizer.output_proj.weight": yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_OUT), data_torch.squeeze(-1)) return if "_transformer.transformer.layers." in name: yield from self._mimi_tfm_tensor(name, data_torch) return if name.startswith("mimi.encoder.model.") or name.startswith("mimi.decoder.model."): yield from self._seanet_tensor(name, data_torch) return return def _flow_net_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: T = gguf.MODEL_TENSOR key = name.split("flow_lm.flow_net.", 1)[1] suffix = "." + key.rsplit(".", 1)[1] simple = { "input_proj": T.A_GEN_FLOW_INPUT_PROJ, "cond_embed": T.A_GEN_FLOW_COND_EMBD, "final_layer.linear": T.A_GEN_FLOW_FINAL_PROJ, "final_layer.adaLN_modulation.1": T.A_GEN_FLOW_FINAL_ADA, } tensor = simple.get(key.rsplit(".", 1)[0]) if tensor is not None: yield (self.format_tensor_name(tensor, suffix=suffix), data_torch) return if key.startswith("time_embed."): bid = int(key.split(".")[1]) rest = key.split(f"time_embed.{bid}.", 1)[1] time_map = { "freqs": (T.A_GEN_FLOW_TIME_FREQS, ""), "mlp.0": (T.A_GEN_FLOW_TIME_UP, suffix), "mlp.2": (T.A_GEN_FLOW_TIME_DOWN, suffix), "mlp.3.alpha": (T.A_GEN_FLOW_TIME_NORM, ""), } entry = time_map.get(rest) or time_map.get(rest.rsplit(".", 1)[0]) if entry is not None: yield (self.format_tensor_name(entry[0], bid, suffix=entry[1]), data_torch) return if key.startswith("res_blocks."): bid = int(key.split(".")[1]) rest = key.split(f"res_blocks.{bid}.", 1)[1].rsplit(".", 1)[0] blk_map = { "in_ln": T.A_GEN_FLOW_BLK_NORM, "mlp.0": T.A_GEN_FLOW_BLK_UP, "mlp.2": T.A_GEN_FLOW_BLK_DOWN, "adaLN_modulation.1": T.A_GEN_FLOW_BLK_ADA, } tensor = blk_map.get(rest) if tensor is not None: yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) return def _mimi_tfm_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: is_decoder = name.startswith("mimi.decoder_transformer.") bid = int(name.split("_transformer.transformer.layers.", 1)[1].split(".")[0]) key_with_suffix = name.split(f".layers.{bid}.", 1)[1] if key_with_suffix == "self_attn.in_proj.weight": q, k, v = data_torch.chunk(3, dim=0) names = self._MIMI_TFM_QKV[1 if is_decoder else 0] for tensor, part in zip(names, (q, k, v)): yield (self.format_tensor_name(tensor, bid), part) return key, suffix = key_with_suffix.rsplit(".", 1) entry = self._MIMI_TFM_MAP.get(key) or self._MIMI_TFM_MAP.get(key_with_suffix) if entry is None: return tensor = entry[1 if is_decoder else 0] suffix = ".weight" if key_with_suffix.endswith(".scale") else "." + suffix yield (self.format_tensor_name(tensor, bid, suffix=suffix), data_torch) def _seanet_tensor(self, name: str, data_torch: Tensor) -> Iterable[tuple[str, Tensor]]: T = gguf.MODEL_TENSOR is_decoder = name.startswith("mimi.decoder.") idx = int(name.split(".model.", 1)[1].split(".")[0]) suffix = "." + name.rsplit(".", 1)[1] conv_in, conv_out, res1, res2, scale = ( (T.A_GEN_WAV_SEANET_CONV_IN, T.A_GEN_WAV_SEANET_CONV_OUT, T.A_GEN_WAV_SEANET_RES_CONV1, T.A_GEN_WAV_SEANET_RES_CONV2, T.A_GEN_WAV_SEANET_SCALE_CONV) if is_decoder else (T.A_ENC_SEANET_CONV_IN, T.A_ENC_SEANET_CONV_OUT, T.A_ENC_SEANET_RES_CONV1, T.A_ENC_SEANET_RES_CONV2, T.A_ENC_SEANET_SCALE_CONV) ) if idx == 0: yield (self.format_tensor_name(conv_in, suffix=suffix), data_torch) return if idx == 3 * _N_SEANET_STAGES + 2: yield (self.format_tensor_name(conv_out, suffix=suffix), data_torch) return for stage in range(_N_SEANET_STAGES): res_idx = _DEC_RES_IDX(stage) if is_decoder else _ENC_RES_IDX(stage) scale_idx = _DEC_SCALE_IDX(stage) if is_decoder else _ENC_SCALE_IDX(stage) if idx == scale_idx: yield (self.format_tensor_name(scale, stage, suffix=suffix), data_torch) return if idx == res_idx: # block.1 is the dilated conv, block.3 the pointwise one (0 and 2 are ELU) inner = int(name.split(".block.", 1)[1].split(".")[0]) tensor = res1 if inner == 1 else res2 yield (self.format_tensor_name(tensor, stage, suffix=suffix), data_torch) return