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")
# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here
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")
# [TAG_HF_EXAMPLE_MISSING] model is gated, and the checkpoint requires cd to subdir, not supported here
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