From 7717dd0d02804ffeff08d9b40796a64afb7c3100 Mon Sep 17 00:00:00 2001 From: Hans Date: Tue, 8 Sep 2026 22:39:41 +0800 Subject: [PATCH 1/4] mtmd : add NeMo Nano Codec decoder Assisted-by: Codex --- conversion/__init__.py | 1 + conversion/nemo_nano_codec.py | 133 ++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 31 ++++++ tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 2 + tools/mtmd/clip-model.h | 19 ++++ tools/mtmd/clip.cpp | 98 +++++++++++++++++-- tools/mtmd/models/models.h | 10 ++ tools/mtmd/models/nemo-nano-codec.cpp | 80 ++++++++++++++++ tools/mtmd/mtmd.cpp | 27 +++++- tools/mtmd/mtmd.h | 3 + tools/tts/README.md | 16 ++++ 12 files changed, 410 insertions(+), 11 deletions(-) create mode 100644 conversion/nemo_nano_codec.py create mode 100644 tools/mtmd/models/nemo-nano-codec.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 4d58bcd1060e..6d4e3f886e35 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -285,6 +285,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "NemoNanoCodecModel": "nemo_nano_codec", "AudioFlamingo3ForConditionalGeneration": "ultravox", "CogVLMForCausalLM": "cogvlm", "DeepseekOCR2ForCausalLM": "deepseek", diff --git a/conversion/nemo_nano_codec.py b/conversion/nemo_nano_codec.py new file mode 100644 index 000000000000..2f75e65d8e7a --- /dev/null +++ b/conversion/nemo_nano_codec.py @@ -0,0 +1,133 @@ +# Copyright (c) 2026 codec.cpp contributors +# SPDX-License-Identifier: MIT + +from __future__ import annotations + +import io +from pathlib import Path +import re +import tarfile + +import torch + +from .base import ModelBase, MmprojModel, gguf + + +_ARCHIVE = "nemo-nano-codec-22khz-0.6kbps-12.5fps.nemo" + + +def _read_member(archive: tarfile.TarFile, name: str) -> bytes: + for member in archive.getmembers(): + if member.name.removeprefix("./") == name and member.isfile(): + with archive.extractfile(member) as f: + return f.read() + raise ValueError(f"NeMo archive is missing {name}") + + +@ModelBase.register_hparams_loader(lambda path: (path / _ARCHIVE).is_file()) +def _load_hparams(path: Path) -> dict: + import yaml + + with tarfile.open(path / _ARCHIVE) as archive: + config = yaml.safe_load(_read_member(archive, "model_config.yaml")) + decoder = config["audio_decoder"] + quantizer = config["vector_quantizer"] + expected = { + "up_sample_rates": [7, 7, 6, 3, 2], "input_dim": 16, "base_channels": 864, + "activation": "half_snake", "output_activation": "clamp", "pad_mode": "zeros", + "n_groups_equal_to_out_channels": True, + } + defaults = {"in_kernel_size": 7, "out_kernel_size": 3, + "resblock_kernel_sizes": [3, 7, 11], "resblock_dilation_sizes": [1, 3, 5]} + if (not decoder.get("_target_", "").endswith(".CausalHiFiGANDecoder") + or any(decoder.get(k) != v for k, v in expected.items()) + or any(decoder.get(k, v) != v for k, v in defaults.items()) + or not quantizer.get("_target_", "").endswith(".GroupFiniteScalarQuantizer") + or config.get("sample_rate") != 22050 or config.get("samples_per_frame") != 1764 + or quantizer.get("num_groups") != 4 or quantizer.get("num_levels_per_group") != [9, 8, 8, 7]): + raise ValueError("Only NeMo Nano Codec 22 kHz / 12.5 fps is supported") + return { + "architectures": ["NemoNanoCodecModel"], "hidden_size": 16, + "audio_config": {"num_hidden_layers": 5}, + } + + +@ModelBase.register("NemoNanoCodecModel") +@ModelBase.example("nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps") +class NemoNanoCodecModel(MmprojModel): + has_vision_encoder = False + has_audio_encoder = False + + def set_gguf_parameters(self): + self.gguf_writer.add_file_type(self.ftype) + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.NEMO_NANO_CODEC) + self.gguf_writer.add_gen_audio_projection_dim(16) + self.gguf_writer.add_gen_audio_embedding_length(864) + self.gguf_writer.add_gen_audio_feed_forward_length(864) + self.gguf_writer.add_gen_audio_block_count(5) + self.gguf_writer.add_gen_audio_head_count(1) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(1e-5) + + def get_tensors(self): + with tarfile.open(self.dir_model / _ARCHIVE) as archive: + state = torch.load(io.BytesIO(_read_member(archive, "model_weights.ckpt")), map_location="cpu", weights_only=True) + state = state.get("state_dict", state) + for name, value in state.items(): + if not name.startswith("audio_decoder.") or name.endswith(".weight_v"): + continue + if name.endswith(".weight_g"): + weight = state[name.removesuffix("weight_g") + "weight_v"].float() + norm = torch.linalg.vector_norm(weight.flatten(1), dim=1).reshape(-1, 1, 1) + value = weight * (value.float().reshape(-1, 1, 1) / norm) + name = name.removesuffix("weight_g") + "weight" + yield name, value + + # All four FSQ groups use the same fixed codebook. + levels = torch.tensor([9, 8, 8, 7]) + bases = torch.tensor([1, 9, 72, 576]) + scale = levels // 2 + codes = (torch.arange(4032)[:, None] // bases) % levels + yield "fsq_codebook", (codes - scale).float() / scale + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if new_name.endswith(".alpha") or name == "fsq_codebook": + return gguf.GGMLQuantizationType.F32 + # Convolution uses the standard ggml F16 im2col path. + if new_name.endswith(".weight") and n_dims == 3: + return gguf.GGMLQuantizationType.F16 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch, name, bid): + T = gguf.MODEL_TENSOR + suffix = "." + name.rsplit(".", 1)[-1] + name = name.removeprefix("audio_decoder.") + if name == "fsq_codebook": + yield self.format_tensor_name(T.A_GEN_WAV_HIFIGAN_CODEBOOK), data_torch + return + if name.startswith("pre_conv."): + tensor = T.A_GEN_WAV_HIFIGAN_PRE + elif name.startswith("post_conv."): + tensor = T.A_GEN_WAV_HIFIGAN_POST + elif name.startswith("post_activation."): + tensor = T.A_GEN_WAV_HIFIGAN_POST_ACT + elif match := re.fullmatch(r"activations\.(\d+)\.activation\.snake_act\.alpha", name): + tensor, bid = T.A_GEN_WAV_HIFIGAN_UP_ACT, int(match[1]) + elif match := re.fullmatch(r"up_sample_conv_layers\.(\d+)\.conv\.(weight|bias)", name): + tensor, bid = T.A_GEN_WAV_HIFIGAN_UP, int(match[1]) + if suffix == ".weight": + # Grouped transposed convolution: [IC, 1, K] -> [OC, K, 2]. + data_torch = data_torch.reshape(-1, 2, data_torch.shape[-1]).transpose(1, 2).contiguous() + elif match := re.fullmatch(r"res_layers\.(\d+)\.res_blocks\.(\d+)\.res_blocks\.(\d+)\.(input_conv|skip_conv|input_activation|skip_activation)\..+", name): + bid = int(match[1]) * 9 + int(match[2]) * 3 + int(match[3]) + tensor = { + "input_conv": T.A_GEN_WAV_HIFIGAN_RES_CONV1, + "skip_conv": T.A_GEN_WAV_HIFIGAN_RES_CONV2, + "input_activation": T.A_GEN_WAV_HIFIGAN_RES_ACT1, + "skip_activation": T.A_GEN_WAV_HIFIGAN_RES_ACT2, + }[match[4]] + else: + raise ValueError(f"Unexpected NeMo decoder tensor: {name}") + if suffix == ".alpha": + data_torch = data_torch.flatten() + yield self.format_tensor_name(tensor, bid, suffix), data_torch diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d3a639f374c0..d497bd1abea2 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1113,6 +1113,16 @@ class MODEL_TENSOR(IntEnum): A_GEN_WAV_UP_PW1 = auto() # ConvNeXt pointwise conv 1 (expand) A_GEN_WAV_UP_PW2 = auto() # ConvNeXt pointwise conv 2 (project) A_GEN_WAV_UP_GAMMA = auto() # ConvNeXt layer scale + A_GEN_WAV_HIFIGAN_CODEBOOK = auto() + A_GEN_WAV_HIFIGAN_PRE = auto() + A_GEN_WAV_HIFIGAN_POST = auto() + A_GEN_WAV_HIFIGAN_POST_ACT = auto() + A_GEN_WAV_HIFIGAN_UP_ACT = auto() + A_GEN_WAV_HIFIGAN_UP = auto() + A_GEN_WAV_HIFIGAN_RES_CONV1 = auto() + A_GEN_WAV_HIFIGAN_RES_CONV2 = auto() + A_GEN_WAV_HIFIGAN_RES_ACT1 = auto() + A_GEN_WAV_HIFIGAN_RES_ACT2 = auto() A_GEN_WAV_DAC_ENTRY = auto() # DAC conv_pre A_GEN_WAV_DAC_UP_SNAKE = auto() # DAC per-block SnakeBeta before the upsample conv A_GEN_WAV_DAC_UP_CONV = auto() # DAC per-block causal ConvTranspose1d @@ -1863,6 +1873,16 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_UP_PW1: "a.gen.wav.up.blk.{bid}.pw1", MODEL_TENSOR.A_GEN_WAV_UP_PW2: "a.gen.wav.up.blk.{bid}.pw2", MODEL_TENSOR.A_GEN_WAV_UP_GAMMA: "a.gen.wav.up.blk.{bid}.gamma", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_CODEBOOK: "a.gen.wav.hifigan.codebook", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_PRE: "a.gen.wav.hifigan.pre", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_POST: "a.gen.wav.hifigan.post", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_POST_ACT: "a.gen.wav.hifigan.post_act", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_UP_ACT: "a.gen.wav.hifigan.up.{bid}.act", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_UP: "a.gen.wav.hifigan.up.{bid}.conv", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_CONV1: "a.gen.wav.hifigan.res.{bid}.conv1", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_CONV2: "a.gen.wav.hifigan.res.{bid}.conv2", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_ACT1: "a.gen.wav.hifigan.res.{bid}.act1", + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_ACT2: "a.gen.wav.hifigan.res.{bid}.act2", MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY: "a.gen.wav.dac.entry", MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE: "a.gen.wav.dac.blk.{bid}.snake", MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV: "a.gen.wav.dac.blk.{bid}.conv", @@ -2220,6 +2240,16 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.A_GEN_WAV_UP_PW1, MODEL_TENSOR.A_GEN_WAV_UP_PW2, MODEL_TENSOR.A_GEN_WAV_UP_GAMMA, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_CODEBOOK, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_PRE, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_POST, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_POST_ACT, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_UP_ACT, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_UP, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_CONV1, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_CONV2, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_ACT1, + MODEL_TENSOR.A_GEN_WAV_HIFIGAN_RES_ACT2, MODEL_TENSOR.A_GEN_WAV_DAC_ENTRY, MODEL_TENSOR.A_GEN_WAV_DAC_UP_SNAKE, MODEL_TENSOR.A_GEN_WAV_DAC_UP_CONV, @@ -5812,6 +5842,7 @@ class VisionProjectorType: 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 + NEMO_NANO_CODEC = "nemo_nano_codec" # audio generation: causal HiFiGAN decoder POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder HUNYUANVL = "hunyuanvl" PARAKEET = "parakeet" # audio diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 907468e87ec7..87e707f97759 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -62,6 +62,7 @@ add_library(mtmd models/qwen3tts-gen.cpp models/pockettts-seanet.cpp models/pockettts-spkenc.cpp + models/nemo-nano-codec.cpp models/pockettts-gen.cpp models/step3vl.cpp models/siglip.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 72148a4d9a9b..d3973ed6c049 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -502,6 +502,7 @@ enum projector_type { PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, PROJECTOR_TYPE_POCKETTTS_SPKENC, + PROJECTOR_TYPE_NEMO_NANO_CODEC, PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_UNKNOWN, @@ -567,6 +568,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, + { PROJECTOR_TYPE_NEMO_NANO_CODEC, "nemo_nano_codec"}, { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index f737ccc24527..3bb8f4b4b0fa 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -485,6 +485,24 @@ struct clip_flow_net { std::vector blocks; }; +// NeMo Nano Codec 22 kHz / 12.5 fps: grouped FSQ codes -> raw PCM. +struct clip_nemo_nano_codec { + static constexpr int n_groups = 4; + static constexpr int codebook_size = 4032; + struct conv { + ggml_tensor * w = nullptr; + ggml_tensor * b = nullptr; + ggml_tensor * alpha = nullptr; + }; + struct stage { + conv up; + conv res[3][3][2]; + }; + ggml_tensor * codebook = nullptr; + conv pre, post; + stage stages[5]; +}; + // qwen3tts code2wav: RVQ codes -> raw PCM struct clip_code2wav { // "upsample" stage: one ConvNeXt block plus the causal ConvTranspose1d before it @@ -786,6 +804,7 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; + clip_nemo_nano_codec nemo; // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path) clip_seanet seanet; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index cd6421def528..b8ab70871604 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1116,6 +1116,11 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const const int n_frames = params && params->feats ? (int) (params->feats->size() / n_latent) : 1; builder = std::make_unique(ctx, img, gen_process, n_step, n_frames); } break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + { + const int n_frames = params && params->codes ? (int) (params->codes->size() / clip_nemo_nano_codec::n_groups) : 1; + builder = std::make_unique(ctx, img, n_frames); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; @@ -1831,6 +1836,12 @@ struct clip_model_loader { hparams.audio_window_len = 1024; hparams.audio_hop_len = 256; } break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + { + if (hparams.n_layer != 5 || hparams.n_embd != 864 || hparams.projection_dim != 16 || hparams.n_head != 1) { + throw std::runtime_error("invalid NeMo Nano Codec dimensions"); + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // TODO: hardcoded for now, read from code_predictor_config instead @@ -2239,7 +2250,8 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && - model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN); + model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN && + model.proj_type != PROJECTOR_TYPE_NEMO_NANO_CODEC); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -3026,6 +3038,57 @@ struct clip_model_loader { layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight")); } } break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + { + auto tensor = [&](const std::string & name, std::initializer_list dims, ggml_type type = GGML_TYPE_F32) { + auto * t = get_tensor("a.gen.wav.hifigan." + name); + if (t->type != type) { + throw std::runtime_error("invalid NeMo tensor type: " + name); + } + int d = 0; + for (auto dim : dims) { + if (t->ne[d++] != dim) { + throw std::runtime_error(string_format("invalid NeMo tensor %s dimension %d: expected %lld, got %lld", name.c_str(), d - 1, (long long) dim, (long long) t->ne[d - 1])); + } + } + for (; d < GGML_MAX_DIMS; ++d) { + if (t->ne[d] != 1) { + throw std::runtime_error("invalid NeMo tensor rank: " + name); + } + } + return t; + }; + auto conv = [&](clip_nemo_nano_codec::conv & c, const std::string & name, int k, int ic, int oc) { + c.w = tensor(name + ".weight", {k, ic, oc}, GGML_TYPE_F16); + c.b = tensor(name + ".bias", {oc}); + }; + auto & m = model.nemo; + m.codebook = tensor("codebook.weight", {4, m.codebook_size}); + conv(m.pre, "pre", 7, 16, 864); + const int rates[] = {7, 7, 6, 3, 2}; + const int kernels[] = {3, 7, 11}; + int channels = 864; + for (int stage = 0; stage < 5; ++stage) { + auto & s = m.stages[stage]; + const std::string prefix = "up." + std::to_string(stage); + s.up.alpha = tensor(prefix + ".act.alpha", {channels / 2}); + s.up.w = tensor(prefix + ".conv.weight", {2, 2 * rates[stage], channels / 2}, GGML_TYPE_F16); + channels /= 2; + s.up.b = tensor(prefix + ".conv.bias", {channels}); + for (int block = 0; block < 3; ++block) { + for (int layer = 0; layer < 3; ++layer) { + const std::string res = "res." + std::to_string(stage * 9 + block * 3 + layer); + for (int j = 0; j < 2; ++j) { + auto & c = s.res[block][layer][j]; + conv(c, res + ".conv" + std::to_string(j + 1), kernels[block], channels, channels); + c.alpha = tensor(res + ".act" + std::to_string(j + 1) + ".alpha", {channels / 2}); + } + } + } + } + conv(m.post, "post", 3, channels, 1); + m.post.alpha = tensor("post_act.alpha", {channels / 2}); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // code_predictor, proj_in is absent when the talker and the predictor share the hidden size @@ -4345,6 +4408,10 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { // pooling gives one speaker embedding, whatever the clip length is n_patches = 1; } break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + { + n_patches = 1; // GEN_WAV sizes the graph from codes, not the placeholder batch. + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // one hidden-state vector fed back to the talker per call @@ -5266,6 +5333,17 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { { // do nothing } break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + { + const int n_frames = (int) (params->codes->size() / clip_nemo_nano_codec::n_groups); + std::vector codes(params->codes->size()); + for (int f = 0; f < n_frames; ++f) { + for (int g = 0; g < clip_nemo_nano_codec::n_groups; ++g) { + codes[g * n_frames + f] = (*params->codes)[f * clip_nemo_nano_codec::n_groups + g]; + } + } + set_input_i32("inp_codes", codes); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { @@ -5820,13 +5898,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { out_audio.resize(ggml_nelements(audio)); ggml_backend_tensor_get(audio, out_audio.data(), 0, ggml_nbytes(audio)); - // drop the tail audio that comes from the code-0 rear padding - const int64_t n_codes = params->codes ? model.gen_code_head_w->ne[2] + 1 : 0; - const int64_t n_frames_w = hparams.wav_tfm_swa; - const int64_t n_frames = params->codes ? (int64_t) params->codes->size() / n_codes : n_frames_w; - if (n_frames < n_frames_w) { - const size_t hop = out_audio.size() / n_frames_w; - out_audio.resize((size_t) n_frames * hop); + // Qwen3-TTS rear-pads codes to one window; drop the audio from that padding. + if (model.proj_type == PROJECTOR_TYPE_QWEN3TTS_GEN && params->codes) { + const int64_t n_codes = model.gen_code_head_w->ne[2] + 1; + const int64_t n_frames_w = hparams.wav_tfm_swa; + const int64_t n_frames = (int64_t) params->codes->size() / n_codes; + if (n_frames < n_frames_w) { + const size_t hop = out_audio.size() / n_frames_w; + out_audio.resize((size_t) n_frames * hop); + } } } if (params->state_out != nullptr) { @@ -5999,6 +6079,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_QWEN3TTS_SPKENC: return ctx->model.mm_fc_w->ne[2]; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + return ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_QWEN3TTS_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; case PROJECTOR_TYPE_POCKETTTS_SPKENC: diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 5945c6d92cb7..b35baec8ecd7 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -12,6 +12,16 @@ * We encourage human contributors to ensure the quality and reliability of the codebase. */ +struct clip_graph_nemo_nano_codec : clip_graph { + clip_graph_nemo_nano_codec(clip_ctx * ctx, const clip_image_f32 & img, int n_frames) + : clip_graph(ctx, img), n_frames(n_frames) {} + int n_frames; + ggml_tensor * half_snake(ggml_tensor * x, ggml_tensor * alpha) const; + ggml_tensor * conv1d(ggml_tensor * x, const clip_nemo_nano_codec::conv & c, int dilation = 1) const; + ggml_tensor * upsample(ggml_tensor * x, const clip_nemo_nano_codec::conv & c, int stride) const; + ggml_cgraph * build() override; +}; + struct clip_graph_siglip : clip_graph { clip_graph_siglip(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/nemo-nano-codec.cpp b/tools/mtmd/models/nemo-nano-codec.cpp new file mode 100644 index 000000000000..768ad995c22d --- /dev/null +++ b/tools/mtmd/models/nemo-nano-codec.cpp @@ -0,0 +1,80 @@ +// Copyright (c) 2026 codec.cpp contributors +// SPDX-License-Identifier: MIT + +#include "models.h" + +ggml_tensor * clip_graph_nemo_nano_codec::half_snake(ggml_tensor * x, ggml_tensor * alpha) const { + const int64_t half = x->ne[1] / 2; + ggml_tensor * left = ggml_view_2d(ctx0, x, x->ne[0], half, x->nb[1], 0); + ggml_tensor * right = ggml_view_2d(ctx0, x, x->ne[0], x->ne[1] - half, x->nb[1], half * x->nb[1]); + ggml_tensor * a = ggml_reshape_2d(ctx0, alpha, 1, half); + ggml_tensor * sine = ggml_sin(ctx0, ggml_mul(ctx0, left, a)); + left = ggml_add(ctx0, left, ggml_div(ctx0, ggml_sqr(ctx0, sine), ggml_scale_bias(ctx0, a, 1.0f, 1e-9f))); + right = ggml_leaky_relu(ctx0, right, 0.01f, false); + return ggml_concat(ctx0, left, right, 1); +} + +ggml_tensor * clip_graph_nemo_nano_codec::conv1d(ggml_tensor * x, const clip_nemo_nano_codec::conv & c, int dilation) const { + const int64_t padding = (c.w->ne[0] - 1) * dilation; + x = ggml_pad_ext(ctx0, x, padding, 0, 0, 0, 0, 0, 0, 0); + x = ggml_conv_1d(ctx0, c.w, x, 1, 0, dilation); + x = ggml_reshape_2d(ctx0, x, x->ne[0], x->ne[1]); + return ggml_add(ctx0, x, ggml_reshape_2d(ctx0, c.b, 1, c.b->ne[0])); +} + +ggml_tensor * clip_graph_nemo_nano_codec::upsample(ggml_tensor * x, const clip_nemo_nano_codec::conv & c, int stride) const { + const int64_t frames = x->ne[0]; + const int64_t channels = x->ne[1] / 2; + const int64_t kernel = c.w->ne[1]; + // Each output channel sums a pair of input channels: [2, K, OC] x [2, T, OC]. + x = ggml_reshape_3d(ctx0, x, frames, 2, channels); + x = ggml_cont(ctx0, ggml_permute(ctx0, x, 1, 0, 2, 3)); + x = ggml_mul_mat(ctx0, c.w, x); + x = ggml_cont(ctx0, ggml_permute(ctx0, x, 0, 2, 1, 3)); + x = ggml_reshape_2d(ctx0, x, kernel * channels, frames); + x = ggml_col2im_1d(ctx0, x, stride, channels, 0); + // Causal transposed convolution discards the right overlap tail. + x = ggml_cont(ctx0, ggml_view_2d(ctx0, x, frames * stride, channels, x->nb[1], 0)); + return ggml_add(ctx0, x, ggml_reshape_2d(ctx0, c.b, 1, channels)); +} + +ggml_cgraph * clip_graph_nemo_nano_codec::build() { + const auto & m = model.nemo; + ggml_tensor * codes = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_frames, m.n_groups); + ggml_set_name(codes, "inp_codes"); + ggml_set_input(codes); + ggml_tensor * cur = nullptr; + for (int group = 0; group < m.n_groups; ++group) { + auto * indices = ggml_view_1d(ctx0, codes, n_frames, group * codes->nb[1]); + auto * embd = ggml_get_rows(ctx0, m.codebook, indices); + cur = cur ? ggml_concat(ctx0, cur, embd, 0) : embd; + } + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = conv1d(cur, m.pre); + const int rates[] = {7, 7, 6, 3, 2}; + const int dilations[] = {1, 3, 5}; + for (int stage = 0; stage < 5; ++stage) { + const auto & s = m.stages[stage]; + cur = upsample(half_snake(cur, s.up.alpha), s.up, rates[stage]); + ggml_tensor * sum = nullptr; + for (int block = 0; block < 3; ++block) { + ggml_tensor * residual = cur; + for (int layer = 0; layer < 3; ++layer) { + const auto & first = s.res[block][layer][0]; + const auto & second = s.res[block][layer][1]; + auto * h = conv1d(half_snake(residual, first.alpha), first, dilations[layer]); + h = conv1d(half_snake(h, second.alpha), second); + residual = ggml_add(ctx0, residual, h); + } + sum = sum ? ggml_add(ctx0, sum, residual) : residual; + } + cur = ggml_scale(ctx0, sum, 1.0f / 3.0f); + cb(cur, "nemo_stage", stage); + } + cur = conv1d(half_snake(cur, m.post.alpha), m.post); + cur = ggml_clamp(ctx0, cur, -1.0f, 1.0f); + ggml_set_name(cur, "out_audio"); + ggml_set_output(cur); + ggml_build_forward_expand(gf, cur); + return gf; +} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 00ecadcf4dfe..c370b9ae6500 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -494,7 +494,7 @@ struct mtmd_context { std::string media_marker; const int n_embd_text = -1; // -1 means llm context not provided, skip checking this const llama_vocab * vocab = nullptr; // can be nullptr if text_model is not provided - mtmd_pos_type pos_type; + mtmd_pos_type pos_type = MTMD_POS_TYPE_NORMAL; // these are not token, but strings used to mark the beginning and end of image/audio embeddings std::string img_beg; @@ -583,7 +583,7 @@ struct mtmd_context { ctx_v = res.ctx_v; ctx_a = res.ctx_a; ctx_gen_a = res.ctx_gen_a; - if (!ctx_v && !ctx_a) { + if (!ctx_v && !ctx_a && !ctx_gen_a) { throw std::runtime_error(string_format("Failed to load CLIP model from %s\n", mmproj_fname)); } @@ -600,7 +600,7 @@ struct mtmd_context { // since we already validate n_embd of vision and audio mmproj, // we can safely assume that they are the same - int n_embd_clip = clip_n_mmproj_embd(ctx_v ? ctx_v : ctx_a); + int n_embd_clip = (ctx_v || ctx_a) ? clip_n_mmproj_embd(ctx_v ? ctx_v : ctx_a) : n_embd_text; if (n_embd_text > 0 && n_embd_text != n_embd_clip) { throw std::runtime_error(string_format( "mismatch between text model (n_embd = %d) and mmproj (n_embd = %d)\n" @@ -1881,6 +1881,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { info.type = MTMD_GEN_AUDIO_TYPE_QWEN3TTS; info.sample_rate = 24000; break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + info.type = MTMD_GEN_AUDIO_TYPE_NEMO_NANO_CODEC; + info.sample_rate = 22050; + break; case PROJECTOR_TYPE_POCKETTTS_GEN: info.type = MTMD_GEN_AUDIO_TYPE_POCKETTTS; info.sample_rate = 24000; @@ -1907,6 +1911,9 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { inp.top_p = 1.0f; inp.temp = 0.9f; // TODO: handle this on graph break; + case PROJECTOR_TYPE_NEMO_NANO_CODEC: + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; + break; case PROJECTOR_TYPE_POCKETTTS_GEN: // https://github.com/kyutai-labs/pocket-tts/blob/main/pocket_tts/default_parameters.py inp.top_k = 50; @@ -1928,6 +1935,20 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, const mtmd_gen_in *out = {}; + if (clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_NEMO_NANO_CODEC) { + if (inp->type != MTMD_GEN_PROCESS_TYPE_GEN_WAV || !inp->codes || inp->n_codes == 0 || + inp->n_codes % 4 != 0 || inp->n_codes > 4 * 128 || inp->feats || inp->n_feats || inp->state_data || inp->state_size) { + LOG_ERR("%s: NeMo Nano Codec requires 1 to 128 frames of four codes, without features or state\n", __func__); + return 1; + } + for (size_t i = 0; i < inp->n_codes; ++i) { + if (inp->codes[i] < 0 || inp->codes[i] >= 4032) { + LOG_ERR("%s: NeMo Nano Codec code out of range at %zu\n", __func__, i); + return 1; + } + } + } + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index c2de26eeee2c..bf813e91608d 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -371,6 +371,7 @@ enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_NONE, // not supported MTMD_GEN_AUDIO_TYPE_QWEN3TTS, MTMD_GEN_AUDIO_TYPE_POCKETTTS, + MTMD_GEN_AUDIO_TYPE_NEMO_NANO_CODEC, }; struct mtmd_gen_audio_info { @@ -402,6 +403,8 @@ struct mtmd_gen_inp { // for MTMD_GEN_PROCESS_TYPE_GEN_WAV // pass either codes (discrete) or feats (continuous), depending on the pipeline + // NeMo Nano Codec: frame-major [n_frames, 4], 1 to 128 frames, codes in [0, 4031]. + // It accepts no feats or state and returns n_frames * 1764 samples at 22050 Hz. int32_t * codes; size_t n_codes; const float * feats; diff --git a/tools/tts/README.md b/tools/tts/README.md index 1b08d5ef3218..b2f55add7f03 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -57,3 +57,19 @@ The [upstream repository](https://huggingface.co/kyutai/pocket-tts) holds one co python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pocket-tts.gguf python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf ``` + +## NeMo Nano Codec decoder (MTMD API) + +The [22 kHz / 0.6 kbps / 12.5 fps variant](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps) is supported as a standalone codes-to-audio decoder. It is not a text-to-speech model and cannot be used directly with `llama-tts -p`. + +Place `nemo-nano-codec-22khz-0.6kbps-12.5fps.nemo` in a local directory and convert it: + +```sh +python convert_hf_to_gguf.py path/to/nemo-nano-codec --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec.gguf +``` + +The converter reads the architecture from the archive's `model_config.yaml`; there is no text backbone or tokenizer to convert. Convolution weights are stored as F16 even with `--outtype f32`, to use the existing ggml convolution path. FSQ codebook and activation parameters remain F32. + +Load the mmproj with `mtmd_init_from_file(path, nullptr, params)`. Pass `MTMD_GEN_PROCESS_TYPE_GEN_WAV` to `mtmd_gen_audio_process`, with `codes` laid out as `[frame][group]`: four codes per frame, each in `[0, 4031]`. NeMo's `[group, batch, frame]` tokens must be transposed for this interface. Each call accepts 1 to 128 complete frames and returns `n_frames * 1764` mono float samples at 22050 Hz. Copy the output before the next call, and release the context with `mtmd_free`. + +This initial implementation decodes a complete sequence with zero initial context. It does not accept continuous features, persistent state or `GEN_CODE`. Independent chunks do not preserve convolution history. Audio encoding, other Nano Codec variants and a text-generation pipeline are not included. NVIDIA describes this particular variant as intended for fine-tuning with a limited set of speakers, rather than general-purpose audio reconstruction. From 89e116c6788b66bb65852fd484b78e8ab24abe69 Mon Sep 17 00:00:00 2001 From: Hans Date: Tue, 8 Sep 2026 23:15:01 +0800 Subject: [PATCH 2/4] mtmd : add KaniTTS-2 text-to-speech generation Assisted-by: Codex --- conversion/__init__.py | 1 + conversion/lfm2.py | 28 +++++ gguf-py/gguf/constants.py | 3 + src/llama-arch.cpp | 2 + src/llama-arch.h | 1 + src/llama-model.cpp | 4 +- src/models/lfm2.cpp | 24 +++- tools/mtmd/mtmd-helper-gen.cpp | 213 +++++++++++++++++++++++++++++++++ tools/mtmd/mtmd.cpp | 3 +- tools/tts/README.md | 32 +++-- tools/tts/tts.cpp | 18 +-- 11 files changed, 306 insertions(+), 23 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 6d4e3f886e35..40ca7e353130 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -148,6 +148,7 @@ "LLaMAForCausalLM": "llama", "Lfm25AudioTokenizer": "lfm2", "Lfm2BidirectionalModel": "lfm2", + "KaniTTS2ForCausalLM": "lfm2", "Lfm2ForCausalLM": "lfm2", "Lfm2Model": "lfm2", "Lfm2MoeForCausalLM": "lfm2", diff --git a/conversion/lfm2.py b/conversion/lfm2.py index 984f44480649..4b53eeeb88b5 100644 --- a/conversion/lfm2.py +++ b/conversion/lfm2.py @@ -65,6 +65,34 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("KaniTTS2ForCausalLM") +@ModelBase.example("nineninesix/kani-tts-2-en") +class KaniTTS2Model(LFM2Model): + model_arch = gguf.MODEL_ARCH.LFM2 + + def set_gguf_parameters(self): + if (self.hparams.get("tokens_per_frame") != 4 or self.hparams.get("audio_step") != 1.0 + or self.hparams.get("text_vocab_size") != 64400 or self.hparams.get("vocab_size") != 80538 + or not self.hparams.get("use_learnable_rope")): + raise ValueError("Unsupported KaniTTS-2 audio configuration") + super().set_gguf_parameters() + head_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_sections([0, head_dim // 2, 0, 0]) + self.gguf_writer.add_string("lfm2.tts.model", "kani-tts-2") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("model.learnable_rope_layers."): + layer = int(name.split(".")[2]) + alpha = self.hparams["alpha_min"] + (self.hparams["alpha_max"] - self.hparams["alpha_min"]) * torch.sigmoid(data_torch.float()) + head_dim = self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + yield f"blk.{layer}.attn_rope_freqs.weight", (1.0 / alpha).expand(head_dim // 2).clone() + return + if name == "model.speaker_emb_projection.weight": + # Speaker embedding extraction is not part of the unconditional TTS pipeline. + return + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("Lfm2Model", "Lfm2BidirectionalModel") @ModelBase.example("LiquidAI/LFM2.5-ColBERT-350M", "LiquidAI/LFM2.5-Embedding-350M") class LFM2ColBertModel(LFM2Model): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d497bd1abea2..288b74174976 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -663,6 +663,7 @@ class MODEL_TENSOR(IntEnum): HC_HEAD_NORM = auto() # qwen4exp HC_HEAD_DOWN = auto() # qwen4exp HC_HEAD_UP = auto() # qwen4exp + ATTN_ROPE_FREQS = auto() ROPE_FREQS = auto() ROPE_FACTORS_LONG = auto() ROPE_FACTORS_SHORT = auto() @@ -1426,6 +1427,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_HEAD_NORM: "output_hc_norm", # qwen4exp MODEL_TENSOR.HC_HEAD_DOWN: "output_hc_down", # qwen4exp MODEL_TENSOR.HC_HEAD_UP: "output_hc_up", # qwen4exp + MODEL_TENSOR.ATTN_ROPE_FREQS: "blk.{bid}.attn_rope_freqs", MODEL_TENSOR.ROPE_FREQS: "rope_freqs", MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long", MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short", @@ -4917,6 +4919,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_EXP, ], MODEL_ARCH.LFM2: [ + MODEL_TENSOR.ATTN_ROPE_FREQS, MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.TOKEN_EMBD_NORM, MODEL_TENSOR.SHORTCONV_CONV, diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index b5efb7206565..5d453f9e8682 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -429,6 +429,7 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_OUTPUT_NORM, "output_norm" }, { LLM_TENSOR_OUTPUT_NORM_LFM2, "token_embd_norm" }, // fix for wrong tensor name { LLM_TENSOR_OUTPUT, "output" }, + { LLM_TENSOR_ATTN_ROPE_FREQS, "blk.%d.attn_rope_freqs" }, { LLM_TENSOR_ROPE_FREQS, "rope_freqs" }, { LLM_TENSOR_ATTN_NORM, "blk.%d.attn_norm" }, { LLM_TENSOR_ATTN_Q, "blk.%d.attn_q" }, @@ -724,6 +725,7 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_OUTPUT_NORM_LFM2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, {LLM_TENSOR_DEC_OUTPUT_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, {LLM_TENSOR_ENC_OUTPUT_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_ATTN_ROPE_FREQS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ROPE}}, {LLM_TENSOR_ROPE_FREQS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ROPE}}, {LLM_TENSOR_ROPE_FACTORS_LONG, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ROPE}}, {LLM_TENSOR_ROPE_FACTORS_SHORT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ROPE}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index f1d173a57556..5555d3d39e05 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -439,6 +439,7 @@ enum llm_tensor { LLM_TENSOR_OUTPUT, LLM_TENSOR_OUTPUT_NORM, LLM_TENSOR_OUTPUT_NORM_LFM2, // fix for wrong tensor name + LLM_TENSOR_ATTN_ROPE_FREQS, LLM_TENSOR_ROPE_FREQS, LLM_TENSOR_ROPE_FACTORS_LONG, LLM_TENSOR_ROPE_FACTORS_SHORT, diff --git a/src/llama-model.cpp b/src/llama-model.cpp index ffedf89e6718..998fa796c795 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2985,7 +2985,6 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_OPENAI_MOE: case LLM_ARCH_HUNYUAN_DENSE: case LLM_ARCH_HY_V3: - case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_SMALLTHINKER: case LLM_ARCH_SEED_OSS: @@ -3006,6 +3005,9 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_MELLUM: return LLAMA_ROPE_TYPE_NEOX; + case LLM_ARCH_LFM2: + return model->hparams.rope_sections[1] > 0 ? LLAMA_ROPE_TYPE_MROPE : LLAMA_ROPE_TYPE_NEOX; + case LLM_ARCH_DFLASH: // drafts for M-RoPE targets carry rope sections and follow the target's temporal dim if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) { diff --git a/src/models/lfm2.cpp b/src/models/lfm2.cpp index 07b71ccd3a60..8a5a0ad8d0d0 100644 --- a/src/models/lfm2.cpp +++ b/src/models/lfm2.cpp @@ -5,6 +5,7 @@ #include void llama_model_lfm2::load_arch_hparams(llama_model_loader & ml) { + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); ml.get_key(LLM_KV_SHORTCONV_L_CACHE, hparams.n_shortconv_l_cache); ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -67,6 +68,7 @@ void llama_model_lfm2::load_arch_tensors(llama_model_loader &) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); if (!hparams.is_recr(i)) { + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ATTN_ROPE_FREQS, "weight", i), {hparams.n_rot()/2}, TENSOR_NOT_REQUIRED); layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); GGML_ASSERT(n_embd_v_gqa == n_embd_k_gqa); @@ -143,10 +145,21 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ cb(k, "model.layers.{}.self_attn.k_layernorm", il); // RoPE - q = ggml_rope_ext(ctx0, q, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow); - k = ggml_rope_ext(ctx0, k, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, - attn_factor, beta_fast, beta_slow); + auto * factors = model.layers[il].rope_freqs; + if (rope_type == LLAMA_ROPE_TYPE_MROPE) { + // KaniTTS-2 uses separate cache and audio-frame positions. + int sections[4]; + std::copy(hparams.rope_sections.begin(), hparams.rope_sections.end(), sections); + q = ggml_rope_multi(ctx0, q, inp_pos, factors, n_rot, sections, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + k = ggml_rope_multi(ctx0, k, inp_pos, factors, n_rot, sections, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + } else { + q = ggml_rope_ext(ctx0, q, inp_pos, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow); + k = ggml_rope_ext(ctx0, k, inp_pos, factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, + attn_factor, beta_fast, beta_slow); + } cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, @@ -284,7 +297,8 @@ llama_model_lfm2::graph::graph(const llama_model & model, const llm_graph_ cb(cur, "result_norm", -1); res->t_embd = cur; - if (!cparams.embeddings) { + // Audio generation needs both the hidden state and next-token logits. + if (!cparams.embeddings || rope_type == LLAMA_ROPE_TYPE_MROPE) { cur = build_lora_mm(model.output, cur, model.output_s); cb(cur, "result_output", -1); diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 1c58d3ae1959..6ad9dc01f929 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -993,10 +993,223 @@ class pockettts_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector out_buf; }; +// KaniTTS-2 predicts four offset FSQ codes per frame for NeMo Nano Codec. +class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + prompt.clear(); + codes.clear(); + pcm.clear(); + wav.clear(); + pos = 0; + first_audio_pos = -1; + started = false; + stopped = false; + ready = false; + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + char kind[32]; + if (llama_model_meta_val_str(model, "lfm2.tts.model", kind, sizeof(kind)) < 0 || + std::strcmp(kind, "kani-tts-2") != 0 || llama_vocab_n_tokens(vocab) != 80538 || + llama_model_rope_type(model) != LLAMA_ROPE_TYPE_MROPE) { + LOG_ERR("KaniTTS-2 requires a backbone converted from KaniTTS2ForCausalLM\n"); + return 1; + } + if (!inp->prompt || !inp->prompt_len || inp->prompt_len > INT32_MAX - 64 || inp->speaker_ref) { + LOG_ERR("KaniTTS-2 requires text; speaker reference audio is not supported\n"); + return 1; + } + if (inp->out_type != MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM && inp->out_type != MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV) { + return 1; + } + std::string lang = inp->lang ? inp->lang : ""; + if (lang == "en") { + lang = "en_us"; + } + if (!lang.empty() && lang != "en_us" && lang != "en_nyork" && lang != "en_oakl" && + lang != "en_glasg" && lang != "en_bost" && lang != "en_scou") { + LOG_ERR("KaniTTS-2: unsupported language tag '%s'\n", lang.c_str()); + return 1; + } + std::string text(inp->prompt, inp->prompt_len); + if (!lang.empty()) { + text = lang + ": " + text; + } + int n = -llama_tokenize(vocab, text.data(), text.size(), nullptr, 0, true, false); + if (n <= 0 || n + 3 >= (int) llama_n_ctx_seq(lctx)) { + return 1; + } + prompt.resize(n + 3); + prompt[0] = 64403; // start_of_human + if (llama_tokenize(vocab, text.data(), text.size(), prompt.data() + 1, n, true, false) != n) { + return 1; + } + prompt[n + 1] = 2; // end_of_text + prompt[n + 2] = 64404; // end_of_human + seq_id = inp->seq_id; + out_type = inp->out_type; + if (!llama_memory_seq_rm(llama_get_memory(lctx), seq_id, -1, -1)) { + return 1; + } + if (tok_embd.empty()) { + const size_t count = llama_model_get_tok_embd(model, nullptr); + if (count != (size_t) n_embd * 80538) { + return 1; + } + tok_embd.resize(count); + if (llama_model_get_tok_embd(model, tok_embd.data()) != count) { + tok_embd.clear(); + return 1; + } + } + ready = true; + return 0; + } + + int32_t step_prompt(int32_t n_batch) override { + if (!ready || n_batch <= 0 || prompt.empty()) { + return -1; + } + int n = std::min({n_batch, (int) llama_n_batch(lctx), (int) prompt.size() - pos}); + if (n > 0 && decode(prompt.data() + pos, n, pos) != 0) { + return -1; + } + return (int) prompt.size() - pos; + } + + int32_t step_gen(llama_token sampled, const float *, const float ** h_state_out, bool * out_stop) override { + *h_state_out = nullptr; + *out_stop = stopped; + if (stopped) { + return 0; + } + if (!ready || pos < (int) prompt.size() || sampled < 0 || sampled >= 80538) { + return 1; + } + if (sampled == 64402) { // end_of_speech + if (!started || codes.empty() || codes.size() % 4) { + LOG_ERR("KaniTTS-2: incomplete audio frame at end of speech\n"); + return 1; + } + stopped = *out_stop = true; + return 0; + } + int rope_pos = pos; + if (sampled == 64401 && !started) { + started = true; + } else if (started) { + const int code = sampled - 64410 - (int) (codes.size() % 4) * 4032; + if (code < 0 || code >= 4032) { + LOG_ERR("KaniTTS-2: invalid codebook token %d at offset %zu\n", sampled, codes.size()); + return 1; + } + if (first_audio_pos < 0) { + first_audio_pos = pos; + } + rope_pos = first_audio_pos + (int) (codes.size() / 4); + codes.push_back(code); + pcm.clear(); + wav.clear(); + } else if (sampled != 64405) { // start_of_ai + LOG_ERR("KaniTTS-2: unexpected token before speech: %d\n", sampled); + return 1; + } + if (decode(&sampled, 1, rope_pos) != 0) { + return 1; + } + *h_state_out = llama_get_embeddings_ith(lctx, -1); + return *h_state_out ? 0 : 1; + } + + int32_t get_output(int32_t * rate, const char ** data, size_t * len, int64_t * samples) override { + if (!ready || codes.empty() || codes.size() % 4) { + LOG_ERR("KaniTTS-2: generation ended with an incomplete audio frame; increase -n\n"); + return 1; + } + if (pcm.empty()) { + // The causal decoder needs at most 28 preceding latent frames. + // Use 32 for alignment and decode at most 128 frames per graph. + std::vector decoded; + const size_t frames = codes.size() / 4; + for (size_t offset = 0; offset < frames;) { + size_t begin = offset > 32 ? offset - 32 : 0; + size_t end = std::min(frames, begin + 128); + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.codes = codes.data() + begin * 4; + inp.n_codes = (end - begin) * 4; + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, &inp, &out) != 0 || out.n_samples != (end - begin) * 1764) { + return 1; + } + decoded.insert(decoded.end(), out.audio + (offset - begin) * 1764, out.audio + out.n_samples); + offset = end; + } + pcm = std::move(decoded); + } + *rate = info.sample_rate; + if (samples) { + *samples = pcm.size(); + } + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *data = (const char *) pcm.data(); + *len = pcm.size() * sizeof(float); + } else { + if (wav.empty() && !write_wav16(wav, pcm, info.sample_rate)) { + return 1; + } + *data = wav.data(); + *len = wav.size(); + } + return 0; + } + +private: + int decode(const llama_token * tokens, int n, int rope_pos) { + if (pos + n > (int) llama_n_ctx_seq(lctx)) { + LOG_ERR("KaniTTS-2: context exhausted; increase -c\n"); + return 1; + } + std::vector embd((size_t) n * n_embd); + for (int i = 0; i < n; ++i) { + std::copy_n(tok_embd.data() + (size_t) tokens[i] * n_embd, n_embd, embd.data() + (size_t) i * n_embd); + } + decode_embd_batch batch(embd.data(), n, 4, n_embd); + batch.set_position_mrope_1d(pos, seq_id); + for (int i = 0; i < n; ++i) { + batch.pos[n + i] = rope_pos + i; + } + batch.logits[n - 1] = true; + if (llama_decode(lctx, batch.batch) != 0) { + return 1; + } + pos += n; + return 0; + } + + llama_seq_id seq_id = 0; + int pos = 0; + int first_audio_pos = -1; + bool started = false; + bool stopped = false; + bool ready = false; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector prompt; + std::vector codes; + std::vector tok_embd; + std::vector pcm; + std::vector wav; +}; + static std::unique_ptr make_pipeline(llama_context * lctx, mtmd_context * mctx) { switch (mtmd_gen_audio_get_info(mctx).type) { case MTMD_GEN_AUDIO_TYPE_QWEN3TTS: return std::unique_ptr(new qwen3tts_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_NEMO_NANO_CODEC: + return std::unique_ptr(new kani_tts2_gen_audio_pipeline(lctx, mctx)); case MTMD_GEN_AUDIO_TYPE_POCKETTTS: return std::unique_ptr(new pockettts_gen_audio_pipeline(lctx, mctx)); default: diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index c370b9ae6500..0e9432f7cf69 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -609,7 +609,8 @@ struct mtmd_context { } if (ctx_gen_a) { int n_embd_gen = clip_n_mmproj_embd(ctx_gen_a); - if (n_embd_text > 0 && n_embd_text != n_embd_gen) { + if (n_embd_text > 0 && n_embd_text != n_embd_gen && + clip_get_projector_type(ctx_gen_a) != PROJECTOR_TYPE_NEMO_NANO_CODEC) { throw std::runtime_error(string_format( "mismatch between text model (n_embd = %d) and gen-audio mmproj (n_embd = %d)\n" "hint: you may be using wrong mmproj\n", diff --git a/tools/tts/README.md b/tools/tts/README.md index b2f55add7f03..812dc1f5dc81 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -58,18 +58,36 @@ python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --outfile pock python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --outfile mmproj-pocket-tts.gguf ``` -## NeMo Nano Codec decoder (MTMD API) +## KaniTTS-2 -The [22 kHz / 0.6 kbps / 12.5 fps variant](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps) is supported as a standalone codes-to-audio decoder. It is not a text-to-speech model and cannot be used directly with `llama-tts -p`. +[KaniTTS-2 English](https://huggingface.co/nineninesix/kani-tts-2-en) uses an LFM2 backbone with learned RoPE frequencies and four audio tokens per frame. Its audio decoder is [NeMo Nano Codec 22 kHz / 0.6 kbps / 12.5 fps](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps). -Place `nemo-nano-codec-22khz-0.6kbps-12.5fps.nemo` in a local directory and convert it: +Download the backbone and the matching codec into separate directories, then convert both: ```sh -python convert_hf_to_gguf.py path/to/nemo-nano-codec --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec.gguf +hf download nineninesix/kani-tts-2-en --local-dir models/kani-tts-2-en +hf download nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps \ + nemo-nano-codec-22khz-0.6kbps-12.5fps.nemo --local-dir models/nemo-nano-codec + +python convert_hf_to_gguf.py models/kani-tts-2-en \ + --outtype f16 --outfile kani-tts-2-f16.gguf +python convert_hf_to_gguf.py models/nemo-nano-codec \ + --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec-f16.gguf + +./build/bin/llama-tts -m kani-tts-2-f16.gguf -mm mmproj-nemo-nano-codec-f16.gguf \ + -p "The weather is beautiful today. Let us take a walk in the park, enjoy the sunshine, and listen to the birds singing in the trees." --tts-lang en_us -o output.wav \ + -c 4096 -n 3000 --temp 1 --top-p 0.95 --top-k 0 --min-p 0.05 \ + --repeat-penalty 1.1 --repeat-last-n 4096 --seed 42 ``` -The converter reads the architecture from the archive's `model_config.yaml`; there is no text backbone or tokenizer to convert. Convolution weights are stored as F16 even with `--outtype f32`, to use the existing ggml convolution path. FSQ codebook and activation parameters remain F32. +Only `KaniTTS2ForCausalLM` selects the Kani converter. Generic LFM2 models keep their existing conversion and inference behavior. Use `llama-tts` for Kani generation: ordinary text generation does not assign the required frame positions. Each token advances the cache, while all four tokens in an audio frame share a RoPE position. The learned per-layer frequency scales are preserved in GGUF. -Load the mmproj with `mtmd_init_from_file(path, nullptr, params)`. Pass `MTMD_GEN_PROCESS_TYPE_GEN_WAV` to `mtmd_gen_audio_process`, with `codes` laid out as `[frame][group]`: four codes per frame, each in `[0, 4031]`. NeMo's `[group, batch, frame]` tokens must be transposed for this interface. Each call accepts 1 to 128 complete frames and returns `n_frames * 1764` mono float samples at 22050 Hz. Copy the output before the next call, and release the context with `mtmd_free`. +The supported language tags are `en_us`, `en_nyork`, `en_oakl`, `en_glasg`, `en_bost`, and `en_scou`; `en` is an alias for `en_us`. Omitting the tag leaves the prompt untagged, as in the reference implementation. Speaker reference audio / voice cloning is not supported; the optional speaker projection is excluded from the backbone conversion. -This initial implementation decodes a complete sequence with zero initial context. It does not accept continuous features, persistent state or `GEN_CODE`. Independent chunks do not preserve convolution history. Audio encoding, other Nano Codec variants and a text-generation pipeline are not included. NVIDIA describes this particular variant as intended for fine-tuning with a limited set of speakers, rather than general-purpose audio reconstruction. +`-n` limits generation steps (tokens for Kani), not audio frames. Four audio tokens produce 1764 samples at 22050 Hz (12.5 frames/s). Increase `-c` for longer prompts and output. The helper validates codebook offsets and end-of-speech boundaries; it reports an error if a token limit interrupts a frame. Greedy decoding can repeat audio codes; use the sampling parameters above. Waveform decoding uses overlapping chunks to preserve the causal convolution history on long outputs. + +### NeMo decoder API + +The same mmproj can be loaded with `mtmd_init_from_file(path, nullptr, params)` for codes-to-audio use. Pass `MTMD_GEN_PROCESS_TYPE_GEN_WAV` to `mtmd_gen_audio_process`, with `codes` laid out as `[frame][group]`: four codes per frame, each in `[0, 4031]`. NeMo's `[group, batch, frame]` tokens must be transposed. Each call accepts 1 to 128 complete frames and returns `n_frames * 1764` mono float samples. Copy the output before the next call and release the context with `mtmd_free`. + +Convolution weights are F16, including with `--outtype f32`, to use the existing ggml convolution path. FSQ codebook and activation parameters remain F32. Each API call has zero initial context; the Kani helper supplies 32 preceding frames when splitting a long sequence. The decoder does not accept continuous features, persistent state or `GEN_CODE`. Audio encoding and other Nano Codec variants are not supported. diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 6fd1936324ce..d0282f0fdf7b 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -20,15 +20,15 @@ struct tts_timings { int64_t t_start_us = ggml_time_us(); int64_t t_last_us = t_start_us; - void report(int n_frames) { + void report(int n_steps) { const int64_t t_now_us = ggml_time_us(); if (t_now_us - t_last_us < 2000000) { return; } t_last_us = t_now_us; const double t_elapsed_s = (t_now_us - t_start_us) / 1e6; - const double fps = t_elapsed_s > 0 ? n_frames / t_elapsed_s : 0.0; - LOG_INF("frames generated: %d, speed: %.2f frames/s\n", n_frames, fps); + const double fps = t_elapsed_s > 0 ? n_steps / t_elapsed_s : 0.0; + LOG_INF("generation steps: %d, speed: %.2f steps/s\n", n_steps, fps); } }; @@ -153,7 +153,7 @@ int main(int argc, char ** argv) { }; const int max_new = params.n_predict > 0 ? params.n_predict : 512; - int n_frames = 0; + int n_steps = 0; llama_token sampled = sample_semantic_code(); const float * h_state = llama_get_embeddings_ith(lctx, -1); @@ -161,23 +161,23 @@ int main(int argc, char ** argv) { const int64_t t_gen_start_us = ggml_time_us(); bool stop = false; - while (!stop && n_frames < max_new) { + while (!stop && n_steps < max_new) { const float * h_next = nullptr; // stage 2+3: semantic --> acoustic details --> audio waveform // step_gen() runs both stages and returns new h_state for next step if (gen.step_gen(sampled, h_state, &h_next, &stop) != 0) { - LOG_ERR("step_gen failed at frame %d\n", n_frames); + LOG_ERR("step_gen failed at step %d\n", n_steps); return 1; } if (!h_next) { break; // stopped without generating a frame } - n_frames++; + n_steps++; h_state = h_next; sampled = sample_semantic_code(); - timings.report(n_frames); + timings.report(n_steps); } const double t_gen_s = (ggml_time_us() - t_gen_start_us) / 1e6; @@ -192,7 +192,7 @@ int main(int argc, char ** argv) { } const double t_wav_s = (ggml_time_us() - t_wav_start_us) / 1e6; - LOG_INF("generated %d frames, %zu bytes of WAV audio (%d Hz)\n", n_frames, data_len, sample_rate); + LOG_INF("generated %d steps, %zu bytes of WAV audio (%d Hz)\n", n_steps, data_len, sample_rate); const double t_prompt_s = (t_gen_start_us - t_prompt_start_us) / 1e6; const double t_total_s = t_prompt_s + t_gen_s + t_wav_s; From b74b0177eb7dc0b3def277fb07cd339b9b7abf42 Mon Sep 17 00:00:00 2001 From: Hans Date: Tue, 8 Sep 2026 23:32:06 +0800 Subject: [PATCH 3/4] docs : match KaniTTS-2 usage to existing TTS sections Assisted-by: Codex --- tools/tts/README.md | 42 ++++++++++++++++++------------------------ 1 file changed, 18 insertions(+), 24 deletions(-) diff --git a/tools/tts/README.md b/tools/tts/README.md index 812dc1f5dc81..8a0b96fea7dd 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -60,34 +60,28 @@ python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --out ## KaniTTS-2 -[KaniTTS-2 English](https://huggingface.co/nineninesix/kani-tts-2-en) uses an LFM2 backbone with learned RoPE frequencies and four audio tokens per frame. Its audio decoder is [NeMo Nano Codec 22 kHz / 0.6 kbps / 12.5 fps](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps). +Available params: +- `--tts-lang` can be `en_us`, `en_nyork`, `en_oakl`, `en_glasg`, `en_bost`, `en_scou`; `en` is an alias for `en_us` +- `--tts-speaker-file` is not supported +- `-n` limits generated tokens; four audio tokens represent one frame -Download the backbone and the matching codec into separate directories, then convert both: +Example usage: ```sh -hf download nineninesix/kani-tts-2-en --local-dir models/kani-tts-2-en -hf download nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps \ - nemo-nano-codec-22khz-0.6kbps-12.5fps.nemo --local-dir models/nemo-nano-codec - -python convert_hf_to_gguf.py models/kani-tts-2-en \ - --outtype f16 --outfile kani-tts-2-f16.gguf -python convert_hf_to_gguf.py models/nemo-nano-codec \ - --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec-f16.gguf - -./build/bin/llama-tts -m kani-tts-2-f16.gguf -mm mmproj-nemo-nano-codec-f16.gguf \ - -p "The weather is beautiful today. Let us take a walk in the park, enjoy the sunshine, and listen to the birds singing in the trees." --tts-lang en_us -o output.wav \ - -c 4096 -n 3000 --temp 1 --top-p 0.95 --top-k 0 --min-p 0.05 \ - --repeat-penalty 1.1 --repeat-last-n 4096 --seed 42 +llama-tts -m kani-tts-2.gguf \ + -mm mmproj-nemo-nano-codec.gguf \ + -p "Hello world" \ + --tts-lang en_us \ + -c 4096 -n 3000 --temp 1 --top-p 0.95 --top-k 0 \ + --repeat-penalty 1.1 --repeat-last-n 4096 \ + --output out.wav ``` -Only `KaniTTS2ForCausalLM` selects the Kani converter. Generic LFM2 models keep their existing conversion and inference behavior. Use `llama-tts` for Kani generation: ordinary text generation does not assign the required frame positions. Each token advances the cache, while all four tokens in an audio frame share a RoPE position. The learned per-layer frequency scales are preserved in GGUF. - -The supported language tags are `en_us`, `en_nyork`, `en_oakl`, `en_glasg`, `en_bost`, and `en_scou`; `en` is an alias for `en_us`. Omitting the tag leaves the prompt untagged, as in the reference implementation. Speaker reference audio / voice cloning is not supported; the optional speaker projection is excluded from the backbone conversion. - -`-n` limits generation steps (tokens for Kani), not audio frames. Four audio tokens produce 1764 samples at 22050 Hz (12.5 frames/s). Increase `-c` for longer prompts and output. The helper validates codebook offsets and end-of-speech boundaries; it reports an error if a token limit interrupts a frame. Greedy decoding can repeat audio codes; use the sampling parameters above. Waveform decoding uses overlapping chunks to preserve the causal convolution history on long outputs. - -### NeMo decoder API +**Note for GGUF conversion:** -The same mmproj can be loaded with `mtmd_init_from_file(path, nullptr, params)` for codes-to-audio use. Pass `MTMD_GEN_PROCESS_TYPE_GEN_WAV` to `mtmd_gen_audio_process`, with `codes` laid out as `[frame][group]`: four codes per frame, each in `[0, 4031]`. NeMo's `[group, batch, frame]` tokens must be transposed. Each call accepts 1 to 128 complete frames and returns `n_frames * 1764` mono float samples. Copy the output before the next call and release the context with `mtmd_free`. +Download [KaniTTS-2 English](https://huggingface.co/nineninesix/kani-tts-2-en) and the `.nemo` archive from [NeMo Nano Codec 22 kHz / 0.6 kbps / 12.5 fps](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps) into separate directories: -Convolution weights are F16, including with `--outtype f32`, to use the existing ggml convolution path. FSQ codebook and activation parameters remain F32. Each API call has zero initial context; the Kani helper supplies 32 preceding frames when splitting a long sequence. The decoder does not accept continuous features, persistent state or `GEN_CODE`. Audio encoding and other Nano Codec variants are not supported. +```sh +python convert_hf_to_gguf.py path/to/kani-tts-2-en --outtype f16 --outfile kani-tts-2.gguf +python convert_hf_to_gguf.py path/to/nemo-nano-codec --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec.gguf +``` From fb7e3386e6e58af40fa10393f2508f6b9171c877 Mon Sep 17 00:00:00 2001 From: Hans Date: Wed, 9 Sep 2026 00:05:43 +0800 Subject: [PATCH 4/4] mtmd: support KaniTTS-2 speaker reference audio Convert the WavLM speaker encoder and backbone speaker projection into the NeMo mmproj, encode reference audio through MTMD, and insert the projected embedding into the Kani prompt. Assisted-by: Codex --- conversion/__init__.py | 1 + conversion/lfm2.py | 2 +- conversion/nemo_nano_codec.py | 127 ++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 16 ++++ tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 2 + tools/mtmd/clip-model.h | 16 ++++ tools/mtmd/clip.cpp | 95 +++++++++++++++++++++- tools/mtmd/models/kani-spkenc.cpp | 89 +++++++++++++++++++++ tools/mtmd/models/models.h | 5 ++ tools/mtmd/mtmd-audio.cpp | 21 +++++ tools/mtmd/mtmd-audio.h | 6 ++ tools/mtmd/mtmd-helper-gen.cpp | 40 +++++++++- tools/mtmd/mtmd.cpp | 4 + tools/tts/README.md | 9 ++- 15 files changed, 425 insertions(+), 9 deletions(-) create mode 100644 tools/mtmd/models/kani-spkenc.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 40ca7e353130..b0d5b6692b19 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -286,6 +286,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "KaniTTS2ForCausalLM": "nemo_nano_codec", "NemoNanoCodecModel": "nemo_nano_codec", "AudioFlamingo3ForConditionalGeneration": "ultravox", "CogVLMForCausalLM": "cogvlm", diff --git a/conversion/lfm2.py b/conversion/lfm2.py index 4b53eeeb88b5..05d5d033cc21 100644 --- a/conversion/lfm2.py +++ b/conversion/lfm2.py @@ -88,7 +88,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield f"blk.{layer}.attn_rope_freqs.weight", (1.0 / alpha).expand(head_dim // 2).clone() return if name == "model.speaker_emb_projection.weight": - # Speaker embedding extraction is not part of the unconditional TTS pipeline. + # The speaker projection is stored in the mmproj. return yield from super().modify_tensors(data_torch, name, bid) diff --git a/conversion/nemo_nano_codec.py b/conversion/nemo_nano_codec.py index 2f75e65d8e7a..69d65f06f536 100644 --- a/conversion/nemo_nano_codec.py +++ b/conversion/nemo_nano_codec.py @@ -131,3 +131,130 @@ def modify_tensors(self, data_torch, name, bid): if suffix == ".alpha": data_torch = data_torch.flatten() yield self.format_tensor_name(tensor, bid, suffix), data_torch + + +@ModelBase.register("KaniTTS2ForCausalLM") +class KaniTTS2MmprojModel(NemoNanoCodecModel): + has_audio_encoder = True + + def get_audio_config(self): + import json + path = self.dir_model / "speaker_encoder" / "config.json" + with path.open() as f: + config = json.load(f) + expected = {"model_type": "wavlm", "hidden_size": 1024, "num_hidden_layers": 24, + "num_attention_heads": 16, "intermediate_size": 4096, "do_stable_layer_norm": True, + "feat_extract_norm": "layer", "layer_norm_eps": 1e-5, "conv_bias": False, "embd_size": 128, + "top_interm_size": 512, "num_conv_pos_embeddings": 128, + "num_conv_pos_embedding_groups": 16, "num_buckets": 320, "max_bucket_distance": 800, + "conv_dim": [512] * 7, "conv_kernel": [10, 3, 3, 3, 3, 2, 2], + "conv_stride": [5, 2, 2, 2, 2, 2, 2], "hidden_act": "gelu", "feat_extract_activation": "gelu"} + if any(config.get(k) != v for k, v in expected.items()) or config.get("add_adapter"): + raise ValueError("Unsupported Kani WavLM speaker encoder configuration") + return config + + def set_gguf_parameters(self): + _load_hparams(self.dir_model) # validate the matching NeMo archive + super().set_gguf_parameters() + self.gguf_writer.add_clip_has_audio_encoder(True) + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.KANI_SPKENC) + self.gguf_writer.add_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_audio_embedding_length(1024) + self.gguf_writer.add_audio_feed_forward_length(4096) + self.gguf_writer.add_audio_block_count(24) + self.gguf_writer.add_audio_head_count(16) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio["layer_norm_eps"]) + self.gguf_writer.add_audio_num_mel_bins(1) + + def get_tensors(self): + from safetensors.torch import load_file + yield from super().get_tensors() + speaker = self.dir_model / "speaker_encoder" + if (speaker / "model.safetensors").is_file(): + state = load_file(speaker / "model.safetensors") + else: + state = torch.load(speaker / "pytorch_model.bin", map_location="cpu", weights_only=True) + pos = "wavlm.encoder.pos_conv_embed.conv." + wv = state.get(pos + "parametrizations.weight.original1", state.get(pos + "weight_v")) + wg = state.get(pos + "parametrizations.weight.original0", state.get(pos + "weight_g")) + if wv is None or wg is None: + raise ValueError("Missing WavLM positional convolution weight normalization") + weight = wv.float() * wg.float() / torch.linalg.vector_norm(wv.float(), dim=(0, 1), keepdim=True) + yield "a.spk.pos.weight", weight.reshape(16, 64, 64, 128).contiguous() + yield "a.spk.pos.bias", state[pos + "bias"] + for name, value in state.items(): + if name.startswith(pos) or name == "wavlm.masked_spec_embed" or name.endswith("num_batches_tracked"): + continue + name = name.replace("wavlm.feature_extractor.conv_layers.", "a.spk.conv.") + name = name.replace("wavlm.feature_projection.", "a.spk.feature.") + name = name.replace("wavlm.encoder.layers.", "a.spk.blk.") + name = name.replace("wavlm.encoder.layer_norm.", "a.spk.output_norm.") + name = name.replace("top_layers.", "a.spk.top.") + if not name.startswith("a.spk."): + raise ValueError(f"Unexpected speaker encoder tensor: {name}") + if "gru_rel_pos_linear" in name: + value = value.float().reshape(2, 4, -1).sum(1) + if name.endswith("bias"): + value = value.flatten() + elif "gru_rel_pos_const" in name: + value = value.reshape(16, 1, 1) + elif ".affine" in name and name.endswith("weight"): + value = value.squeeze(-1) + if ".batchnorm" in name: + if name.endswith("running_var"): + continue + if not name.endswith("running_mean"): + raise ValueError(f"Unexpected speaker batch norm tensor: {name}") + key = name.replace("a.spk.top.", "top_layers.").replace("running_mean", "running_var") + scale = torch.rsqrt(state[key].float() + 1e-3) + yield name.replace("running_mean", "weight"), scale + yield name.replace("running_mean", "bias"), -value.float() * scale + continue + yield name, value + from safetensors import safe_open + with safe_open(self.dir_model / "model.safetensors", framework="pt", device="cpu") as backbone: + yield "a.spk.projection.weight", backbone.get_tensor("model.speaker_emb_projection.weight") + + def modify_tensors(self, data_torch, name, bid): + if not name.startswith("a.spk."): + yield from super().modify_tensors(data_torch, name, bid) + return + T = gguf.MODEL_TENSOR + name = name.removeprefix("a.spk.") + suffix = name.rsplit(".", 1)[-1] + stem = name.rsplit(".", 1)[0] + direct = {"pos": T.A_ENC_POSITION_CONV, "feature.layer_norm": T.A_PRE_NORM, + "feature.projection": T.A_ENC_INP_PROJ, "output_norm": T.A_POST_NORM, + "projection": T.A_ENC_SPEAKER_PROJ} + if stem in direct: + tensor = direct[stem] + elif match := re.fullmatch(r"conv\.(\d+)\.(conv|layer_norm)", stem): + bid = int(match[1]) + tensor = T.A_ENC_CONV1D if match[2] == "conv" else T.A_ENC_CONV1D_NORM + elif match := re.fullmatch(r"top\.(affine|batchnorm)([12])", stem): + bid = int(match[2]) - 1 + tensor = T.A_ENC_SPK_FC if match[1] == "affine" else T.A_ENC_SPK_FC_NORM + elif match := re.fullmatch(r"blk\.(\d+)\.(.+)", name): + bid = int(match[1]) + tail = match[2] + if tail == "attention.gru_rel_pos_const": + tensor, suffix = T.A_ENC_ATTN_REL_GATE_CONST, "weight" + else: + tensor = {"layer_norm": T.A_ENC_INPUT_NORM, "final_layer_norm": T.A_ENC_FFN_NORM, + "attention.q_proj": T.A_ENC_ATTN_Q, "attention.k_proj": T.A_ENC_ATTN_K, + "attention.v_proj": T.A_ENC_ATTN_V, "attention.out_proj": T.A_ENC_OUTPUT, + "attention.gru_rel_pos_linear": T.A_ENC_ATTN_REL_GATE, + "attention.rel_attn_embed": T.A_ENC_ATTN_REL_POS_EMB, + "feed_forward.intermediate_dense": T.A_ENC_FFN_UP, + "feed_forward.output_dense": T.A_ENC_FFN_DOWN}[tail.rsplit(".", 1)[0]] + else: + raise ValueError(f"Unexpected speaker tensor: {name}") + yield self.format_tensor_name(tensor, bid, "." + suffix), data_torch + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if name.startswith("a.spk."): + if ".conv.weight" in name or name == "a.spk.pos.weight": + return gguf.GGMLQuantizationType.F16 + if "gru_rel_pos" in name or "rel_attn_embed" in name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 288b74174976..287a1a660c28 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1028,6 +1028,11 @@ class MODEL_TENSOR(IntEnum): V_MULTI_PROJ_POST_NORM = auto() # audio (mtmd) + A_ENC_POSITION_CONV = auto() + A_ENC_ATTN_REL_GATE = auto() + A_ENC_ATTN_REL_GATE_CONST = auto() + A_ENC_SPK_FC = auto() + A_ENC_SPK_FC_NORM = auto() A_ENC_EMBD_POS = auto() A_ENC_EMBD_NORM = auto() A_ENC_EMBD_TO_LOGITS = auto() # lfm2 @@ -1791,6 +1796,11 @@ class MODEL_TENSOR(IntEnum): # audio (mtmd) # note: all audio tensor names must use prefix "a." or "mm.a." + MODEL_TENSOR.A_ENC_POSITION_CONV: "a.position_conv", + MODEL_TENSOR.A_ENC_ATTN_REL_GATE: "a.blk.{bid}.attn_rel_gate", + MODEL_TENSOR.A_ENC_ATTN_REL_GATE_CONST: "a.blk.{bid}.attn_rel_gate_const", + MODEL_TENSOR.A_ENC_SPK_FC: "a.spk_fc.{bid}", + MODEL_TENSOR.A_ENC_SPK_FC_NORM: "a.spk_fc.{bid}.norm", MODEL_TENSOR.A_ENC_EMBD_POS: "a.position_embd", MODEL_TENSOR.A_ENC_EMBD_NORM: "a.position_embd_norm", MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS: "a.embd_to_logits", @@ -2134,6 +2144,11 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.V_MULTI_PROJ_NORM, MODEL_TENSOR.V_MULTI_PROJ_POST_NORM, # audio + MODEL_TENSOR.A_ENC_POSITION_CONV, + MODEL_TENSOR.A_ENC_ATTN_REL_GATE, + MODEL_TENSOR.A_ENC_ATTN_REL_GATE_CONST, + MODEL_TENSOR.A_ENC_SPK_FC, + MODEL_TENSOR.A_ENC_SPK_FC_NORM, MODEL_TENSOR.A_ENC_EMBD_POS, MODEL_TENSOR.A_ENC_EMBD_NORM, MODEL_TENSOR.A_ENC_EMBD_TO_LOGITS, @@ -5845,6 +5860,7 @@ class VisionProjectorType: 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 + KANI_SPKENC = "kani_spkenc" # audio: WavLM speaker encoder NEMO_NANO_CODEC = "nemo_nano_codec" # audio generation: causal HiFiGAN decoder POCKETTTS_GEN = "pockettts_gen" # audio generation: flow-matching decoder + mimi decoder HUNYUANVL = "hunyuanvl" diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 87e707f97759..a3a5bad6c64f 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -62,6 +62,7 @@ add_library(mtmd models/qwen3tts-gen.cpp models/pockettts-seanet.cpp models/pockettts-spkenc.cpp + models/kani-spkenc.cpp models/nemo-nano-codec.cpp models/pockettts-gen.cpp models/step3vl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index d3973ed6c049..82a7a3f59647 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -502,6 +502,7 @@ enum projector_type { PROJECTOR_TYPE_QWEN3TTS_SPKENC, PROJECTOR_TYPE_QWEN3TTS_GEN, PROJECTOR_TYPE_POCKETTTS_SPKENC, + PROJECTOR_TYPE_KANI_SPKENC, PROJECTOR_TYPE_NEMO_NANO_CODEC, PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, @@ -568,6 +569,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_QWEN3TTS_SPKENC, "qwen3tts_spkenc"}, { PROJECTOR_TYPE_QWEN3TTS_GEN, "qwen3tts_gen"}, { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, + { PROJECTOR_TYPE_KANI_SPKENC, "kani_spkenc"}, { PROJECTOR_TYPE_NEMO_NANO_CODEC, "nemo_nano_codec"}, { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 3bb8f4b4b0fa..8e6bab81798e 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -485,6 +485,21 @@ struct clip_flow_net { std::vector blocks; }; +// WavLM speaker embedding and Kani backbone projection. +struct clip_kani_speaker { + struct linear { ggml_tensor * w = nullptr; ggml_tensor * b = nullptr; }; + struct conv { ggml_tensor * w = nullptr; linear norm; }; + struct layer { + linear norm, ffn_norm, q, k, v, o, up, down, gate; + ggml_tensor * gate_const = nullptr; + }; + conv convs[7]; + linear feature_norm, feature_proj, pos, output_norm, top[2], top_norm[2]; + layer layers[24]; + ggml_tensor * relative = nullptr; + ggml_tensor * projection = nullptr; +}; + // NeMo Nano Codec 22 kHz / 12.5 fps: grouped FSQ codes -> raw PCM. struct clip_nemo_nano_codec { static constexpr int n_groups = 4; @@ -805,6 +820,7 @@ struct clip_model { // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; clip_nemo_nano_codec nemo; + clip_kani_speaker kani_speaker; // pocket-tts: SEANet stack, shared by the encoder (speaker path) and the decoder (gen path) clip_seanet seanet; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index b8ab70871604..8eaa165ae016 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1093,6 +1093,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_KANI_SPKENC: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_QWEN3TTS_SPKENC: { builder = std::make_unique(ctx, img); @@ -1836,6 +1840,13 @@ struct clip_model_loader { hparams.audio_window_len = 1024; hparams.audio_hop_len = 256; } break; + case PROJECTOR_TYPE_KANI_SPKENC: + { + if (hparams.n_layer != 24 || hparams.n_embd != 1024 || hparams.n_head != 16 || hparams.projection_dim != 1024) { + throw std::runtime_error("invalid Kani speaker encoder dimensions"); + } + hparams.audio_sample_rate = 16000; + } break; case PROJECTOR_TYPE_NEMO_NANO_CODEC: { if (hparams.n_layer != 5 || hparams.n_embd != 864 || hparams.projection_dim != 16 || hparams.n_head != 1) { @@ -2089,7 +2100,8 @@ struct clip_model_loader { // GEMMA4UA is encoder-free: it uses n_mel_bins as a raw-waveform frame size (640) and has no FFT/filterbank, so the mel-range and FFT // checks below do not apply to it. // pocket-tts is encoder-free in the same sense: mimi convolves the raw waveform - const bool fft_based = model.proj_type != PROJECTOR_TYPE_GEMMA4UA && + const bool fft_based = model.proj_type != PROJECTOR_TYPE_KANI_SPKENC && + model.proj_type != PROJECTOR_TYPE_GEMMA4UA && model.proj_type != PROJECTOR_TYPE_POCKETTTS_SPKENC; // Validate audio hparams loaded from GGUF metadata @@ -2249,6 +2261,7 @@ struct clip_model_loader { const bool has_standard_layers = ( model.proj_type != PROJECTOR_TYPE_GEMMA3NV && + model.proj_type != PROJECTOR_TYPE_KANI_SPKENC && model.proj_type != PROJECTOR_TYPE_QWEN3TTS_SPKENC && model.proj_type != PROJECTOR_TYPE_POCKETTTS_GEN && model.proj_type != PROJECTOR_TYPE_NEMO_NANO_CODEC); @@ -2922,6 +2935,67 @@ struct clip_model_loader { model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight")); model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight")); } break; + case PROJECTOR_TYPE_KANI_SPKENC: + { + auto & m = model.kani_speaker; + auto tensor = [&](const std::string & name, std::initializer_list shape) { + auto * t = get_tensor("a." + name); + int d = 0; + for (auto n : shape) { + if (t->ne[d++] != n) { throw std::runtime_error("invalid Kani speaker tensor: " + name); } + } + for (; d < 4; ++d) { + if (t->ne[d] != 1) { throw std::runtime_error("invalid Kani speaker tensor rank: " + name); } + } + return t; + }; + auto linear = [&](clip_kani_speaker::linear & l, const std::string & name, int in, int out) { + l.w = tensor(name + ".weight", {in, out}); + l.b = tensor(name + ".bias", {out}); + }; + auto norm = [&](clip_kani_speaker::linear & l, const std::string & name, int dim) { + l.w = tensor(name + ".weight", {dim}); + l.b = tensor(name + ".bias", {dim}); + }; + const int kernels[] = {10, 3, 3, 3, 3, 2, 2}; + for (int i = 0; i < 7; ++i) { + auto p = "conv1d." + std::to_string(i); + m.convs[i].w = tensor(p + ".weight", {kernels[i], i ? 512 : 1, 512}); + norm(m.convs[i].norm, p + ".norm", 512); + } + m.feature_norm = {model.pre_ln_w, model.pre_ln_b}; + if (!m.feature_norm.w || !m.feature_norm.b || m.feature_norm.w->ne[0] != 512 || m.feature_norm.b->ne[0] != 512) { + throw std::runtime_error("invalid Kani feature normalization"); + } + linear(m.feature_proj, "input_projection", 512, 1024); + m.pos.w = tensor("position_conv.weight", {128, 64, 64, 16}); + m.pos.b = tensor("position_conv.bias", {1024}); + m.output_norm = {model.post_ln_w, model.post_ln_b}; + if (!m.output_norm.w || !m.output_norm.b || m.output_norm.w->ne[0] != 1024 || m.output_norm.b->ne[0] != 1024) { + throw std::runtime_error("invalid Kani output normalization"); + } + m.relative = tensor("blk.0.attn_rel_pos_emb.weight", {16, 320}); + for (int i = 0; i < 24; ++i) { + auto & l = m.layers[i]; + auto p = "blk." + std::to_string(i) + "."; + norm(l.norm, p + "ln1", 1024); + norm(l.ffn_norm, p + "ffn_norm", 1024); + linear(l.q, p + "attn_q", 1024, 1024); + linear(l.k, p + "attn_k", 1024, 1024); + linear(l.v, p + "attn_v", 1024, 1024); + linear(l.o, p + "attn_out", 1024, 1024); + linear(l.up, p + "ffn_up", 1024, 4096); + linear(l.down, p + "ffn_down", 4096, 1024); + linear(l.gate, p + "attn_rel_gate", 64, 2); + l.gate_const = tensor(p + "attn_rel_gate_const.weight", {1, 1, 16}); + } + for (int i = 0; i < 2; ++i) { + auto p = "spk_fc." + std::to_string(i); + linear(m.top[i], p, i ? 512 : 2048, i ? 128 : 512); + norm(m.top_norm[i], p + ".norm", i ? 128 : 512); + } + m.projection = tensor("speaker_proj.weight", {128, hparams.projection_dim}); + } break; case PROJECTOR_TYPE_QWEN3TTS_SPKENC: { // stem TDNN (block 0) @@ -4403,6 +4477,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { const int ds = ctx->model.hparams.audio_proj_downsample_rate; n_patches = ((img->nx() + ws - 1) / ws) * (ws / ds); } break; + case PROJECTOR_TYPE_KANI_SPKENC: case PROJECTOR_TYPE_QWEN3TTS_SPKENC: { // pooling gives one speaker embedding, whatever the clip length is @@ -5333,6 +5408,22 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { { // do nothing } break; + case PROJECTOR_TYPE_KANI_SPKENC: + { + int64_t n = imgs.entries[0].nx(); + const int kernels[] = {10, 3, 3, 3, 3, 2, 2}; + const int strides[] = {5, 2, 2, 2, 2, 2, 2}; + for (int i = 0; i < 7; ++i) { n = (n - kernels[i]) / strides[i] + 1; } + std::vector buckets(n * n); + for (int64_t q = 0; q < n; ++q) { + for (int64_t k = 0; k < n; ++k) { + const int distance = (int) std::abs(k - q); + const int bucket = distance < 80 ? distance : std::min(159, 80 + (int) (std::log((float) distance / 80) / std::log(10.0f) * 80)); + buckets[q * n + k] = (k > q ? 160 : 0) + bucket; + } + } + set_input_i32("kani_buckets", buckets); + } break; case PROJECTOR_TYPE_NEMO_NANO_CODEC: { const int n_frames = (int) (params->codes->size() / clip_nemo_nano_codec::n_groups); @@ -6077,6 +6168,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_MIMO_AUDIO: return ctx->model.mm_2_w->ne[1]; + case PROJECTOR_TYPE_KANI_SPKENC: + return ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_QWEN3TTS_SPKENC: return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_NEMO_NANO_CODEC: diff --git a/tools/mtmd/models/kani-spkenc.cpp b/tools/mtmd/models/kani-spkenc.cpp new file mode 100644 index 000000000000..ae96a2ee4e45 --- /dev/null +++ b/tools/mtmd/models/kani-spkenc.cpp @@ -0,0 +1,89 @@ +#include "models.h" + +ggml_cgraph * clip_graph_kani_spkenc::build() { + const auto & m = model.kani_speaker; + auto linear = [&](ggml_tensor * x, const clip_kani_speaker::linear & l) { + return ggml_add(ctx0, build_mm(l.w, x), l.b); + }; + auto norm = [&](ggml_tensor * x, const clip_kani_speaker::linear & l) { + return ggml_add(ctx0, ggml_mul(ctx0, ggml_norm(ctx0, x, 1e-5f), l.w), l.b); + }; + auto transpose = [&](ggml_tensor * x) { return ggml_cont(ctx0, ggml_transpose(ctx0, x)); }; + auto * raw = build_inp_raw(1); + auto * x = ggml_reshape_1d(ctx0, raw, img.nx()); + // PyTorch std() uses the sample variance (correction=1). + auto * centered = ggml_sub(ctx0, x, ggml_mean(ctx0, x)); + auto * variance = ggml_scale(ctx0, ggml_mean(ctx0, ggml_sqr(ctx0, centered)), (float) img.nx() / (img.nx() - 1)); + x = ggml_div(ctx0, centered, ggml_scale_bias(ctx0, ggml_sqrt(ctx0, variance), 1.0f, 1e-10f)); + x = ggml_reshape_2d(ctx0, x, img.nx(), 1); + const int strides[] = {5, 2, 2, 2, 2, 2, 2}; + for (int i = 0; i < 7; ++i) { + x = ggml_conv_1d(ctx0, m.convs[i].w, x, strides[i], 0, 1); + x = transpose(ggml_reshape_2d(ctx0, x, x->ne[0], x->ne[1])); + x = ggml_gelu_erf(ctx0, norm(x, m.convs[i].norm)); + cb(x, "kani_conv", i); + x = transpose(x); + } + x = linear(norm(transpose(x), m.feature_norm), m.feature_proj); + cb(x, "kani_feature", -1); + const int64_t nt = x->ne[1]; + auto * xt = transpose(x); + ggml_tensor * pos = nullptr; + for (int group = 0; group < 16; ++group) { + auto * w = ggml_view_3d(ctx0, m.pos.w, 128, 64, 64, m.pos.w->nb[1], m.pos.w->nb[2], group * m.pos.w->nb[3]); + auto * chunk = ggml_cont(ctx0, ggml_view_2d(ctx0, xt, nt, 64, xt->nb[1], group * 64 * xt->nb[1])); + chunk = ggml_conv_1d(ctx0, w, chunk, 1, 64, 1); + chunk = ggml_cont(ctx0, ggml_view_2d(ctx0, chunk, nt, 64, chunk->nb[1], 0)); + pos = pos ? ggml_concat(ctx0, pos, chunk, 1) : chunk; + } + pos = ggml_gelu_erf(ctx0, ggml_add(ctx0, transpose(pos), m.pos.b)); + x = ggml_add(ctx0, x, pos); + cb(x, "kani_position", -1); + + auto * buckets = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, nt * nt); + ggml_set_name(buckets, "kani_buckets"); + ggml_set_input(buckets); + auto * bias = transpose(ggml_get_rows(ctx0, m.relative, buckets)); + bias = ggml_reshape_3d(ctx0, bias, nt, nt, 16); + for (int i = 0; i < 24; ++i) { + const auto & l = m.layers[i]; + auto * h = norm(x, l.norm); + auto * gh = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, h, 64, 16, nt), 0, 2, 1, 3)); + auto * gates = ggml_sigmoid(ctx0, linear(gh, l.gate)); + auto * ga = ggml_view_3d(ctx0, gates, 1, nt, 16, gates->nb[1], gates->nb[2], 0); + auto * gb = ggml_view_3d(ctx0, gates, 1, nt, 16, gates->nb[1], gates->nb[2], sizeof(float)); + auto * gate = ggml_scale_bias(ctx0, ggml_mul(ctx0, ga, + ggml_scale_bias(ctx0, ggml_mul(ctx0, gb, l.gate_const), 1.0f, -1.0f)), 1.0f, 2.0f); + auto * gated_bias = ggml_mul(ctx0, bias, gate); + auto heads = [&](const clip_kani_speaker::linear & proj) { + return ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, linear(h, proj), 64, 16, nt), 0, 2, 1, 3)); + }; + auto * q = heads(l.q); + auto * k = heads(l.k); + auto * v = heads(l.v); + auto * scores = ggml_mul_mat(ctx0, k, q); + scores = ggml_soft_max_ext(ctx0, scores, gated_bias, 0.125f, 0.0f); + auto * attn = ggml_mul_mat(ctx0, transpose(v), scores); + attn = ggml_cont(ctx0, ggml_permute(ctx0, attn, 0, 2, 1, 3)); + x = ggml_add(ctx0, x, linear(ggml_reshape_2d(ctx0, attn, 1024, nt), l.o)); + h = ggml_gelu_erf(ctx0, linear(norm(x, l.ffn_norm), l.up)); + x = ggml_add(ctx0, x, linear(h, l.down)); + cb(x, "kani_layer", i); + } + x = transpose(norm(x, m.output_norm)); + auto * mean = ggml_mean(ctx0, x); + auto * delta = ggml_sub(ctx0, x, mean); + auto * var = ggml_scale(ctx0, ggml_mean(ctx0, ggml_sqr(ctx0, delta)), (float) nt / (nt - 1)); + auto * std = ggml_sqrt(ctx0, ggml_clamp(ctx0, var, 1e-10f, INFINITY)); + x = ggml_concat(ctx0, transpose(mean), transpose(std), 0); + for (int i = 0; i < 2; ++i) { + x = ggml_relu(ctx0, linear(x, m.top[i])); + x = ggml_add(ctx0, ggml_mul(ctx0, x, m.top_norm[i].w), m.top_norm[i].b); + } + x = ggml_l2_norm(ctx0, x, 1e-12f); + cb(x, "kani_speaker_embedding", -1); + x = build_mm(m.projection, x); + cb(x, "kani_speaker_projection", -1); + ggml_build_forward_expand(gf, x); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index b35baec8ecd7..2ffc2525da04 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -12,6 +12,11 @@ * We encourage human contributors to ensure the quality and reliability of the codebase. */ +struct clip_graph_kani_spkenc : clip_graph { + clip_graph_kani_spkenc(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_nemo_nano_codec : clip_graph { clip_graph_nemo_nano_codec(clip_ctx * ctx, const clip_image_f32 & img, int n_frames) : clip_graph(ctx, img), n_frames(n_frames) {} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index c25bf4ec8967..3334a95816f8 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1555,3 +1555,24 @@ bool mtmd_audio_preprocessor_pockettts::preprocess(const float * output.push_back(std::move(out)); return true; } + +bool mtmd_audio_preprocessor_kani::preprocess(const float * samples, size_t n_samples, std::vector & output) const { + if (n_samples < 720) { + LOG_ERR("%s: Kani speaker reference must contain at least 45 ms of audio\n", __func__); + return false; + } + const size_t max_samples = 30 * 16000; + if (n_samples > max_samples) { + LOG_WRN("%s: truncating Kani speaker reference to 30 seconds\n", __func__); + n_samples = max_samples; + } + mtmd_audio_mel out; + out.n_mel = 1; + out.n_len = out.n_len_org = n_samples; + out.data.assign(samples, samples + n_samples); + for (float sample : out.data) { + if (!std::isfinite(sample)) { return false; } + } + output.push_back(std::move(out)); + return true; +} diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 4a15fe6d4a74..781a8d9e6306 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -194,3 +194,9 @@ struct mtmd_audio_streaming_istft { std::vector ifft_in; std::vector ifft_out; }; + +struct mtmd_audio_preprocessor_kani : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_kani(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} + void initialize() override {} + bool preprocess(const float * samples, size_t n_samples, std::vector & output) const override; +}; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 6ad9dc01f929..6371470972b7 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -1000,6 +1000,7 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { void reset() override { prompt.clear(); + speaker_embd.clear(); codes.clear(); pcm.clear(); wav.clear(); @@ -1019,8 +1020,8 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { LOG_ERR("KaniTTS-2 requires a backbone converted from KaniTTS2ForCausalLM\n"); return 1; } - if (!inp->prompt || !inp->prompt_len || inp->prompt_len > INT32_MAX - 64 || inp->speaker_ref) { - LOG_ERR("KaniTTS-2 requires text; speaker reference audio is not supported\n"); + if (!inp->prompt || !inp->prompt_len || inp->prompt_len > INT32_MAX - 64) { + LOG_ERR("KaniTTS-2 requires text\n"); return 1; } if (inp->out_type != MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM && inp->out_type != MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV) { @@ -1050,6 +1051,11 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } prompt[n + 1] = 2; // end_of_text prompt[n + 2] = 64404; // end_of_human + if (inp->speaker_ref) { + if (!encode_speaker(inp->speaker_ref)) { return 1; } + prompt.insert(prompt.begin() + 1, LLAMA_TOKEN_NULL); + if (prompt.size() >= llama_n_ctx_seq(lctx)) { return 1; } + } seq_id = inp->seq_id; out_type = inp->out_type; if (!llama_memory_seq_rm(llama_get_memory(lctx), seq_id, -1, -1)) { @@ -1168,6 +1174,32 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } private: + bool encode_speaker(mtmd_bitmap * bitmap) { + if (!mtmd_support_audio(mctx)) { + LOG_ERR("KaniTTS-2: mmproj has no speaker encoder; reconvert with speaker_encoder weights\n"); + return false; + } + const std::string marker = mtmd_default_marker(); + mtmd_input_text text{marker.c_str(), marker.size(), false, true}; + mtmd_input_chunks * chunks = mtmd_input_chunks_init(); + const mtmd_bitmap * bptr = bitmap; + bool ok = mtmd_tokenize(mctx, chunks, &text, &bptr, 1) == 0; + if (ok) { + ok = false; + for (size_t i = 0; i < mtmd_input_chunks_size(chunks); ++i) { + const auto * chunk = mtmd_input_chunks_get(chunks, i); + if (mtmd_input_chunk_get_type(chunk) != MTMD_INPUT_CHUNK_TYPE_AUDIO) { continue; } + if (mtmd_input_chunk_get_n_tokens(chunk) != 1 || mtmd_encode_chunk(mctx, chunk) != 0) { break; } + const float * embd = mtmd_get_output_embd(mctx); + speaker_embd.assign(embd, embd + n_embd); + ok = true; + break; + } + } + mtmd_input_chunks_free(chunks); + return ok; + } + int decode(const llama_token * tokens, int n, int rope_pos) { if (pos + n > (int) llama_n_ctx_seq(lctx)) { LOG_ERR("KaniTTS-2: context exhausted; increase -c\n"); @@ -1175,7 +1207,8 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { } std::vector embd((size_t) n * n_embd); for (int i = 0; i < n; ++i) { - std::copy_n(tok_embd.data() + (size_t) tokens[i] * n_embd, n_embd, embd.data() + (size_t) i * n_embd); + const float * row = tokens[i] == LLAMA_TOKEN_NULL ? speaker_embd.data() : tok_embd.data() + (size_t) tokens[i] * n_embd; + std::copy_n(row, n_embd, embd.data() + (size_t) i * n_embd); } decode_embd_batch batch(embd.data(), n, 4, n_embd); batch.set_position_mrope_1d(pos, seq_id); @@ -1200,6 +1233,7 @@ class kani_tts2_gen_audio_pipeline : public mtmd_gen_audio_pipeline { std::vector prompt; std::vector codes; std::vector tok_embd; + std::vector speaker_embd; std::vector pcm; std::vector wav; }; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 0e9432f7cf69..14e43591fe70 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1002,6 +1002,10 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_KANI_SPKENC: + { + audio_preproc = std::make_unique(ctx_a); + } break; case PROJECTOR_TYPE_POCKETTTS_SPKENC: { audio_preproc = std::make_unique(ctx_a); diff --git a/tools/tts/README.md b/tools/tts/README.md index 8a0b96fea7dd..e9201156864e 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -62,14 +62,14 @@ python convert_hf_to_gguf.py path/to/pocket-tts/languages/english --mmproj --out Available params: - `--tts-lang` can be `en_us`, `en_nyork`, `en_oakl`, `en_glasg`, `en_bost`, `en_scou`; `en` is an alias for `en_us` -- `--tts-speaker-file` is not supported +- `--tts-speaker-file` optionally provides reference audio for voice cloning - `-n` limits generated tokens; four audio tokens represent one frame Example usage: ```sh llama-tts -m kani-tts-2.gguf \ - -mm mmproj-nemo-nano-codec.gguf \ + -mm mmproj-kani-tts-2.gguf \ -p "Hello world" \ --tts-lang en_us \ -c 4096 -n 3000 --temp 1 --top-p 0.95 --top-k 0 \ @@ -79,9 +79,10 @@ llama-tts -m kani-tts-2.gguf \ **Note for GGUF conversion:** -Download [KaniTTS-2 English](https://huggingface.co/nineninesix/kani-tts-2-en) and the `.nemo` archive from [NeMo Nano Codec 22 kHz / 0.6 kbps / 12.5 fps](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps) into separate directories: +Download [KaniTTS-2 English](https://huggingface.co/nineninesix/kani-tts-2-en), place the [NeMo Nano Codec archive](https://huggingface.co/nvidia/nemo-nano-codec-22khz-0.6kbps-12.5fps) in the same directory, and download the [speaker encoder](https://huggingface.co/nineninesix/speaker-emb-tbr) into its `speaker_encoder` subdirectory: ```sh +hf download nineninesix/speaker-emb-tbr --local-dir path/to/kani-tts-2-en/speaker_encoder python convert_hf_to_gguf.py path/to/kani-tts-2-en --outtype f16 --outfile kani-tts-2.gguf -python convert_hf_to_gguf.py path/to/nemo-nano-codec --mmproj --outtype f16 --outfile mmproj-nemo-nano-codec.gguf +python convert_hf_to_gguf.py path/to/kani-tts-2-en --mmproj --outtype f16 --outfile mmproj-kani-tts-2.gguf ```