diff --git a/conversion/__init__.py b/conversion/__init__.py index 4d58bcd1060e..b0d5b6692b19 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", @@ -285,6 +286,8 @@ MMPROJ_MODEL_MAP: dict[str, str] = { + "KaniTTS2ForCausalLM": "nemo_nano_codec", + "NemoNanoCodecModel": "nemo_nano_codec", "AudioFlamingo3ForConditionalGeneration": "ultravox", "CogVLMForCausalLM": "cogvlm", "DeepseekOCR2ForCausalLM": "deepseek", diff --git a/conversion/lfm2.py b/conversion/lfm2.py index 984f44480649..05d5d033cc21 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": + # The speaker projection is stored in the mmproj. + 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/conversion/nemo_nano_codec.py b/conversion/nemo_nano_codec.py new file mode 100644 index 000000000000..69d65f06f536 --- /dev/null +++ b/conversion/nemo_nano_codec.py @@ -0,0 +1,260 @@ +# 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 + + +@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 d3a639f374c0..287a1a660c28 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() @@ -1027,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 @@ -1113,6 +1119,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 @@ -1416,6 +1432,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", @@ -1779,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", @@ -1863,6 +1885,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", @@ -2112,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, @@ -2220,6 +2257,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, @@ -4887,6 +4934,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, @@ -5812,6 +5860,8 @@ 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" PARAKEET = "parakeet" # audio 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/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 907468e87ec7..a3a5bad6c64f 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -62,6 +62,8 @@ 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 models/siglip.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 72148a4d9a9b..82a7a3f59647 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -502,6 +502,8 @@ 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, PROJECTOR_TYPE_UNKNOWN, @@ -567,6 +569,8 @@ 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 f737ccc24527..8e6bab81798e 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -485,6 +485,39 @@ 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; + 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 +819,8 @@ 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 cd6421def528..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); @@ -1116,6 +1120,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 +1840,19 @@ 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) { + 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 @@ -2078,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 @@ -2238,8 +2261,10 @@ 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_POCKETTTS_GEN && + model.proj_type != PROJECTOR_TYPE_NEMO_NANO_CODEC); // layers const int n_layers_to_load = has_standard_layers ? hparams.n_layer : 0; @@ -2910,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) @@ -3026,6 +3112,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 @@ -4340,11 +4477,16 @@ 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 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 +5408,33 @@ 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); + 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 +5989,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) { @@ -5997,8 +6168,12 @@ 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: + 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/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 5945c6d92cb7..2ffc2525da04 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -12,6 +12,21 @@ * 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) {} + 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-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 1c58d3ae1959..6371470972b7 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -993,10 +993,257 @@ 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(); + speaker_embd.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) { + 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) { + 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 + 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)) { + 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: + 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"); + return 1; + } + std::vector embd((size_t) n * n_embd); + for (int i = 0; i < n; ++i) { + 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); + 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 speaker_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 00ecadcf4dfe..14e43591fe70 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" @@ -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", @@ -1001,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); @@ -1881,6 +1886,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 +1916,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 +1940,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..e9201156864e 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -57,3 +57,32 @@ 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 ``` + +## KaniTTS-2 + +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` 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-kani-tts-2.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 +``` + +**Note for GGUF conversion:** + +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/kani-tts-2-en --mmproj --outtype f16 --outfile mmproj-kani-tts-2.gguf +``` 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;