diff --git a/common/arg.cpp b/common/arg.cpp index f733c2470..e6d6450f4 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4512,6 +4512,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.tts_lang = value; } ).set_examples({LLAMA_EXAMPLE_TTS})); + add_opt(common_arg( + {"--tts-voice"}, "VOICE", + "voice for audio generation (model-dependent)\n" + "see tts/README.md for per-model usage notes", + [](common_params & params, const std::string & value) { + params.tts_voice = value; + } + ).set_examples({LLAMA_EXAMPLE_TTS})); add_opt(common_arg( {"--tts-speaker-file"}, "FNAME", "speaker file path for audio generation", diff --git a/common/chat.cpp b/common/chat.cpp index 46b98bdbe..ad0571790 100644 --- a/common/chat.cpp +++ b/common/chat.cpp @@ -316,6 +316,7 @@ struct common_chat_templates { bool has_explicit_template; // Model had builtin template or template overridden was specified. std::unique_ptr template_default; // always set (defaults to chatml) std::unique_ptr template_tool_use; + std::unique_ptr template_tts; // the prompt of a text to speak }; common_chat_tool_choice common_chat_tool_choice_parse_oaicompat(const std::string & tool_choice) { @@ -803,19 +804,70 @@ bool common_chat_templates_was_explicit(const struct common_chat_templates * tmp return tmpls->has_explicit_template; } -common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates) { +// renders a task template (e.g. "tts") with the given variables +static std::string common_chat_template_render(const common_chat_template & tmpl, const json & vars) { + jinja::context ctx(tmpl.source()); + jinja::global_from_json(ctx, vars, false); + ctx.set_val("bos_token", jinja::mk_val(tmpl.bos_token())); + ctx.set_val("eos_token", jinja::mk_val(tmpl.eos_token())); + jinja::runtime runtime(ctx); + return jinja::runtime::gather_string_parts(runtime.execute(tmpl.prog))->as_string().str(); +} + +// LFM3 format detection: template renders deferred tools inside ... +static bool is_lfm3_template(const std::string & src) { + return src.find("") != std::string::npos; +} + +// LFM3-Audio: a known language code maps to its name, no language lets the model detect it +static std::string lfm3_asr_system_prompt(const std::string & language) { + static const std::map names = { + { "ar", "Arabic" }, { "da", "Danish" }, { "de", "German" }, { "el", "Greek" }, { "en", "English" }, + { "es", "Spanish" }, { "fi", "Finnish" }, { "fr", "French" }, { "he", "Hebrew" }, { "it", "Italian" }, + { "ja", "Japanese" }, { "ko", "Korean" }, { "nl", "Dutch" }, { "pl", "Polish" }, { "pt", "Portuguese" }, + { "sv", "Swedish" }, { "tr", "Turkish" }, { "zh", "Chinese" }, + }; + if (language.empty()) { + return "Perform ASR."; + } + std::string code = language; + std::transform(code.begin(), code.end(), code.begin(), [](unsigned char c) { return (char) std::tolower(c); }); + const auto it = names.find(code); + return "Perform ASR in " + (it != names.end() ? it->second : language) + "."; +} + +common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates, const std::string & language) { common_chat_prompt_preset asr_preset; asr_preset.system = ""; asr_preset.user = "Transcribe audio to text"; - if (chat_templates && chat_templates->template_default && is_lfm2_template(chat_templates->template_default->source())) { - asr_preset.system = "Perform ASR."; - asr_preset.user = ""; + if (chat_templates && chat_templates->template_default) { + const std::string src = chat_templates->template_default->source(); + if (is_lfm2_template(src)) { + asr_preset.system = "Perform ASR."; + asr_preset.user = ""; + } else if (is_lfm3_template(src)) { + asr_preset.system = lfm3_asr_system_prompt(language); + asr_preset.user = ""; + asr_preset.has_language = true; + } } return asr_preset; } +std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templates, const std::string & text, + const std::string & language, const std::string & voice) { + if (!chat_templates || !chat_templates->template_tts) { + return ""; + } + return common_chat_template_render(*chat_templates->template_tts, { + {"messages", json::array({{{"role", "user"}, {"content", text}}})}, + {"language", language}, + {"voice", voice}, + }); +} + std::string common_chat_templates_source(const struct common_chat_templates * tmpls, const std::string & variant) { if (!variant.empty()) { if (variant == "tool_use") { @@ -926,6 +978,16 @@ common_chat_templates_ptr common_chat_templates_init(const struct llama_model * LOG_ERR("%s: failed to parse tool use chat template (ignoring it): %s\n", __func__, e.what()); } } + // task templates of the model, kept with a chat template override + if (model) { + if (const char * src = llama_model_chat_template(model, "tts")) { + try { + tmpls->template_tts = std::make_unique(src, token_bos, token_eos); + } catch (const std::exception & e) { + LOG_ERR("%s: failed to parse the tts chat template (ignoring it): %s\n", __func__, e.what()); + } + } + } return tmpls; } diff --git a/common/chat.h b/common/chat.h index 647e56cc4..86ecc1315 100644 --- a/common/chat.h +++ b/common/chat.h @@ -491,6 +491,12 @@ std::optional common_chat_try_specialized_template( struct common_chat_prompt_preset { std::string system; std::string user; + bool has_language = false; // the preset already names the requested language }; -common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates); +common_chat_prompt_preset common_chat_get_asr_prompt(const common_chat_templates * chat_templates, const std::string & language = ""); + +// prompt of the model's "tts" chat template for the text to speak, empty if the model has none +// language and voice are optional, the template defines their defaults +std::string common_chat_get_tts_prompt(const common_chat_templates * chat_templates, const std::string & text, + const std::string & language = "", const std::string & voice = ""); diff --git a/common/common.h b/common/common.h index a0d960c5e..acdbd66ec 100644 --- a/common/common.h +++ b/common/common.h @@ -772,6 +772,7 @@ struct common_params { // TTS params std::string tts_lang = ""; + std::string tts_voice = ""; std::string tts_speaker_file = ""; bool is_gen_docs = false; // whether we are running inside llama-gen-docs diff --git a/conversion/__init__.py b/conversion/__init__.py index 4e5c5f4c6..c2a5f0b5e 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -173,6 +173,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "Lfm2Model": "lfm2", "Lfm2MoeForCausalLM": "lfm2", "Lfm3ForCausalLM": "lfm3", + "Lfm3OmniForConditionalGeneration": "lfm3", "Lfm3VlForConditionalGeneration": "lfm3", "Llama4ForCausalLM": "llama", "Llama4ForConditionalGeneration": "llama", @@ -360,6 +361,8 @@ MMPROJ_MODEL_MAP: dict[str, str] = { "Lfm2AudioForConditionalGeneration": "lfm2", "D1OmniModel": "lfm2", "Lfm2VlForConditionalGeneration": "lfm2", + "Lfm3AudioForConditionalGeneration": "lfm3", + "Lfm3OmniForConditionalGeneration": "lfm3", "Lfm3VlForConditionalGeneration": "lfm3", "LightOnOCRForConditionalGeneration": "lighton_ocr", "Llama4ForConditionalGeneration": "llama4", diff --git a/conversion/lfm2.py b/conversion/lfm2.py index b1f0d1d32..7f4f792e7 100644 --- a/conversion/lfm2.py +++ b/conversion/lfm2.py @@ -15,25 +15,30 @@ from .base import MmprojModel, ModelBase, TextModel, gguf, jinja_str_or_json, lo from .gemma import ConformerAudioModel +def lfm2_feed_forward_length(hparams: dict[str, Any]) -> int: + ff_dim = hparams.get("block_ff_dim", hparams.get("intermediate_size")) + assert ff_dim is not None + auto_adjust_ff_dim = hparams["block_auto_adjust_ff_dim"] + ffn_dim_multiplier = hparams["block_ffn_dim_multiplier"] + multiple_of = hparams["block_multiple_of"] + + if auto_adjust_ff_dim: + ff_dim = int(2 * ff_dim / 3) + # custom dim factor multiplier + if ffn_dim_multiplier is not None: + ff_dim = int(ffn_dim_multiplier * ff_dim) + ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of) + + return ff_dim + + @ModelBase.register("Lfm2ForCausalLM", "LFM2ForCausalLM") @ModelBase.example("LiquidAI/LFM2-1.2B", "LiquidAI/LFM2.5-350M") class LFM2Model(TextModel): model_arch = gguf.MODEL_ARCH.LFM2 def _add_feed_forward_length(self): - ff_dim = self.find_hparam(["block_ff_dim", "intermediate_size"]) - auto_adjust_ff_dim = self.hparams["block_auto_adjust_ff_dim"] - ffn_dim_multiplier = self.hparams["block_ffn_dim_multiplier"] - multiple_of = self.hparams["block_multiple_of"] - - if auto_adjust_ff_dim: - ff_dim = int(2 * ff_dim / 3) - # custom dim factor multiplier - if ffn_dim_multiplier is not None: - ff_dim = int(ffn_dim_multiplier * ff_dim) - ff_dim = multiple_of * ((ff_dim + multiple_of - 1) // multiple_of) - - self.gguf_writer.add_feed_forward_length(ff_dim) + self.gguf_writer.add_feed_forward_length(lfm2_feed_forward_length(self.hparams)) def set_gguf_parameters(self): # set num_key_value_heads only for attention layers diff --git a/conversion/lfm3.py b/conversion/lfm3.py index 79ff6a196..6f80efabc 100644 --- a/conversion/lfm3.py +++ b/conversion/lfm3.py @@ -1,6 +1,10 @@ from __future__ import annotations -from typing import Iterable, TYPE_CHECKING +import json +import math + +from pathlib import Path +from typing import Any, Callable, Iterable, TYPE_CHECKING import torch @@ -8,13 +12,73 @@ if TYPE_CHECKING: from torch import Tensor from .base import MmprojModel, ModelBase, TextModel, gguf +from .lfm2 import lfm2_feed_forward_length -@ModelBase.register("Lfm3ForCausalLM", "Lfm3VlForConditionalGeneration") +@ModelBase.register("Lfm3ForCausalLM", "Lfm3VlForConditionalGeneration", "Lfm3OmniForConditionalGeneration") @ModelBase.example("LiquidAI/LFM3-1.7B", "LiquidAI/LFM3-1.7B-Base") class LFM3Model(TextModel): model_arch = gguf.MODEL_ARCH.LFM3 + # LFM3-Audio/Omni: encoder, projector and audio-output heads live next to the text model + _media_prefixes = ("encoder.", "mlp.", "codebook_embedding.", "depth_embeddings.", "depth_linear.", "depthformer.", + "mask_embedding", "model.codebook_embedding.", "model.depth_decoder.") + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + if name.startswith(cls._media_prefixes): + return None + + name = name.removeprefix("lfm.") + + return super().filter_tensors((name, gen)) + + # LFM3-Audio live ASR prompt of the model card, the audio follows in the assistant turn + _asr_template = ( + "{{- bos_token -}}" + "{{- '<|im_start|>system\\nPerform live ASR.<|im_end|>\\n' -}}" + "{%- if add_generation_prompt -%}{{- '<|im_start|>assistant\\n' -}}{%- endif -%}" + ) + + # TTS prompt of the model card: a language code is looked up, anything else is taken as the language name + _tts_langs = { + "ar": "Arabic", "da": "Danish", "de": "German", "el": "Greek", "en": "American English", "en-us": "American English", + "en-gb": "British English", "fi": "Finnish", "fr": "French", "he": "Hebrew", "it": "Italian", "ja": "Japanese", + "ko": "Korean", "nl": "Dutch", "pl": "Polish", "pt": "Portuguese", "es": "Spanish", "sv": "Swedish", "tr": "Turkish", + "zh": "Chinese", + } + + # "language" and "voice" are optional, the text to speak is the first message + def _tts_template(self, think: bool) -> str: + return ( + "{%- set names = " + json.dumps(self._tts_langs) + " -%}" + "{%- set code = (language | lower) if language is defined and language else 'en' -%}" + "{%- set lang = names[code] if code in names else language -%}" + "{%- set v = (voice | lower) if voice is defined and voice and (voice | lower) in ['female', 'male'] else 'female' -%}" + "{{- bos_token -}}" + "{{- '<|im_start|>system\\nPerform TTS in ' + lang + ' using a ' + v + ' voice.<|im_end|>\\n' -}}" + "{{- '<|im_start|>user\\n' + messages[0]['content'] + '<|im_end|>\\n' -}}" + "{{- '<|im_start|>assistant\\n" + ("" if think else "") + "<|audio_start|>' -}}" + ) + + def set_vocab(self): + super().set_vocab() + # hparams hold the merged text config, the checkpoint's own model_type is in config.json + with open(self.dir_model / "config.json", encoding="utf-8") as f: + model_type = json.load(f).get("model_type") + is_audio = model_type == "lfm3_audio" + is_omni = model_type == "lfm3_omni" + templates = [] + if is_audio: + templates.append({"name": "asr", "template": self._asr_template}) + if is_audio or is_omni: + # LFM3-Omni opens every assistant turn with an empty think block + templates.append({"name": "tts", "template": self._tts_template(think=is_omni)}) + if templates: + self.gguf_writer.add_chat_template(templates) + def set_gguf_parameters(self): layer_types = self.hparams["layer_types"] n_kv_heads = self.hparams["num_key_value_heads"] @@ -55,28 +119,19 @@ class LFM3Model(TextModel): yield from super().modify_tensors(data_torch, name, bid) -@ModelBase.register("Lfm3VlForConditionalGeneration") -@ModelBase.example("LiquidAI/LFM3-VL-2B") -class LFM3VLMmprojModel(MmprojModel): - has_vision_encoder = True - has_audio_encoder = False +class LFM3VLVisionMixin(MmprojModel): + """LFM3-VL vision tower and projector, shared with LFM3-Omni.""" - def __init__(self, *args, **kwargs): - super().__init__(*args, **kwargs) + def _set_lfm3vl_gguf_parameters(self): assert self.hparams_vision is not None - # dynamic resolution, image_size is only set for MmprojModel compatibility - self.hparams_vision["image_size"] = 512 - - def set_gguf_parameters(self): - super().set_gguf_parameters() - assert self.hparams_vision is not None - # clip.cpp hardcodes the vision rope base, the RGB patchifier, the 3x3 spatial conv and the GELU projector - if (self.hparams_vision["rope_parameters"]["rope_theta"] != 1e6 or self.hparams_vision["in_channels"] != 3 - or self.hparams_vision["conv_kernel_size"] != 3): - raise ValueError("unsupported vision rope_theta / in_channels / conv_kernel_size") + # clip.cpp hardcodes the GELU projector and its input norm eps if (self.global_config["projector_hidden_act"] != "gelu" or not self.global_config["projector_bias"] or self.global_config["projector_layer_norm_eps"] != 1e-6): raise ValueError("unsupported projector_hidden_act / projector_bias / projector_layer_norm_eps") + # clip.cpp hardcodes the vision rope base, the RGB patchifier and the 3x3 spatial conv + if (self.hparams_vision["rope_parameters"]["rope_theta"] != 1e6 or self.hparams_vision["in_channels"] != 3 + or self.hparams_vision["conv_kernel_size"] != 3): + raise ValueError("unsupported vision rope_theta / in_channels / conv_kernel_size") self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.LFM3VL) self.gguf_writer.add_vision_head_count_kv(self.hparams_vision["num_key_value_heads"]) self.gguf_writer.add_vision_spatial_merge_size(self.global_config["downsample_factor"]) @@ -100,10 +155,7 @@ class LFM3VLMmprojModel(MmprojModel): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) - def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - if not name.startswith(("model.vision_tower.", "model.multi_modal_projector.")): - return - + def _lfm3vl_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.vision_tower.patchifier.weight": # qwen2vl-style layout [channel, temporal, ph, pw], one conv kernel per temporal slice assert self.hparams_vision is not None @@ -137,3 +189,368 @@ class LFM3VLMmprojModel(MmprojModel): return yield (self.map_tensor_name(name), data_torch) + + +@ModelBase.register("Lfm3VlForConditionalGeneration") +@ModelBase.example("LiquidAI/LFM3-VL-2B") +class LFM3VLMmprojModel(LFM3VLVisionMixin): + has_vision_encoder = True + has_audio_encoder = False + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams_vision is not None + # dynamic resolution, image_size is only set for MmprojModel compatibility + self.hparams_vision["image_size"] = 512 + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self._set_lfm3vl_gguf_parameters() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith(("model.vision_tower.", "model.multi_modal_projector.")): + yield from self._lfm3vl_tensors(data_torch, name, bid) + + +def _is_lfm3_audio_checkpoint(dir_model: Path) -> bool: + if not (dir_model / "config.json").is_file(): + return False + with open(dir_model / "config.json", encoding="utf-8") as f: + config = json.load(f) + return config.get("model_type") == "lfm3_audio" and "lfm" in config + + +# read the release config.json as is: a transformers that knows lfm3_audio rewrites it into another layout +@ModelBase.register_hparams_loader(_is_lfm3_audio_checkpoint) +def _load_lfm3_audio_hparams(dir_model: Path) -> dict[str, Any]: + with open(dir_model / "config.json", encoding="utf-8") as f: + config = json.load(f) + config["text_config"] = config["lfm"] + return config + + +@ModelBase.register("Lfm3AudioForConditionalGeneration") +@ModelBase.example("LiquidAI/LFM3-Audio-2B") +class LFM3AudioMmprojModel(MmprojModel): + has_vision_encoder = False + has_audio_encoder = True + + # the downsampler has the LFM2-Audio pre-encoder layout: conv, dw conv, pw conv, dw conv, pw conv + _downsample_idx = {"conv_in": 0, "conv_mid": 2, "lin_mid": 3, "conv_end": 5, "lin_end": 6} + + # audio output: MaskGIT depthformer (a.gen.code.*) and the codec decoder (a.gen.wav.*) + _depth_layer_map = { + "operator_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_NORM, + "self_attn.q_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q, + "self_attn.q_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_Q_NORM, + "self_attn.k_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K, + "self_attn.k_norm": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_K_NORM, + "self_attn.v_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_V, + "self_attn.o_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_OUT, + "self_attn.gate_proj": gguf.MODEL_TENSOR.A_GEN_CODE_ATTN_GATE, + "ffn_norm": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_NORM, + "feed_forward.w1": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_GATE, + "feed_forward.w3": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_UP, + "feed_forward.w2": gguf.MODEL_TENSOR.A_GEN_CODE_FFN_DOWN, + } + _codec_layer_map = { + "operator_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_NORM, + "post_operator_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_POST_NORM, + "self_attn.q_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q, + "self_attn.q_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q_NORM, + "self_attn.k_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K, + "self_attn.k_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K_NORM, + "self_attn.v_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_V, + "self_attn.o_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_OUT, + "self_attn.gate_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_GATE, + "conv.in_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW1, + "conv.conv": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_DW, + "conv.out_proj": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW2, + "ffn_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_NORM, + "post_ffn_norm": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_POST_NORM, + "feed_forward.w1": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_GATE, + "feed_forward.w3": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP, + "feed_forward.w2": gguf.MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN, + } + # the codec graph in clip.cpp hardcodes the liquid-codec FSQ front end and decoder shape, conversion fails on a mismatch + _codec = { + "model_type": "liquid_codec", "fsq_levels": [8, 7, 6, 6], "n_q": 12, "quantizer_dim": 256, + "encoder_dim": 512, "stft_n_fft": 1280, "stft_hop_length": 320, "stft_win_length": 1280, "stft_frame_stack": 6, + "decoder_window": 30, "sampling_rate": 24000, + } + # the lfm3a pipeline in mtmd-helper-gen.cpp hardcodes these + _generation = {"codebooks": 12, "maskgit_schedule": [4, 4, 4], "interleaved_n_text": 6, "interleaved_n_audio": 12} + + # speech-code sampling of the model card, generation_config.json overrides it + _audio_sampling = { + "audio_temperature": 1.0, "audio_top_k": 128, "audio_repetition_penalty": 1.1, "audio_repetition_window": 4096, + } + + # mtmd_audio_preprocessor_lfm3a hardcodes this front end, conversion fails on a mismatch + _feature_extractor = { + "sampling_rate": 16000, "feature_size": 128, "n_fft": 400, "win_length": 400, "hop_length": 160, + "log_floor": -12.0, "normalization_offset": 4.0, "normalization_divisor": 4.0, "encoder_stride": 8, + } + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # audio output is exported when the checkpoint ships the codec + self._audio_output = (self.dir_model / "audio_tokenizer").is_dir() + gen = self.get_generation_config() + self._n_codebooks: int = gen["codebooks"] + self._audio_vocab: int = gen["audio_vocab_size"] + self._depth_embd: dict[int, Tensor] = {} + self._depth_norm: dict[int, Tensor] = {} + + def get_audio_config(self) -> dict[str, Any] | None: + enc = self.global_config["encoder"] + cfg = dict(enc["lfm"]) + cfg["num_mel_bins"] = enc["dim_in"] + cfg["sliding_window"] = enc["sliding_window"] + return cfg + + def get_generation_config(self) -> dict[str, Any]: + c = self.global_config + steps = c["maskgit_steps"] + return { + "codebooks": c["codebooks"], "audio_vocab_size": c["audio_vocab_size"], + "maskgit_schedule": [c["codebooks"] // steps] * steps, + "interleaved_n_text": c["interleaved_n_text"], "interleaved_n_audio": c["interleaved_n_audio"], + } + + def get_depth_config(self) -> dict[str, Any]: + return self.global_config["depthformer"] + + def set_gguf_parameters(self): + assert self.hparams_audio is not None + fe = self.preprocessor_config["feature_extractor"] + if fe["feature_extractor_type"] not in ("Lfm3AudioFeatureExtractor", "Lfm3OmniFeatureExtractor"): + raise ValueError(f"unsupported feature_extractor_type {fe['feature_extractor_type']!r}") + for k, v in self._feature_extractor.items(): + if fe[k] != v: + raise ValueError(f"unsupported feature_extractor.{k}={fe[k]!r}, expected {v!r}") + + # clip.cpp hardcodes the encoder rope base + if self.hparams_audio["rope_parameters"]["rope_theta"] != 1e6: + raise ValueError("unsupported audio encoder rope_theta") + + self.hparams_audio["intermediate_size"] = lfm2_feed_forward_length(self.hparams_audio) + super().set_gguf_parameters() + self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.LFM3A) + self.gguf_writer.add_audio_head_count_kv(self.hparams_audio["num_key_value_heads"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio["norm_eps"]) + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_window_size(self.hparams_audio["sliding_window"]) + + if not self._audio_output: + return + + # audio generation: depthformer hparams, the codec ones are constant and live in clip.cpp + gen = self.get_generation_config() + for k, v in self._generation.items(): + if gen[k] != v: + raise ValueError(f"unsupported {k}={gen[k]!r}, expected {v!r}") + depth = self.get_depth_config() + if depth.get("is_causal", False) or depth["use_sandwich_norm"] or set(depth["layer_types"]) != {"full_attention"}: + raise ValueError("unsupported depthformer, expected a bidirectional all-attention stack") + if depth["rope_parameters"]["full_attention"]["rope_theta"] != 1e6: + raise ValueError("unsupported depthformer rope_theta") + self.gguf_writer.add_clip_has_gen_audio_encoder(True) + self.gguf_writer.add_clip_gen_audio_projector_type(gguf.VisionProjectorType.LFM3A_GEN) + self.gguf_writer.add_gen_audio_projection_dim(self.n_embd_text) + self.gguf_writer.add_gen_audio_embedding_length(depth["hidden_size"]) + self.gguf_writer.add_gen_audio_feed_forward_length(depth["intermediate_size"]) + self.gguf_writer.add_gen_audio_block_count(depth["num_hidden_layers"]) + self.gguf_writer.add_gen_audio_head_count(depth["num_attention_heads"]) + self.gguf_writer.add_gen_audio_head_count_kv(depth["num_key_value_heads"]) + self.gguf_writer.add_gen_audio_attention_layernorm_eps(depth["norm_eps"]) + + # sampling defaults of the speech codes, next to the text ones that become general.sampling.* + gen = gguf.Metadata.load_generation_config(self.dir_model) + for key, add in (("audio_temperature", self.gguf_writer.add_gen_audio_sampling_temp), + ("audio_top_k", self.gguf_writer.add_gen_audio_sampling_top_k), + ("audio_repetition_penalty", self.gguf_writer.add_gen_audio_sampling_penalty_repeat), + ("audio_repetition_window", self.gguf_writer.add_gen_audio_sampling_penalty_last_n)): + value = gen.get(key, self._audio_sampling.get(key)) + if value is not None: + add(value) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if new_name.startswith("a.conv1d.") or "conv_dw" in new_name or new_name.startswith("a.gen.wav.quant."): + return gguf.GGMLQuantizationType.F32 + # the codec decoder is small and its log-magnitude/phase output does not survive quantization + if new_name.startswith("a.gen.wav."): + return gguf.GGMLQuantizationType.F16 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + if self._audio_output: + yield from self._codec_tensors() + + def _codec_tensors(self) -> Iterable[tuple[str, Tensor]]: + import json + from safetensors.torch import load_file + + T = gguf.MODEL_TENSOR + codec_dir = self.dir_model / "audio_tokenizer" + with open(codec_dir / "config.json", encoding="utf-8") as f: + cfg = json.load(f) + for k, v in self._codec.items(): + if cfg.get(k) != v: + raise ValueError(f"unsupported audio_tokenizer {k}={cfg.get(k)!r}, expected {v!r}") + if math.prod(cfg["fsq_levels"]) != self._audio_vocab - 1: + raise ValueError(f"audio_vocab_size {self._audio_vocab} does not match the codec, expected codebook size + 1") + + state_dict = load_file(codec_dir / "model.safetensors") + + # FSQ code table: index -> the len(levels) quantized values, see FiniteScalarQuantization._indices_to_codes + levels = torch.tensor(cfg["fsq_levels"], dtype=torch.int64) + basis = torch.cumprod(torch.cat([torch.ones(1, dtype=torch.int64), levels[:-1]]), dim=0) + half = levels // 2 + idx = torch.arange(int(levels.prod()), dtype=torch.int64)[:, None] + codes = ((idx // basis) % levels - half) / half + yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_CB), codes.to(torch.float32)) + + # one projection over the concatenated code vectors of all codebooks, under quantizer.fsq. in newer exports + quantizer = "quantizer.fsq." if "quantizer.fsq.project_out.weight" in state_dict else "quantizer." + for suffix in ("weight", "bias"): + yield (self.format_tensor_name(T.A_GEN_WAV_QUANT_FIRST_OUT, suffix=f".{suffix}"), state_dict[f"{quantizer}project_out.{suffix}"]) + yield (self.format_tensor_name(T.A_GEN_WAV_TFM_IN_PROJ, suffix=f".{suffix}"), state_dict[f"from_quantizer.{suffix}"]) + yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUT_PROJ, suffix=f".{suffix}"), state_dict[f"decoder.lin.{suffix}"]) + yield (self.format_tensor_name(T.A_GEN_WAV_TFM_OUTPUT_NORM), state_dict["decoder.lfm.embedding_norm.weight"]) + + for name, data_torch in state_dict.items(): + if not name.startswith("decoder.lfm.layers."): + continue + bid_str, key = name.removeprefix("decoder.lfm.layers.").rsplit(".weight", 1)[0].split(".", 1) + if key == "conv.conv": + data_torch = data_torch.squeeze(1) + yield (self.format_tensor_name(self._codec_layer_map[key], int(bid_str)), data_torch) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith("a.gen.wav."): + yield (name, data_torch) + return + + if name.startswith(("depthformer.", "depth_", "codebook_embedding.", "mask_embedding")): + if self._audio_output: + yield from self._depthformer_tensor(data_torch, name, bid) + return + + if not name.startswith(("encoder.", "mlp.")): + return + + if name.startswith("encoder.downsample."): + layer, suffix = name.removeprefix("encoder.downsample.").split(".") + if layer == "lin_out": + if suffix == "weight": + data_torch = data_torch.squeeze(-1) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_OUT, suffix=f".{suffix}"), data_torch) + return + if suffix == "bias": + data_torch = data_torch[:, None, None] + yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_CONV1D, self._downsample_idx[layer], f".{suffix}"), data_torch) + return + + # conv op requires 2d tensor + if name.endswith("conv.conv.weight"): + data_torch = data_torch.squeeze(1) + + yield (self.map_tensor_name(name), data_torch) + + def _depthformer_tensor(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + T = gguf.MODEL_TENSOR + + if name.startswith("depthformer.model.layers."): + key = name.removeprefix("depthformer.model.layers.").split(".", 1)[1].rsplit(".weight", 1)[0] + yield (self.format_tensor_name(self._depth_layer_map[key], bid), data_torch) + return + if name == "depthformer.model.embedding_norm.weight": + yield (self.format_tensor_name(T.A_GEN_CODE_OUTPUT_NORM), data_torch) + return + if name.startswith("depth_linear."): + yield (self.format_tensor_name(T.A_GEN_CODE_PROJ_IN, suffix="." + name.rsplit(".", 1)[1]), data_torch) + return + if name == "mask_embedding": + yield (self.format_tensor_name(T.A_GEN_CODE_MASK_EMBD, suffix=""), data_torch) + return + if name == "codebook_embedding.weight": + assert data_torch.shape[0] == self._n_codebooks * self._audio_vocab + yield (self.format_tensor_name(T.A_GEN_CODE_OUT_EMBD), data_torch) + return + + # per-codebook heads: the embedding table is tied to to_logits, merge the 12 tables into one tensor + assert name.startswith("depth_embeddings.") + idx = int(name.split(".")[1]) + key = name.split(".", 2)[2] + if key == "embedding.weight": + self._depth_embd[idx] = data_torch + if len(self._depth_embd) == self._n_codebooks: + yield (self.format_tensor_name(T.A_GEN_CODE_EMBD), torch.stack([self._depth_embd[i] for i in range(self._n_codebooks)], dim=0)) + elif key == "embedding_norm.weight": + self._depth_norm[idx] = data_torch + if len(self._depth_norm) == self._n_codebooks: + yield (self.format_tensor_name(T.A_GEN_CODE_HEAD_NORM), torch.stack([self._depth_norm[i] for i in range(self._n_codebooks)], dim=0)) + elif key == "to_logits.weight": + if not torch.equal(data_torch, self._depth_embd[idx]): + raise ValueError(f"depth_embeddings.{idx} is not tied") + else: + raise ValueError(f"unexpected tensor {name}") + + +@ModelBase.register("Lfm3OmniForConditionalGeneration") +@ModelBase.example("LiquidAI/LFM3-Omni-2B") +class LFM3OmniMmprojModel(LFM3VLVisionMixin, LFM3AudioMmprojModel): + has_vision_encoder = True + has_audio_encoder = True + + _audio_sampling = {"audio_temperature": 0.4, "audio_repetition_penalty": 1.2, "audio_repetition_window": 64} + + # the audio side is the LFM3-Audio model under the HF-native names + _audio_renames = ( + ("model.audio_encoder.lfm.", "encoder.lfm.lfm."), + ("model.audio_encoder.", "encoder."), + ("model.audio_projector.layer_norm.", "mlp.model.0."), + ("model.audio_projector.linear_1.", "mlp.model.1."), + ("model.audio_projector.linear_2.", "mlp.model.3."), + ("model.depth_decoder.model.", "depthformer.model."), + ("model.depth_decoder.", ""), + ("model.codebook_embedding.", "codebook_embedding."), + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + assert self.hparams_vision is not None + # dynamic resolution, image_size is only set for MmprojModel compatibility + self.hparams_vision["image_size"] = 512 + + def get_audio_config(self) -> dict[str, Any] | None: + enc = self.global_config["audio_encoder_config"] + cfg = dict(enc["text_config"]) + cfg["num_mel_bins"] = enc["num_mel_bins"] + cfg["sliding_window"] = enc["sliding_window"] + return cfg + + def get_generation_config(self) -> dict[str, Any]: + c = self.global_config + return { + "codebooks": c["num_codebooks"], "audio_vocab_size": c["codebook_size"], "maskgit_schedule": c["maskgit_schedule"], + "interleaved_n_text": c["interleaved_n_text"], "interleaved_n_audio": c["interleaved_n_audio"], + } + + def get_depth_config(self) -> dict[str, Any]: + return self.global_config["depth_config"] + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self._set_lfm3vl_gguf_parameters() + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name.startswith(("model.vision_tower.", "model.multi_modal_projector.")): + yield from self._lfm3vl_tensors(data_torch, name, bid) + return + for old, new in self._audio_renames: + if name.startswith(old): + yield from super().modify_tensors(data_torch, new + name.removeprefix(old), bid) + return + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d94551a4f..51474c29e 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -501,6 +501,7 @@ class Keys: class Attention: HEAD_COUNT = "clip.audio.attention.head_count" + HEAD_COUNT_KV = "clip.audio.attention.head_count_kv" HEAD_DIM = "clip.audio.attention.head_dim" # set when qkv width != n_embd LAYERNORM_EPS = "clip.audio.attention.layer_norm_epsilon" @@ -518,6 +519,11 @@ class Keys: FEED_FORWARD_LENGTH = "clip.gen.audio.feed_forward_length" BLOCK_COUNT = "clip.gen.audio.block_count" PROJECTION_DIM = "clip.gen.audio.projection_dim" + # sampling defaults of the generated audio, like general.sampling.* for text + SAMPLING_TEMP = "clip.gen.audio.sampling.temp" + SAMPLING_TOP_K = "clip.gen.audio.sampling.top_k" + SAMPLING_PENALTY_REPEAT = "clip.gen.audio.sampling.penalty_repeat" + SAMPLING_PENALTY_LAST_N = "clip.gen.audio.sampling.penalty_last_n" class Attention: HEAD_COUNT = "clip.gen.audio.attention.head_count" @@ -1176,6 +1182,8 @@ class MODEL_TENSOR(IntEnum): A_ENC_ATTN_Q = auto() A_ENC_ATTN_K = auto() A_ENC_ATTN_V = auto() + A_ENC_ATTN_Q_NORM = auto() # lfm3a + A_ENC_ATTN_K_NORM = auto() # lfm3a A_ENC_ATTN_POST_NORM = auto() A_ENC_ATTN_PRE_NORM = auto() A_ENC_ATTN_K_REL = auto() # gemma4 @@ -1220,6 +1228,9 @@ class MODEL_TENSOR(IntEnum): A_GEN_CODE_FFN_UP = auto() A_GEN_CODE_FFN_DOWN = auto() A_GEN_CODE_OUTPUT_NORM = auto() + A_GEN_CODE_ATTN_GATE = auto() # lfm3a: per-head sigmoid gate on the attention output + A_GEN_CODE_HEAD_NORM = auto() # lfm3a: per-codebook norm before the output head, merged 2D [n_codebooks, dim] + A_GEN_CODE_MASK_EMBD = auto() # lfm3a: MaskGIT mask embedding # qwen3tts code2wav: RVQ codes -> raw PCM A_GEN_WAV_QUANT_FIRST_IN = auto() # semantic RVQ, in_proj (1x1 conv, loaded as 2D) A_GEN_WAV_QUANT_FIRST_OUT = auto() # semantic RVQ, out_proj @@ -1242,6 +1253,14 @@ class MODEL_TENSOR(IntEnum): A_GEN_WAV_TFM_FFN_UP = auto() A_GEN_WAV_TFM_FFN_DOWN = auto() A_GEN_WAV_TFM_FFN_SCALE = auto() # layer scale (gamma) on the FFN output + A_GEN_WAV_TFM_ATTN_Q_NORM = auto() # lfm3a + A_GEN_WAV_TFM_ATTN_K_NORM = auto() # lfm3a + A_GEN_WAV_TFM_CONV_PW1 = auto() # lfm3a: short conv in_proj + A_GEN_WAV_TFM_CONV_DW = auto() # lfm3a: short conv depthwise kernel + A_GEN_WAV_TFM_CONV_PW2 = auto() # lfm3a: short conv out_proj + A_GEN_WAV_TFM_ATTN_GATE = auto() # lfm3a: per-head sigmoid gate on the attention output + A_GEN_WAV_TFM_ATTN_POST_NORM = auto() # lfm3a: sandwich norm after the operator + A_GEN_WAV_TFM_FFN_POST_NORM = auto() # lfm3a: sandwich norm after the FFN A_GEN_WAV_UP_CONV = auto() # causal ConvTranspose1d, 2x upsample A_GEN_WAV_UP_DWCONV = auto() # ConvNeXt depthwise conv A_GEN_WAV_UP_NORM = auto() # ConvNeXt LayerNorm @@ -2004,6 +2023,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.A_ENC_ATTN_Q: "a.blk.{bid}.attn_q", MODEL_TENSOR.A_ENC_ATTN_K: "a.blk.{bid}.attn_k", MODEL_TENSOR.A_ENC_ATTN_V: "a.blk.{bid}.attn_v", + MODEL_TENSOR.A_ENC_ATTN_Q_NORM: "a.blk.{bid}.attn_q_norm", + MODEL_TENSOR.A_ENC_ATTN_K_NORM: "a.blk.{bid}.attn_k_norm", MODEL_TENSOR.A_ENC_ATTN_POST_NORM: "a.blk.{bid}.attn_post_norm", MODEL_TENSOR.A_ENC_ATTN_PRE_NORM: "a.blk.{bid}.attn_pre_norm", MODEL_TENSOR.A_ENC_ATTN_K_REL: "a.blk.{bid}.attn_k_rel", @@ -2047,6 +2068,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.A_GEN_CODE_FFN_UP: "a.gen.code.blk.{bid}.ffn_up", MODEL_TENSOR.A_GEN_CODE_FFN_DOWN: "a.gen.code.blk.{bid}.ffn_down", MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM: "a.gen.code.output_norm", + MODEL_TENSOR.A_GEN_CODE_ATTN_GATE: "a.gen.code.blk.{bid}.attn_gate", + MODEL_TENSOR.A_GEN_CODE_HEAD_NORM: "a.gen.code.head_norm", + MODEL_TENSOR.A_GEN_CODE_MASK_EMBD: "a.gen.code.mask_embd", MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN: "a.gen.wav.quant.first.in_proj", MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT: "a.gen.wav.quant.first.out_proj", MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB: "a.gen.wav.quant.first.codebook", @@ -2068,6 +2092,14 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP: "a.gen.wav.tfm.blk.{bid}.ffn_up", MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN: "a.gen.wav.tfm.blk.{bid}.ffn_down", MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE: "a.gen.wav.tfm.blk.{bid}.ls2", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q_NORM: "a.gen.wav.tfm.blk.{bid}.attn_q_norm", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K_NORM: "a.gen.wav.tfm.blk.{bid}.attn_k_norm", + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW1: "a.gen.wav.tfm.blk.{bid}.conv_pw1", + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_DW: "a.gen.wav.tfm.blk.{bid}.conv_dw", + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW2: "a.gen.wav.tfm.blk.{bid}.conv_pw2", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_GATE: "a.gen.wav.tfm.blk.{bid}.attn_gate", + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_POST_NORM: "a.gen.wav.tfm.blk.{bid}.attn_post_norm", + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_POST_NORM: "a.gen.wav.tfm.blk.{bid}.ffn_post_norm", MODEL_TENSOR.A_GEN_WAV_UP_CONV: "a.gen.wav.up.blk.{bid}.conv", MODEL_TENSOR.A_GEN_WAV_UP_DWCONV: "a.gen.wav.up.blk.{bid}.dwconv", MODEL_TENSOR.A_GEN_WAV_UP_NORM: "a.gen.wav.up.blk.{bid}.norm", @@ -2364,6 +2396,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.A_ENC_ATTN_Q, MODEL_TENSOR.A_ENC_ATTN_K, MODEL_TENSOR.A_ENC_ATTN_V, + MODEL_TENSOR.A_ENC_ATTN_Q_NORM, + MODEL_TENSOR.A_ENC_ATTN_K_NORM, MODEL_TENSOR.A_ENC_ATTN_POST_NORM, MODEL_TENSOR.A_ENC_ATTN_PRE_NORM, MODEL_TENSOR.A_ENC_ATTN_K_REL, @@ -2431,6 +2465,9 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.A_GEN_CODE_FFN_UP, MODEL_TENSOR.A_GEN_CODE_FFN_DOWN, MODEL_TENSOR.A_GEN_CODE_OUTPUT_NORM, + MODEL_TENSOR.A_GEN_CODE_ATTN_GATE, + MODEL_TENSOR.A_GEN_CODE_HEAD_NORM, + MODEL_TENSOR.A_GEN_CODE_MASK_EMBD, MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_IN, MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_OUT, MODEL_TENSOR.A_GEN_WAV_QUANT_FIRST_CB, @@ -2452,6 +2489,14 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.A_GEN_WAV_TFM_FFN_UP, MODEL_TENSOR.A_GEN_WAV_TFM_FFN_DOWN, MODEL_TENSOR.A_GEN_WAV_TFM_FFN_SCALE, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_Q_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_K_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW1, + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_DW, + MODEL_TENSOR.A_GEN_WAV_TFM_CONV_PW2, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_GATE, + MODEL_TENSOR.A_GEN_WAV_TFM_ATTN_POST_NORM, + MODEL_TENSOR.A_GEN_WAV_TFM_FFN_POST_NORM, MODEL_TENSOR.A_GEN_WAV_UP_CONV, MODEL_TENSOR.A_GEN_WAV_UP_DWCONV, MODEL_TENSOR.A_GEN_WAV_UP_NORM, @@ -6547,6 +6592,8 @@ class VisionProjectorType: DEEPSEEK4V = "deepseek4v" DEEPSEEK41V = "deepseek41v" LFM2A = "lfm2a" # audio + LFM3A = "lfm3a" # audio + LFM3A_GEN = "lfm3a_gen" # audio generation: MaskGIT depthformer + liquid-codec decoder MUSIC_FLAMINGO = "musicflamingo" # audio GLM4V = "glm4v" GLM5V = "glm5v" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 605013610..efdf8f09c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1628,6 +1628,9 @@ class GGUFWriter: def add_audio_head_count(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Attention.HEAD_COUNT, value) + def add_audio_head_count_kv(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.Attention.HEAD_COUNT_KV, value) + def add_audio_head_dim(self, value: int) -> None: self.add_uint32(Keys.ClipAudio.Attention.HEAD_DIM, value) @@ -1708,6 +1711,18 @@ class GGUFWriter: def add_gen_audio_block_count(self, value: int) -> None: self.add_uint32(Keys.ClipGenAudio.BLOCK_COUNT, value) + def add_gen_audio_sampling_temp(self, value: float) -> None: + self.add_float32(Keys.ClipGenAudio.SAMPLING_TEMP, value) + + def add_gen_audio_sampling_top_k(self, value: int) -> None: + self.add_int32(Keys.ClipGenAudio.SAMPLING_TOP_K, value) + + def add_gen_audio_sampling_penalty_repeat(self, value: float) -> None: + self.add_float32(Keys.ClipGenAudio.SAMPLING_PENALTY_REPEAT, value) + + def add_gen_audio_sampling_penalty_last_n(self, value: int) -> None: + self.add_int32(Keys.ClipGenAudio.SAMPLING_PENALTY_LAST_N, value) + def add_gen_audio_head_count(self, value: int) -> None: self.add_uint32(Keys.ClipGenAudio.Attention.HEAD_COUNT, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index beacc5a67..d06dc4ec5 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2554,12 +2554,14 @@ class TensorNameMap: "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note "audio_tower.ln_post", # qwen2omni "encoder.layer_norm", # mimo-audio-tokenizer + "encoder.lfm.lfm.embedding_norm", # lfm3a ), MODEL_TENSOR.A_ENC_ATTN_Q: ( "audio_tower.layers.{bid}.self_attn.q_proj", # ultravox "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.q_proj", # dots3note "conformer.layers.{bid}.self_attn.linear_q", # lfm2 + "encoder.lfm.lfm.layers.{bid}.self_attn.q_proj", # lfm3a "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n "conformer.layers.{bid}.self_attn.q_proj", # gemma4 "sound_encoder.encoder.layers.{bid}.self_attn.q_proj", # parakeet @@ -2571,6 +2573,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn.k_proj", # ultravox "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.k_proj", # dots3note "conformer.layers.{bid}.self_attn.linear_k", # lfm2 + "encoder.lfm.lfm.layers.{bid}.self_attn.k_proj", # lfm3a "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n "conformer.layers.{bid}.self_attn.k_proj", # gemma4 "sound_encoder.encoder.layers.{bid}.self_attn.k_proj", # parakeet @@ -2582,6 +2585,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn.v_proj", # ultravox "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.v_proj", # dots3note "conformer.layers.{bid}.self_attn.linear_v", # lfm2 + "encoder.lfm.lfm.layers.{bid}.self_attn.v_proj", # lfm3a "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n "conformer.layers.{bid}.self_attn.v_proj", # gemma4 "sound_encoder.encoder.layers.{bid}.self_attn.v_proj", # parakeet @@ -2589,6 +2593,14 @@ class TensorNameMap: "encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer ), + MODEL_TENSOR.A_ENC_ATTN_Q_NORM: ( + "encoder.lfm.lfm.layers.{bid}.self_attn.q_layernorm", # lfm3a + ), + + MODEL_TENSOR.A_ENC_ATTN_K_NORM: ( + "encoder.lfm.lfm.layers.{bid}.self_attn.k_layernorm", # lfm3a + ), + MODEL_TENSOR.A_ENC_ATTN_K_REL: ( "conformer.layers.{bid}.self_attn.relative_k_proj", # gemma4 ), @@ -2614,6 +2626,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn_layer_norm", # ultravox "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn_layer_norm", # dots3note "conformer.layers.{bid}.norm_self_att", # lfm2 + "encoder.lfm.lfm.layers.{bid}.operator_norm", # lfm3a "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n "sound_encoder.encoder.layers.{bid}.norm_self_att", # parakeet "encoder.layers.{bid}.attn.pre_norm", # granite_speech @@ -2625,6 +2638,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.self_attn.o_proj", # voxtral realtime "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.self_attn.out_proj", # dots3note "conformer.layers.{bid}.self_attn.linear_out", # lfm2 + "encoder.lfm.lfm.layers.{bid}.self_attn.out_proj", # lfm3a "conformer.layers.{bid}.attention.post", # gemma3n "conformer.layers.{bid}.self_attn.post", # gemma4 "sound_encoder.encoder.layers.{bid}.self_attn.o_proj", # parakeet @@ -2636,6 +2650,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.final_layer_norm", # ultravox "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.final_layer_norm", # dots3note "conformer.layers.{bid}.norm_out", # lfm2 + "encoder.lfm.lfm.layers.{bid}.ffn_norm", # lfm3a "conformer.layers.{bid}.attention.post_norm", # gemma3n "sound_encoder.encoder.layers.{bid}.norm_out", # parakeet "encoder.layers.{bid}.post_norm", # granite_speech @@ -2664,6 +2679,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.fc1", # ultravox "audio_tower.layers.{bid}.mlp.up_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear1", # lfm2 + "encoder.lfm.lfm.layers.{bid}.feed_forward.w3", # lfm3a "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 "sound_encoder.encoder.layers.{bid}.feed_forward1.linear1", # parakeet @@ -2673,6 +2689,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_FFN_GATE: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_gate", # dots3note (split from fc1 in conversion code) + "encoder.lfm.lfm.layers.{bid}.feed_forward.w1", # lfm3a "audio_tower.layers.{bid}.mlp.gate_proj", # voxtral realtime ), @@ -2681,6 +2698,7 @@ class TensorNameMap: "audio_tower.layers.{bid}.fc2", # ultravox "audio_tower.layers.{bid}.mlp.down_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear2", # lfm2 + "encoder.lfm.lfm.layers.{bid}.feed_forward.w2", # lfm3a "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 "sound_encoder.encoder.layers.{bid}.feed_forward1.linear2", # parakeet @@ -2764,6 +2782,7 @@ class TensorNameMap: "audio.multi_modal_projector.linear_{bid}", # ultravox, meralion "audio_encoder.audio_adapter.proj.{bid}", # dots3note (proj.1, proj.3) "audio_adapter.model.{bid}", # lfm2 + "mlp.model.{bid}", # lfm3a "audio_tower.proj{bid}", # qwen3omni "sound_projection.linear{bid}", # parakeet (linear1, linear2) ), @@ -2823,6 +2842,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_DW: ( "conformer.layers.{bid}.conv.depthwise_conv", # lfm2 + "encoder.lfm.lfm.layers.{bid}.conv.conv", # lfm3a "conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n "sound_encoder.encoder.layers.{bid}.conv.depthwise_conv", # parakeet "encoder.layers.{bid}.conv.depth_conv.conv", # granite_speech @@ -2845,6 +2865,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_PW1: ( "conformer.layers.{bid}.conv.pointwise_conv1", # lfm2 + "encoder.lfm.lfm.layers.{bid}.conv.in_proj", # lfm3a "conformer.layers.{bid}.lconv1d.linear_start", # gemma3n "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv1", # parakeet "encoder.layers.{bid}.conv.up_conv", # granite_speech @@ -2853,6 +2874,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV_PW2: ( "conformer.layers.{bid}.conv.pointwise_conv2", # lfm2 + "encoder.lfm.lfm.layers.{bid}.conv.out_proj", # lfm3a "conformer.layers.{bid}.lconv1d.linear_end", # gemma3n "sound_encoder.encoder.layers.{bid}.conv.pointwise_conv2", # parakeet "encoder.layers.{bid}.conv.down_conv", # granite_speech diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 5ebd57807..9ddd6a63a 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -345,6 +345,11 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize LFM3-VL's depthwise conv kernels, CONV_2D_DW takes F32 kernels quantize &= name.find("conv2d_spatial.weight") == std::string::npos; + // do not quantize LFM3-Audio's encoder convs and codec, the codec output does not survive quantization + quantize &= name.rfind("a.conv1d.", 0) != 0; + quantize &= name.find("conv_dw") == std::string::npos; + quantize &= name.rfind("a.gen.wav.", 0) != 0; + // do not quantize MiniMax's indexer projection weights, they are tiny quantize &= name.find("indexer.k_proj.weight") == std::string::npos; quantize &= name.find("indexer.q_proj.weight") == std::string::npos; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index 7c054c0c6..f1ec3c456 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -59,6 +59,8 @@ add_library(mtmd models/minimax-m3.cpp models/qwen3vl.cpp models/lfm3vl.cpp + models/lfm3a.cpp + models/lfm3a-gen.cpp models/ling3vl.cpp models/mimovl.cpp models/qwen3a.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 141954f10..0504e8fdc 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -106,6 +106,11 @@ #define KEY_GEN_AUDIO_PROJ_TYPE "clip.gen.audio.projector_type" // for models with mixed modalities // name of the weight variant, for settings that are not in the checkpoint #define KEY_GEN_AUDIO_VARIANT "clip.gen.audio.model_variant" +// sampling defaults of the generated audio, like general.sampling.* for text +#define KEY_GEN_AUDIO_TEMP "clip.gen.audio.sampling.temp" +#define KEY_GEN_AUDIO_TOP_K "clip.gen.audio.sampling.top_k" +#define KEY_GEN_AUDIO_PENALTY_REPEAT "clip.gen.audio.sampling.penalty_repeat" +#define KEY_GEN_AUDIO_PENALTY_LAST_N "clip.gen.audio.sampling.penalty_last_n" #define KEY_AUDIO_SUBSMPL_FACTOR "clip.audio.subsampling_factor" // @@ -246,6 +251,11 @@ #define TN_A_GEN_CODE_OUT_EMBD "a.gen.code.out_embd.%s" #define TN_A_GEN_CODE_NORM "a.gen.code.output_norm.%s" +// lfm3a depthformer (MaskGIT), the codec decoder uses the generic layer names with prefix "a.gen.wav.tfm" +#define TN_A_GEN_CODE_HEAD_NORM "a.gen.code.head_norm.%s" +#define TN_A_GEN_CODE_MASK_EMBD "a.gen.code.mask_embd" +#define TN_ATTN_GATE "%s.blk.%d.attn_gate.%s" + // qwen3tts code2wav (RVQ codes -> raw PCM) // pre_transformer layers use the generic TN_ATTN_*/TN_FFN_*/TN_LN_*/TN_LS_* macros, prefix "a.gen.wav.tfm" #define TN_A_GEN_WAV_QUANT_FIRST_IN "a.gen.wav.quant.first.in_proj.%s" @@ -474,6 +484,9 @@ struct clip_enc_state { ggml_backend_buffer_ptr buf; std::vector k; // per layer [n_embd_kv, n_ctx] std::vector v; + + // lfm3a: same order as list_lfm3a_state_slots(), F32 [ne0, ne1], updated in place + std::vector slots; }; // per-sequence state of audio generation, see clip_gen_state_init() @@ -539,6 +552,7 @@ enum projector_type { PROJECTOR_TYPE_DEEPSEEK41V, PROJECTOR_TYPE_LFM2A, PROJECTOR_TYPE_D1OMNI_A, + PROJECTOR_TYPE_LFM3A, PROJECTOR_TYPE_GLM4V, PROJECTOR_TYPE_GLM5V, PROJECTOR_TYPE_YOUTUVL, @@ -562,6 +576,7 @@ enum projector_type { PROJECTOR_TYPE_POCKETTTS_SPKENC, PROJECTOR_TYPE_POCKETTTS_GEN, PROJECTOR_TYPE_SOPRANO, + PROJECTOR_TYPE_LFM3A_GEN, PROJECTOR_TYPE_MUSE_GLIMMER, PROJECTOR_TYPE_SENSENOVA_U1, PROJECTOR_TYPE_COHERE2V, @@ -618,6 +633,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_DEEPSEEK41V, "deepseek41v"}, { PROJECTOR_TYPE_LFM2A, "lfm2a"}, { PROJECTOR_TYPE_D1OMNI_A, "d1omni_a"}, + { PROJECTOR_TYPE_LFM3A, "lfm3a"}, { PROJECTOR_TYPE_GLM4V, "glm4v"}, { PROJECTOR_TYPE_GLM5V, "glm5v"}, { PROJECTOR_TYPE_YOUTUVL, "youtuvl"}, @@ -641,6 +657,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_POCKETTTS_SPKENC, "pockettts_spkenc"}, { PROJECTOR_TYPE_POCKETTTS_GEN, "pockettts_gen"}, { PROJECTOR_TYPE_SOPRANO, "soprano"}, + { PROJECTOR_TYPE_LFM3A_GEN, "lfm3a_gen"}, { PROJECTOR_TYPE_MUSE_GLIMMER, "muse-glimmer"}, { PROJECTOR_TYPE_SENSENOVA_U1, "sensenova_u1"}, { PROJECTOR_TYPE_COHERE2V, "cohere2v"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 2c5cf1fed..88aa20de1 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -175,6 +175,12 @@ struct clip_hparams { // name of the weight variant, some pipelines tune themselves on it std::string gen_model_variant; + // sampling defaults of the generated audio: greedy, no top-k, no repetition penalty + float gen_temp = 0.0f; + int32_t gen_top_k = 0; + float gen_penalty_repeat = 1.0f; + int32_t gen_penalty_last_n = 0; + // pocket-tts static constexpr int32_t pockettts_max_spk_seconds = 30; int32_t seanet_n_stage = 0; @@ -195,6 +201,7 @@ struct clip_hparams { int32_t wav_dac_n_block = 0; int32_t wav_dac_n_res = 0; int32_t wav_tfm_swa = 0; // pre_transformer's KV cache size, in frames + int32_t wav_frame_stack = 0; // lfm3a codec: decoder positions per frame // mimo-v2.5: LLM-side connector (input_local_transformer) int32_t audio_local_n_layer = 0; @@ -286,6 +293,9 @@ struct clip_layer { ggml_tensor * k_norm = nullptr; ggml_tensor * q_norm = nullptr; + // lfm3: per-head sigmoid gate on the attention output + ggml_tensor * attn_gate_w = nullptr; + ggml_tensor * attn_post_norm_w = nullptr; ggml_tensor * ff_up_w = nullptr; @@ -566,6 +576,7 @@ struct clip_code2wav { // quantizer: RVQ codebook decode ggml_tensor * quant_first_in_w = nullptr; // semantic RVQ, in_proj (1x1 conv, loaded as 2D) ggml_tensor * quant_first_out_w = nullptr; + ggml_tensor * quant_first_out_b = nullptr; ggml_tensor * quant_first_cb_w = nullptr; // codebook (1 layer) ggml_tensor * quant_rest_in_w = nullptr; // acoustic RVQ ggml_tensor * quant_rest_out_w = nullptr; @@ -852,6 +863,10 @@ struct clip_model { ggml_tensor * gen_code_out_embd_w = nullptr; // codebook-0 embedding, fed back into the talker ggml_tensor * gen_code_norm_w = nullptr; // final norm + // lfm3a depthformer (MaskGIT) + ggml_tensor * gen_code_head_norm_w = nullptr; // per-codebook norm before the tied head, merged 2D + ggml_tensor * gen_code_mask_embd = nullptr; + // qwen3tts code2wav: RVQ codes -> raw PCM clip_code2wav c2w; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 0ec65e4c1..26aa05e25 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1148,6 +1148,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_LFM3A: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_GEMMA4A: { builder = std::make_unique(ctx, img); @@ -1196,6 +1200,24 @@ 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_LFM3A_GEN: + { + const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; + const int64_t n_cb = ctx->model.gen_code_embd_w->ne[2]; + bool all_known = false; + int n_frames = 1; + if (params && params->codes) { + GGML_ASSERT(params->codes->size() % (size_t) n_cb == 0); + if (gen_process == CLIP_GEN_PROCESS_GEN_WAV) { + n_frames = (int) (params->codes->size() / n_cb); + GGML_ASSERT(n_frames > 0); + } else { + // -1 marks a masked codebook, none masked means only the feedback embedding is needed + all_known = std::all_of(params->codes->begin(), params->codes->end(), [](int32_t c) { return c >= 0; }); + } + } + builder = std::make_unique(ctx, img, gen_process, all_known, n_frames); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { const auto gen_process = params ? params->gen_process : CLIP_GEN_PROCESS_GEN_CODE; @@ -2136,6 +2158,33 @@ struct clip_model_loader { hparams.audio_window_len = 1024; hparams.audio_hop_len = 256; } break; + case PROJECTOR_TYPE_LFM3A_GEN: + { + // fixed by LiquidCodecConfig, head counts come from the tensors + hparams.rope_theta = 1000000.0f; + + hparams.wav_tfm_n_layer = 8; + hparams.wav_tfm_n_embd = 512; + hparams.wav_tfm_eps = 1e-5f; + hparams.wav_tfm_rope_theta = 1000000.0f; + hparams.wav_tfm_swa = 30; // decoder_window + hparams.wav_frame_stack = 6; // stft_frame_stack, STFT frames per codec frame + // the ISTFT runs on the host, see mtmd_gen_audio_process() + hparams.audio_sample_rate = 24000; + hparams.audio_n_fft = 1280; + hparams.audio_window_len = 1280; + hparams.audio_hop_len = 320; + + // speech-code sampling, the converter writes the model card values; else the LFM3-Omni card + hparams.gen_temp = 0.4f; + hparams.gen_top_k = 0; + hparams.gen_penalty_repeat = 1.2f; + hparams.gen_penalty_last_n = 64; + get_f32(KEY_GEN_AUDIO_TEMP, hparams.gen_temp, false); + get_i32(KEY_GEN_AUDIO_TOP_K, hparams.gen_top_k, false); + get_f32(KEY_GEN_AUDIO_PENALTY_REPEAT, hparams.gen_penalty_repeat, false); + get_i32(KEY_GEN_AUDIO_PENALTY_LAST_N, hparams.gen_penalty_last_n, false); + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // TODO: hardcoded for now, read from code_predictor_config instead @@ -2241,6 +2290,20 @@ struct clip_model_loader { hparams.audio_window_len = 400; hparams.audio_hop_len = 160; } break; + case PROJECTOR_TYPE_LFM3A: + { + get_u32(KEY_A_ATTN_WINDOW_SIZE, hparams.attn_window_size); + if (hparams.attn_window_size <= 0) { + throw std::runtime_error(string_format("%s: invalid audio attention window %d\n", __func__, hparams.attn_window_size)); + } + hparams.rope_theta = 1000000.0f; // checked by the converter + // audio preprocessing params, see mtmd_audio_preprocessor_lfm3a + hparams.audio_chunk_len = 0; + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 400; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + } break; case PROJECTOR_TYPE_EXAONE4_5: { hparams.n_merge = 2; @@ -2567,8 +2630,9 @@ struct clip_model_loader { model.layers.resize(n_layers_to_load); for (int il = 0; il < n_layers_to_load; ++il) { auto & layer = model.layers[il]; - // lfm3vl conv layers have no attention tensors - const bool o_w_required = model.proj_type != PROJECTOR_TYPE_LFM3VL; + // lfm3vl / lfm3a conv layers have no attention tensors + const bool o_w_required = model.proj_type != PROJECTOR_TYPE_LFM3VL + && model.proj_type != PROJECTOR_TYPE_LFM3A; layer.k_w = get_tensor(string_format(TN_ATTN_K, prefix, il, "weight"), false); layer.q_w = get_tensor(string_format(TN_ATTN_Q, prefix, il, "weight"), false); layer.v_w = get_tensor(string_format(TN_ATTN_V, prefix, il, "weight"), false); @@ -3438,6 +3502,65 @@ struct clip_model_loader { layer.ls_2_w = get_tensor(string_format(TN_LS_2, p, il, "weight")); } } break; + case PROJECTOR_TYPE_LFM3A_GEN: + { + // depthformer, the generic loop above loaded the a.gen.code.blk.* tensors + for (int il = 0; il < hparams.n_layer; ++il) { + model.layers[il].attn_gate_w = get_tensor(string_format(TN_ATTN_GATE, prefix, il, "weight")); + } + model.gen_code_proj_in_w = get_tensor(string_format(TN_A_GEN_CODE_PROJ_IN, "weight")); + model.gen_code_proj_in_b = get_tensor(string_format(TN_A_GEN_CODE_PROJ_IN, "bias")); + model.gen_code_embd_w = get_tensor(string_format(TN_A_GEN_CODE_EMBD, "weight")); // tied head + model.gen_code_out_embd_w = get_tensor(string_format(TN_A_GEN_CODE_OUT_EMBD, "weight")); + model.gen_code_norm_w = get_tensor(string_format(TN_A_GEN_CODE_NORM, "weight")); + model.gen_code_head_norm_w = get_tensor(string_format(TN_A_GEN_CODE_HEAD_NORM, "weight")); + model.gen_code_mask_embd = get_tensor(TN_A_GEN_CODE_MASK_EMBD); + + // codec decoder: FSQ level table, one projection over all codebooks, LFM3 decoder -> polar STFT frames + auto & c2w = model.c2w; + c2w.quant_first_cb_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_CB, "weight")); + c2w.quant_first_out_w = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_OUT, "weight")); + c2w.quant_first_out_b = get_tensor(string_format(TN_A_GEN_WAV_QUANT_FIRST_OUT, "bias")); + c2w.tfm_in_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "weight")); + c2w.tfm_in_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_IN_PROJ, "bias")); + c2w.tfm_out_proj_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "weight")); + c2w.tfm_out_proj_b = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_PROJ, "bias")); + c2w.tfm_output_norm_w = get_tensor(string_format(TN_A_GEN_WAV_TFM_OUT_NORM, "weight")); + + // each layer is a short conv block or a gated attention block, both with sandwich norms + c2w.tfm_layers.resize(hparams.wav_tfm_n_layer); + for (int il = 0; il < hparams.wav_tfm_n_layer; il++) { + auto & layer = c2w.tfm_layers[il]; + const char * p = "a.gen.wav.tfm"; + layer.ln_1_w = get_tensor(string_format(TN_LN_1, p, il, "weight")); + layer.attn_post_norm_w = get_tensor(string_format(TN_ATTN_POST_NORM, p, il, "weight")); + layer.ln_2_w = get_tensor(string_format(TN_LN_2, p, il, "weight")); + layer.ff_post_norm_w = get_tensor(string_format(TN_FFN_POST_NORM, p, il, "weight")); + layer.ff_gate_w = get_tensor(string_format(TN_FFN_GATE, p, il, "weight")); + layer.ff_up_w = get_tensor(string_format(TN_FFN_UP, p, il, "weight")); + layer.ff_down_w = get_tensor(string_format(TN_FFN_DOWN, p, il, "weight")); + layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, p, il, "weight"), false); + layer.conv_dw_w = get_tensor(string_format(TN_CONV_DW, p, il, "weight"), false); + layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, p, il, "weight"), false); + layer.q_w = get_tensor(string_format(TN_ATTN_Q, p, il, "weight"), false); + layer.k_w = get_tensor(string_format(TN_ATTN_K, p, il, "weight"), false); + layer.v_w = get_tensor(string_format(TN_ATTN_V, p, il, "weight"), false); + layer.o_w = get_tensor(string_format(TN_ATTN_OUTPUT, p, il, "weight"), false); + layer.q_norm = get_tensor(string_format(TN_ATTN_Q_NORM, p, il, "weight"), false); + layer.k_norm = get_tensor(string_format(TN_ATTN_K_NORM, p, il, "weight"), false); + layer.attn_gate_w = get_tensor(string_format(TN_ATTN_GATE, p, il, "weight"), false); + const bool is_conv = layer.conv_pw1_w && layer.conv_dw_w && layer.conv_pw2_w; + const bool is_attn = layer.q_w && layer.k_w && layer.v_w && layer.o_w && layer.q_norm && layer.k_norm && layer.attn_gate_w; + if (is_conv == is_attn) { + throw std::runtime_error(string_format("%s: codec layer %d is neither a conv nor an attention block\n", __func__, il)); + } + if (is_attn) { + const int d_head = layer.q_norm->ne[0]; + hparams.wav_tfm_n_head = layer.q_w->ne[1] / d_head; + hparams.wav_tfm_n_head_kv = layer.k_w->ne[1] / d_head; + } + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { // code_predictor, proj_in is absent when the talker and the predictor share the hidden size @@ -3830,6 +3953,39 @@ struct clip_model_loader { layer.conv_pw2_b = get_tensor(string_format(TN_CONV_PW2, prefix, il, "bias")); } } break; + case PROJECTOR_TYPE_LFM3A: + { + for (int i : {0, 2, 3, 5, 6}) { + model.pre_encode_conv_X_w[i] = get_tensor(string_format(TN_CONV1D, i, "weight")); + model.pre_encode_conv_X_b[i] = get_tensor(string_format(TN_CONV1D, i, "bias")); + } + model.pre_encode_out_w = get_tensor(string_format(TN_PRE_ENCODE_OUT, "weight")); + model.pre_encode_out_b = get_tensor(string_format(TN_PRE_ENCODE_OUT, "bias")); + + model.mm_0_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 0, "weight")); + model.mm_0_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 0, "bias")); + model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight")); + model.mm_1_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "bias")); + model.mm_3_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "weight")); + model.mm_3_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 3, "bias")); + + if (!model.post_ln_w) { + throw std::runtime_error(string_format("%s: missing tensor %s\n", __func__, string_format(TN_LN_POST, prefix, "weight").c_str())); + } + + // each layer is either a short conv block or a complete attention block + for (int il = 0; il < hparams.n_layer; ++il) { + auto & layer = model.layers[il]; + layer.conv_pw1_w = get_tensor(string_format(TN_CONV_PW1, prefix, il, "weight"), false); + layer.conv_dw_w = get_tensor(string_format(TN_CONV_DW, prefix, il, "weight"), false); + layer.conv_pw2_w = get_tensor(string_format(TN_CONV_PW2, prefix, il, "weight"), false); + const bool is_conv = layer.conv_pw1_w && layer.conv_dw_w && layer.conv_pw2_w; + const bool is_attn = layer.q_w && layer.k_w && layer.v_w && layer.o_w && layer.q_norm && layer.k_norm; + if (is_conv == is_attn) { + throw std::runtime_error(string_format("%s: lfm3a layer %d is neither a complete conv nor a complete attention block\n", __func__, il)); + } + } + } break; case PROJECTOR_TYPE_PARAKEET: { @@ -4449,6 +4605,10 @@ struct clip_init_result clip_init(const char * fname, struct clip_context_params if (loader.has_gen_audio) { ctx_gen_audio = new clip_ctx(ctx_params); loader.load_hparams(ctx_gen_audio->model, CLIP_MODALITY_GEN_AUDIO); + // TODO: CUDA has no flash attention kernel for the depthformer head dim 32, enable it once it does + if (ctx_gen_audio->model.proj_type == PROJECTOR_TYPE_LFM3A_GEN && ctx_gen_audio->flash_attn_type == CLIP_FLASH_ATTN_TYPE_AUTO) { + ctx_gen_audio->flash_attn_type = CLIP_FLASH_ATTN_TYPE_DISABLED; + } loader.load_tensors(*ctx_gen_audio); // TODO: fix warmup ctx_gen_audio->buf_compute_meta.resize(ctx_gen_audio->max_nodes * ggml_tensor_overhead() + ggml_graph_overhead()); @@ -4814,6 +4974,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { } break; case PROJECTOR_TYPE_LFM2A: case PROJECTOR_TYPE_D1OMNI_A: + case PROJECTOR_TYPE_LFM3A: { n_patches = ((((img->nx() + 1) / 2) + 1) / 2 + 1) / 2; } break; @@ -4856,6 +5017,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { n_patches = 1; } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: + case PROJECTOR_TYPE_LFM3A_GEN: { // one hidden-state vector fed back to the talker per call n_patches = 1; @@ -4916,11 +5078,12 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 return clip_encode(ctx, ¶ms); } -// persisted state slots of the gen-audio decoder, per pipeline +// persisted state slots of the gen-audio decoders and the streaming audio encoders static std::vector list_gen_state_slots(const clip_hparams & hparams, const clip_model & model) { switch (model.proj_type) { case PROJECTOR_TYPE_QWEN3TTS_GEN: return list_c2w_state_slots(hparams, model); case PROJECTOR_TYPE_POCKETTTS_GEN: return list_pockettts_state_slots(hparams, model); + case PROJECTOR_TYPE_LFM3A_GEN: return list_lfm3a_gen_state_slots(hparams, model); default: return {}; } } @@ -4980,6 +5143,30 @@ void clip_gen_state_reset(clip_gen_state * state) { clip_enc_state * clip_enc_state_init(clip_ctx * ctx, int32_t n_frames_max) { const auto & hparams = ctx->model.hparams; + if (ctx->proj_type() == PROJECTOR_TYPE_LFM3A) { + // downsampler context, short-conv taps and attention window, see list_lfm3a_state_slots() + const auto slots = list_lfm3a_state_slots(hparams, ctx->model); + auto st = std::make_unique(); + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead() * slots.size(), + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + st->ctx_data.reset(ggml_init(params)); + for (const auto & slot : slots) { + ggml_tensor * t = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_format_name(t, "enc_state_%s", slot.name.c_str()); + st->slots.push_back(t); + } + st->buf.reset(ggml_backend_alloc_ctx_tensors(st->ctx_data.get(), ctx->backend)); + if (!st->buf) { + LOG_ERR("%s: failed to allocate the encoder state\n", __func__); + return nullptr; + } + // a zero state is the cold start + ggml_backend_buffer_clear(st->buf.get(), 0); + return st.release(); + } if (ctx->proj_type() != PROJECTOR_TYPE_VOXTRAL_RT || hparams.attn_window_size <= 0) { LOG_ERR("%s: the encoder does not support streaming input\n", __func__); return nullptr; @@ -5051,6 +5238,11 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // build the inference graph ggml_backend_sched_reset(ctx->sched.get()); ggml_cgraph * gf = clip_get_graph_builder(ctx, imgs, params)->build(); + + // the last node is the embedding tensor, code2wav has no out_embd + // read it before scheduling: backends may reorder the nodes of graphs with several outputs (Metal) + ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr; + if (!ggml_backend_sched_alloc_graph(ctx->sched.get(), gf)) { LOG_ERR("%s: failed to allocate compute graph\n", __func__); return false; @@ -5119,6 +5311,28 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } }; + // lfm3a encoder and codec: causal sliding window over [cached window | this call] + auto set_lfm3a_window_inputs = [&](int base, int n_pos, int window) { + std::vector positions(n_pos); + for (int i = 0; i < n_pos; i++) { + positions[i] = base + i; + } + set_input_i32("positions", positions); + + // key j holds position base - prefix + j + const int prefix = window - 1; + const int n_kv = prefix + n_pos; + std::vector mask((size_t) n_kv * n_pos, -INFINITY); + for (int q = 0; q < n_pos; q++) { + for (int j = std::max(0, q + prefix - window + 1); j <= q + prefix; j++) { + if (base - prefix + j >= 0) { + mask[(size_t) q * n_kv + j] = 0.0f; + } + } + } + set_input_attn_mask(mask); + }; + // rope positions and attention mask of the mimi transformers (pocket-tts). // the mask is causal with a sliding window, see _build_attention_mask() in the reference auto set_pockettts_tfm_inputs = [&]() { @@ -5190,8 +5404,9 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("inp_raw", inp_raw); - } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV) { - // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below + } else if (params->gen_process != CLIP_GEN_PROCESS_GEN_WAV && ggml_graph_get_tensor(gf, "inp_raw") != nullptr) { + // audio input. GEN_WAV is not here: it takes codes or feats, set in the switch below. + // a GEN_CODE graph may have no hidden-state input (lfm3a_gen feedback embedding) GGML_ASSERT(imgs.entries.size() == 1); const auto & mel_inp = imgs.entries[0]; @@ -6003,6 +6218,53 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { { // do nothing } break; + case PROJECTOR_TYPE_LFM3A_GEN: + { + const int64_t n_cb = model.gen_code_embd_w->ne[2]; + const int64_t n_vocab = model.gen_code_embd_w->ne[1]; + if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { + GGML_ASSERT(params->codes != nullptr); + // codes index the FSQ table via ggml_get_rows, the end-of-audio code has no waveform + const int64_t n_codes = model.c2w.quant_first_cb_w->ne[1]; + for (int32_t c : *params->codes) { + if (c < 0 || c >= n_codes) { + LOG_ERR("%s: code out of range (%d, codebook size %lld)\n", __func__, c, (long long) n_codes); + return false; + } + } + std::vector codes = *params->codes; + set_input_i32("inp_codes", codes); + set_gen_state_in(); + // the "pos" slot leads the state + float base = 0.0f; + ggml_backend_tensor_get(params->gen_state->slots[0], &base, 0, sizeof(float)); + set_lfm3a_window_inputs((int) base, (int) ggml_nelements(get_inp_tensor("positions")), hparams.wav_tfm_swa); + } else { + // rows of the merged per-codebook tables; a masked codebook (-1) takes row 0, the mask drops it + std::vector codes(n_cb, -1); + if (params->codes) { + GGML_ASSERT(params->codes->size() == (size_t) n_cb); + codes = *params->codes; + } + std::vector mask(n_cb, 0.0f); + for (int64_t i = 0; i < n_cb; i++) { + if (codes[i] >= n_vocab) { + LOG_ERR("%s: code out of range (%d, vocab %lld)\n", __func__, codes[i], (long long) n_vocab); + return false; + } + if (codes[i] < 0) { + codes[i] = 0; + } else { + mask[i] = 1.0f; + } + codes[i] += (int32_t) (i * n_vocab); + } + set_input_i32("inp_codes", codes); + if (ggml_graph_get_tensor(gf, "inp_mask") != nullptr) { + set_input_f32("inp_mask", mask); + } + } + } break; case PROJECTOR_TYPE_QWEN3TTS_GEN: { if (params->gen_process == CLIP_GEN_PROCESS_GEN_WAV) { @@ -6271,6 +6533,31 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { } set_input_f32("pos_emb", pos_emb); } break; + case PROJECTOR_TYPE_LFM3A: + { + GGML_ASSERT(imgs.entries.size() == 1); + const int n_pos = clip_n_output_tokens(ctx, &imgs.entries.front()); + + // the "pos" slot leads the state; without a stream the state is zero, the clip is encoded from its start + int base = 0; + if (params->enc_state != nullptr) { + const auto & slots = params->enc_state->slots; + float base_f = 0.0f; + ggml_backend_tensor_get(slots[0], &base_f, 0, sizeof(float)); + base = (int) base_f; + size_t i = 0; + for (const auto & slot : list_lfm3a_state_slots(hparams, model)) { + ggml_backend_tensor_copy(slots[i++], get_inp_tensor(("state_in_" + slot.name).c_str())); + } + } else { + for (const auto & slot : list_lfm3a_state_slots(hparams, model)) { + ggml_tensor * t = get_inp_tensor(("state_in_" + slot.name).c_str()); + ggml_backend_tensor_memset(t, 0, 0, ggml_nbytes(t)); + } + } + + set_lfm3a_window_inputs(base, n_pos, hparams.attn_window_size); + } break; case PROJECTOR_TYPE_PARAKEET: { GGML_ASSERT(imgs.entries.size() == 1); @@ -6490,9 +6777,6 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { return false; } - // the last node is the embedding tensor, code2wav has no out_embd - ggml_tensor * embeddings = params->out_embd ? ggml_graph_node(gf, -1) : nullptr; - if (embeddings != nullptr) { // sanity check (assuming that all images in batch have the same number of tokens, so we only check the first one) const int n_tokens_out = embeddings->ne[1]; @@ -6572,7 +6856,15 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { out_audio.resize((size_t) n_frames * hop); } } - if (params->enc_state != nullptr) { + if (params->enc_state != nullptr && model.proj_type == PROJECTOR_TYPE_LFM3A) { + const auto & slots = params->enc_state->slots; + size_t i = 0; + for (const auto & slot : list_lfm3a_state_slots(hparams, model)) { + ggml_tensor * t = ggml_graph_get_tensor(gf, ("state_out_" + slot.name).c_str()); + GGML_ASSERT(t != nullptr); + copy_gen_state_out(t, slots[i++]); + } + } else if (params->enc_state != nullptr) { clip_enc_state * st = params->enc_state; ggml_tensor * tail = ggml_graph_get_tensor(gf, "out_stem_tail"); GGML_ASSERT(tail != nullptr); @@ -6754,6 +7046,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_LFM2A: return ctx->model.position_embeddings->ne[0]; case PROJECTOR_TYPE_D1OMNI_A: + case PROJECTOR_TYPE_LFM3A: return ctx->model.mm_3_w->ne[1]; case PROJECTOR_TYPE_GRANITE_SPEECH: return ctx->model.qf_proj_blocks[0].qf_proj_linear_w->ne[1]; @@ -6767,6 +7060,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_QWEN3TTS_SPKENC: return ctx->model.mm_fc_w->ne[2]; case PROJECTOR_TYPE_QWEN3TTS_GEN: + case PROJECTOR_TYPE_LFM3A_GEN: return ctx->model.gen_code_out_embd_w->ne[0]; case PROJECTOR_TYPE_POCKETTTS_SPKENC: return ctx->model.spk_proj_w->ne[1]; diff --git a/tools/mtmd/models/lfm3a-gen.cpp b/tools/mtmd/models/lfm3a-gen.cpp new file mode 100644 index 000000000..998605711 --- /dev/null +++ b/tools/mtmd/models/lfm3a-gen.cpp @@ -0,0 +1,257 @@ +#include "models.h" + +#include + +// LFM3-Audio generation +// +// GEN_CODE: one MaskGIT pass of the depthformer, hidden state + known codes (-1 = masked) -> logits per codebook +// once all codes are known, only the feedback embedding is built +// GEN_WAV : a window of frames -> polar STFT frames of the liquid-codec decoder, the ISTFT runs on the host + +// state carried between GEN_WAV calls: rope offset, short-conv taps and attention window of the codec decoder +std::vector list_lfm3a_gen_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector slots; + if (model.c2w.tfm_layers.empty()) { + return slots; + } + const int64_t prefix = hparams.wav_tfm_swa - 1; + slots.push_back({"pos", 1, 1}); + for (size_t il = 0; il < model.c2w.tfm_layers.size(); il++) { + const auto & layer = model.c2w.tfm_layers[il]; + if (layer.conv_dw_w) { + slots.push_back({"conv_" + std::to_string(il), layer.conv_dw_w->ne[0] - 1, hparams.wav_tfm_n_embd}); + } else { + slots.push_back({"k_" + std::to_string(il), layer.k_w->ne[1], prefix}); + slots.push_back({"v_" + std::to_string(il), layer.v_w->ne[1], prefix}); + } + } + return slots; +} + +// pre-norm LFM block, with the LFM3 sandwich norms when the layer has them +ggml_tensor * clip_graph_lfm3a_gen::build_lfm_layer(ggml_tensor * inpL, const clip_layer & layer, float eps, int il, + const std::function & op) { + ggml_tensor * cur = build_norm(inpL, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il); + cur = op(cur); + if (layer.attn_post_norm_w) { + cur = build_norm(cur, layer.attn_post_norm_w, nullptr, NORM_TYPE_RMS, eps, il); + } + inpL = ggml_add(ctx0, cur, inpL); + + cur = build_norm(inpL, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il); + cur = build_ffn(cur, layer.ff_up_w, nullptr, layer.ff_gate_w, nullptr, layer.ff_down_w, nullptr, FFN_SILU, il); + if (layer.ff_post_norm_w) { + cur = build_norm(cur, layer.ff_post_norm_w, nullptr, NORM_TYPE_RMS, eps, il); + } + return ggml_add(ctx0, inpL, cur); +} + +// GQA with q/k RMSNorm, NEOX rope, an optional cached K/V prefix and the LFM3 per-head sigmoid gate +ggml_tensor * clip_graph_lfm3a_gen::build_lfm_attn(ggml_tensor * cur, const clip_layer & layer, int n_head, int n_head_kv, + float eps, float rope_theta, int il, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, ggml_tensor * k_cache, ggml_tensor * v_cache, + ggml_tensor ** k_all, ggml_tensor ** v_all) { + const int64_t n_pos = cur->ne[1]; + const int64_t head_dim = layer.q_norm->ne[0]; + + ggml_tensor * Qcur = ggml_reshape_3d(ctx0, build_mm(layer.q_w, cur), head_dim, n_head, n_pos); + ggml_tensor * Kcur = ggml_reshape_3d(ctx0, build_mm(layer.k_w, cur), head_dim, n_head_kv, n_pos); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + Qcur = build_norm(Qcur, layer.q_norm, nullptr, NORM_TYPE_RMS, eps, il); + Kcur = build_norm(Kcur, layer.k_norm, nullptr, NORM_TYPE_RMS, eps, il); + Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, nullptr, head_dim, GGML_ROPE_TYPE_NEOX, 0, rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, inp_pos, nullptr, head_dim, GGML_ROPE_TYPE_NEOX, 0, rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + ggml_tensor * k = ggml_reshape_2d(ctx0, Kcur, head_dim * n_head_kv, n_pos); + ggml_tensor * v = Vcur; + if (k_cache) { + k = ggml_concat(ctx0, k_cache, k, 1); + v = ggml_concat(ctx0, v_cache, v, 1); + } + if (k_all) { + *k_all = k; + *v_all = v; + } + const int64_t n_kv = k->ne[1]; + + ggml_tensor * attn = build_attn(layer.attn_gate_w ? nullptr : layer.o_w, nullptr, + Qcur, + ggml_reshape_3d(ctx0, k, head_dim, n_head_kv, n_kv), + ggml_reshape_3d(ctx0, v, head_dim, n_head_kv, n_kv), + kq_mask, 1.0f / sqrtf((float) head_dim), il); + if (layer.attn_gate_w) { + ggml_tensor * gate = ggml_sigmoid(ctx0, build_mm(layer.attn_gate_w, cur)); + attn = ggml_mul(ctx0, ggml_reshape_3d(ctx0, attn, head_dim, n_head, n_pos), ggml_reshape_3d(ctx0, gate, 1, n_head, n_pos)); + attn = build_mm(layer.o_w, ggml_reshape_2d(ctx0, attn, head_dim * n_head, n_pos)); + } + return attn; +} + +// short conv: in_proj -> b * x -> causal depthwise conv over [conv_state | x] -> c * conv -> out_proj +ggml_tensor * clip_graph_lfm3a_gen::build_lfm_conv(ggml_tensor * cur, const clip_layer & layer, int il, + ggml_tensor * conv_state, ggml_tensor ** conv_state_out) { + const int64_t n_embd_c = cur->ne[0]; + const int64_t n_pos = cur->ne[1]; + + ggml_tensor * bcx = build_mm(layer.conv_pw1_w, cur); + ggml_tensor * b = ggml_view_2d(ctx0, bcx, n_embd_c, n_pos, bcx->nb[1], 0); + ggml_tensor * c = ggml_view_2d(ctx0, bcx, n_embd_c, n_pos, bcx->nb[1], ggml_row_size(bcx->type, n_embd_c)); + ggml_tensor * x = ggml_view_2d(ctx0, bcx, n_embd_c, n_pos, bcx->nb[1], ggml_row_size(bcx->type, 2 * n_embd_c)); + + ggml_tensor * bx = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mul(ctx0, b, x))); + bx = ggml_concat(ctx0, conv_state, bx, 0); + const int64_t d_conv = conv_state->ne[0]; + *conv_state_out = ggml_cont(ctx0, ggml_view_2d(ctx0, bx, d_conv, n_embd_c, bx->nb[1], (size_t) n_pos * bx->nb[0])); + + cur = ggml_ssm_conv(ctx0, ggml_reshape_3d(ctx0, bx, bx->ne[0], n_embd_c, 1), layer.conv_dw_w); + cur = ggml_mul(ctx0, c, ggml_reshape_2d(ctx0, cur, n_embd_c, n_pos)); + cb(cur, "conv", il); + return build_mm(layer.conv_pw2_w, cur); +} + +ggml_tensor * clip_graph_lfm3a_gen::build_gen_code() { + ggml_tensor * embd_w = model.gen_code_embd_w; // [n_embd, n_vocab, n_codebooks], tied to the output head + const int64_t n_cb = embd_w->ne[2]; + const int64_t n_vocab = embd_w->ne[1]; + + // rows of the merged per-codebook tables: code + codebook * n_vocab + ggml_tensor * inp_codes = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_cb); + ggml_set_name(inp_codes, "inp_codes"); + ggml_set_input(inp_codes); + + // feedback embedding for the backbone + ggml_tensor * fb = ggml_get_rows(ctx0, model.gen_code_out_embd_w, inp_codes); // [n_mmproj_embd, n_cb] + fb = ggml_sum_rows(ctx0, ggml_cont(ctx0, ggml_transpose(ctx0, fb))); + fb = ggml_reshape_2d(ctx0, fb, n_mmproj_embd, 1); + cb(fb, "gen_embd", -1); + if (all_known) { + return fb; + } + + ggml_tensor * h_state = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_mmproj_embd); + ggml_set_name(h_state, "inp_raw"); + ggml_set_input(h_state); + + // 1 for a revealed codebook, 0 for a masked one + ggml_tensor * inp_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_cb); + ggml_set_name(inp_mask, "inp_mask"); + ggml_set_input(inp_mask); + + // one position per codebook: projected hidden state + token embedding or the mask embedding + ggml_tensor * cur = ggml_add(ctx0, build_mm(model.gen_code_proj_in_w, h_state), model.gen_code_proj_in_b); + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_cb); + + ggml_tensor * tok = ggml_get_rows(ctx0, ggml_reshape_2d(ctx0, embd_w, n_embd, n_vocab * n_cb), inp_codes); + tok = ggml_mul(ctx0, tok, inp_mask); + ggml_tensor * msk = ggml_mul(ctx0, ggml_repeat(ctx0, model.gen_code_mask_embd, cur), ggml_scale_bias(ctx0, inp_mask, -1.0f, 1.0f)); + cur = ggml_add(ctx0, cur, ggml_add(ctx0, tok, msk)); + cb(cur, "depth_inp", -1); + + // bidirectional: no mask, no state + ggml_tensor * inp_pos = ggml_cast(ctx0, ggml_arange(ctx0, 0.0f, (float) n_cb, 1.0f), GGML_TYPE_I32); + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.layers[il]; + cur = build_lfm_layer(cur, layer, eps, il, [&](ggml_tensor * x) { + return build_lfm_attn(x, layer, hparams.n_head, hparams.n_head_kv, eps, hparams.rope_theta, il, inp_pos, + nullptr, nullptr, nullptr, nullptr, nullptr); + }); + } + + cur = build_norm(cur, model.gen_code_norm_w, nullptr, NORM_TYPE_RMS, eps, -1); + cur = build_norm(cur, model.gen_code_head_norm_w, nullptr, NORM_TYPE_RMS, eps, -1); // per codebook + ggml_tensor * logits = ggml_mul_mat(ctx0, embd_w, ggml_reshape_3d(ctx0, cur, n_embd, 1, n_cb)); // [n_vocab, 1, n_cb] + logits = ggml_reshape_2d(ctx0, logits, n_vocab, n_cb); + ggml_set_name(logits, "out_feats"); + ggml_set_output(logits); + ggml_build_forward_expand(gf, logits); + + return fb; +} + +ggml_tensor * clip_graph_lfm3a_gen::build_gen_wav() { + const auto & c2w = model.c2w; + const int64_t n_cb = model.gen_code_embd_w->ne[2]; + const int64_t n_fs = hparams.wav_frame_stack; + const int64_t n_pos = n_fs * n_frames; + const int64_t prefix = hparams.wav_tfm_swa - 1; + + // frame-major codes + ggml_tensor * inp_codes = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_cb * n_frames); + ggml_set_name(inp_codes, "inp_codes"); + ggml_set_input(inp_codes); + + std::map state_in; + std::vector> state_out; + for (const auto & slot : list_lfm3a_gen_state_slots(hparams, model)) { + ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_set_name(t, ("state_in_" + slot.name).c_str()); + ggml_set_input(t); + state_in[slot.name] = t; + } + + // FSQ decode: every codebook shares the level table, one projection over the n_cb code vectors of a frame + ggml_tensor * q = ggml_get_rows(ctx0, c2w.quant_first_cb_w, inp_codes); // [n_lvl, n_cb * n_frames] + q = ggml_reshape_2d(ctx0, q, q->ne[0] * n_cb, n_frames); + ggml_tensor * cur = ggml_add(ctx0, build_mm(c2w.quant_first_out_w, q), c2w.quant_first_out_b); + cur = ggml_add(ctx0, build_mm(c2w.tfm_in_proj_w, cur), c2w.tfm_in_proj_b); // [n_embd_dec, n_frames] + const int64_t n_embd_dec = cur->ne[0]; + + // nearest-exact upsampling: every frame becomes n_fs decoder positions + cur = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, cur, n_embd_dec, 1, n_frames), n_embd_dec, n_fs, n_frames, 1); + cur = ggml_reshape_2d(ctx0, cur, n_embd_dec, n_pos); + + // positions continue across calls from the "pos" slot, see clip_encode() + ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(inp_pos, "positions"); + ggml_set_input(inp_pos); + state_out.push_back({"pos", ggml_scale_bias(ctx0, state_in.at("pos"), 1.0f, (float) n_pos)}); + + // sliding-window causal mask over [cached window | this chunk] + ggml_tensor * kq_mask = build_inp_attn_mask(prefix + n_pos, n_pos); + + const float eps_dec = hparams.wav_tfm_eps; + for (size_t il = 0; il < c2w.tfm_layers.size(); il++) { + const auto & layer = c2w.tfm_layers[il]; + cur = build_lfm_layer(cur, layer, eps_dec, (int) il, [&](ggml_tensor * x) { + if (layer.conv_dw_w) { + const std::string name = "conv_" + std::to_string(il); + ggml_tensor * st = nullptr; + ggml_tensor * out = build_lfm_conv(x, layer, (int) il, state_in.at(name), &st); + state_out.push_back({name, st}); + return out; + } + // this chunk's tail of the window is kept for the next call + const std::string k_name = "k_" + std::to_string(il); + const std::string v_name = "v_" + std::to_string(il); + ggml_tensor * k = nullptr; + ggml_tensor * v = nullptr; + ggml_tensor * out = build_lfm_attn(x, layer, hparams.wav_tfm_n_head, hparams.wav_tfm_n_head_kv, eps_dec, + hparams.wav_tfm_rope_theta, (int) il, inp_pos, kq_mask, + state_in.at(k_name), state_in.at(v_name), &k, &v); + state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, k, k->ne[0], prefix, k->nb[1], (size_t) n_pos * k->nb[1]))}); + state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, v, v->ne[0], prefix, v->nb[1], (size_t) n_pos * v->nb[1]))}); + return out; + }); + } + + cur = build_norm(cur, c2w.tfm_output_norm_w, nullptr, NORM_TYPE_RMS, eps_dec, -1); + cur = ggml_add(ctx0, build_mm(c2w.tfm_out_proj_w, cur), c2w.tfm_out_proj_b); // [2 * n_bins, n_pos] + cb(cur, "polar_out", -1); + + for (const auto & s : state_out) { + ggml_set_name(s.second, ("state_out_" + s.first).c_str()); + ggml_set_output(s.second); + ggml_build_forward_expand(gf, s.second); + } + + ggml_set_name(cur, "out_feats"); + ggml_set_output(cur); + return cur; +} + +ggml_cgraph * clip_graph_lfm3a_gen::build() { + ggml_tensor * out = gen_process == CLIP_GEN_PROCESS_GEN_WAV ? build_gen_wav() : build_gen_code(); + ggml_build_forward_expand(gf, out); + return gf; +} diff --git a/tools/mtmd/models/lfm3a.cpp b/tools/mtmd/models/lfm3a.cpp new file mode 100644 index 000000000..4882fc42f --- /dev/null +++ b/tools/mtmd/models/lfm3a.cpp @@ -0,0 +1,189 @@ +#include "models.h" + +// streaming encoder state: base position, conv left context, short-conv taps and attention window +// a zero state gives the same result as encoding the whole clip +std::vector list_lfm3a_state_slots(const clip_hparams & hparams, const clip_model & model) { + std::vector slots; + + const int64_t n_mel = hparams.n_mel_bins; + const int64_t n_ch = model.pre_encode_conv_X_w[0]->ne[3]; + const int64_t prefix = hparams.attn_window_size - 1; + + slots.push_back({"pos", 1, 1}); + slots.push_back({"ds0", 1, n_mel}); + slots.push_back({"ds1", 1, (n_mel / 2) * n_ch}); + slots.push_back({"ds2", 1, (n_mel / 4) * n_ch}); + for (size_t il = 0; il < model.layers.size(); il++) { + const auto & layer = model.layers[il]; + if (layer.conv_dw_w) { + slots.push_back({"conv_" + std::to_string(il), layer.conv_dw_w->ne[0] - 1, hparams.n_embd}); + } else { + slots.push_back({"k_" + std::to_string(il), layer.k_w->ne[1], prefix}); + slots.push_back({"v_" + std::to_string(il), layer.v_w->ne[1], prefix}); + } + } + + return slots; +} + +// LFM3-Audio encoder: conv downsampler (8x in time) -> causal LFM2 stack -> MLP projector +ggml_cgraph * clip_graph_lfm3a::build() { + const norm_type norm_t = NORM_TYPE_RMS; + const float rope_theta = hparams.rope_theta; + const int64_t prefix = hparams.attn_window_size - 1; + + // the graph always streams, a zero state encodes a clip from its start + std::map state_in; + std::vector> state_out; + for (const auto & slot : list_lfm3a_state_slots(hparams, model)) { + ggml_tensor * t = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, slot.ne0, slot.ne1); + ggml_set_name(t, ("state_in_" + slot.name).c_str()); + ggml_set_input(t); + state_in[slot.name] = t; + } + + // stride-2 conv over [time, freq, ch]; the left column comes from the state and the right edge is zero-padded + auto conv_down = [&](ggml_tensor * x, const std::string & name, ggml_tensor * w, ggml_tensor * b, bool depthwise) { + ggml_tensor * left = ggml_reshape_4d(ctx0, state_in.at(name), 1, x->ne[1], x->ne[2], 1); + x = ggml_concat(ctx0, left, x, 0); + ggml_tensor * last = ggml_view_4d(ctx0, x, 1, x->ne[1], x->ne[2], 1, x->nb[1], x->nb[2], x->nb[3], (x->ne[0] - 1) * x->nb[0]); + state_out.push_back({name, ggml_reshape_2d(ctx0, ggml_cont(ctx0, last), 1, x->ne[1] * x->ne[2])}); + x = ggml_pad_ext(ctx0, x, 0, 1, 0, 0, 0, 0, 0, 0); + x = depthwise ? ggml_conv_2d_dw_direct(ctx0, w, x, 2, 2, 0, 1, 1, 1) + : ggml_conv_2d(ctx0, w, x, 2, 2, 0, 1, 1, 1); + return ggml_add(ctx0, x, b); + }; + + // inp: [n_frames, n_mel, 1] + ggml_tensor * cur = build_inp_raw(1); + + // downsampler, same layout as the lfm2a pre-encoder but with SiLU + cur = conv_down(cur, "ds0", model.pre_encode_conv_X_w[0], model.pre_encode_conv_X_b[0], false); + cur = ggml_silu_inplace(ctx0, cur); + + cur = conv_down(cur, "ds1", model.pre_encode_conv_X_w[2], model.pre_encode_conv_X_b[2], true); + cur = ggml_conv_2d_direct(ctx0, model.pre_encode_conv_X_w[3], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[3]); + cur = ggml_silu_inplace(ctx0, cur); + + cur = conv_down(cur, "ds2", model.pre_encode_conv_X_w[5], model.pre_encode_conv_X_b[5], true); + cur = ggml_conv_2d_direct(ctx0, model.pre_encode_conv_X_w[6], cur, 1, 1, 0, 0, 1, 1); + cur = ggml_add(ctx0, cur, model.pre_encode_conv_X_b[6]); + cur = ggml_silu_inplace(ctx0, cur); + cb(cur, "downsample.conv", -1); + + // [n_pos, n_mel/8, n_ch] -> [n_ch * n_mel/8, n_pos] + const int n_pos = cur->ne[0]; + cur = ggml_reshape_2d(ctx0, cur, n_pos, cur->ne[1] * cur->ne[2]); + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); + cur = build_mm(model.pre_encode_out_w, cur); + cur = ggml_add(ctx0, cur, model.pre_encode_out_b); + cb(cur, "downsample.out", -1); + + ggml_tensor * positions = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(positions, "positions"); + ggml_set_input(positions); + state_out.push_back({"pos", ggml_scale_bias(ctx0, state_in.at("pos"), 1.0f, (float) n_pos)}); + + // causal mask with a sliding window, shared by all attention layers + ggml_tensor * kq_mask = build_inp_attn_mask(prefix + n_pos, n_pos); + + ggml_tensor * inpL = cur; + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.layers[il]; + + cur = build_norm(inpL, layer.ln_1_w, nullptr, norm_t, eps, il); + cb(cur, "operator_norm", il); + + if (layer.conv_dw_w) { + // short conv: in_proj -> b * x -> causal depthwise conv -> c * conv -> out_proj + ggml_tensor * bcx = build_mm(layer.conv_pw1_w, cur); + cb(bcx, "conv_in_proj", il); + + ggml_tensor * b = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], 0); + ggml_tensor * c = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], ggml_row_size(bcx->type, n_embd)); + ggml_tensor * x = ggml_view_2d(ctx0, bcx, n_embd, n_pos, bcx->nb[1], ggml_row_size(bcx->type, 2 * n_embd)); + + ggml_tensor * bx = ggml_cont(ctx0, ggml_transpose(ctx0, ggml_mul(ctx0, b, x))); + const int d_conv = layer.conv_dw_w->ne[0] - 1; + const std::string name = "conv_" + std::to_string(il); + bx = ggml_concat(ctx0, state_in.at(name), bx, 0); + state_out.push_back({name, ggml_cont(ctx0, ggml_view_2d(ctx0, bx, d_conv, n_embd, bx->nb[1], (size_t) n_pos * bx->nb[0]))}); + bx = ggml_reshape_3d(ctx0, bx, bx->ne[0], n_embd, 1); + + cur = ggml_ssm_conv(ctx0, bx, layer.conv_dw_w); + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_pos); + cb(cur, "conv_conv", il); + + cur = ggml_mul(ctx0, c, cur); + cur = build_mm(layer.conv_pw2_w, cur); + cb(cur, "conv_out_proj", il); + } else { + ggml_tensor * Qcur = build_mm(layer.q_w, cur); + ggml_tensor * Kcur = build_mm(layer.k_w, cur); + ggml_tensor * Vcur = build_mm(layer.v_w, cur); + + Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); + Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head_kv, n_pos); + Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head_kv, n_pos); + + Qcur = build_norm(Qcur, layer.q_norm, nullptr, norm_t, eps, il); + Kcur = build_norm(Kcur, layer.k_norm, nullptr, norm_t, eps, il); + + Qcur = ggml_rope_ext(ctx0, Qcur, positions, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + Kcur = ggml_rope_ext(ctx0, Kcur, positions, nullptr, d_head, GGML_ROPE_TYPE_NEOX, 0, + rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + + // the cached window goes in front, this chunk's tail is kept for the next call + const std::string k_name = "k_" + std::to_string(il); + const std::string v_name = "v_" + std::to_string(il); + const int64_t n_embd_kv = d_head * n_head_kv; + ggml_tensor * K = ggml_concat(ctx0, state_in.at(k_name), ggml_reshape_2d(ctx0, Kcur, n_embd_kv, n_pos), 1); + ggml_tensor * V = ggml_concat(ctx0, state_in.at(v_name), ggml_reshape_2d(ctx0, Vcur, n_embd_kv, n_pos), 1); + state_out.push_back({k_name, ggml_cont(ctx0, ggml_view_2d(ctx0, K, n_embd_kv, prefix, K->nb[1], (size_t) n_pos * K->nb[1]))}); + state_out.push_back({v_name, ggml_cont(ctx0, ggml_view_2d(ctx0, V, n_embd_kv, prefix, V->nb[1], (size_t) n_pos * V->nb[1]))}); + Kcur = ggml_reshape_3d(ctx0, K, d_head, n_head_kv, prefix + n_pos); + Vcur = ggml_reshape_3d(ctx0, V, d_head, n_head_kv, prefix + n_pos); + + cur = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, kq_mask, kq_scale, il); + cb(cur, "attn_out", il); + } + + cur = ggml_add(ctx0, cur, inpL); + inpL = cur; + + cur = build_norm(cur, layer.ln_2_w, nullptr, norm_t, eps, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + layer.ff_up_w, nullptr, + layer.ff_gate_w, nullptr, + layer.ff_down_w, nullptr, + FFN_SILU, il); + cb(cur, "ffn_out", il); + + cur = ggml_add(ctx0, inpL, cur); + inpL = cur; + } + + cur = build_norm(inpL, model.post_ln_w, nullptr, norm_t, eps, -1); + cb(cur, "embedding_norm", -1); + + // projector: LayerNorm -> Linear -> GELU -> Linear + cur = build_norm(cur, model.mm_0_w, model.mm_0_b, NORM_TYPE_NORMAL, 1e-5, -1); + cur = build_ffn(cur, model.mm_1_w, model.mm_1_b, nullptr, nullptr, model.mm_3_w, model.mm_3_b, FFN_GELU_ERF, -1); + cb(cur, "projected", -1); + + // the state outputs go first so that the embeddings stay the last node of the graph + for (const auto & s : state_out) { + ggml_set_name(s.second, ("state_out_" + s.first).c_str()); + ggml_set_output(s.second); + ggml_build_forward_expand(gf, s.second); + } + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index ba5769c76..c75004d43 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -2,6 +2,7 @@ #include "../clip-graph.h" +#include #include #include #include @@ -189,6 +190,11 @@ struct clip_graph_conformer : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_lfm3a : clip_graph { + clip_graph_lfm3a(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_granite_speech : clip_graph { clip_graph_granite_speech(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; @@ -433,6 +439,39 @@ struct clip_graph_soprano : clip_graph { int n_frames; }; +// +// LFM3-Audio generation: +// GEN_CODE = one MaskGIT pass of the depthformer, or only the feedback embedding once all codes are known +// GEN_WAV = liquid-codec decoder, a window of frames -> polar STFT frames +// +struct clip_graph_lfm3a_gen : clip_graph { + clip_graph_lfm3a_gen(clip_ctx * ctx, const clip_image_f32 & img, clip_gen_process_type gen_process, bool all_known, int n_frames) + : clip_graph(ctx, img), gen_process(gen_process), all_known(all_known), n_frames(n_frames) {} + ggml_cgraph * build() override; + + clip_gen_process_type gen_process; + bool all_known; // GEN_CODE only + int n_frames; // GEN_WAV only + + ggml_tensor * build_gen_code(); + ggml_tensor * build_gen_wav(); + + ggml_tensor * build_lfm_layer(ggml_tensor * inpL, const clip_layer & layer, float eps, int il, + const std::function & op); + ggml_tensor * build_lfm_attn(ggml_tensor * cur, const clip_layer & layer, int n_head, int n_head_kv, + float eps, float rope_theta, int il, ggml_tensor * inp_pos, + ggml_tensor * kq_mask, ggml_tensor * k_cache, ggml_tensor * v_cache, + ggml_tensor ** k_all, ggml_tensor ** v_all); + ggml_tensor * build_lfm_conv(ggml_tensor * cur, const clip_layer & layer, int il, + ggml_tensor * conv_state, ggml_tensor ** conv_state_out); +}; + +// state slots of the streaming codec decoder (lfm3a_gen GEN_WAV) +std::vector list_lfm3a_gen_state_slots(const clip_hparams & hparams, const clip_model & model); + +// state slots of the streaming LFM3-Audio encoder (lfm3a) +std::vector list_lfm3a_state_slots(const clip_hparams & hparams, const clip_model & model); + struct clip_graph_kimik25 : clip_graph { clip_graph_kimik25(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index 558475458..14b5a7213 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -1055,6 +1055,61 @@ bool mtmd_audio_preprocessor_d1omni::preprocess(const float * sa return true; } +// +// mtmd_audio_preprocessor_lfm3a +// + +void mtmd_audio_preprocessor_lfm3a::initialize() { + cache.fill_sin_cos_table(hparams.audio_n_fft); + cache.fill_hann_window(hparams.audio_window_len, true); + cache.fill_mel_filterbank_matrix(hparams.n_mel_bins, hparams.audio_n_fft, hparams.audio_sample_rate); +} + +bool mtmd_audio_preprocessor_lfm3a::preprocess(const float * samples, + size_t n_samples, + std::vector & output) const { + if (n_samples == 0) { + return false; + } + + const size_t win = hparams.audio_window_len; + const size_t hop = hparams.audio_hop_len; + + // pad on the right so that no sample is dropped by the last frame + size_t n_padded = win; + if (n_samples > win) { + n_padded = n_samples + (hop - (n_samples - win) % hop) % hop; + } + std::vector padded(samples, samples + n_samples); + padded.resize(n_padded, 0.0f); + + filter_params params; + params.n_mel = hparams.n_mel_bins; + params.n_fft_bins = 1 + (hparams.audio_n_fft / 2); + params.hann_window_size = hparams.audio_window_len; + params.hop_length = hparams.audio_hop_len; + params.sample_rate = hparams.audio_sample_rate; + params.no_padding = true; + params.mel_floor = 1e-12f; // log10 floor of -12 + + GGML_ASSERT(!cache.sin_vals.empty()); + GGML_ASSERT(!cache.cos_vals.empty()); + GGML_ASSERT(!cache.filters.data.empty()); + + mtmd_audio_mel out; + bool ok = log_mel_spectrogram(padded.data(), (int) padded.size(), 4, params, cache, out); + if (!ok) { + return false; + } + + for (auto & v : out.data) { + v = (v + 4.0f) / 4.0f; + } + + output.push_back(std::move(out)); + return true; +} + // // mtmd_audio_preprocessor_granite_speech // diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 26ace2143..e9194ad5e 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -95,6 +95,15 @@ struct mtmd_audio_preprocessor_d1omni : mtmd_audio_preprocessor_conformer { bool preprocess(const float * samples, size_t n_samples, std::vector & output) const override; }; +struct mtmd_audio_preprocessor_lfm3a : mtmd_audio_preprocessor { + mtmd_audio_preprocessor_lfm3a(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; + + private: + mtmd_audio_cache cache; +}; + struct mtmd_audio_preprocessor_granite_speech : mtmd_audio_preprocessor { mtmd_audio_preprocessor_granite_speech(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {} void initialize() override; diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index 72e8ae657..ba156f1c5 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -120,7 +120,8 @@ struct mtmd_cli_context { init_vision_context(params); mtmd_cap = mtmd_helper_get_capabilities(model, ctx_vision.get()); - if (!mtmd_cap.chat) { + // a realtime transcription model has no chat, it streams the audio below + if (!mtmd_cap.chat && !mtmd_cap.audio_stream) { LOG_ERR("Model does not support chat mode\n"); LOG_ERR("Hint: for TTS models, please use llama-tts\n"); exit(1); @@ -567,7 +568,10 @@ int main(int argc, char ** argv) { if (g_is_interrupted) return 130; - if (ctx.mtmd_cap.audio_stream) { + // realtime transcription models (no chat) transcribe while they listen, the others answer after the audio + const bool transcribes = ctx.mtmd_cap.audio_stream && !ctx.mtmd_cap.chat; + + if (transcribes) { if (params.image.size() != 1) { LOG_ERR("This model transcribes audio in realtime, pass exactly one --audio file\n"); return 1; diff --git a/tools/mtmd/mtmd-helper-gen.cpp b/tools/mtmd/mtmd-helper-gen.cpp index 44bee60bd..d0342b9bd 100644 --- a/tools/mtmd/mtmd-helper-gen.cpp +++ b/tools/mtmd/mtmd-helper-gen.cpp @@ -8,7 +8,9 @@ #include #include #include +#include #include +#include #include #include #include @@ -103,6 +105,11 @@ public: virtual int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) = 0; // forces any buffered codes through code2wav now, regardless of window_frames virtual int32_t flush() = 0; + // continuing a sequence (inp->n_past > 0) with the model's own text/audio switching + virtual bool supports_continue() const { return false; } + virtual bool next_is_text() const { return false; } + + std::vector logit_bias; // see mtmd_helper_gen_audio_get_logit_bias() protected: // one state per generation, so a fixed seed always gives the same output @@ -1261,6 +1268,455 @@ private: bool sent = false; }; +// LFM3-Audio TTS, see generate() in the reference: each step samples a frame of codes by MaskGIT +// and feeds its embedding back, the end-of-audio code stops +class lfm3a_gen_audio_pipeline : public mtmd_gen_audio_pipeline { +public: + using mtmd_gen_audio_pipeline::mtmd_gen_audio_pipeline; + + void reset() override { + seq_id = 0; + pos = 0; + continuing = false; + mode = MODE_AUDIO; + slots_left = 0; + speaking = false; + text_done = false; + prompt_tokens.clear(); + prompt_pos = 0; + hist.clear(); + gen_state.reset(); + pending_tail = false; + pcm_sent = 0; + wav_header_sent = false; + audio_pcm.clear(); + h_state_buf.clear(); + out_buf.clear(); + } + + int32_t set_input(const mtmd_helper_gen_audio_inp * inp) override { + reset(); + seq_id = inp->seq_id; + + if (inp->speaker_ref) { + LOG_WRN("mtmd_helper_gen_audio: the model has no voice conditioning, the speaker file is ignored\n"); + } + + // the prompt is rendered by the caller: the "tts" chat template of the model, or the next turn of a conversation + continuing = inp->n_past > 0; + pos = continuing ? inp->n_past : 0; + const std::string full(inp->prompt, inp->prompt_len); + prompt_tokens.resize(full.size() + 16); + const int n_ids = llama_tokenize(vocab, full.c_str(), (int32_t) full.size(), prompt_tokens.data(), + (int32_t) prompt_tokens.size(), false, true); + if (n_ids <= 0) { + LOG_ERR("mtmd_helper_gen_audio: tokenization failed\n"); + return 1; + } + prompt_tokens.resize((size_t) n_ids); + + if (!find_tokens()) { + return 1; + } + // a continued response opens with the token of the asked mode, a text-only one cannot switch to audio + if (continuing) { + switch (inp->out_mode) { + case MTMD_HELPER_GEN_AUDIO_OUT_AUDIO: prompt_tokens.push_back(tok_audio_start); break; + case MTMD_HELPER_GEN_AUDIO_OUT_MIXED: prompt_tokens.push_back(tok_mixed_start); break; + case MTMD_HELPER_GEN_AUDIO_OUT_TEXT: + logit_bias = { { tok_audio_start, -INFINITY }, { tok_mixed_start, -INFINITY } }; + break; + default: break; + } + } + // the last prompt token picks the mode, see generate() of the reference + const llama_token last = prompt_tokens.back(); + mode = last == tok_audio_start ? MODE_AUDIO : last == tok_mixed_start ? MODE_MIXED : MODE_TEXT; + slots_left = n_text_block; + if (!continuing && mode != MODE_AUDIO) { + LOG_ERR("mtmd_helper_gen_audio: the prompt must come from the \"tts\" chat template of the model\n"); + return 1; + } + + // the sampling defaults of the model, the caller's top-k/top-p are not used + const mtmd_gen_inp def = mtmd_gen_inp_default(mctx); + temp = def.temp; + top_k = def.top_k; + rep_penalty = def.penalty_repeat; + rep_window = def.penalty_last_n; + + out_type = inp->out_type; + stream = inp->stream; + rng.seed(inp->seed == UINT32_MAX ? std::random_device{}() : inp->seed); + return init_gen_state(inp->seed) ? 0 : 1; + } + + int32_t step_prompt(int32_t n_batch) override { + GGML_ASSERT(n_batch > 0); + const int n_prompt = (int) prompt_tokens.size(); + if (prompt_pos >= n_prompt) { + return 0; + } + const int n = std::min(n_batch, n_prompt - prompt_pos); + llama_batch batch = llama_batch_init(n, 0, 1); + for (int i = 0; i < n; i++) { + batch.token [i] = prompt_tokens[(size_t) (prompt_pos + i)]; + batch.pos [i] = pos + i; + batch.n_seq_id[i] = 1; + batch.seq_id [i][0] = seq_id; + batch.logits [i] = prompt_pos + i == n_prompt - 1; + } + batch.n_tokens = n; + const int ret = llama_decode(lctx, batch); + llama_batch_free(batch); + if (ret != 0) { + LOG_ERR("mtmd_helper_gen_audio: prompt decode failed\n"); + return -1; + } + pos += n; + prompt_pos += n; + return n_prompt - prompt_pos; + } + + bool supports_continue() const override { + return true; + } + + bool next_is_text() const override { + return mode == MODE_TEXT || (mode == MODE_MIXED && !speaking); + } + + // one generate() step: a text token from the caller or an audio frame from the hidden state + int32_t step_gen(llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) override { + const bool interleaving = mode == MODE_MIXED; + if (next_is_text()) { + // <|im_end|> ends the response in any text slot, <|text_end|> outside an interleaved block too + if (sampled == tok_im_end || (!interleaving && sampled == tok_text_end)) { + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + // the switching tokens are fed back like any text token + if (!interleaving) { + if (sampled == tok_audio_start) { + mode = MODE_AUDIO; + } else if (sampled == tok_mixed_start) { + mode = MODE_MIXED; + slots_left = n_text_block; + speaking = false; + text_done = false; + } + } else if (sampled == tok_text_end) { + text_done = true; + speaking = true; + } + llama_batch batch = llama_batch_get_one(&sampled, 1); + if (!decode(batch)) { + return 1; + } + } else { + std::vector codes; + const float * fb = nullptr; + bool is_eos = false; + if (!sample_frame(h_state_in, codes, &fb, &is_eos)) { + return 1; + } + if (is_eos) { + if (!continuing) { + // TTS: the end of audio ends the generation + *out_stop = true; + *h_state_out = nullptr; + return 0; + } + // the segment ends, the codec tail goes out and the response continues in text + if (flush() != 0) { + return 1; + } + mode = MODE_TEXT; + } else { + hist.insert(hist.end(), codes.begin(), codes.end()); + if (!decode_wav(codes.data(), codes.size())) { + return 1; + } + } + decode_embd_batch batch_embd(const_cast(fb), 1, 1, n_embd); + batch_embd.set_position_normal(pos, seq_id); + batch_embd.logits[0] = 1; + if (llama_process(lctx, LLAMA_PROCESS_TYPE_DECODE, batch_embd.render(lctx, 0, 1)) != 0) { + LOG_ERR("mtmd_helper_gen_audio: decode failed\n"); + return 1; + } + pos++; + } + + // a slot emitted while interleaving counts down the block, the switching token itself does not + if (interleaving && !text_done && --slots_left == 0) { + speaking = !speaking; + slots_left = speaking ? n_audio_block : n_text_block; + } + + const float * he = llama_get_embeddings_ith(lctx, -1); + h_state_buf.assign(he, he + n_embd); + *h_state_out = h_state_buf.data(); + return 0; + } + + int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples) override { + *out_sample_rate = info.sample_rate; + + if (!stream) { + if (flush() != 0) { + return 1; + } + if (out_n_samples) { + *out_n_samples = (int64_t) audio_pcm.size(); + } + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *out_data = (const char *) audio_pcm.data(); + *out_data_len = audio_pcm.size() * sizeof(float); + return 0; + } + out_buf.clear(); + if (!write_wav16(out_buf, audio_pcm, info.sample_rate)) { + LOG_ERR("mtmd_helper_gen_audio: output too large for WAV\n"); + return 1; + } + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + + // streaming: only the audio produced since the previous call, every frame is decoded as it is sampled + const size_t n_new = audio_pcm.size() - pcm_sent; + if (out_n_samples) { + *out_n_samples = (int64_t) n_new; + } + if (out_type == MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM) { + *out_data = (const char *) (audio_pcm.data() + pcm_sent); + *out_data_len = n_new * sizeof(float); + pcm_sent = audio_pcm.size(); + return 0; + } + out_buf.clear(); + if (!wav_header_sent) { + write_wav16_header(out_buf, UINT32_MAX, info.sample_rate); + wav_header_sent = true; + } + append_wav16_pcm(out_buf, audio_pcm.data() + pcm_sent, n_new); + pcm_sent = audio_pcm.size(); + *out_data = out_buf.data(); + *out_data_len = out_buf.size(); + return 0; + } + + // the ISTFT holds back the overlap of the last frame + int32_t flush() override { + if (pending_tail) { + if (!decode_wav(nullptr, 0)) { + return 1; + } + pending_tail = false; + } + return 0; + } + +private: + // checked by the converter + static constexpr int n_cb = 12; + static constexpr int schedule[] = { 4, 4, 4 }; // codebooks revealed per depthformer pass + + // a random order of the codebooks, revealed by the schedule; fb is the feedback embedding of the codes. + // at the end of audio every code is the end code, its feedback embedding is only built when continuing + bool sample_frame(const float * h_state, std::vector & codes, const float ** fb, bool * is_eos) { + std::vector order(n_cb); + for (int i = 0; i < n_cb; i++) { + order[(size_t) i] = i; + } + std::shuffle(order.begin(), order.end(), rng); + + codes.assign(n_cb, -1); + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_CODE; + inp.embd = const_cast(h_state); + inp.codes = codes.data(); + inp.n_codes = codes.size(); + mtmd_gen_out out{}; + + int k = 0; + for (int n_pick : schedule) { + if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0 || out.n_feats == 0 || out.n_feats % n_cb != 0) { + LOG_ERR("mtmd_helper_gen_audio: depthformer pass failed\n"); + return false; + } + const size_t n_vocab = out.n_feats / n_cb; + logits.assign(out.feats, out.feats + out.n_feats); + for (int j = 0; j < n_pick; j++, k++) { + const int i = order[(size_t) k]; + float * l = logits.data() + (size_t) i * n_vocab; + apply_rep_penalty(l, n_vocab, i); + codes[(size_t) i] = sample(l, n_vocab); + } + } + + // the last code of every codebook ends the audio + const int32_t eos = (int32_t) (logits.size() / n_cb) - 1; + *is_eos = std::find(codes.begin(), codes.end(), eos) != codes.end(); + if (*is_eos) { + codes.assign(n_cb, eos); + if (!continuing) { + *fb = nullptr; + return true; + } + } + + // all codes known: only the feedback embedding + if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: codebook embedding failed\n"); + return false; + } + fb_buf.assign(out.embd, out.embd + n_embd); + *fb = fb_buf.data(); + return true; + } + + // CTRL penalty over the codes codebook cb took in the last rep_window frames + void apply_rep_penalty(float * l, size_t n_vocab, int cb) const { + if (rep_penalty <= 1.0f || rep_window <= 0) { + return; + } + const size_t n_hist = hist.size() / n_cb; + const size_t n_win = std::min(n_hist, (size_t) rep_window); + std::vector seen(n_vocab, false); + for (size_t f = n_hist - n_win; f < n_hist; f++) { + const int32_t c = hist[f * n_cb + (size_t) cb]; + if (c >= 0 && (size_t) c < n_vocab && !seen[(size_t) c]) { + seen[(size_t) c] = true; + l[c] = l[c] > 0.0f ? l[c] / rep_penalty : l[c] * rep_penalty; + } + } + } + + // TODO: reuse the samplers of common (top-k, temperature, penalties) instead of this one + int32_t sample(const float * l, size_t n_vocab) { + if (temp <= 0.0f) { + return (int32_t) (std::max_element(l, l + n_vocab) - l); + } + // top-k keeps the codes scoring at least the k-th best, ties included + float lmin = -INFINITY; + if (top_k > 0 && (size_t) top_k < n_vocab) { + std::vector sorted(l, l + n_vocab); + std::nth_element(sorted.begin(), sorted.begin() + (top_k - 1), sorted.end(), std::greater()); + lmin = sorted[(size_t) top_k - 1]; + } + const float lmax = *std::max_element(l, l + n_vocab); + std::vector prob(n_vocab); + double sum = 0.0; + for (size_t v = 0; v < n_vocab; v++) { + prob[v] = l[v] < lmin ? 0.0 : std::exp((double) (l[v] - lmax) / temp); + sum += prob[v]; + } + double r = std::uniform_real_distribution(0.0, 1.0)(rng) * sum; + for (size_t v = 0; v < n_vocab; v++) { + r -= prob[v]; + if (r <= 0.0) { + return (int32_t) v; + } + } + return (int32_t) n_vocab - 1; + } + + // n_codes == 0 flushes the ISTFT tail + bool decode_wav(const int32_t * codes, size_t n_codes) { + mtmd_gen_inp inp = mtmd_gen_inp_default(mctx); + inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; + inp.codes = const_cast(codes); + inp.n_codes = n_codes; + mtmd_gen_out out{}; + if (mtmd_gen_audio_process(mctx, gen_state.get(), &inp, &out) != 0) { + LOG_ERR("mtmd_helper_gen_audio: codec decode failed\n"); + return false; + } + audio_pcm.insert(audio_pcm.end(), out.audio, out.audio + out.n_samples); + pending_tail = n_codes > 0; + return true; + } + + bool find_tokens() { + tok_audio_start = find_special_token(vocab, "<|audio_start|>"); + tok_mixed_start = find_special_token(vocab, "<|mixed_start|>"); + tok_text_end = find_special_token(vocab, "<|text_end|>"); + tok_im_end = find_special_token(vocab, "<|im_end|>"); + if (tok_audio_start == LLAMA_TOKEN_NULL || tok_mixed_start == LLAMA_TOKEN_NULL || + tok_text_end == LLAMA_TOKEN_NULL || tok_im_end == LLAMA_TOKEN_NULL) { + LOG_ERR("mtmd_helper_gen_audio: missing LFM3-Audio control tokens in vocab\n"); + return false; + } + return true; + } + + bool decode(llama_batch batch) { + std::vector positions(batch.n_tokens); + std::vector n_seq(batch.n_tokens, 1); + std::vector seq(batch.n_tokens, seq_id); + std::vector seq_ptr(batch.n_tokens); + std::vector logits(batch.n_tokens, 1); + for (int i = 0; i < batch.n_tokens; i++) { + positions[i] = pos + i; + seq_ptr[i] = &seq[i]; + } + batch.pos = positions.data(); + batch.n_seq_id = n_seq.data(); + batch.seq_id = seq_ptr.data(); + batch.logits = logits.data(); + if (llama_decode(lctx, batch) != 0) { + LOG_ERR("mtmd_helper_gen_audio: decode failed\n"); + return false; + } + pos += batch.n_tokens; + return true; + } + + // generate() modes: text, audio until end of audio, interleaved until <|text_end|> + enum gen_mode { MODE_TEXT, MODE_AUDIO, MODE_MIXED }; + // checked by the converter + static constexpr int n_text_block = 6; // interleaved_n_text + static constexpr int n_audio_block = 12; // interleaved_n_audio + + llama_token tok_audio_start = LLAMA_TOKEN_NULL; + llama_token tok_mixed_start = LLAMA_TOKEN_NULL; + llama_token tok_text_end = LLAMA_TOKEN_NULL; + llama_token tok_im_end = LLAMA_TOKEN_NULL; + + bool continuing = false; // the prompt continues a sequence, the response follows the model's switching + gen_mode mode = MODE_AUDIO; + int slots_left = 0; + bool speaking = false; // interleaving: the current block is audio + bool text_done = false; // interleaving: <|text_end|> was emitted, audio until the end-of-audio frame + + llama_seq_id seq_id = 0; + int pos = 0; + std::vector prompt_tokens; + int prompt_pos = 0; + + float temp = 0.0f; + int32_t top_k = 0; + float rep_penalty = 1.0f; + int32_t rep_window = 0; // frames + std::mt19937 rng; + std::vector logits; + std::vector fb_buf; + std::vector hist; // frame-major codes of the frames so far + + bool pending_tail = false; // the ISTFT holds back the overlap of the last frame + std::vector audio_pcm; + std::vector h_state_buf; + mtmd_helper_gen_audio_outtype out_type = MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV; + std::vector out_buf; + bool stream = false; + size_t pcm_sent = 0; // samples already returned by get_output() + bool wav_header_sent = false; +}; + 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: @@ -1269,6 +1725,8 @@ static std::unique_ptr make_pipeline(llama_context * lc return std::unique_ptr(new pockettts_gen_audio_pipeline(lctx, mctx)); case MTMD_GEN_AUDIO_TYPE_SOPRANO: return std::unique_ptr(new soprano_gen_audio_pipeline(lctx, mctx)); + case MTMD_GEN_AUDIO_TYPE_LFM3A: + return std::unique_ptr(new lfm3a_gen_audio_pipeline(lctx, mctx)); default: return nullptr; } @@ -1308,6 +1766,11 @@ int32_t mtmd_helper_gen_audio_set_input(mtmd_helper_gen_audio * ctx, const mtmd_ LOG_ERR("mtmd_helper_gen_audio: unsupported or missing gen-audio pipeline\n"); return 1; } + if (inp->n_past > 0 && !ctx->pipeline->supports_continue()) { + LOG_ERR("mtmd_helper_gen_audio: the pipeline cannot continue a sequence\n"); + return 1; + } + ctx->pipeline->logit_bias.clear(); return ctx->pipeline->set_input(inp); } @@ -1318,6 +1781,34 @@ int32_t mtmd_helper_gen_audio_step_prompt(mtmd_helper_gen_audio * ctx, int32_t n return ctx->pipeline->step_prompt(n_batch); } +bool mtmd_helper_gen_audio_next_is_text(mtmd_helper_gen_audio * ctx) { + return ctx->pipeline && ctx->pipeline->next_is_text(); +} + +const llama_logit_bias * mtmd_helper_gen_audio_get_logit_bias(mtmd_helper_gen_audio * ctx, size_t * n_bias) { + *n_bias = ctx->pipeline ? ctx->pipeline->logit_bias.size() : 0; + return *n_bias ? ctx->pipeline->logit_bias.data() : nullptr; +} + +int32_t mtmd_helper_audio_stream_get_logit_bias(const mtmd_context * mctx, const llama_model * model, + llama_logit_bias * out, int32_t n_max) { + std::vector bias; + if (mtmd_gen_audio_get_info(mctx).type == MTMD_GEN_AUDIO_TYPE_LFM3A) { + // LFM3-Audio live ASR: only <|text_end|> ends the text of a chunk, the stream adds the final <|mixed_end|> + const llama_vocab * vocab = llama_model_get_vocab(model); + for (const char * piece : { "<|mixed_end|>", "<|mixed_start|>", "<|audio_start|>", "<|im_start|>", "<|im_end|>", "<|endoftext|>" }) { + const llama_token tok = find_special_token(vocab, piece); + if (tok != LLAMA_TOKEN_NULL) { + bias.push_back({ tok, -INFINITY }); + } + } + } + for (int32_t i = 0; i < n_max && i < (int32_t) bias.size(); i++) { + out[i] = bias[i]; + } + return (int32_t) bias.size(); +} + int32_t mtmd_helper_gen_audio_step_gen(mtmd_helper_gen_audio * ctx, llama_token sampled, const float * h_state_in, const float ** h_state_out, bool * out_stop) { diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index d06af8f02..28f1efdca 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -11,6 +11,7 @@ #include "mtmd-helper.h" #include "mtmd-helper-common.h" #include "llama.h" +#include "../src/llama-ext.h" #include "hash/hash.h" @@ -1043,8 +1044,10 @@ mtmd_helper_cap mtmd_helper_get_capabilities(const llama_model * model, const mt cap.audio_gen_spk_ref = gen_audio_type == MTMD_GEN_AUDIO_TYPE_QWEN3TTS || gen_audio_type == MTMD_GEN_AUDIO_TYPE_POCKETTTS; - // tts-only model cannot be used for chat (no chat template) - cap.chat = !cap.audio_gen || llama_model_chat_template(model, nullptr) != nullptr; + // a tts model without a chat template only outputs audio, it cannot be used for chat + // neither can a realtime transcription model (e.g. Voxtral Realtime): its text follows the audio by a token delay + cap.text_gen = !cap.audio_gen || llama_model_chat_template(model, nullptr) != nullptr; + cap.chat = cap.text_gen && llama_model_n_tokens_delay(model) == 0; return cap; } diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index e46b5d36a..c9f7ad75d 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -184,7 +184,7 @@ MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, char ** out_text); struct mtmd_helper_cap { - bool chat; // false for tts-only model (no chat template) + bool chat; // false for tts-only models (no chat template) and realtime transcription models (token delay) // input capabilities bool vision; @@ -195,7 +195,7 @@ struct mtmd_helper_cap { bool audio_stream; // realtime transcription, see mtmd_audio_stream_init() // output capabilities - bool text_gen; // always true for now + bool text_gen; // false for an audio-only model (tts-only, no chat template) bool audio_gen; bool audio_gen_spk_ref; // accepts a speaker reference audio for voice cloning }; @@ -223,15 +223,28 @@ enum mtmd_helper_gen_audio_outtype { MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM, // raw PCM MTMD_HELPER_GEN_AUDIO_OUTTYPE_WAV, // WAV PCM 16-bit LE, mono }; +// what a response that continues a sequence (n_past > 0) outputs +enum mtmd_helper_gen_audio_out_mode { + MTMD_HELPER_GEN_AUDIO_OUT_AUTO, // the model switches between text and audio itself + MTMD_HELPER_GEN_AUDIO_OUT_TEXT, + MTMD_HELPER_GEN_AUDIO_OUT_AUDIO, + MTMD_HELPER_GEN_AUDIO_OUT_MIXED, // interleaved text and audio +}; struct mtmd_helper_gen_audio_inp { bool stream; // if true, output() must be called after each step_gen() llama_seq_id seq_id; + // the text to speak; for a model with a "tts" chat template (e.g. LFM3-Audio), the prompt rendered from it const char * prompt; size_t prompt_len; + // > 0: continue a sequence of n_past positions with raw text (LFM3-Audio only) + int32_t n_past; + enum mtmd_helper_gen_audio_out_mode out_mode; // n_past > 0 only (default: auto) + mtmd_bitmap * speaker_ref; // optional, can be NULL const char * lang; // optional, can be NULL + const char * voice; // optional, can be NULL, model-dependent int32_t top_k; float top_p; @@ -271,6 +284,18 @@ MTMD_API int32_t mtmd_helper_gen_audio_step_gen( const float ** h_state_out, bool * out_stop); +// true if the next step_gen() takes a text token: the caller samples it from the logits and passes it as sampled +// false if the next step generates an audio frame from h_state_in +MTMD_API bool mtmd_helper_gen_audio_next_is_text(mtmd_helper_gen_audio * ctx); + +// bias for the text tokens the caller samples (e.g. no switch to audio in a text-only response), n_bias is set to its size +MTMD_API const llama_logit_bias * mtmd_helper_gen_audio_get_logit_bias(mtmd_helper_gen_audio * ctx, size_t * n_bias); + +// bias for the text the caller samples while it transcribes an audio stream opened with a prompt (see mtmd_audio_stream_params) +// writes up to n_max entries to out, returns the size of the bias +MTMD_API int32_t mtmd_helper_audio_stream_get_logit_bias(const mtmd_context * mctx, const struct llama_model * model, + llama_logit_bias * out, int32_t n_max); + // out_data valid until next get_output() or reset() call // out_n_samples (optional, can be NULL) receives the number of generated PCM samples // if inp->stream is true: returns only audio produced since the previous call, and @@ -293,6 +318,7 @@ MTMD_API int32_t mtmd_helper_gen_audio_flush(mtmd_helper_gen_audio * ctx); #ifdef __cplusplus #include #include +#include namespace mtmd_helper { @@ -318,6 +344,7 @@ struct gen_audio { mtmd_helper_gen_audio_inp data = mtmd_helper_gen_audio_inp_default(); std::string prompt_str; std::string lang_str; + std::string voice_str; mtmd::bitmap_ptr speaker_ref_ptr; inp() = default; @@ -328,6 +355,7 @@ struct gen_audio { void set_prompt (std::string p) { prompt_str = std::move(p); } void set_lang (std::string l) { lang_str = std::move(l); } + void set_voice (std::string v) { voice_str = std::move(v); } void set_speaker_ref(mtmd::bitmap_ptr bmp) { speaker_ref_ptr = std::move(bmp); } // pointers are only valid as long as *this is alive @@ -335,6 +363,7 @@ struct gen_audio { data.prompt = prompt_str.c_str(); data.prompt_len = prompt_str.size(); data.lang = lang_str.empty() ? nullptr : lang_str.c_str(); + data.voice = voice_str.empty() ? nullptr : voice_str.c_str(); data.speaker_ref = speaker_ref_ptr.get(); return &data; } @@ -359,6 +388,14 @@ struct gen_audio { int32_t step_gen(llama_token sampled, const float * h_state, const float ** h_state_out, bool * out_stop = nullptr) { return mtmd_helper_gen_audio_step_gen(ctx.get(), sampled, h_state, h_state_out, out_stop); } + bool next_is_text() { + return mtmd_helper_gen_audio_next_is_text(ctx.get()); + } + std::vector logit_bias() { + size_t n = 0; + const llama_logit_bias * bias = mtmd_helper_gen_audio_get_logit_bias(ctx.get(), &n); + return std::vector(bias, bias + n); + } int32_t get_output(int32_t * out_sample_rate, const char ** out_data, size_t * out_data_len, int64_t * out_n_samples = nullptr) { return mtmd_helper_gen_audio_get_output(ctx.get(), out_sample_rate, out_data, out_data_len, out_n_samples); } @@ -367,6 +404,12 @@ struct gen_audio { } }; +inline std::vector audio_stream_logit_bias(const mtmd_context * mctx, const llama_model * model) { + std::vector bias(mtmd_helper_audio_stream_get_logit_bias(mctx, model, nullptr, 0)); + mtmd_helper_audio_stream_get_logit_bias(mctx, model, bias.data(), (int32_t) bias.size()); + return bias; +} + } // namespace mtmd_helper #endif diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 37d5eb21d..7e0da513a 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1080,6 +1080,15 @@ struct mtmd_context { { audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_LFM3A: + { + // LFM3-Omni wraps the audio in the image span, LFM3-Audio (no vision tower) places it bare + if (ctx_v) { + aud_beg = img_beg; + aud_end = img_end; + } + audio_preproc = std::make_unique(ctx_a); + } break; case PROJECTOR_TYPE_GRANITE_SPEECH: { audio_preproc = std::make_unique(ctx_a); @@ -2101,6 +2110,10 @@ mtmd_gen_audio_info mtmd_gen_audio_get_info(const mtmd_context * ctx) { info.type = MTMD_GEN_AUDIO_TYPE_SOPRANO; info.sample_rate = 32000; break; + case PROJECTOR_TYPE_LFM3A_GEN: + info.type = MTMD_GEN_AUDIO_TYPE_LFM3A; + info.sample_rate = 24000; + break; default: info.type = MTMD_GEN_AUDIO_TYPE_NONE; break; @@ -2131,6 +2144,16 @@ mtmd_gen_inp mtmd_gen_inp_default(const mtmd_context * ctx) { case PROJECTOR_TYPE_SOPRANO: inp.type = MTMD_GEN_PROCESS_TYPE_GEN_WAV; break; + case PROJECTOR_TYPE_LFM3A_GEN: + { + // MaskGIT sampling runs in the helper, with the settings of the checkpoint + const auto * hparams = clip_get_hparams(ctx->ctx_gen_a); + inp.top_k = hparams->gen_top_k; + inp.top_p = 1.0f; + inp.temp = hparams->gen_temp; + inp.penalty_repeat = hparams->gen_penalty_repeat; + inp.penalty_last_n = hparams->gen_penalty_last_n; + } break; default: break; } @@ -2146,6 +2169,9 @@ struct mtmd_gen_state { std::vector out_embd; // next-step hidden state fed back to backbone (GEN_CODE) std::vector out_audio; // decoded PCM samples for the current frame (GEN_WAV) + // lfm3a: the codec graph yields STFT frames, the ISTFT runs on the host + std::unique_ptr istft; + ~mtmd_gen_state() { clip_gen_state_free(clip_state); } @@ -2168,6 +2194,10 @@ mtmd_gen_state * mtmd_gen_state_init(mtmd_context * ctx, mtmd_gen_state_params p if (!st->clip_state) { return nullptr; } + if (clip_get_projector_type(ctx->ctx_gen_a) == PROJECTOR_TYPE_LFM3A_GEN) { + const auto * hparams = clip_get_hparams(ctx->ctx_gen_a); + st->istft = std::make_unique(hparams->audio_n_fft, hparams->audio_hop_len); + } return st.release(); } @@ -2177,6 +2207,72 @@ void mtmd_gen_state_free(mtmd_gen_state * state) { void mtmd_gen_state_reset(mtmd_gen_state * state) { clip_gen_state_reset(state->clip_state); + if (state->istft) { + const auto * hparams = clip_get_hparams(state->ctx_clip); + state->istft = std::make_unique(hparams->audio_n_fft, hparams->audio_hop_len); + } +} + +// lfm3a GEN_WAV: ISTFT of the codec output, n_codes == 0 flushes and resets the state +static int32_t mtmd_gen_wav_lfm3a(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) { + clip_ctx * ctx_clip = ctx->ctx_gen_a; + const auto & hparams = *clip_get_hparams(ctx_clip); + + const int n_bins = hparams.audio_n_fft / 2 + 1; + const int n_trim = (hparams.audio_n_fft - hparams.audio_hop_len) / 2; + + auto & audio = state->out_audio; + audio.clear(); + + if (!inp->codes || inp->n_codes == 0) { + // the tail beyond the last frame is padding, like the centered ISTFT of the reference + audio = state->istft->flush(); + audio.resize(audio.size() - std::min(audio.size(), (size_t) n_trim)); + mtmd_gen_state_reset(state); + } else { + std::vector in_codes(inp->codes, inp->codes + inp->n_codes); + std::vector feats; + + // no hidden-state input, the batch entry is an unused placeholder + clip_image_f32 dummy; + dummy.set_size({1, 1}, false, true); + dummy.cpy_buf(std::vector(1, 0.0f)); + + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(dummy)); + + clip_encode_params params; + params.imgs = &batch; + params.n_threads = ctx->n_threads; + params.gen_process = CLIP_GEN_PROCESS_GEN_WAV; + params.codes = &in_codes; + params.out_feats = &feats; + params.gen_state = state->clip_state; + if (!clip_encode(ctx_clip, ¶ms)) { + LOG_ERR("%s: clip_encode failed (codec)\n", __func__); + return 1; + } + + // [2 * n_bins, n_stft_frames]: log magnitude, then phase + GGML_ASSERT(feats.size() % (2 * (size_t) n_bins) == 0); + const size_t n_frames = feats.size() / (2 * (size_t) n_bins); + std::vector spec(2 * (size_t) n_bins); + for (size_t f = 0; f < n_frames; f++) { + const float * fr = feats.data() + f * 2 * n_bins; + for (int j = 0; j < n_bins; j++) { + const float mag = std::exp(fr[j]); + spec[2 * j + 0] = mag * std::cos(fr[n_bins + j]); + spec[2 * j + 1] = mag * std::sin(fr[n_bins + j]); + } + const std::vector pcm = state->istft->process_frame(spec.data()); + audio.insert(audio.end(), pcm.begin(), pcm.end()); + } + } + + out->audio = audio.data(); + out->n_samples = audio.size(); + return 0; } static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * state, const mtmd_gen_inp * inp, mtmd_gen_out * out) { @@ -2200,12 +2296,22 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * } } + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_WAV && clip_get_projector_type(ctx_clip) == PROJECTOR_TYPE_LFM3A_GEN) { + return mtmd_gen_wav_lfm3a(ctx, state, inp, out); + } + if (inp->type == MTMD_GEN_PROCESS_TYPE_GEN_CODE) { const size_t n_embd = (size_t) clip_n_mmproj_embd(ctx_clip); + // the hidden state may be omitted when the graph does not use it clip_image_f32 hidden_state; hidden_state.set_size({(int) n_embd, 1}, false, true); - hidden_state.cpy_buf(std::vector(inp->embd, inp->embd + n_embd)); + hidden_state.cpy_buf(inp->embd ? std::vector(inp->embd, inp->embd + n_embd) : std::vector(n_embd, 0.0f)); + + std::vector in_codes; + if (inp->codes) { + in_codes.assign(inp->codes, inp->codes + inp->n_codes); + } clip_image_f32_batch batch; batch.is_audio = true; @@ -2227,6 +2333,7 @@ static int32_t mtmd_gen_audio_process_impl(mtmd_context * ctx, mtmd_gen_state * params.top_k = inp->top_k; params.top_p = inp->top_p; params.temp = inp->temp; + params.codes = inp->codes ? &in_codes : nullptr; params.out_is_eos = &is_eos; params.gen_state = state->clip_state; @@ -3095,26 +3202,18 @@ mtmd_memory_usage mtmd_get_memory_usage(const char * mmproj_fname, // realtime audio input // +// the model-independent part of a stream, the model-specific part is in the subclasses below struct mtmd_audio_stream { mtmd_context * ctx; clip_ctx * ctx_clip; - clip_enc_state * enc_state = nullptr; // model constants int n_mel; - int n_fft; int hop; int n_embd; // decoder embedding size int frames_per_tok; // mel frames per decoder row int samples_per_tok; - int n_frames_max; // encoder frames per clip_encode() call int n_step_mel; // mel frames per encoder call - int n_right_pad; // silence tokens added by flush() - float log_mel_max; // global_log_mel_max in params.json - - std::vector prefix; // text tokens paired with the first rows - - mtmd_audio_cache cache; // samples not yet turned into mel frames, buf[0] is at absolute index buf_pos std::vector buf; @@ -3123,103 +3222,143 @@ struct mtmd_audio_stream { int64_t n_frames = 0; // mel frames computed so far std::vector mel; // pending mel frames, frame-major - std::vector mel_ctx; // last 2 mel frames, the left context of the conv stem std::vector embd; // pending audio embeddings size_t embd_off = 0; // rows of embd already consumed std::vector out; // rows of the last step - std::vector tok_embd; int64_t n_rows = 0; // rows produced so far - int64_t n_rows_flushed = 0; // rows produced once the audio up to the last flush is consumed + int64_t n_rows_done = 0; // audio rows whose text has been sampled, see mtmd_audio_stream_get_n_rows() + int64_t n_rows_flushed = 0; // audio rows up to the end of the last flush bool pending_flush = false; bool ended = false; + bool seq_start = false; // the rows of the last step start a new sequence - ~mtmd_audio_stream() { - clip_enc_state_free(enc_state); + + virtual ~mtmd_audio_stream() = default; + + // model constants and initial state, false on failure + virtual bool init(const mtmd_audio_stream_params & params) = 0; + virtual int32_t flush(bool end) = 0; + virtual int32_t step(llama_token last_token) = 0; + + void add_silence(int64_t n) { + buf.insert(buf.end(), n, 0.0f); + n_samples += n; } }; +// Voxtral Realtime: a row is an audio embedding plus the embedding of the previous token +struct mtmd_audio_stream_voxtral_rt : mtmd_audio_stream { + clip_enc_state * enc_state = nullptr; + + int n_fft; + int n_frames_max; // encoder frames per clip_encode() call + int n_right_pad; // silence tokens added by flush() + float log_mel_max; // global_log_mel_max in params.json + + std::vector prefix; // text tokens paired with the first rows + + mtmd_audio_cache cache; + + std::vector mel_ctx; // last 2 mel frames, the left context of the conv stem + std::vector tok_embd; + + ~mtmd_audio_stream_voxtral_rt() override { + clip_enc_state_free(enc_state); + } + + bool init(const mtmd_audio_stream_params & params) override; + int32_t flush(bool end) override; + int32_t step(llama_token last_token) override; +}; + +// LFM3-Audio: the rows are plain audio embeddings of a chat turn, the model answers after the audio +// live ASR (with a prompt), trained layout: [A x 12] text<|text_end|> ... [A x r] <|mixed_end|> text<|text_end|> +// training samples are max 30 s: restart the sequence at a silent chunk after 15 s, always after 30 s; the encoder keeps its state +struct mtmd_audio_stream_lfm3a : mtmd_audio_stream { + int win; // samples per mel frame + clip_enc_state * enc_state = nullptr; + + // live ASR + static constexpr int64_t chunk_rows = 12; + static constexpr int32_t max_chunk_tokens = 48; // text tokens per chunk, then the chunk is closed + static constexpr int32_t max_final_tokens = 96; + + std::vector prompt; // empty: the rows of a chat turn + llama_token tok_text_end = LLAMA_TOKEN_NULL; + llama_token tok_mixed_end = LLAMA_TOKEN_NULL; + int64_t soft_seq_rows = 0; + int64_t max_seq_rows = 0; + + bool started = false; + bool done = false; // the text of the final chunk is complete + bool in_chunk = false; // the text of the last chunk is being sampled + bool final_chunk = false; + bool chunk_has_text = false; + int32_t n_chunk_text = 0; + int64_t n_chunk_rows = 0; + int64_t n_seq_rows = 0; // audio rows since the sequence started + std::vector closing; // the tokens that close the last chunk, they go in before the next one + std::vector tok_embd; + // the audio markers of the model (LFM3-Omni), as rows around the audio + std::vector beg_rows; + std::vector end_rows; + bool beg_done = false; + bool end_done = false; + + ~mtmd_audio_stream_lfm3a() override { + clip_enc_state_free(enc_state); + } + + bool init(const mtmd_audio_stream_params & params) override; + int32_t flush(bool end) override; + int32_t step(llama_token last_token) override; + + bool add_tokens(const llama_token * tokens, size_t n); + int32_t step_live(llama_token last_token); +}; + bool mtmd_support_audio_stream(const mtmd_context * ctx) { - return ctx->ctx_a && ctx->text_model && clip_get_projector_type(ctx->ctx_a) == PROJECTOR_TYPE_VOXTRAL_RT; + if (!ctx->ctx_a || !ctx->text_model) { + return false; + } + const projector_type proj = clip_get_projector_type(ctx->ctx_a); + return proj == PROJECTOR_TYPE_VOXTRAL_RT || proj == PROJECTOR_TYPE_LFM3A; } mtmd_audio_stream_params mtmd_audio_stream_params_default(void) { mtmd_audio_stream_params params; - params.step_ms = 320; + params.step_ms = 320; + params.prompt = nullptr; + params.n_prompt = 0; return params; } -static void mtmd_audio_stream_add_silence(mtmd_audio_stream * s, int64_t n) { - s->buf.insert(s->buf.end(), n, 0.0f); - s->n_samples += n; -} - mtmd_audio_stream * mtmd_audio_stream_init(mtmd_context * ctx, mtmd_audio_stream_params params) { if (!mtmd_support_audio_stream(ctx)) { LOG_ERR("%s: the model does not support streaming audio input\n", __func__); return nullptr; } const clip_hparams * hparams = clip_get_hparams(ctx->ctx_a); - const llama_vocab * vocab = ctx->vocab; - auto s = std::make_unique(); + std::unique_ptr s; + switch (clip_get_projector_type(ctx->ctx_a)) { + case PROJECTOR_TYPE_VOXTRAL_RT: s = std::make_unique(); break; + case PROJECTOR_TYPE_LFM3A: s = std::make_unique(); break; + default: GGML_ABORT("unsupported projector"); + } s->ctx = ctx; s->ctx_clip = ctx->ctx_a; s->n_mel = hparams->n_mel_bins; - s->n_fft = hparams->audio_n_fft; s->hop = hparams->audio_hop_len; s->n_embd = clip_n_mmproj_embd(ctx->ctx_a); if (s->n_embd != ctx->n_embd_text) { LOG_ERR("%s: mmproj embedding size %d does not match the text model (%d)\n", __func__, s->n_embd, ctx->n_embd_text); return nullptr; } - - // one decoder row per 80 ms of audio: the conv stem has stride 2, then frames are stacked - s->frames_per_tok = 2 * hparams->proj_stack_factor; - s->samples_per_tok = s->frames_per_tok * s->hop; - s->n_frames_max = (hparams->warmup_audio_size - 2) / 2; - s->log_mel_max = 1.5f; - - // the encoder works in whole decoder rows - int n_step = (int) ((int64_t) params.step_ms * hparams->audio_sample_rate / 1000 / s->hop); - n_step = std::max(n_step, s->frames_per_tok); - n_step = CLIP_ALIGN(n_step, s->frames_per_tok); - s->n_step_mel = std::min(n_step, 2 * s->n_frames_max); - - // prefix: BOS then one pad token per silence token and per delay token - // the output lags behind the audio by n_delay rows, the text model is conditioned on the same value - const int n_delay = llama_model_n_tokens_delay(ctx->text_model); - const int n_left_pad = 32; - s->n_right_pad = n_delay + 1 + 10; - { - const std::string pad_str = "[STREAMING_PAD]"; - std::vector pad(8); - const int n_pad = llama_tokenize(vocab, pad_str.c_str(), pad_str.size(), pad.data(), pad.size(), false, true); - if (n_pad != 1) { - LOG_ERR("%s: cannot find the %s token\n", __func__, pad_str.c_str()); - return nullptr; - } - s->prefix.push_back(llama_vocab_bos(vocab)); - s->prefix.insert(s->prefix.end(), n_left_pad + n_delay, pad[0]); - } - - s->cache.fill_sin_cos_table(s->n_fft); - s->cache.fill_hann_window(hparams->audio_window_len, true); - s->cache.fill_mel_filterbank_matrix(s->n_mel, s->n_fft, hparams->audio_sample_rate); - - s->enc_state = clip_enc_state_init(s->ctx_clip, s->n_frames_max); - if (!s->enc_state) { + if (!s->init(params)) { return nullptr; } - - // the first mel frame is centered on the first sample, so its window starts before the audio - s->buf.assign(s->n_fft / 2, 0.0f); - s->buf_pos = -s->n_fft / 2; - s->mel_ctx.assign(2 * s->n_mel, 0.0f); - - // silence before the audio - mtmd_audio_stream_add_silence(s.get(), (int64_t) n_left_pad * s->samples_per_tok); - return s.release(); } @@ -3240,18 +3379,99 @@ int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) { if (s->ended) { return 1; } + return s->flush(end); +} + +int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) { + return s->step(last_token); +} + +const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * s) { + return s->out.data(); +} + +bool mtmd_audio_stream_is_seq_start(const mtmd_audio_stream * s) { + return s->seq_start; +} + + +// +// Voxtral Realtime +// + +bool mtmd_audio_stream_voxtral_rt::init(const mtmd_audio_stream_params & params) { + const clip_hparams * hparams = clip_get_hparams(ctx_clip); + const llama_vocab * vocab = ctx->vocab; + + n_fft = hparams->audio_n_fft; + + // one decoder row per 80 ms of audio: the conv stem has stride 2, then frames are stacked + frames_per_tok = 2 * hparams->proj_stack_factor; + samples_per_tok = frames_per_tok * hop; + n_frames_max = (hparams->warmup_audio_size - 2) / 2; + log_mel_max = 1.5f; + + // the encoder works in whole decoder rows + int n_step = (int) ((int64_t) params.step_ms * hparams->audio_sample_rate / 1000 / hop); + n_step = std::max(n_step, frames_per_tok); + n_step = CLIP_ALIGN(n_step, frames_per_tok); + n_step_mel = std::min(n_step, 2 * n_frames_max); + + if (params.n_prompt > 0) { + LOG_ERR("%s: Voxtral Realtime takes no prompt\n", __func__); + return false; + } + + // prefix: BOS then one pad token per silence token and per delay token + // the output lags behind the audio by n_delay rows, the text model is conditioned on the same value + const int n_delay = llama_model_n_tokens_delay(ctx->text_model); + const int n_left_pad = 32; + n_right_pad = n_delay + 1 + 10; + { + const std::string pad_str = "[STREAMING_PAD]"; + std::vector pad(8); + const int n_pad = llama_tokenize(vocab, pad_str.c_str(), pad_str.size(), pad.data(), pad.size(), false, true); + if (n_pad != 1) { + LOG_ERR("%s: cannot find the %s token\n", __func__, pad_str.c_str()); + return false; + } + prefix.push_back(llama_vocab_bos(vocab)); + prefix.insert(prefix.end(), n_left_pad + n_delay, pad[0]); + } + + cache.fill_sin_cos_table(n_fft); + cache.fill_hann_window(hparams->audio_window_len, true); + cache.fill_mel_filterbank_matrix(n_mel, n_fft, hparams->audio_sample_rate); + + enc_state = clip_enc_state_init(ctx_clip, n_frames_max); + if (!enc_state) { + return false; + } + + // the first mel frame is centered on the first sample, so its window starts before the audio + buf.assign(n_fft / 2, 0.0f); + buf_pos = -n_fft / 2; + mel_ctx.assign(2 * n_mel, 0.0f); + + // silence before the audio + add_silence((int64_t) n_left_pad * samples_per_tok); + + return true; +} + +int32_t mtmd_audio_stream_voxtral_rt::flush(bool end) { // align to a whole row, then add the silence the model expects after the audio - const int64_t n_align = (s->samples_per_tok - s->n_samples % s->samples_per_tok) % s->samples_per_tok; - mtmd_audio_stream_add_silence(s, n_align + (int64_t) s->n_right_pad * s->samples_per_tok); + const int64_t n_align = (samples_per_tok - n_samples % samples_per_tok) % samples_per_tok; + add_silence(n_align + (int64_t) n_right_pad * samples_per_tok); // one row per samples_per_tok, the left padding included - s->n_rows_flushed = s->n_samples / s->samples_per_tok; - s->pending_flush = true; - s->ended = end; + n_rows_flushed = n_samples / samples_per_tok; + pending_flush = true; + ended = end; return 0; } // mel frame k is centered on sample k * hop, with force the samples past the end are taken as silence -static void mtmd_audio_stream_compute_mel(mtmd_audio_stream * s, bool force) { +static void mtmd_audio_stream_compute_mel(mtmd_audio_stream_voxtral_rt * s, bool force) { const int half = s->n_fft / 2; std::vector window(s->n_fft); std::vector frame(s->n_mel); @@ -3281,7 +3501,7 @@ static void mtmd_audio_stream_compute_mel(mtmd_audio_stream * s, bool force) { } } -static bool mtmd_audio_stream_encode(mtmd_audio_stream * s, bool force) { +static bool mtmd_audio_stream_encode(mtmd_audio_stream_voxtral_rt * s, bool force) { while (true) { const int n_avail = (int) (s->mel.size() / s->n_mel); int n = n_avail - n_avail % s->frames_per_tok; @@ -3333,7 +3553,7 @@ static bool mtmd_audio_stream_encode(mtmd_audio_stream * s, bool force) { } // out rows = audio embeddings + token embeddings -static bool mtmd_audio_stream_add_rows(mtmd_audio_stream * s, const llama_token * tokens, size_t n) { +static bool mtmd_audio_stream_add_rows(mtmd_audio_stream_voxtral_rt * s, const llama_token * tokens, size_t n) { s->tok_embd.resize(n * s->n_embd); if (llama_model_get_tok_embd_rows(s->ctx->text_model, tokens, n, s->tok_embd.data()) == 0) { LOG_ERR("%s: invalid token\n", __func__); @@ -3348,55 +3568,330 @@ static bool mtmd_audio_stream_add_rows(mtmd_audio_stream * s, const llama_token return true; } -int32_t mtmd_audio_stream_step(mtmd_audio_stream * s, llama_token last_token) { - const bool force = s->pending_flush; - s->pending_flush = false; +int32_t mtmd_audio_stream_voxtral_rt::step(llama_token last_token) { + const bool force = pending_flush; + pending_flush = false; - mtmd_audio_stream_compute_mel(s, force); - if (!mtmd_audio_stream_encode(s, force)) { + // last_token was sampled from the last row, the text of every row so far is out + n_rows_done = n_rows; + + mtmd_audio_stream_compute_mel(this, force); + if (!mtmd_audio_stream_encode(this, force)) { return -2; } // drop the consumed rows once in a while - if (s->embd_off >= 1024) { - s->embd.erase(s->embd.begin(), s->embd.begin() + s->embd_off * s->n_embd); - s->embd_off = 0; + if (embd_off >= 1024) { + embd.erase(embd.begin(), embd.begin() + embd_off * n_embd); + embd_off = 0; } - const int64_t n_avail = (int64_t) (s->embd.size() / s->n_embd) - s->embd_off; - s->out.clear(); + const int64_t n_avail = (int64_t) (embd.size() / n_embd) - embd_off; + out.clear(); - if (s->n_rows == 0) { + if (n_rows == 0) { // the prefix rows go out together, the caller samples from the last one - const int64_t n_prefix = (int64_t) s->prefix.size(); + const int64_t n_prefix = (int64_t) prefix.size(); if (n_avail < n_prefix) { - return s->ended ? -1 : 0; + return ended ? -1 : 0; } - if (!mtmd_audio_stream_add_rows(s, s->prefix.data(), s->prefix.size())) { + if (!mtmd_audio_stream_add_rows(this, prefix.data(), prefix.size())) { return -2; } return (int32_t) n_prefix; } if (n_avail == 0) { - return s->ended ? -1 : 0; + return ended ? -1 : 0; } if (last_token == LLAMA_TOKEN_NULL) { LOG_ERR("%s: the token sampled from the previous row is required\n", __func__); return -2; } - if (!mtmd_audio_stream_add_rows(s, &last_token, 1)) { + if (!mtmd_audio_stream_add_rows(this, &last_token, 1)) { return -2; } return 1; } -const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * s) { - return s->out.data(); +// +// LFM3-Audio +// + +bool mtmd_audio_stream_lfm3a::init(const mtmd_audio_stream_params & params) { + const clip_hparams * hparams = clip_get_hparams(ctx_clip); + + // one row per 8 mel frames (three stride-2 convs), mel frames are not centered + win = hparams->audio_window_len; + frames_per_tok = 8; + samples_per_tok = frames_per_tok * hop; + + const int n_step = (int) ((int64_t) params.step_ms * hparams->audio_sample_rate / 1000 / hop); + n_step_mel = CLIP_ALIGN(std::max(n_step, frames_per_tok), frames_per_tok); + + if (params.n_prompt > 0) { + // live ASR, a chunk per encoder call + n_step_mel = (int) chunk_rows * frames_per_tok; + prompt.assign(params.prompt, params.prompt + params.n_prompt); + + const llama_vocab * vocab = llama_model_get_vocab(ctx->text_model); + auto find_token = [&](const char * piece) { + llama_token tok[2]; + const int n = llama_tokenize(vocab, piece, (int32_t) strlen(piece), tok, 2, false, true); + return n == 1 ? tok[0] : LLAMA_TOKEN_NULL; + }; + tok_text_end = find_token("<|text_end|>"); + tok_mixed_end = find_token("<|mixed_end|>"); + if (tok_text_end == LLAMA_TOKEN_NULL || tok_mixed_end == LLAMA_TOKEN_NULL) { + LOG_ERR("%s: the vocab lacks the tokens of LFM3-Audio live ASR\n", __func__); + return false; + } + soft_seq_rows = std::lround(15.0 * hparams->audio_sample_rate / samples_per_tok); + max_seq_rows = std::lround(30.0 * hparams->audio_sample_rate / samples_per_tok); + } + + enc_state = clip_enc_state_init(ctx_clip, n_step_mel); + if (!enc_state) { + return false; + } + + // the markers the tokenizer puts around an audio clip go around the stream as well + const llama_vocab * vocab = llama_model_get_vocab(ctx->text_model); + auto marker_rows = [&](const std::string & text, std::vector & rows) { + if (text.empty()) { + return true; + } + std::vector tokens(text.size() + 4); + const int32_t n = llama_tokenize(vocab, text.c_str(), (int32_t) text.size(), tokens.data(), (int32_t) tokens.size(), false, true); + if (n <= 0) { + return false; + } + rows.resize((size_t) n * n_embd); + return llama_model_get_tok_embd_rows(ctx->text_model, tokens.data(), n, rows.data()) != 0; + }; + return marker_rows(ctx->aud_beg, beg_rows) && marker_rows(ctx->aud_end, end_rows); +} + +// end: pad the tail like mtmd_audio_preprocessor_lfm3a, so the rows match encoding the clip at once +// otherwise only whole rows are encoded, a padded partial row would change the encoder state +// live ASR: pad the last chunk with silence, so its text comes out +int32_t mtmd_audio_stream_lfm3a::flush(bool end) { + if (end) { + const int64_t n_padded = n_samples <= win ? win : n_samples + (hop - (n_samples - win) % hop) % hop; + add_silence(n_padded - n_samples); + } else if (!prompt.empty() && n_samples > 0) { + const int64_t frames_per_chunk = chunk_rows * frames_per_tok; + const int64_t n_frames_audio = n_samples <= win ? 1 : (n_samples - win + hop - 1) / hop + 1; + const int64_t n_frames_padded = (n_frames_audio + frames_per_chunk - 1) / frames_per_chunk * frames_per_chunk; + add_silence(win + (n_frames_padded - 1) * hop - n_samples); + } + n_rows_flushed = (n_samples >= win ? (n_samples - win) / hop + 1 : 0) / frames_per_tok; + pending_flush = true; + ended = end; + return 0; +} + +// mel frame k covers samples [k * hop, k * hop + win), computed with the offline preprocessor +static bool mtmd_audio_stream_compute_mel(mtmd_audio_stream_lfm3a * s) { + const int64_t n_total = s->n_samples >= s->win ? (s->n_samples - s->win) / s->hop + 1 : 0; + const int64_t n_new = n_total - s->n_frames; + if (n_new <= 0) { + return true; + } + const int64_t start = s->n_frames * s->hop - s->buf_pos; + const size_t len = (size_t) (s->win + (n_new - 1) * s->hop); + std::vector mel_chunks; + if (!s->ctx->audio_preproc->preprocess(s->buf.data() + start, len, mel_chunks) || mel_chunks.size() != 1) { + LOG_ERR("%s: unable to preprocess audio\n", __func__); + return false; + } + const auto & mel = mel_chunks[0]; + GGML_ASSERT(mel.n_len == n_new && mel.n_mel == s->n_mel); + for (int64_t t = 0; t < n_new; t++) { + for (int m = 0; m < s->n_mel; m++) { + s->mel.push_back(mel.data[(size_t) m * mel.n_len + t]); + } + } + s->n_frames = n_total; + + const int64_t keep_from = s->n_frames * s->hop; + if (keep_from > s->buf_pos) { + s->buf.erase(s->buf.begin(), s->buf.begin() + (keep_from - s->buf_pos)); + s->buf_pos = keep_from; + } + return true; +} + +// whole rows only, except for the last call of an ended stream; the encoder state carries the rest +static bool mtmd_audio_stream_encode(mtmd_audio_stream_lfm3a * s, bool flush) { + const int n_avail = (int) (s->mel.size() / s->n_mel); + int n = n_avail - n_avail % s->frames_per_tok; + if (flush && s->ended) { + n = n_avail; + } else if (!flush) { + n -= n % s->n_step_mel; + } + if (n <= 0) { + return true; + } + + std::vector data((size_t) s->n_mel * n); + for (int m = 0; m < s->n_mel; m++) { + for (int t = 0; t < n; t++) { + data[(size_t) m * n + t] = s->mel[(size_t) t * s->n_mel + m]; + } + } + clip_image_f32 img; + img.set_size({n, s->n_mel}, false, /* is_audio */ true); + img.cpy_buf(data); + const int n_rows = clip_n_output_tokens(s->ctx_clip, &img); + + clip_image_f32_batch batch; + batch.is_audio = true; + batch.entries.push_back(std::move(img)); + + std::vector res((size_t) s->n_embd * n_rows); + clip_encode_params params; + params.n_threads = s->ctx->n_threads; + params.imgs = &batch; + params.out_embd = &res; + params.enc_state = s->enc_state; + if (!clip_encode(s->ctx_clip, ¶ms)) { + LOG_ERR("%s: failed to encode audio\n", __func__); + return false; + } + s->embd.insert(s->embd.end(), res.begin(), res.end()); + s->mel.erase(s->mel.begin(), s->mel.begin() + (size_t) n * s->n_mel); + return true; +} + +int32_t mtmd_audio_stream_lfm3a::step(llama_token last_token) { + const bool flush = pending_flush; + pending_flush = false; + + if (!mtmd_audio_stream_compute_mel(this) || !mtmd_audio_stream_encode(this, flush)) { + return -2; + } + if (!prompt.empty()) { + return step_live(last_token); + } + + // rows carry no text, all of them go out at once and last_token is not used + const int64_t n_avail = (int64_t) (embd.size() / n_embd) - embd_off; + out.clear(); + if (n_avail > 0 || ended) { + if (!beg_done) { + beg_done = true; + out = beg_rows; + } + out.insert(out.end(), embd.begin() + embd_off * n_embd, embd.end()); + embd.clear(); + embd_off = 0; + if (ended && !end_done) { + end_done = true; + out.insert(out.end(), end_rows.begin(), end_rows.end()); + } + } + const int64_t n_out = (int64_t) (out.size() / n_embd); + if (n_out == 0) { + return ended ? -1 : 0; + } + n_rows += n_out; + n_rows_done += n_avail; // the markers are not audio rows + return (int32_t) n_out; +} + +bool mtmd_audio_stream_lfm3a::add_tokens(const llama_token * tokens, size_t n) { + if (n == 0) { + return true; + } + tok_embd.resize(n * n_embd); + if (llama_model_get_tok_embd_rows(ctx->text_model, tokens, n, tok_embd.data()) == 0) { + LOG_ERR("%s: invalid token\n", __func__); + return false; + } + out.insert(out.end(), tok_embd.begin(), tok_embd.end()); + return true; +} + +// last_token continues the text of the chunk, or <|text_end|> closes it and the next 12 rows go out (the rest at the end) +int32_t mtmd_audio_stream_lfm3a::step_live(llama_token last_token) { + out.clear(); + seq_start = false; + + // drop the consumed rows once in a while + if (embd_off >= 1024) { + embd.erase(embd.begin(), embd.begin() + embd_off * n_embd); + embd_off = 0; + } + + if (in_chunk) { + if (last_token == LLAMA_TOKEN_NULL) { + LOG_ERR("%s: the token sampled from the previous row is required\n", __func__); + return -2; + } + if (last_token != tok_text_end) { + chunk_has_text = true; + if (n_chunk_text + 1 < (final_chunk ? max_final_tokens : max_chunk_tokens)) { + n_chunk_text++; + return add_tokens(&last_token, 1) ? 1 : -2; + } + // close a runaway chunk like the model should have + closing.push_back(last_token); + } + // the text of the chunk is complete, its closing tokens go in before the next chunk + in_chunk = false; + n_rows_done += n_chunk_rows; + if (final_chunk) { + done = true; + } + closing.push_back(tok_text_end); + } + if (done) { + return -1; + } + + const int64_t n_avail = (int64_t) (embd.size() / n_embd) - embd_off; + const bool drained = ended && mel.empty(); + const bool final = drained && n_avail <= chunk_rows; + if (!final && n_avail < chunk_rows) { + return 0; + } + if (final && n_avail == 0 && !started) { + done = true; // no audio at all + return -1; + } + + if (!started || n_seq_rows >= max_seq_rows || (n_seq_rows >= soft_seq_rows && !chunk_has_text)) { + closing = prompt; // the closing tokens go away with the sequence + seq_start = true; + n_seq_rows = 0; + } + const int64_t n_feed = final ? n_avail : chunk_rows; + if (!add_tokens(closing.data(), closing.size())) { + return -2; + } + closing.clear(); + const float * rows = embd.data() + embd_off * n_embd; + out.insert(out.end(), rows, rows + (size_t) n_feed * n_embd); + embd_off += n_feed; + if (final && !add_tokens(&tok_mixed_end, 1)) { + return -2; + } + + started = true; + in_chunk = true; + final_chunk = final; + chunk_has_text = false; + n_chunk_text = 0; + n_chunk_rows = n_feed; + n_seq_rows += n_feed; + n_rows += n_feed; + return (int32_t) (out.size() / n_embd); } int64_t mtmd_audio_stream_get_n_rows(const mtmd_audio_stream * s) { - return s->n_rows; + return s->n_rows_done; } int64_t mtmd_audio_stream_get_n_rows_flushed(const mtmd_audio_stream * s) { diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 0a6bf4c4c..f552dddf9 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -383,6 +383,11 @@ typedef struct mtmd_audio_stream mtmd_audio_stream; struct mtmd_audio_stream_params { int32_t step_ms; // audio encoded per encoder call, rounded up to the model step (default: 320) + + // tokens that open the stream, they go out with the first rows (default: none) + // with the prompt of its "asr" chat template, LFM3-Audio transcribes live + const llama_token * prompt; + size_t n_prompt; }; MTMD_API struct mtmd_audio_stream_params mtmd_audio_stream_params_default(void); @@ -412,11 +417,14 @@ MTMD_API int32_t mtmd_audio_stream_step(mtmd_audio_stream * stream, llama_token // rows of the last step, the reading size (in bytes) is n_rows * llama_model_n_embd_inp(model) * sizeof(float) MTMD_API const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * stream); -// rows produced so far, the prefix rows included +// true if the rows of the last step start a new sequence: the caller drops its sequence before decoding them +MTMD_API bool mtmd_audio_stream_is_seq_start(const mtmd_audio_stream * stream); + +// audio rows whose text has been sampled (prefix rows included), step() updates it from last_token MTMD_API int64_t mtmd_audio_stream_get_n_rows(const mtmd_audio_stream * stream); -// rows the stream produces up to the end of the last flush -// once get_n_rows() reaches it, every row of the flushed audio has been produced +// audio rows up to the end of the last flush +// once get_n_rows() reaches it, the text of the flushed audio has been sampled MTMD_API int64_t mtmd_audio_stream_get_n_rows_flushed(const mtmd_audio_stream * stream); ///////////////////////////////////////// @@ -428,6 +436,7 @@ enum mtmd_gen_audio_type { MTMD_GEN_AUDIO_TYPE_QWEN3TTS, MTMD_GEN_AUDIO_TYPE_POCKETTTS, MTMD_GEN_AUDIO_TYPE_SOPRANO, + MTMD_GEN_AUDIO_TYPE_LFM3A, }; struct mtmd_gen_audio_info { @@ -444,6 +453,7 @@ enum mtmd_gen_process_type { MTMD_GEN_PROCESS_TYPE_GEN_WAV, // convert semantic to PCM audio // for qwen3tts, this is code2wav // for pocket-tts, this is mimi decoder + // for lfm3a, this is the codec decoder + ISTFT, n_codes == 0 flushes the tail }; struct mtmd_gen_inp { @@ -455,9 +465,12 @@ struct mtmd_gen_inp { int32_t top_k; float top_p; float temp; // sampling temperature, or noise scale for flow-matching decoders + float penalty_repeat; // repetition penalty of the sampled codes, <= 1 is off + int32_t penalty_last_n; // frames the repetition penalty looks back on // for MTMD_GEN_PROCESS_TYPE_GEN_WAV // pass either codes (discrete) or feats (continuous), depending on the pipeline + // lfm3a GEN_CODE also takes codes: the codes revealed so far, -1 for a masked codebook int32_t * codes; size_t n_codes; const float * feats; diff --git a/tools/server/README.md b/tools/server/README.md index fe4012c4a..b6202bd36 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -857,6 +857,8 @@ Returns raw audio bytes (`audio/wav` by default) rather than JSON. For more info `lang`: Language code for the utterance (model-dependent, e.g. `en`, `zh`). Optional. +`voice`: Voice for the utterance (model-dependent, e.g. `female` or `male` for LFM3-Audio). Optional. + `speaker_ref_b64`: Base64-encoded reference audio to clone the speaker's voice. Optional. Alternatively, upload the reference audio as a `speaker_ref` file field via `multipart/form-data` (see example below) - if both are provided, the uploaded file takes precedence. `top_k`, `top_p`: Sampling params for the acoustic code predictor. Optional, model-dependent defaults apply. diff --git a/tools/server/server-chat.cpp b/tools/server/server-chat.cpp index b0245786c..6d1c56c0d 100644 --- a/tools/server/server-chat.cpp +++ b/tools/server/server-chat.cpp @@ -687,11 +687,11 @@ json convert_transcriptions_to_chatcmpl( if (response_format != "json") { throw std::invalid_argument("Only 'json' response_format is supported for transcription"); } - const common_chat_prompt_preset preset = common_chat_get_asr_prompt(tmpls); + const common_chat_prompt_preset preset = common_chat_get_asr_prompt(tmpls, language); if (prompt.empty()) { prompt = preset.user; } - if (!language.empty()) { + if (!language.empty() && !preset.has_language) { prompt += string_format(" (language: %s)", language.c_str()); } prompt += get_media_marker(); diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index da0d58fb4..0ec9b548f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -3424,7 +3424,9 @@ private: return true; // rows from a previous step, not decoded yet } - // every produced row is decoded and its delta queued, the commits ending there are complete + const int32_t n = mtmd_audio_stream_step(slot.audio_stream.get(), slot.sampled); + + // complete the commits whose rows have their text sampled, before the new rows: their text goes to the next item const int64_t n_rows = mtmd_audio_stream_get_n_rows(slot.audio_stream.get()); while (!slot.audio_commit_rows.empty() && slot.audio_commit_rows.front() <= n_rows) { slot.audio_commit_rows.pop_front(); @@ -3440,7 +3442,6 @@ private: queue_results.send(std::move(res)); } - const int32_t n = mtmd_audio_stream_step(slot.audio_stream.get(), slot.sampled); if (n > 0) { const float * rows = mtmd_audio_stream_get_output_embd(slot.audio_stream.get()); slot.inp_embd.assign(rows, rows + (size_t) n * batch.n_embd); @@ -5142,7 +5143,8 @@ server_context_meta server_context::get_meta() const { /* has_inp_video */ impl->chat_params.allow_video, /* has_cap_chat */ impl->has_cap_chat(), /* has_cap_tts */ impl->has_cap_tts(), - /* has_audio_stream */ impl->mtmd_cap.audio_stream, + // only realtime transcription models (no chat) stream without a prompt + /* has_audio_stream */ impl->mtmd_cap.audio_stream && !impl->mtmd_cap.chat, /* json_ui_settings */ impl->json_ui_settings, /* slot_n_ctx */ impl->n_ctx_slot(), /* pooling_type */ llama_pooling_type(impl->ctx_tgt), @@ -6367,8 +6369,13 @@ void server_routes::init_routes() { server_task task(SERVER_TASK_TYPE_TTS); server_schema::eval_tts_schema(params, task, body); const bool stream = task.params.stream; - task.tts_inp.set_prompt(prompt); - task.tts_inp.set_lang(json_value(body, "lang", std::string())); + // a model with a "tts" chat template (e.g. LFM3-Audio) speaks the prompt rendered from it + const std::string lang = json_value(body, "lang", std::string()); + const std::string voice = json_value(body, "voice", std::string()); + const std::string tts_prompt = common_chat_get_tts_prompt(meta->chat_params.tmpls.get(), prompt, lang, voice); + task.tts_inp.set_prompt(tts_prompt.empty() ? prompt : tts_prompt); + task.tts_inp.set_lang(lang); + task.tts_inp.set_voice(voice); task.tts_inp.data.stream = stream; task.tts_inp.data.out_type = response_format == "pcm" ? MTMD_HELPER_GEN_AUDIO_OUTTYPE_PCM diff --git a/tools/tts/README.md b/tools/tts/README.md index 41397b82e..1793b7992 100644 --- a/tools/tts/README.md +++ b/tools/tts/README.md @@ -84,3 +84,29 @@ Both [Soprano-80M](https://huggingface.co/ekwek/Soprano-80M) and [Soprano-1.1-80 python convert_hf_to_gguf.py path/to/Soprano-1.1-80M --outtype f32 --outfile soprano.gguf python convert_hf_to_gguf.py path/to/Soprano-1.1-80M --mmproj --outtype f16 --outfile mmproj-soprano.gguf ``` + +## LFM3-Audio + +Available params: +- `--tts-lang` can be `ar`, `da`, `de`, `el`, `en`, `en-gb`, `es`, `fi`, `fr`, `he`, `it`, `ja`, `ko`, `nl`, `pl`, `pt`, `sv`, `tr`, `zh` (default: `en`), other values are used as the language name +- `--tts-voice` can be `female` (default) or `male` +- Note: `--tts-speaker-file` is not used, the model has no voice cloning +- Note: `--top-k` and `--top-p` are not used, audio frames are sampled with the settings of the model card (LFM3-Audio: temperature 1.0, top-k 128, repetition penalty 1.1 over 4096 frames; LFM3-Omni: temperature 0.4, repetition penalty 1.2 over 64 frames), the `audio_*` keys of `generation_config.json` override them at conversion +- Each frame is 80 ms, generation stops at the end-of-audio frame + +Example usage: + +```sh +llama-tts -m lfm3-audio.gguf \ + -mm mmproj-lfm3-audio.gguf \ + -p "Hello world" \ + --tts-voice male \ + --output out.wav +``` + +The mmproj holds the audio encoder (for ASR), the depthformer and the codec decoder, the latter two are exported when the checkpoint has an `audio_tokenizer` directory: + +```sh +python convert_hf_to_gguf.py path/to/lfm3-audio --outfile lfm3-audio.gguf +python convert_hf_to_gguf.py path/to/lfm3-audio --mmproj --outfile mmproj-lfm3-audio.gguf +``` diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index aea794e65..8949373cb 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -1,4 +1,5 @@ #include "arg.h" +#include "chat.h" #include "common.h" #include "sampling.h" #include "log.h" @@ -36,7 +37,7 @@ static void print_usage(int, char ** argv) { LOG("\nexample usage:\n"); LOG("\n %s -m backbone.gguf -mm mmproj.gguf -p \"text to speak\" -o output.wav", argv[0]); LOG("\n %s -hf user/model -p \"text to speak\" -o output.wav\n", argv[0]); - LOG("\nnote: --tts-lang and --tts-speaker-file may not be supported in all models"); + LOG("\nnote: --tts-lang, --tts-voice and --tts-speaker-file may not be supported in all models"); LOG("\n use -n to limit the output length"); LOG("\n see tts/README.md for per-model usage notes"); LOG("\n\n"); @@ -111,14 +112,22 @@ int main(int argc, char ** argv) { speaker_bitmap.reset(wrapper.bitmap); } + // a model with a "tts" chat template (e.g. LFM3-Audio) speaks the prompt rendered from it + auto tmpls = common_chat_templates_init(model, ""); + std::string prompt = common_chat_get_tts_prompt(tmpls.get(), params.prompt, params.tts_lang, params.tts_voice); + if (prompt.empty()) { + prompt = params.prompt; + } + mtmd_helper::gen_audio gen; gen.init(lctx, mctx.get()); mtmd_helper_gen_audio_inp inp = mtmd_helper_gen_audio_inp_default(); inp.seq_id = 0; - inp.prompt = params.prompt.c_str(); - inp.prompt_len = params.prompt.size(); + inp.prompt = prompt.c_str(); + inp.prompt_len = prompt.size(); inp.speaker_ref = speaker_bitmap.get(); inp.lang = params.tts_lang.c_str(); + inp.voice = params.tts_voice.c_str(); inp.top_k = params.sampling.top_k; inp.top_p = params.sampling.top_p; inp.seed = params.sampling.seed;