Public/llama.cpp
1
0
Fork 0
mirror of https://github.com/ggml-org/llama.cpp.git synced 2026-09-19 01:04:55 +02:00
Code Issues Packages Projects Releases Wiki Activity

text model ok

Browse Source
This commit is contained in:
Xuan Son Nguyen
2026-08-05 02:41:43 +02:00
parent f55a236349
commit a0c869f043
9 changed files with 632 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
conversion/__init__.py
+2
View File
@@ -210,6 +210,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Qwen3MoeForCausalLM": "qwen",
"Qwen3NextForCausalLM": "qwen",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
@@ -305,6 +306,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"Qwen2_5_VLForConditionalGeneration": "qwenvl",
"Qwen3ASRForConditionalGeneration": "qwen3vl",
"Qwen3OmniMoeForConditionalGeneration": "qwen3vl",
"PocketTTSModel": "pockettts",
"Qwen3TTSForConditionalGeneration": "qwen3tts",
"Qwen3VLForConditionalGeneration": "qwen3vl",
"Qwen3VLMoeForConditionalGeneration": "qwen3vl",
conversion/base.py
+47
View File
@@ -1053,6 +1053,10 @@ class ModelBase:
config = AutoConfig.from_pretrained(dir_model, trust_remote_code=False).to_dict()
except Exception as e:
logger.warning(f"Failed to load model config from {dir_model}: {e}")
if not (dir_model / "config.json").is_file():
config = load_hparams_non_hf(dir_model)
if config is not None:
return config
logger.warning("Trying to load config.json instead")
with open(dir_model / "config.json", "r", encoding="utf-8") as f:
config = json.load(f)
@@ -2618,6 +2622,49 @@ else:
LazyTorchTensor._dtype_str_map["F8_E8M0"] = torch.uint8
def load_hparams_non_hf(dir_model: Path) -> dict[str, Any] | None:
# some models ship no config.json at all, their hparams are derived from the checkpoint
part_names = ModelBase.get_model_part_names(dir_model, "model", ".safetensors")
if len(part_names) != 1:
return None
with gguf.utility.SafetensorsLocal(dir_model / part_names[0]) as part:
shapes = {name: tuple(part[name].shape) for name in part.keys()}
if "flow_lm.bos_emb" in shapes:
return _load_hparams_pockettts(shapes)
return None
def _load_hparams_pockettts(shapes: dict[str, tuple[int, ...]]) -> dict[str, Any]:
logger.info("gguf: detected pocket-tts checkpoint, deriving hparams from tensor shapes")
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 stored anywhere in the checkpoint, but every released variant uses head_dim 64
"num_attention_heads": n_embd // 64,
# 2 learned vectors are appended to the embedding table as extra tokens, see pockettts.py
"vocab_size": n_vocab + 2,
"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,
},
}
def get_model_architecture(hparams: dict[str, Any], model_type: ModelType) -> str:
# TODO @ngxson : this won't work correctly if the model has both audio & vision encoders
# maybe we should fallback to text model's arch in that case, since not many models have both
conversion/pockettts.py
+312
View File
@@ -0,0 +1,312 @@
from __future__ import annotations
from typing import Iterable, TYPE_CHECKING
import torch
if TYPE_CHECKING:
from torch import Tensor
from .base import ModelBase, MmprojModel, SentencePieceTokenTypes, TextModel, gguf
# Pocket TTS is a CALM: an autoregressive backbone conditions a flow-matching decoder that
# generates one continuous 32-d latent per frame. There is no codebook anywhere in this model.
#
# The checkpoint ships no config.json, hparams are derived in base.load_hparams_non_hf().
#
# 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 so the helper can look them 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 so that a sampler
# can run over 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
@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):
tokens, scores, toktypes = self._create_vocab_sentencepiece()
# the last 3 rows of the embedding table are not sentencepiece pieces: the conditioner's
# padding row, then the two learned vectors appended by generate_extra_tensors()
extra = ["<|pad|>", "<|bos_before_voice|>", "<|audio_bos|>"]
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("llama")
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_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 _embd_table(self, embed: Tensor) -> Tensor:
bos_before_voice = self.model_tensors["flow_lm.bos_before_voice"]().reshape(1, -1)
# 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)
return torch.cat([embed, bos_before_voice.to(embed.dtype), audio_bos.to(embed.dtype)], 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, 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)
# generation: flow-matching decoder + mimi decoder
# note: the SEANet and flow net hparams are hardcoded on the clip.cpp side for now
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)
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]
# layer_scale is stored without a .weight/.bias suffix
suffix = "" 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
gguf-py/gguf/constants.py
+107
View File
@@ -572,6 +572,7 @@ class MODEL_ARCH(IntEnum):
MELLUM = auto()
NANBEIGE = auto()
QWEN3TTS = auto()
POCKETTTS = auto()
class VISION_PROJECTOR_TYPE(IntEnum):
@@ -1031,6 +1032,37 @@ class MODEL_TENSOR(IntEnum):
A_GEN_WAV_DAC_RES_CONV2 = auto() # DAC residual unit, pointwise causal conv
A_GEN_WAV_DAC_POST_SNAKE = auto() # DAC final SnakeBeta
A_GEN_WAV_DAC_POST_CONV = auto() # DAC conv_post -> 1-channel PCM
# pocket-tts: SEANet encoder (speaker path) and decoder (a.gen.wav path)
A_ENC_SEANET_CONV_IN = auto()
A_ENC_SEANET_CONV_OUT = auto()
A_ENC_SEANET_RES_CONV1 = auto() # residual unit, dilated conv
A_ENC_SEANET_RES_CONV2 = auto() # residual unit, pointwise conv
A_ENC_SEANET_SCALE_CONV = auto() # strided downsample conv
A_ENC_ATTN_SCALE = auto() # layer scale (gamma) on the attn output
A_ENC_SPEAKER_PROJ = auto() # voice latent -> backbone embd
A_GEN_FLOW_INPUT_PROJ = auto()
A_GEN_FLOW_COND_EMBD = auto()
A_GEN_FLOW_TIME_FREQS = auto() # timestep embedder, stored cos/sin frequencies
A_GEN_FLOW_TIME_UP = auto()
A_GEN_FLOW_TIME_DOWN = auto()
A_GEN_FLOW_TIME_NORM = auto() # RMSNorm alpha
A_GEN_FLOW_BLK_NORM = auto() # AdaLN res block, in_ln
A_GEN_FLOW_BLK_UP = auto()
A_GEN_FLOW_BLK_DOWN = auto()
A_GEN_FLOW_BLK_ADA = auto() # AdaLN modulation, -> shift/scale/gate
A_GEN_FLOW_FINAL_ADA = auto() # final layer AdaLN modulation, -> shift/scale
A_GEN_FLOW_FINAL_PROJ = auto()
A_GEN_OUT_EOS = auto() # end-of-speech head on the backbone hidden state
A_GEN_INPUT_LINEAR = auto() # generated latent -> backbone embd
A_GEN_EMB_MEAN = auto() # latent denormalization stats
A_GEN_EMB_STD = auto()
A_GEN_WAV_QUANT_OUT = auto() # DummyQuantizer output_proj, latent -> decoder dim
A_GEN_WAV_UPSAMPLE = auto() # frame rate -> encoder frame rate, depthwise convtr
A_GEN_WAV_SEANET_CONV_IN = auto()
A_GEN_WAV_SEANET_CONV_OUT = auto() # -> 1-channel PCM
A_GEN_WAV_SEANET_RES_CONV1 = auto()
A_GEN_WAV_SEANET_RES_CONV2 = auto()
A_GEN_WAV_SEANET_SCALE_CONV = auto() # strided upsample convtr
A_MMPROJ = auto()
A_MMPROJ_FC = auto()
A_MM_NORM_PRE = auto()
@@ -1244,6 +1276,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.MELLUM: "mellum",
MODEL_ARCH.NANBEIGE: "nanbeige",
MODEL_ARCH.QWEN3TTS: "qwen3tts",
MODEL_ARCH.POCKETTTS: "pockettts",
}
VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = {
@@ -1698,6 +1731,36 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2: "a.gen.wav.dac.blk.{bid}.res.{xid}.conv2",
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE: "a.gen.wav.dac.post_snake",
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV: "a.gen.wav.dac.post_conv",
MODEL_TENSOR.A_ENC_SEANET_CONV_IN: "a.seanet.conv_in",
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT: "a.seanet.conv_out",
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1: "a.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2: "a.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV: "a.seanet.blk.{bid}.scale_conv",
MODEL_TENSOR.A_ENC_ATTN_SCALE: "a.blk.{bid}.attn_scale",
MODEL_TENSOR.A_ENC_SPEAKER_PROJ: "a.speaker_proj",
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ: "a.gen.flow.input_proj",
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD: "a.gen.flow.cond_embd",
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS: "a.gen.flow.time.{bid}.freqs",
MODEL_TENSOR.A_GEN_FLOW_TIME_UP: "a.gen.flow.time.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN: "a.gen.flow.time.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM: "a.gen.flow.time.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM: "a.gen.flow.blk.{bid}.norm",
MODEL_TENSOR.A_GEN_FLOW_BLK_UP: "a.gen.flow.blk.{bid}.up",
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN: "a.gen.flow.blk.{bid}.down",
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA: "a.gen.flow.blk.{bid}.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA: "a.gen.flow.final.ada",
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ: "a.gen.flow.final.proj",
MODEL_TENSOR.A_GEN_OUT_EOS: "a.gen.out_eos",
MODEL_TENSOR.A_GEN_INPUT_LINEAR: "a.gen.input_linear",
MODEL_TENSOR.A_GEN_EMB_MEAN: "a.gen.emb_mean",
MODEL_TENSOR.A_GEN_EMB_STD: "a.gen.emb_std",
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT: "a.gen.wav.quant_out",
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE: "a.gen.wav.upsample",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN: "a.gen.wav.seanet.conv_in",
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT: "a.gen.wav.seanet.conv_out",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1: "a.gen.wav.seanet.blk.{bid}.res_conv1",
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2: "a.gen.wav.seanet.blk.{bid}.res_conv2",
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV: "a.gen.wav.seanet.blk.{bid}.scale_conv",
MODEL_TENSOR.A_MMPROJ: "mm.a.mlp.{bid}",
MODEL_TENSOR.A_MMPROJ_FC: "mm.a.fc",
MODEL_TENSOR.A_MM_NORM_PRE: "mm.a.norm_pre",
@@ -2009,6 +2072,36 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.A_GEN_WAV_DAC_RES_CONV2,
MODEL_TENSOR.A_GEN_WAV_DAC_POST_SNAKE,
MODEL_TENSOR.A_GEN_WAV_DAC_POST_CONV,
MODEL_TENSOR.A_ENC_SEANET_CONV_IN,
MODEL_TENSOR.A_ENC_SEANET_CONV_OUT,
MODEL_TENSOR.A_ENC_SEANET_RES_CONV1,
MODEL_TENSOR.A_ENC_SEANET_RES_CONV2,
MODEL_TENSOR.A_ENC_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_ATTN_SCALE,
MODEL_TENSOR.A_ENC_SPEAKER_PROJ,
MODEL_TENSOR.A_GEN_FLOW_INPUT_PROJ,
MODEL_TENSOR.A_GEN_FLOW_COND_EMBD,
MODEL_TENSOR.A_GEN_FLOW_TIME_FREQS,
MODEL_TENSOR.A_GEN_FLOW_TIME_UP,
MODEL_TENSOR.A_GEN_FLOW_TIME_DOWN,
MODEL_TENSOR.A_GEN_FLOW_TIME_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_NORM,
MODEL_TENSOR.A_GEN_FLOW_BLK_UP,
MODEL_TENSOR.A_GEN_FLOW_BLK_DOWN,
MODEL_TENSOR.A_GEN_FLOW_BLK_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_ADA,
MODEL_TENSOR.A_GEN_FLOW_FINAL_PROJ,
MODEL_TENSOR.A_GEN_OUT_EOS,
MODEL_TENSOR.A_GEN_INPUT_LINEAR,
MODEL_TENSOR.A_GEN_EMB_MEAN,
MODEL_TENSOR.A_GEN_EMB_STD,
MODEL_TENSOR.A_GEN_WAV_QUANT_OUT,
MODEL_TENSOR.A_GEN_WAV_UPSAMPLE,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_IN,
MODEL_TENSOR.A_GEN_WAV_SEANET_CONV_OUT,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV1,
MODEL_TENSOR.A_GEN_WAV_SEANET_RES_CONV2,
MODEL_TENSOR.A_GEN_WAV_SEANET_SCALE_CONV,
MODEL_TENSOR.A_ENC_CONV_NORM_MEAN,
MODEL_TENSOR.A_ENC_CONV_NORM_VAR,
MODEL_TENSOR.A_ENC_MEL_FILTERS,
@@ -4852,6 +4945,18 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
MODEL_ARCH.POCKETTTS: [
MODEL_TENSOR.TOKEN_EMBD,
MODEL_TENSOR.OUTPUT_NORM,
MODEL_TENSOR.ATTN_NORM,
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
MODEL_TENSOR.ATTN_V,
MODEL_TENSOR.ATTN_OUT,
MODEL_TENSOR.FFN_NORM,
MODEL_TENSOR.FFN_DOWN,
MODEL_TENSOR.FFN_UP,
],
}
# tensors that will not be serialized
@@ -5128,6 +5233,8 @@ class VisionProjectorType:
NEMOTRON_V2_VL = "nemotron_v2_vl"
QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder
QWEN3TTS_GEN = "qwen3tts_gen" # audio generation: code_predictor
POCKETTTS_SPKENC = "pockettts_spkenc" # audio: mimi encoder as voice-prompt encoder
POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder
HUNYUANVL = "hunyuanvl"
PARAKEET = "parakeet" # audio
MINIMAXM3 = "minimax_m3"
src/llama-arch.cpp
+1
View File
@@ -145,6 +145,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
{ LLM_ARCH_QWEN3TTS, "qwen3tts" },
{ LLM_ARCH_POCKETTTS, "pockettts" },
{ LLM_ARCH_UNKNOWN, "(unknown)" },
};
src/llama-arch.h
+1
View File
@@ -150,6 +150,7 @@ enum llm_arch {
LLM_ARCH_DFLASH,
LLM_ARCH_NANBEIGE,
LLM_ARCH_QWEN3TTS,
LLM_ARCH_POCKETTTS,
LLM_ARCH_UNKNOWN,
};
src/llama-model.cpp
+3
View File
@@ -114,6 +114,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_qwen3vlmoe(params);
case LLM_ARCH_QWEN3TTS:
return new llama_model_qwen3tts(params);
case LLM_ARCH_POCKETTTS:
return new llama_model_pockettts(params);
case LLM_ARCH_PHI2:
return new llama_model_phi2(params);
case LLM_ARCH_PHI3:
@@ -2610,6 +2612,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
case LLM_ARCH_NANBEIGE:
case LLM_ARCH_POCKETTTS:
return LLAMA_ROPE_TYPE_NORM;
// the pairs of head values are offset by n_rot/2
src/models/models.h
+13
View File
@@ -697,6 +697,19 @@ struct llama_model_gpt2 : public llama_model_base {
};
struct llama_model_pockettts : public llama_model_base {
llama_model_pockettts(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
void load_arch_tensors(llama_model_loader & ml) override;
struct graph : public llm_graph_context {
graph(const llama_model & model, const llm_graph_params & params);
};
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
struct llama_model_codeshell : public llama_model_base {
llama_model_codeshell(const struct llama_model_params & params) : llama_model_base(params) {}
void load_arch_hparams(llama_model_loader & ml) override;
src/models/pockettts.cpp
+146
View File
@@ -0,0 +1,146 @@
#include "models.h"
// backbone of the pocket-tts CALM pipeline: the "text" side of a flow language model.
// it has no lm_head, the audio latents are produced by the flow net inside the mmproj
void llama_model_pockettts::load_arch_hparams(llama_model_loader & ml) {
ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps);
switch (hparams.n_layer()) {
case 6: type = LLM_TYPE_109M; break;
case 24: type = LLM_TYPE_335M; break;
default: type = LLM_TYPE_UNKNOWN;
}
}
void llama_model_pockettts::load_arch_tensors(llama_model_loader &) {
LLAMA_LOAD_LOCALS;
tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
output_norm_b = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "bias"), {n_embd}, 0);
// no output head, the logits are unused; reuse the embedding table so a sampler can still run
output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED);
for (int i = 0; i < n_layer; ++i) {
auto & layer = layers[i];
layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
layer.attn_norm_b = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "bias", i), {n_embd}, 0);
create_tensor_qkv(layer, i, n_embd, n_embd, n_embd_gqa, n_embd_gqa, TENSOR_NOT_REQUIRED);
layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd, n_embd}, 0);
layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
layer.ffn_norm_b = create_tensor(tn(LLM_TENSOR_FFN_NORM, "bias", i), {n_embd}, 0);
layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
}
}
std::unique_ptr<llm_graph_context> llama_model_pockettts::build_arch_graph(const llm_graph_params & params) const {
return std::make_unique<graph>(*this, params);
}
llama_model_pockettts::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
const int64_t n_embd_head = hparams.n_embd_head_v();
GGML_ASSERT(n_embd_head == hparams.n_embd_head_k());
GGML_ASSERT(n_embd_head == n_rot);
ggml_tensor * cur;
ggml_tensor * inpL;
inpL = build_inp_embd(model.tok_embd);
ggml_tensor * inp_pos = build_inp_pos();
auto * inp_attn = build_attn_inp_kv();
ggml_tensor * inp_out_ids = build_inp_out_ids();
for (int il = 0; il < n_layer; ++il) {
cur = build_norm(inpL,
model.layers[il].attn_norm,
model.layers[il].attn_norm_b,
LLM_NORM, il);
cb(cur, "attn_norm", il);
// self-attention
{
auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur,
n_embd_head, n_head, n_head_kv, il);
Qcur = ggml_rope_ext(
ctx0, Qcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
Kcur = ggml_rope_ext(
ctx0, Kcur, inp_pos, nullptr,
n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
ext_factor, attn_factor, beta_fast, beta_slow
);
cb(Qcur, "Qcur", il);
cb(Kcur, "Kcur", il);
cb(Vcur, "Vcur", il);
cur = build_attn(inp_attn,
model.layers[il].wo, NULL, model.layers[il].wo_s,
Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, 1.0f/sqrtf(float(n_embd_head)), il);
}
if (il == n_layer - 1 && inp_out_ids) {
cur = ggml_get_rows(ctx0, cur, inp_out_ids);
inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
}
ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
cb(ffn_inp, "ffn_inp", il);
// FF
{
cur = build_norm(ffn_inp,
model.layers[il].ffn_norm,
model.layers[il].ffn_norm_b,
LLM_NORM, il);
cb(cur, "ffn_norm", il);
cur = build_ffn(cur,
model.layers[il].ffn_up, NULL, NULL,
NULL, NULL, NULL,
model.layers[il].ffn_down, NULL, NULL,
NULL,
LLM_FFN_GELU, LLM_FFN_SEQ, il);
cb(cur, "ffn_out", il);
}
cur = ggml_add(ctx0, cur, ffn_inp);
cur = build_cvec(cur, il);
cb(cur, "l_out", il);
// input for next layer
inpL = cur;
}
cur = build_norm(inpL,
model.output_norm,
model.output_norm_b,
LLM_NORM, -1);
cb(cur, "result_norm", -1);
res->t_embd = cur;
cur = build_lora_mm(model.output, cur, model.output_s);
cb(cur, "result_output", -1);
res->t_logits = cur;
ggml_build_forward_expand(gf, cur);
}