mtmd: support pocket-tts (#26871)

* adapt the api

* text model ok

* working impl, need verify and clean up

* mtmd: build the pocket-tts transposed convolutions as GEMM + col2im

ggml_conv_transpose_1d has no grouped mode, so the depthwise upsample
was built as one convolution and one concat per channel, which floods
the graph with small nodes and makes kernel launches dominate the
decoder.

Fold both cases into the column form the seanet decoder already needs:
the general case reshapes the kernel to [IC, K * OC] and matmuls it
with the input, the depthwise case batches a matmul over the channels
so a step scales its own kernel. A single col2im_1d then scatter-adds
the columns back to the signal, with the same shape as before, so the
overlap-add tail, the streaming state and the bias are untouched.

Generation time per frame drops by 80% on CUDA and by 50% on CPU. The
output matches the previous implementation sample for sample, with a
correlation of 0.999994 and identical frame counts.

* flow_temp +  frames_after_eos

* chunking

* mtmd: carry the remaining pocket-tts per-pack settings

The language packs also tune the end-of-speech padding and the padding
of short prompts, next to the temperature already carried in the
mmproj: french_24l asks for 8 tail frames instead of the guessed 3,
english_2026-01 asks for short prompts to be padded with spaces.

Write both in the mmproj as clip.gen.audio.frames_after_eos and
clip.gen.audio.pad_short_text, keyed on the pack in the conversion
script like the temperature. The loader keeps them optional, so a
mmproj without them behaves as before. Map semicolons to commas for
every pack instead, the reference only asks for it on three of them and
it costs nothing elsewhere.

Existing mmproj files must be converted again to carry the two keys.

On a long french text the port now lands within 2% of the reference:
22.96s against 23.44s, with the same peak level and the same amount of
silence.

* clip.gen.audio.model_variant

* clean up code comments

* nit: drop the dead flow_temp hparam, the pack table holds the default

* update docs

* address security problems

* less invasive base.py

* lint

* add mtmd_gen_inp_default

* add docs

* rm gen_flow_temp

---------

Co-authored-by: Pascal <admin@serveurperso.com>
This commit is contained in:
Xuan-Son Nguyen
2026-08-11 14:18:30 +02:00
committed by GitHub
co-authored by Pascal
parent 8d274dd7c6
commit 6e62ba5384
29 changed files with 2458 additions and 73 deletions
+112
View File
@@ -407,6 +407,8 @@ class Keys:
class ClipGenAudio:
PROJECTOR_TYPE = "clip.gen.audio.projector_type" # for mixed modality models
# name of the weight variant, for settings that are not in the checkpoint
MODEL_VARIANT = "clip.gen.audio.model_variant"
EMBEDDING_LENGTH = "clip.gen.audio.embedding_length"
FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length"
BLOCK_COUNT = "clip.gen.audio.block_count"
@@ -581,6 +583,7 @@ class MODEL_ARCH(IntEnum):
MELLUM = auto()
NANBEIGE = auto()
QWEN3TTS = auto()
POCKETTTS = auto()
class VISION_PROJECTOR_TYPE(IntEnum):
@@ -1040,6 +1043,38 @@ 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_FFN_SCALE_LS = auto() # layer scale (gamma) on the FFN 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()
@@ -1255,6 +1290,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] = {
@@ -1709,6 +1745,37 @@ 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}.ls1",
MODEL_TENSOR.A_ENC_FFN_SCALE_LS: "a.blk.{bid}.ls2",
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",
@@ -2020,6 +2087,37 @@ 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_FFN_SCALE_LS,
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,
@@ -4903,6 +5001,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
@@ -5179,6 +5289,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"