diff --git a/conversion/__init__.py b/conversion/__init__.py index 6c3914fbd..c2683aec4 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -199,6 +199,7 @@ TEXT_MODEL_MAP: dict[str, str] = { "MiniMaxM3SparseForConditionalGeneration": "minimax", "Ministral3ForCausalLM": "mistral3", "Mistral3ForConditionalGeneration": "mistral3", + "VoxtralRealtimeForConditionalGeneration": "mistral3", "MistralForCausalLM": "llama", "MixtralForCausalLM": "llama", "ModernBertDecisionModel": "bert", @@ -383,6 +384,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = { "UltravoxModel": "ultravox", "UnlimitedOCRForCausalLM": "deepseek", "VoxtralForConditionalGeneration": "ultravox", + "VoxtralRealtimeForConditionalGeneration": "ultravox", "YoutuVLForConditionalGeneration": "youtuvl", } diff --git a/conversion/mistral3.py b/conversion/mistral3.py index fee039b35..661475251 100644 --- a/conversion/mistral3.py +++ b/conversion/mistral3.py @@ -66,3 +66,13 @@ class Mistral3Model(TextModel): def write(self): self.impl.write() + + +@ModelBase.register("VoxtralRealtimeForConditionalGeneration") +@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602") +class VoxtralRealtimeModel(LlamaModel): + model_arch = gguf.MODEL_ARCH.MISTRAL3 + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_realtime_delay_tokens(self.hparams["default_num_delay_tokens"]) diff --git a/conversion/ultravox.py b/conversion/ultravox.py index 62819e574..4a246daa1 100644 --- a/conversion/ultravox.py +++ b/conversion/ultravox.py @@ -6,6 +6,7 @@ if TYPE_CHECKING: from torch import Tensor from .base import MmprojModel, ModelBase, TextModel, gguf +from .llama import LlamaModel @ModelBase.register("UltravoxModel") @@ -208,3 +209,47 @@ class AudioFlamingo3WhisperEncoderModel(WhisperEncoderModel): # Was trained in BF16, being safe, avoiding quantizing to FP16 return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) + + +@ModelBase.register("VoxtralRealtimeForConditionalGeneration") +@ModelBase.example("mistralai/Voxtral-Mini-4B-Realtime-2602") +class VoxtralRealtimeEncoderModel(MmprojModel): + has_vision_encoder = False + has_audio_encoder = True + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_audio is not None + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.VOXTRAL_RT) + self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["num_mel_bins"]) + self.gguf_writer.add_audio_head_dim(self.hparams_audio["head_dim"]) + self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("rms_norm_eps", 1e-5)) + self.gguf_writer.add_audio_window_size(self.hparams_audio["sliding_window"]) + self.gguf_writer.add_audio_stack_factor(self.global_config["downsample_factor"]) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + if ".conv" in name and ".weight" in name: + return gguf.GGMLQuantizationType.F16 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + # prevent clash naming with vision tensors + if name.startswith("multi_modal_projector"): + name = "audio." + name + + return super().filter_tensors((name, gen)) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if "conv1.bias" in name or "conv2.bias" in name: + data_torch = data_torch.unsqueeze(-1) + + # HF stores q/k in rotate-half layout, the encoder uses interleaved (normal) RoPE + if ".self_attn.q_proj." in name or ".self_attn.k_proj." in name: + assert self.hparams_audio is not None + n_head = self.hparams_audio["num_attention_heads"] + data_torch = LlamaModel.permute(data_torch, n_head, n_head) + + yield from super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a3a9c478b..92a69bd57 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -175,6 +175,7 @@ class Keys: HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" + REALTIME_DELAY_TOKENS = "{arch}.realtime_delay_tokens" class Adapters: COUNT = "{arch}.adapters.count" @@ -481,6 +482,7 @@ class Keys: class Attention: HEAD_COUNT = "clip.audio.attention.head_count" + HEAD_DIM = "clip.audio.attention.head_dim" # set when qkv width != n_embd LAYERNORM_EPS = "clip.audio.attention.layer_norm_epsilon" class Projector: @@ -758,6 +760,8 @@ class MODEL_TENSOR(IntEnum): FFN_POST_NORM = auto() FFN_POST_NORM_1 = auto() # gemma4 FFN_POST_NORM_2 = auto() # gemma4 + FFN_NORM_ADA_DOWN = auto() # voxtral realtime + FFN_NORM_ADA_UP = auto() # voxtral realtime FFN_GATE = auto() FFN_DOWN = auto() FFN_UP = auto() @@ -1572,6 +1576,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.FFN_PRE_NORM_2: "blk.{bid}.pre_ffw_norm_2", # gemma4 MODEL_TENSOR.FFN_POST_NORM_1: "blk.{bid}.post_ffw_norm_1", # gemma4 MODEL_TENSOR.FFN_POST_NORM_2: "blk.{bid}.post_ffw_norm_2", # gemma4 + MODEL_TENSOR.FFN_NORM_ADA_DOWN: "blk.{bid}.ffn_norm_ada_down", + MODEL_TENSOR.FFN_NORM_ADA_UP: "blk.{bid}.ffn_norm_ada_up", MODEL_TENSOR.FFN_GATE: "blk.{bid}.ffn_gate", MODEL_TENSOR.FFN_DOWN: "blk.{bid}.ffn_down", MODEL_TENSOR.FFN_UP: "blk.{bid}.ffn_up", @@ -5554,6 +5560,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.ATTN_ROT_EMBD, MODEL_TENSOR.FFN_GATE_INP, MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_NORM_ADA_DOWN, + MODEL_TENSOR.FFN_NORM_ADA_UP, MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, @@ -6297,6 +6305,7 @@ class VisionProjectorType: GLMA = "glma" # audio QWEN25O = "qwen2.5o" # omni VOXTRAL = "voxtral" + VOXTRAL_RT = "voxtral_rt" MERALION = "meralion" # audio: Whisper + gated MLP adaptor LFM2 = "lfm2" KIMIVL = "kimivl" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index a1b1f1263..93ec95eb6 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1138,6 +1138,9 @@ class GGUFWriter: def add_attn_temperature_scale(self, value: float) -> None: self.add_float32(Keys.Attention.TEMPERATURE_SCALE.format(arch=self.arch), value) + def add_realtime_delay_tokens(self, value: int) -> None: + self.add_uint32(Keys.LLM.REALTIME_DELAY_TOKENS.format(arch=self.arch), value) + def add_pooling_type(self, value: PoolingType) -> None: self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value) @@ -1579,6 +1582,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_dim(self, value: int) -> None: + self.add_uint32(Keys.ClipAudio.Attention.HEAD_DIM, value) + def add_audio_attention_layernorm_eps(self, value: float) -> None: self.add_float32(Keys.ClipAudio.Attention.LAYERNORM_EPS, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 01796e978..ab6abf9cb 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -486,6 +486,14 @@ class TensorNameMap: ), # Post feed-forward norm + MODEL_TENSOR.FFN_NORM_ADA_DOWN: ( + "model.layers.{bid}.ada_rms_norm.linear1", # voxtral realtime + ), + + MODEL_TENSOR.FFN_NORM_ADA_UP: ( + "model.layers.{bid}.ada_rms_norm.linear2", # voxtral realtime + ), + MODEL_TENSOR.FFN_POST_NORM: ( "model.layers.{bid}.post_feedforward_layernorm", # gemma2 olmo2 "layers.{bid}.post_feedforward_layernorm", # embeddinggemma @@ -2491,6 +2499,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_CONV1D: ( "audio_tower.conv{bid}", # ultravox + "audio_tower.embedder.conv{bid}", # voxtral realtime "conformer.pre_encode.conv.{bid}", # lfm2 "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4 @@ -2523,6 +2532,7 @@ class TensorNameMap: MODEL_TENSOR.A_POST_NORM: ( "audio_tower.layer_norm", # ultravox + "audio_tower.norm", # voxtral realtime "audio_encoder.dots_encoder.speech_encoder.layer_norm", # dots3note "audio_tower.ln_post", # qwen2omni "encoder.layer_norm", # mimo-audio-tokenizer @@ -2594,6 +2604,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_OUTPUT: ( "audio_tower.layers.{bid}.self_attn.out_proj", # ultravox + "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 "conformer.layers.{bid}.attention.post", # gemma3n @@ -2633,6 +2644,7 @@ class TensorNameMap: MODEL_TENSOR.A_ENC_FFN_UP: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc1_up", # dots3note (split from fc1 in conversion code) "audio_tower.layers.{bid}.fc1", # ultravox + "audio_tower.layers.{bid}.mlp.up_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear1", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4 @@ -2643,11 +2655,13 @@ 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) + "audio_tower.layers.{bid}.mlp.gate_proj", # voxtral realtime ), MODEL_TENSOR.A_ENC_FFN_DOWN: ( "audio_encoder.dots_encoder.speech_encoder.layers.{bid}.fc2", # dots3note "audio_tower.layers.{bid}.fc2", # ultravox + "audio_tower.layers.{bid}.mlp.down_proj", # voxtral realtime "conformer.layers.{bid}.feed_forward1.linear2", # lfm2 "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4 diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 73c6d51d9..1476a079a 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -291,6 +291,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" }, { LLM_KV_ATTENTION_TEMPERATURE_LENGTH, "%s.attention.temperature_length" }, { LLM_KV_ATTENTION_TEMPERATURE_SCALE, "%s.attention.temperature_scale" }, + { LLM_KV_REALTIME_DELAY_TOKENS, "%s.realtime_delay_tokens" }, { LLM_KV_ATTENTION_KEY_LENGTH_MLA, "%s.attention.key_length_mla" }, { LLM_KV_ATTENTION_VALUE_LENGTH_MLA, "%s.attention.value_length_mla" }, { LLM_KV_ATTENTION_KEY_LENGTH_SWA, "%s.attention.key_length_swa" }, @@ -507,6 +508,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_FFN_POST_NORM, "blk.%d.post_ffw_norm" }, { LLM_TENSOR_FFN_POST_NORM_1, "blk.%d.post_ffw_norm_1" }, { LLM_TENSOR_FFN_POST_NORM_2, "blk.%d.post_ffw_norm_2" }, + { LLM_TENSOR_FFN_NORM_ADA_DOWN, "blk.%d.ffn_norm_ada_down" }, + { LLM_TENSOR_FFN_NORM_ADA_UP, "blk.%d.ffn_norm_ada_up" }, { LLM_TENSOR_FFN_PRE_NORM_2, "blk.%d.pre_ffw_norm_2" }, { LLM_TENSOR_FFN_GATE_SHEXP, "blk.%d.ffn_gate_shexp" }, { LLM_TENSOR_FFN_UP_SHEXP, "blk.%d.ffn_up_shexp" }, @@ -991,6 +994,8 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_FFN_PRE_NORM_2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_POST_NORM_1, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_POST_NORM_2, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_FFN_NORM_ADA_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_FFN_NORM_ADA_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_POST_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_NORM_EXPS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4b607fec7..0c842d312 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -266,6 +266,7 @@ enum llm_kv { LLM_KV_HRM_H_CYCLES, LLM_KV_HRM_L_CYCLES, LLM_KV_HRM_PREFIX_LM, + LLM_KV_REALTIME_DELAY_TOKENS, LLM_KV_ATTENTION_HEAD_COUNT, LLM_KV_ATTENTION_HEAD_COUNT_KV, @@ -510,6 +511,8 @@ enum llm_tensor { LLM_TENSOR_FFN_POST_NORM, LLM_TENSOR_FFN_POST_NORM_1, LLM_TENSOR_FFN_POST_NORM_2, + LLM_TENSOR_FFN_NORM_ADA_DOWN, + LLM_TENSOR_FFN_NORM_ADA_UP, LLM_TENSOR_FFN_PRE_NORM_2, LLM_TENSOR_FFN_GATE, LLM_TENSOR_FFN_DOWN, diff --git a/src/llama-ext.h b/src/llama-ext.h index b8c5e8033..8b3932bd7 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -154,3 +154,10 @@ LLAMA_API uint32_t llama_model_target_layer_ids_n(const struct llama_mode // if out is nullptr, returns the number of tokens without writing to out // caller must allocate enough memory for out before calling LLAMA_API uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out); + +// same as above, but only the rows of the given tokens, in order (n_tokens * n_embd elements) +// returns the number of elements written or 0 on error +LLAMA_API uint32_t llama_model_get_tok_embd_rows(const struct llama_model * model, const llama_token * tokens, uint32_t n_tokens, float * out); + +// realtime models: number of tokens the output lags behind the audio input, used by mtmd; 0 if not applicable +LLAMA_API uint32_t llama_model_n_tokens_delay(const struct llama_model * model); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e8e249b21..80be4bb57 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1991,6 +1991,19 @@ ggml_tensor * llm_graph_context::build_ffn( cb(cur, "ffn_act", il); } } break; + case LLM_FFN_GELU_ERF: + if (gate && type_gate == LLM_FFN_PAR) { + cur = ggml_geglu_erf_split(ctx0, cur, tmp); + cb(cur, "ffn_geglu_erf", il); + type_gate = LLM_FFN_SEQ; + } else { + cur = ggml_gelu_erf(ctx0, cur); + cb(cur, "ffn_gelu_erf", il); + if (act_scales != NULL) { + cur = ggml_div(ctx0, cur, act_scales); + cb(cur, "ffn_act", il); + } + } break; case LLM_FFN_RELU: if (gate && type_gate == LLM_FFN_PAR) { cur = ggml_reglu_split(ctx0, cur, tmp); diff --git a/src/llama-graph.h b/src/llama-graph.h index 1a3698cad..816fa07cc 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -59,6 +59,7 @@ enum llm_ffn_op_type : int { LLM_FFN_NONE = 0, // sentinel: unset; archs must assign before use LLM_FFN_SILU, LLM_FFN_GELU, + LLM_FFN_GELU_ERF, LLM_FFN_RELU, LLM_FFN_RELU_SQR, LLM_FFN_SWIGLU, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index bc0d49215..3e025368b 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -281,6 +281,9 @@ struct llama_hparams { float f_attn_temp_scale = 0.0f; float f_attn_temp_offset = 0.0f; // offset position index + // realtime models: number of tokens the output lags behind the audio input + uint32_t n_tokens_delay = 0; + // gemma3n altup uint32_t n_altup = 4; // altup_num_inputs uint32_t i_altup_act = 0; // altup_active_idx diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 2660bf38a..5d48d702c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -3906,6 +3906,32 @@ uint32_t llama_model_target_layer_ids_n(const struct llama_model * model) { return (uint32_t) model->target_layer_ids.size(); } +// read rows [row0, row0 + n_rows) of a 2D tensor as F32 +static void llama_tensor_rows_to_f32(const ggml_tensor * tensor, int64_t row0, int64_t n_rows, float * out) { + const size_t nelements = (size_t) n_rows * tensor->ne[0]; + const size_t offset = (size_t) row0 * tensor->nb[1]; + const size_t nbytes = ggml_row_size(tensor->type, tensor->ne[0]) * n_rows; + + if (tensor->type == GGML_TYPE_F32) { + ggml_backend_tensor_get(tensor, out, offset, nbytes); + return; + } + + std::vector buf(nbytes); + ggml_backend_tensor_get(tensor, buf.data(), offset, nbytes); + + const ggml_type_traits * traits = ggml_get_type_traits(tensor->type); + if (tensor->type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *) buf.data(), out, nelements); + } else if (tensor->type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row((const ggml_bf16_t *) buf.data(), out, nelements); + } else if (ggml_is_quantized(tensor->type) && traits->to_float != nullptr) { + traits->to_float(buf.data(), out, nelements); + } else { + GGML_ABORT("unsupported tensor type for dequantization: %s", ggml_type_name(tensor->type)); + } +} + uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) { if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr) { return 0; @@ -3919,24 +3945,27 @@ uint32_t llama_model_get_tok_embd(const struct llama_model * model, float * out) return (uint32_t) nelements; } - if (tensor->type == GGML_TYPE_F32) { - ggml_backend_tensor_get(tensor, out, 0, nelements * sizeof(float)); - return (uint32_t) nelements; - } - - std::vector buf(ggml_nbytes(tensor)); - ggml_backend_tensor_get(tensor, buf.data(), 0, buf.size()); - - const ggml_type_traits * traits = ggml_get_type_traits(tensor->type); - if (tensor->type == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((const ggml_fp16_t *) buf.data(), out, nelements); - } else if (tensor->type == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((const ggml_bf16_t *) buf.data(), out, nelements); - } else if (ggml_is_quantized(tensor->type) && traits->to_float != nullptr) { - traits->to_float(buf.data(), out, nelements); - } else { - GGML_ABORT("unsupported tensor type for dequantization: %s", ggml_type_name(tensor->type)); - } + llama_tensor_rows_to_f32(tensor, 0, tensor->ne[1], out); return (uint32_t) nelements; } + +uint32_t llama_model_n_tokens_delay(const struct llama_model * model) { + return model->hparams.n_tokens_delay; +} + +uint32_t llama_model_get_tok_embd_rows(const struct llama_model * model, const llama_token * tokens, uint32_t n_tokens, float * out) { + if (model->vocab.n_tokens() == 0 || model->tok_embd == nullptr || tokens == nullptr || out == nullptr) { + return 0; + } + + const ggml_tensor * tensor = model->tok_embd; + for (uint32_t i = 0; i < n_tokens; i++) { + if (tokens[i] < 0 || tokens[i] >= tensor->ne[1]) { + return 0; + } + llama_tensor_rows_to_f32(tensor, tokens[i], 1, out + (size_t) i * tensor->ne[0]); + } + + return n_tokens * (uint32_t) tensor->ne[0]; +} diff --git a/src/llama-model.h b/src/llama-model.h index 66f9129e6..6335dc2bc 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -327,6 +327,8 @@ struct llama_layer { struct ggml_tensor * ffn_post_norm_1 = nullptr; // gemma4 struct ggml_tensor * ffn_post_norm_2 = nullptr; // gemma4 struct ggml_tensor * ffn_pre_norm_2 = nullptr; // gemma4 + struct ggml_tensor * ffn_norm_ada_down = nullptr; // voxtral realtime + struct ggml_tensor * ffn_norm_ada_up = nullptr; // voxtral realtime struct ggml_tensor * layer_out_norm = nullptr; struct ggml_tensor * layer_out_norm_b = nullptr; struct ggml_tensor * ffn_norm_exps = nullptr; diff --git a/src/models/mistral3.cpp b/src/models/mistral3.cpp index 484e1b698..af058d558 100644 --- a/src/models/mistral3.cpp +++ b/src/models/mistral3.cpp @@ -8,6 +8,8 @@ void llama_model_mistral3::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ROPE_SCALING_YARN_BETA_SLOW, hparams.yarn_beta_slow, false); ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false); + ml.get_key(LLM_KV_REALTIME_DELAY_TOKENS, hparams.n_tokens_delay, false); + hparams.f_attn_temp_offset = 0.0f; // TODO: maybe add n_attn_temp_floor_scale as a separate KV? @@ -26,7 +28,7 @@ void llama_model_mistral3::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { +void llama_model_mistral3::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); @@ -53,6 +55,13 @@ void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + // optional ada rms norm (voxtral realtime), the hidden size is taken from the tensor itself + if (const ggml_tensor * meta = ml.get_tensor_meta(tn(LLM_TENSOR_FFN_NORM_ADA_DOWN, "weight", i).str().c_str())) { + const int64_t n_ada = meta->ne[1]; + layer.ffn_norm_ada_down = create_tensor(tn(LLM_TENSOR_FFN_NORM_ADA_DOWN, "weight", i), {n_embd, n_ada}, 0); + layer.ffn_norm_ada_up = create_tensor(tn(LLM_TENSOR_FFN_NORM_ADA_UP, "weight", i), {n_ada, n_embd}, 0); + } + if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) { layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); @@ -86,6 +95,34 @@ void llama_model_mistral3::load_arch_tensors(llama_model_loader &) { } } +// sinusoidal embedding of the output delay, constant for the whole session +class llama_model_mistral3::input_t_cond : public llm_graph_input_i { +public: + input_t_cond(uint32_t n_delay, int64_t n_embd) : data(n_embd) { + const int64_t half = n_embd / 2; + for (int64_t i = 0; i < half; i++) { + const double freq = exp(-log(10000.0) * i / half) * n_delay; + data[i] = cos(freq); + data[half + i] = sin(freq); + } + } + + void set_input(const llama_ubatch * ubatch) override { + GGML_UNUSED(ubatch); + ggml_backend_tensor_set(t_cond, data.data(), 0, ggml_nbytes(t_cond)); + } + + // the input does not depend on the batch + bool can_reuse(const llm_graph_params & params) override { + GGML_UNUSED(params); + return true; + } + + ggml_tensor * t_cond = nullptr; // F32 [n_embd] + + std::vector data; +}; + std::unique_ptr llama_model_mistral3::build_arch_graph(const llm_graph_params & params) const { return std::make_unique(*this, params); } @@ -112,6 +149,17 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa auto * inp_attn = build_attn_inp_kv(); + // (optional) time conditioning of the ada rms norm + ggml_tensor * t_cond = nullptr; + if (model.layers[0].ffn_norm_ada_down) { + auto inp = std::make_unique(hparams.n_tokens_delay, n_embd); + t_cond = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, n_embd); + ggml_set_input(t_cond); + cb(t_cond, "inp_t_cond", -1); + inp->t_cond = t_cond; + res->add_input(std::move(inp)); + } + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; ggml_tensor * inp_out_ids = build_inp_out_ids(); @@ -176,6 +224,18 @@ llama_model_mistral3::graph::graph(const llama_model & model, const llm_graph_pa LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + if (t_cond) { + // ada rms norm: scale the normalized input by (1 + mlp(t_cond)) + ggml_tensor * ada = build_ffn(t_cond, + model.layers[il].ffn_norm_ada_down, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_norm_ada_up, NULL, NULL, + NULL, + LLM_FFN_GELU_ERF, LLM_FFN_SEQ, il); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, cur, ada)); + cb(cur, "ffn_norm_ada", il); + } + cur = build_ffn(cur, model.layers[il].ffn_up, model.layers[il].ffn_up_b, model.layers[il].ffn_up_s, model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, model.layers[il].ffn_gate_s, diff --git a/src/models/models.h b/src/models/models.h index 51fdcfbb0..94a789617 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2744,6 +2744,8 @@ struct llama_model_mistral3 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; + class input_t_cond; // voxtral realtime: time conditioning of the ada rms norm + struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); }; diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index edb5d3d89..398b73904 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -69,6 +69,7 @@ add_library(mtmd models/step3vl.cpp models/siglip.cpp models/whisper-enc.cpp + models/voxtral-rt.cpp models/deepseekocr.cpp models/deepseekocr2.cpp models/mobilenetv5.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index ad291a557..10eea6f74 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -1,6 +1,7 @@ #pragma once #include "ggml.h" +#include "ggml-cpp.h" #include "gguf.h" #include "clip.h" @@ -446,6 +447,20 @@ // align x to upper multiple of n #define CLIP_ALIGN(x, n) ((((x) + (n) - 1) / (n)) * (n)) +// sliding window KV cache of a causal audio encoder, see clip_enc_state_init() +struct clip_enc_state { + int32_t n_ctx = 0; // ring slots: attention window + frames per call + int32_t n_past = 0; // absolute position of the next frame + + std::vector slot_pos; // absolute position held by each slot, -1 if empty + std::vector stem_tail; // last conv stem output frame, empty on a cold start + + ggml_context_ptr ctx_data; + ggml_backend_buffer_ptr buf; + std::vector k; // per layer [n_embd_kv, n_ctx] + std::vector v; +}; + // forward declaration // TODO: improve this later struct clip_ctx; @@ -480,6 +495,7 @@ enum projector_type { PROJECTOR_TYPE_GLMA, PROJECTOR_TYPE_QWEN25O, // will be replaced by QWEN2A or QWEN25VL depending on clip_ctx PROJECTOR_TYPE_VOXTRAL, + PROJECTOR_TYPE_VOXTRAL_RT, PROJECTOR_TYPE_MERALION, PROJECTOR_TYPE_MUSIC_FLAMINGO, PROJECTOR_TYPE_LFM2, @@ -553,6 +569,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_GLMA, "glma"}, { PROJECTOR_TYPE_QWEN25O, "qwen2.5o"}, { PROJECTOR_TYPE_VOXTRAL, "voxtral"}, + { PROJECTOR_TYPE_VOXTRAL_RT, "voxtral_rt"}, { PROJECTOR_TYPE_MERALION, "meralion"}, { PROJECTOR_TYPE_MUSIC_FLAMINGO, "musicflamingo"}, { PROJECTOR_TYPE_LFM2, "lfm2"}, diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 336c266b5..793b8cd1a 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1050,6 +1050,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + builder = std::make_unique(ctx, img, params ? params->enc_state : nullptr); + } break; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_VOXTRAL: case PROJECTOR_TYPE_QWEN2A: @@ -1923,6 +1927,22 @@ struct clip_model_loader { hparams.audio_window_len = 400; hparams.audio_hop_len = 160; } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + get_u32(KEY_A_PROJ_STACK_FACTOR, hparams.proj_stack_factor); + get_u32(KEY_A_ATTN_WINDOW_SIZE, hparams.attn_window_size); + hparams.rope_theta = 1000000.0f; // params.json + hparams.ffn_op = FFN_SILU; + log_ffn_op = "silu"; + + // audio preprocessing params + hparams.audio_sample_rate = 16000; + hparams.audio_n_fft = 400; + hparams.audio_window_len = 400; + hparams.audio_hop_len = 160; + // 2 context frames + 8 s of new frames, see clip_enc_state_init() + hparams.warmup_audio_size = 2 + 2*400; + } break; case PROJECTOR_TYPE_MIMO_AUDIO: { get_u32(KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers, false); @@ -3320,6 +3340,7 @@ struct clip_model_loader { } } break; case PROJECTOR_TYPE_VOXTRAL: + case PROJECTOR_TYPE_VOXTRAL_RT: { model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight")); model.conv1d_1_b = get_tensor(string_format(TN_CONV1D, 1, "bias")); @@ -4459,6 +4480,11 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { n_patches = n_patches_y * n_patches_x; } } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + // 2 context frames, conv stride 2, then stacked frames + n_patches = (img->nx() - 2) / 2 / ctx->model.hparams.proj_stack_factor; + } break; case PROJECTOR_TYPE_MUSE_SPARK_AUDIO: { const int64_t fold = ctx->model.hparams.audio_samples_per_token(); @@ -4656,6 +4682,56 @@ static std::vector list_gen_state_slots(const clip_hparams & hpa } } +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_VOXTRAL_RT || hparams.attn_window_size <= 0) { + LOG_ERR("%s: the encoder does not support streaming input\n", __func__); + return nullptr; + } + if (n_frames_max <= 0) { + LOG_ERR("%s: invalid n_frames_max %d\n", __func__, n_frames_max); + return nullptr; + } + + const int n_layer = hparams.n_layer; + const int d_head = hparams.n_embd_head > 0 ? hparams.n_embd_head : hparams.n_embd / hparams.n_head; + const int n_embd_kv = hparams.n_head_kv * d_head; + + auto st = std::make_unique(); + st->n_ctx = hparams.attn_window_size + n_frames_max; + st->slot_pos.assign(st->n_ctx, -1); + + ggml_init_params params = { + /*.mem_size =*/ ggml_tensor_overhead() * 2 * n_layer, + /*.mem_buffer =*/ nullptr, + /*.no_alloc =*/ true, + }; + st->ctx_data.reset(ggml_init(params)); + for (int il = 0; il < n_layer; il++) { + ggml_tensor * k = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F16, n_embd_kv, st->n_ctx); + ggml_tensor * v = ggml_new_tensor_2d(st->ctx_data.get(), GGML_TYPE_F16, n_embd_kv, st->n_ctx); + ggml_format_name(k, "enc_state_k_%d", il); + ggml_format_name(v, "enc_state_v_%d", il); + st->k.push_back(k); + st->v.push_back(v); + } + 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; + } + // empty slots are masked out, but they must not hold NaN + ggml_backend_buffer_clear(st->buf.get(), 0); + + LOG_INF("%s: encoder state: %d slots, %.2f MiB\n", __func__, st->n_ctx, ggml_backend_buffer_get_size(st->buf.get()) / 1024.0 / 1024.0); + + return st.release(); +} + +void clip_enc_state_free(clip_enc_state * state) { + delete state; +} + bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { const clip_image_f32_batch & imgs = *params->imgs; int n_batch_cur = imgs.entries.size(); @@ -4839,6 +4915,44 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { // set input per projector switch (ctx->model.proj_type) { + case PROJECTOR_TYPE_VOXTRAL_RT: + { + clip_enc_state * st = params->enc_state; + if (!st) { + LOG_ERR("%s: this encoder needs a clip_enc_state\n", __func__); + return false; + } + const int n_window = hparams.attn_window_size; + const int n_new = (imgs.entries[0].nx() - 2) / 2; + GGML_ASSERT(n_new <= st->n_ctx - n_window); + + std::vector pos(n_new); + std::vector idx(n_new); + std::vector slot_pos = st->slot_pos; // including the frames written by this call + for (int i = 0; i < n_new; i++) { + pos[i] = st->n_past + i; + idx[i] = pos[i] % st->n_ctx; + slot_pos[idx[i]] = pos[i]; + } + set_input_i32("inp_pos", pos); + ggml_backend_tensor_set(get_inp_tensor("inp_kv_idx"), idx.data(), 0, idx.size() * sizeof(int64_t)); + + // each frame sees itself and the n_window - 1 frames before it + std::vector mask((size_t) st->n_ctx * n_new, -INFINITY); + for (int i = 0; i < n_new; i++) { + for (int s = 0; s < st->n_ctx; s++) { + const int32_t p = slot_pos[s]; + if (p >= 0 && p <= pos[i] && p > pos[i] - n_window) { + mask[(size_t) i * st->n_ctx + s] = 0.0f; + } + } + } + set_input_attn_mask(mask); + + std::vector tail = st->stem_tail; + tail.resize(hparams.n_embd, 0.0f); + set_input_f32("inp_stem_tail", tail); + } break; case PROJECTOR_TYPE_MUSE_GLIMMER: { const int grid_w = pos_w; // image_size_width / patch_size @@ -6099,6 +6213,20 @@ 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) { + clip_enc_state * st = params->enc_state; + ggml_tensor * tail = ggml_graph_get_tensor(gf, "out_stem_tail"); + GGML_ASSERT(tail != nullptr); + st->stem_tail.resize(ggml_nelements(tail)); + ggml_backend_tensor_get(tail, st->stem_tail.data(), 0, ggml_nbytes(tail)); + + const int n_new = (imgs.entries[0].nx() - 2) / 2; + for (int i = 0; i < n_new; i++) { + st->slot_pos[(st->n_past + i) % st->n_ctx] = st->n_past + i; + } + st->n_past += n_new; + } + if (params->state_out != nullptr) { auto & state_out = *params->state_out; size_t total = 0; @@ -6233,6 +6361,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_fc_w->ne[1]; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_VOXTRAL: + case PROJECTOR_TYPE_VOXTRAL_RT: case PROJECTOR_TYPE_MUSIC_FLAMINGO: return ctx->model.mm_2_w->ne[1]; case PROJECTOR_TYPE_MERALION: diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index 9e12702ce..19db47608 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -13,6 +13,7 @@ #define MTMD_INTERNAL_HEADER struct clip_ctx; +struct clip_enc_state; struct clip_image_size { int width; @@ -119,9 +120,18 @@ struct clip_encode_params { std::vector * out_audio = nullptr; // decoded PCM samples, F32 const std::vector * state_in = nullptr; // state from previous call, null or wrong size means cold start std::vector * state_out = nullptr; // state for the next call + + // streaming audio input, updated in place by clip_encode() + struct clip_enc_state * enc_state = nullptr; }; bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params); +// persistent state of a causal audio encoder, for streaming input (one per stream) +// it holds the sliding window of past keys/values, so each clip_encode() call only takes the new frames +// n_frames_max is the maximum number of encoder frames per call +struct clip_enc_state * clip_enc_state_init(struct clip_ctx * ctx, int32_t n_frames_max); +void clip_enc_state_free(struct clip_enc_state * state); + bool clip_is_llava(const struct clip_ctx * ctx); // note for contributor: this clip_is_(model) pattern is deprecated // do NOT add new functions like this diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 7d7f5ac1d..294448605 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -155,6 +155,13 @@ struct clip_graph_whisper_enc : clip_graph { ggml_cgraph * build() override; }; +// causal whisper-like encoder with a sliding window KV cache (voxtral realtime) +struct clip_graph_voxtral_rt : clip_graph { + clip_enc_state * state; // nullptr during warmup, the cache is then a graph input + clip_graph_voxtral_rt(clip_ctx * ctx, const clip_image_f32 & img, clip_enc_state * state) : clip_graph(ctx, img), state(state) {} + ggml_cgraph * build() override; +}; + struct clip_graph_deepseekocr : clip_graph { clip_graph_deepseekocr(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/models/voxtral-rt.cpp b/tools/mtmd/models/voxtral-rt.cpp new file mode 100644 index 000000000..4a01512d2 --- /dev/null +++ b/tools/mtmd/models/voxtral-rt.cpp @@ -0,0 +1,111 @@ +#include "models.h" + +// the input holds 2 context mel frames followed by the new frames, so the causal conv stem +// sees the same left context as on the full spectrogram +ggml_cgraph * clip_graph_voxtral_rt::build() { + const int n_frames_in = img.nx(); + const int n_new = (n_frames_in - 2) / 2; // encoder frames, the conv stem has stride 2 + const int n_window = hparams.attn_window_size; + const int stack = hparams.proj_stack_factor; + const int n_embd_kv = n_head_kv * d_head; + const int n_ctx = state ? state->n_ctx : n_window + n_new; + + GGML_ASSERT(n_window > 0 && stack > 0); + GGML_ASSERT(n_new > 0 && n_new % stack == 0); + GGML_ASSERT(n_new <= n_ctx - n_window); + + ggml_tensor * inp = build_inp_raw(1); // [n_frames_in, n_mel] + + ggml_tensor * stem_tail = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_embd); + ggml_set_name(stem_tail, "inp_stem_tail"); + ggml_set_input(stem_tail); + + ggml_tensor * pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_new); + ggml_set_name(pos, "inp_pos"); + ggml_set_input(pos); + + ggml_tensor * kv_idx = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new); + ggml_set_name(kv_idx, "inp_kv_idx"); + ggml_set_input(kv_idx); + + ggml_tensor * kq_mask = build_inp_attn_mask(n_ctx, n_new); + + // causal conv stem + ggml_tensor * cur = ggml_conv_1d(ctx0, model.conv1d_1_w, inp, 1, 0, 1); + cur = ggml_add(ctx0, cur, model.conv1d_1_b); + cur = ggml_gelu_erf(ctx0, cur); // [2*n_new, n_embd] + + // the last frame is the left context of the second conv on the next call + ggml_tensor * tail = ggml_view_2d(ctx0, cur, 1, n_embd, cur->nb[1], (2*n_new - 1) * ggml_element_size(cur)); + tail = ggml_cont(ctx0, tail); + ggml_set_name(tail, "out_stem_tail"); + ggml_set_output(tail); + ggml_build_forward_expand(gf, tail); + + cur = ggml_concat(ctx0, stem_tail, cur, 0); // [2*n_new + 1, n_embd] + cur = ggml_conv_1d(ctx0, model.conv1d_2_w, cur, 2, 0, 1); + cur = ggml_add(ctx0, cur, model.conv1d_2_b); + cur = ggml_gelu_erf(ctx0, cur); // [n_new, n_embd] + cur = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_embd, n_new] + cb(cur, "after_conv1d", -1); + + for (int il = 0; il < n_layer; il++) { + const auto & layer = model.layers[il]; + + ggml_tensor * x = build_norm(cur, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il); + + ggml_tensor * q = ggml_add(ctx0, build_mm(layer.q_w, x), layer.q_b); + ggml_tensor * k = build_mm(layer.k_w, x); + ggml_tensor * v = ggml_add(ctx0, build_mm(layer.v_w, x), layer.v_b); + + q = ggml_reshape_3d(ctx0, q, d_head, n_head, n_new); + k = ggml_reshape_3d(ctx0, k, d_head, n_head_kv, n_new); + + // interleaved rope on absolute frame positions + q = ggml_rope_ext(ctx0, q, pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + k = ggml_rope_ext(ctx0, k, pos, nullptr, d_head, GGML_ROPE_TYPE_NORMAL, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + cb(q, "q", il); + cb(k, "k", il); + cb(v, "v", il); + + ggml_tensor * k_cache; + ggml_tensor * v_cache; + if (state) { + k_cache = state->k[il]; + v_cache = state->v[il]; + } else { + k_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_embd_kv, n_ctx); + v_cache = ggml_new_tensor_2d(ctx0, GGML_TYPE_F16, n_embd_kv, n_ctx); + ggml_set_input(k_cache); + ggml_set_input(v_cache); + } + + // write the new frames into the ring, the mask selects the window out of the whole ring + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, k_cache, ggml_reshape_2d(ctx0, k, n_embd_kv, n_new), kv_idx)); + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, v_cache, v, kv_idx)); + + ggml_tensor * k_all = ggml_reshape_3d(ctx0, k_cache, d_head, n_head_kv, n_ctx); + ggml_tensor * v_all = ggml_reshape_3d(ctx0, v_cache, d_head, n_head_kv, n_ctx); + + x = build_attn(layer.o_w, layer.o_b, q, k_all, v_all, kq_mask, kq_scale, il); + cur = ggml_add(ctx0, cur, x); + + x = build_norm(cur, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il); + x = build_ffn(x, layer.ff_up_w, nullptr, layer.ff_gate_w, nullptr, layer.ff_down_w, layer.ff_down_b, hparams.ffn_op, il); + cur = ggml_add(ctx0, cur, x); + cb(cur, "layer_out", il); + } + + cur = build_norm(cur, model.post_ln_w, nullptr, NORM_TYPE_RMS, eps, -1); + cb(cur, "after_transformer", -1); + + cur = build_stack(cur, stack, n_embd); + cur = build_mm(model.mm_1_w, cur); + cur = ggml_gelu_erf(ctx0, cur); + cur = build_mm(model.mm_2_w, cur); + cb(cur, "projected", -1); + + ggml_build_forward_expand(gf, cur); + + return gf; +} diff --git a/tools/mtmd/mtmd-audio.cpp b/tools/mtmd/mtmd-audio.cpp index bbeb57c80..558475458 100644 --- a/tools/mtmd/mtmd-audio.cpp +++ b/tools/mtmd/mtmd-audio.cpp @@ -366,6 +366,31 @@ static void log_mel_spectrogram_worker_thread(int ith, } } +void mtmd_audio_log_mel_frame(const mtmd_audio_cache & cache, const float * samples, int n_fft, int n_mel, float * out) { + const int n_fft_bins = n_fft / 2 + 1; + GGML_ASSERT((int) cache.hann_window.size() == n_fft); + GGML_ASSERT(cache.filters.n_mel == n_mel && cache.filters.n_fft == n_fft); + + std::vector fft_in(n_fft * 2, 0.0f); + std::vector fft_out(n_fft * 2 * 2 * 2); + for (int j = 0; j < n_fft; j++) { + fft_in[j] = cache.hann_window[j] * samples[j]; + } + fft(cache, fft_in.data(), n_fft, fft_out.data()); + + for (int j = 0; j < n_fft_bins; j++) { + fft_out[j] = fft_out[2 * j + 0] * fft_out[2 * j + 0] + fft_out[2 * j + 1] * fft_out[2 * j + 1]; + } + + for (int j = 0; j < n_mel; j++) { + double sum = 0.0; + for (int k = 0; k < n_fft_bins; k++) { + sum += fft_out[k] * cache.filters.data[(size_t) j * n_fft_bins + k]; + } + out[j] = log10(std::max(sum, 1e-10)); + } +} + // ref: https://github.com/openai/whisper/blob/main/whisper/audio.py#L110-L157 static bool log_mel_spectrogram( const float * samples, diff --git a/tools/mtmd/mtmd-audio.h b/tools/mtmd/mtmd-audio.h index 3b66c2758..26ace2143 100644 --- a/tools/mtmd/mtmd-audio.h +++ b/tools/mtmd/mtmd-audio.h @@ -50,6 +50,9 @@ struct mtmd_audio_cache { ); }; +// compute one log10-mel frame from n_fft windowed samples, cache must be initialized for n_fft +void mtmd_audio_log_mel_frame(const mtmd_audio_cache & cache, const float * samples, int n_fft, int n_mel, float * out); + struct mtmd_audio_preprocessor { const clip_hparams & hparams; diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index 99ae5373c..1d7d65ef8 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -257,6 +257,99 @@ static int generate_response(mtmd_cli_context & ctx, int n_predict) { return 0; } +// feed an audio file to a realtime model chunk by chunk, printing the transcript as it is produced +static int run_audio_stream(mtmd_cli_context & ctx, const std::string & fname, int n_predict) { + auto res = mtmd_helper_bitmap_init_from_file(ctx.ctx_vision.get(), fname.c_str(), false, ctx.init_opt); + if (!res.bitmap) { + return 1; // error is already printed by libmtmd + } + mtmd::bitmap bmp(res.bitmap); + if (!mtmd_bitmap_is_audio(bmp.ptr.get())) { + LOG_ERR("%s: %s is not an audio file\n", __func__, fname.c_str()); + return 1; + } + const float * samples = (const float *) bmp.data(); + const size_t n_samples = bmp.n_bytes() / sizeof(float); + const int rate = mtmd_get_audio_sample_rate(ctx.ctx_vision.get()); + + const auto sparams = mtmd_audio_stream_params_default(); + mtmd::audio_stream_ptr stream(mtmd_audio_stream_init(ctx.ctx_vision.get(), sparams)); + if (!stream) { + return 1; + } + + const size_t n_chunk = (size_t) rate * sparams.step_ms / 1000; + const int n_embd = llama_model_n_embd_inp(ctx.model); + + LOG("Transcribing %.1f s of audio, fed in chunks of %d ms\n\n", (double) n_samples / rate, sparams.step_ms); + + size_t n_fed = 0; + int n_gen = 0; + llama_token tok = LLAMA_TOKEN_NULL; + const int64_t t_start = ggml_time_us(); + int64_t t_enc = 0; // time spent in the stream (mel + encoder) + int64_t t_dec = 0; // time spent in the decoder + while (!g_is_interrupted) { + const int64_t t0 = ggml_time_us(); + const int32_t n = mtmd_audio_stream_step(stream.get(), tok); + t_enc += ggml_time_us() - t0; + if (n < 0) { + if (n == -2) { + LOG_ERR("%s: stream error\n", __func__); + return 1; + } + break; // end of stream + } + if (n == 0) { + if (n_fed < n_samples) { + const size_t cnt = std::min(n_chunk, n_samples - n_fed); + mtmd_audio_stream_push(stream.get(), samples + n_fed, cnt); + n_fed += cnt; + } else if (mtmd_audio_stream_flush(stream.get(), true) != 0) { + break; // already ended + } + continue; + } + + const int64_t t1 = ggml_time_us(); + const float * rows = mtmd_audio_stream_get_output_embd(stream.get()); + ctx.batch.clear(); + for (int32_t i = 0; i < n; i++) { + const llama_pos pos[GGML_MROPE_SECTIONS] = { ctx.n_past, ctx.n_past, ctx.n_past, 0 }; + ctx.n_past++; + ctx.batch.add_embd({ rows + (size_t) i * n_embd, 1, (size_t) n_embd }, pos, 0, i == n - 1); + } + if (llama_process(ctx.lctx, LLAMA_PROCESS_TYPE_DECODE, ctx.batch.get())) { + LOG_ERR("%s: failed to decode\n", __func__); + return 1; + } + + tok = common_sampler_sample(ctx.smpl, ctx.lctx, -1); + common_sampler_accept(ctx.smpl, tok, true); + t_dec += ggml_time_us() - t1; + if (llama_vocab_is_eog(ctx.vocab, tok)) { + break; + } + // the model emits pad and word-boundary tokens between words + if (!llama_vocab_is_control(ctx.vocab, tok)) { + LOG("%s", common_token_to_piece(ctx.lctx, tok).c_str()); + fflush(stdout); + } + if (++n_gen >= n_predict) { + break; + } + } + + const double t_total = (ggml_time_us() - t_start) / 1e6; + LOG("\n\n%s: %d tokens, %.2f s audio in %.2f s (%.2fx realtime)\n", __func__, + n_gen, (double) n_samples / rate, t_total, (double) n_samples / rate / t_total); + LOG("%s: encoder %.2f s, decoder %.2f s, %.1f ms per 80 ms step\n", __func__, + t_enc / 1e6, t_dec / 1e6, (t_enc + t_dec) / 1e3 / std::max(n_gen, 1)); + llama_perf_context_print(ctx.lctx); + + return 0; +} + static std::string chat_add_and_format(mtmd_cli_context & ctx, common_chat_msg & new_msg) { LOG_DBG("chat_add_and_format: new_msg.role='%s', new_msg.content='%s'\n", new_msg.role.c_str(), new_msg.content.c_str()); @@ -472,6 +565,14 @@ int main(int argc, char ** argv) { if (g_is_interrupted) return 130; + if (mtmd_support_audio_stream(ctx.ctx_vision.get())) { + if (params.image.size() != 1) { + LOG_ERR("This model transcribes audio in realtime, pass exactly one --audio file\n"); + return 1; + } + return run_audio_stream(ctx, params.image[0], n_predict); + } + auto eval_system_prompt_if_present = [&] { if (params.system_prompt.empty()) { return 0; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 343d29571..26d5ee410 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -7,6 +7,7 @@ #include "debug/mtmd-debug.h" #include "llama.h" +#include "../src/llama-ext.h" // fix problem with std::min and std::max #if defined(_WIN32) @@ -533,6 +534,7 @@ struct mtmd_context { std::string media_marker; const int n_embd_text = -1; // -1 means llm context not provided, skip checking this const llama_vocab * vocab = nullptr; // can be nullptr if text_model is not provided + const llama_model * text_model = nullptr; // for the token embeddings of mtmd_audio_stream mtmd_pos_type pos_type; // these are not token, but strings used to mark the beginning and end of image/audio embeddings @@ -578,6 +580,7 @@ struct mtmd_context { media_marker (ctx_params.media_marker), n_embd_text (text_model ? llama_model_n_embd_inp(text_model) : -1), vocab (text_model ? llama_model_get_vocab(text_model) : nullptr), + text_model (text_model), batch_max_tokens(ctx_params.batch_max_tokens) { if (ctx_params.image_marker != nullptr) { @@ -1033,6 +1036,10 @@ struct mtmd_context { aud_beg = ""; audio_preproc = std::make_unique(ctx_a); } break; + case PROJECTOR_TYPE_VOXTRAL_RT: + { + // audio and text share positions, only usable via mtmd_audio_stream + } break; case PROJECTOR_TYPE_ULTRAVOX: case PROJECTOR_TYPE_GLMA: case PROJECTOR_TYPE_MERALION: @@ -1093,8 +1100,9 @@ struct mtmd_context { } // initialize audio preprocessor - GGML_ASSERT(audio_preproc != nullptr); - audio_preproc->initialize(); + if (audio_preproc) { + audio_preproc->initialize(); + } } // get clip ctx based on chunk type @@ -1690,6 +1698,10 @@ struct mtmd_tokenizer { samples = (const float *)buf.data(); n_samples = buf.size() / sizeof(float); } + if (!ctx->audio_preproc) { + LOG_ERR("%s: this model takes audio input only via mtmd_audio_stream\n", __func__); + return 2; + } bool ok = ctx->audio_preproc->preprocess(samples, n_samples, mel_spec_chunks); if (!ok) { LOG_ERR("Unable to preprocess audio\n"); @@ -2919,7 +2931,7 @@ void mtmd_debug_preprocess_image(mtmd_context * ctx, const std::vector } void mtmd_debug_preprocess_audio(mtmd_context * ctx, const std::vector & samples) { - if (!ctx->ctx_a) { + if (!ctx->ctx_a || !ctx->audio_preproc) { LOG_ERR("%s: model does not support audio input\n", __func__); return; } @@ -2982,3 +2994,304 @@ mtmd_memory_usage mtmd_get_memory_usage(const char * mmproj_fname, return {{}, -1, false}; } } + +// +// realtime audio input +// + +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; + int64_t buf_pos = 0; + int64_t n_samples = 0; // total samples pushed, including padding + 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 + bool pending_flush = false; + bool ended = false; + + ~mtmd_audio_stream() { + clip_enc_state_free(enc_state); + } +}; + +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; +} + +mtmd_audio_stream_params mtmd_audio_stream_params_default(void) { + mtmd_audio_stream_params params; + params.step_ms = 320; + 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(); + 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) { + 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(); +} + +void mtmd_audio_stream_free(mtmd_audio_stream * stream) { + delete stream; +} + +int32_t mtmd_audio_stream_push(mtmd_audio_stream * s, const float * samples, size_t n_samples) { + if (s->ended) { + return 1; + } + s->buf.insert(s->buf.end(), samples, samples + n_samples); + s->n_samples += n_samples; + return 0; +} + +int32_t mtmd_audio_stream_flush(mtmd_audio_stream * s, bool end) { + if (s->ended) { + return 1; + } + // 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); + s->pending_flush = true; + s->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) { + const int half = s->n_fft / 2; + std::vector window(s->n_fft); + std::vector frame(s->n_mel); + + while (true) { + const int64_t start = s->n_frames * s->hop - half; + if (start + s->n_fft > s->n_samples && !(force && start + half < s->n_samples)) { + break; + } + for (int j = 0; j < s->n_fft; j++) { + const int64_t idx = start + j - s->buf_pos; + window[j] = idx < (int64_t) s->buf.size() ? s->buf[idx] : 0.0f; + } + mtmd_audio_log_mel_frame(s->cache, window.data(), s->n_fft, s->n_mel, frame.data()); + for (int m = 0; m < s->n_mel; m++) { + const float v = std::max(frame[m], s->log_mel_max - 8.0f); + s->mel.push_back((v + 4.0f) / 4.0f); + } + s->n_frames++; + } + + // drop the samples no frame will need again + const int64_t keep_from = s->n_frames * s->hop - half; + 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; + } +} + +static bool mtmd_audio_stream_encode(mtmd_audio_stream * 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; + if (!force) { + if (n_avail < s->n_step_mel) { + break; + } + n = std::min(n, s->n_step_mel); + } + n = std::min(n, 2 * s->n_frames_max); + if (n <= 0) { + break; + } + + // 2 context frames, then the new frames, stored mel-major + const int n_in = 2 + n; + std::vector data((size_t) s->n_mel * n_in); + for (int m = 0; m < s->n_mel; m++) { + data[(size_t) m * n_in + 0] = s->mel_ctx[m]; + data[(size_t) m * n_in + 1] = s->mel_ctx[s->n_mel + m]; + for (int t = 0; t < n; t++) { + data[(size_t) m * n_in + 2 + t] = s->mel[(size_t) t * s->n_mel + m]; + } + } + + clip_image_f32 img; + img.set_size({n_in, s->n_mel}, false, /* is_audio */ true); + img.cpy_buf(data); + 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 / s->frames_per_tok)); + 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_ctx.assign(s->mel.begin() + (size_t) (n - 2) * s->n_mel, s->mel.begin() + (size_t) n * s->n_mel); + s->mel.erase(s->mel.begin(), s->mel.begin() + (size_t) n * s->n_mel); + } + return true; +} + +// out rows = audio embeddings + token embeddings +static bool mtmd_audio_stream_add_rows(mtmd_audio_stream * 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__); + return false; + } + const float * audio = s->embd.data() + s->embd_off * s->n_embd; + for (size_t i = 0; i < n * s->n_embd; i++) { + s->out.push_back(audio[i] + s->tok_embd[i]); + } + s->embd_off += n; + s->n_rows += n; + 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; + + mtmd_audio_stream_compute_mel(s, force); + if (!mtmd_audio_stream_encode(s, 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; + } + + const int64_t n_avail = (int64_t) (s->embd.size() / s->n_embd) - s->embd_off; + s->out.clear(); + + if (s->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(); + if (n_avail < n_prefix) { + return s->ended ? -1 : 0; + } + if (!mtmd_audio_stream_add_rows(s, s->prefix.data(), s->prefix.size())) { + return -2; + } + return (int32_t) n_prefix; + } + + if (n_avail == 0) { + return s->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)) { + return -2; + } + return 1; +} + +const float * mtmd_audio_stream_get_output_embd(const mtmd_audio_stream * s) { + return s->out.data(); +} diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index bcd76a9d7..df24c5656 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -363,6 +363,51 @@ struct mtmd_caps { }; MTMD_API struct mtmd_caps mtmd_get_cap_from_file(const char * mmproj_fname); +///////////////////////////////////////// +// EXPERIMENTAL API for realtime audio input, subjected to breaking changes +// +// a stream turns incoming PCM into decoder-ready embedding rows, one row per decoder position +// it keeps the encoder state between calls, so audio can be fed as it arrives +// how text tokens are mixed into the rows is a model detail hidden inside the stream +// the text model given to mtmd_init_from_file() is required + +// whether the current model supports mtmd_audio_stream +MTMD_API bool mtmd_support_audio_stream(const mtmd_context * ctx); + +struct mtmd_audio_stream; +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) +}; +MTMD_API struct mtmd_audio_stream_params mtmd_audio_stream_params_default(void); + +// returns nullptr on failure +MTMD_API mtmd_audio_stream * mtmd_audio_stream_init(mtmd_context * ctx, struct mtmd_audio_stream_params params); +MTMD_API void mtmd_audio_stream_free(mtmd_audio_stream * stream); + +// samples are PCM F32 at mtmd_get_audio_sample_rate() +// returns 0 on success, 1 if the stream has ended +MTMD_API int32_t mtmd_audio_stream_push(mtmd_audio_stream * stream, const float * samples, size_t n_samples); + +// pad the stream so that all pending rows can be produced +// end = false: the stream stays open (use it on a pause, to get a partial transcript) +// end = true: no more push() is accepted +// returns 0 on success, 1 if the stream has ended +MTMD_API int32_t mtmd_audio_stream_flush(mtmd_audio_stream * stream, bool end); + +// produce the next decoder rows +// last_token is the token sampled from the previous row, LLAMA_TOKEN_NULL before the first sample +// return value: +// > 0 number of rows ready, read them with mtmd_audio_stream_get_output_embd() +// 0 more audio is needed +// -1 end of stream, nothing more will be produced +// -2 error +MTMD_API int32_t mtmd_audio_stream_step(mtmd_audio_stream * stream, llama_token last_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); + ///////////////////////////////////////// // EXPERIMENTAL API for audio generation, subjected to breaking changes @@ -493,6 +538,11 @@ struct mtmd_batch_deleter { }; using batch_ptr = std::unique_ptr; +struct mtmd_audio_stream_deleter { + void operator()(mtmd_audio_stream * val) { mtmd_audio_stream_free(val); } +}; +using audio_stream_ptr = std::unique_ptr; + struct bitmap { bitmap_ptr ptr; bitmap() : ptr(nullptr) {}