]> git.djapps.eu Git - pkg/ggml/sources/llama.cpp/commitdiff
mtmd: support MiMo-V2.5 audio input (RVQ-based model) (#26190)
authorXuan-Son Nguyen <redacted>
Mon, 27 Jul 2026 21:17:09 +0000 (23:17 +0200)
committerGitHub <redacted>
Mon, 27 Jul 2026 21:17:09 +0000 (23:17 +0200)
* gguf converter for mimo audio

* fix conv

* cpp impl

* nits

* nits 2

15 files changed:
conversion/mimo.py
gguf-py/gguf/constants.py
gguf-py/gguf/gguf_writer.py
gguf-py/gguf/tensor_mapping.py
src/llama-quant.cpp
tools/mtmd/CMakeLists.txt
tools/mtmd/clip-graph.h
tools/mtmd/clip-impl.h
tools/mtmd/clip-model.h
tools/mtmd/clip.cpp
tools/mtmd/models/mimo-audio.cpp [new file with mode: 0644]
tools/mtmd/models/models.h
tools/mtmd/mtmd-audio.cpp
tools/mtmd/mtmd-audio.h
tools/mtmd/mtmd.cpp

index 11ec2867940a21c4633459cd4ec0c32ddeea9d11..ca2ed28ad39175e0167a8de6926b7bb924f9c6d9 100644 (file)
@@ -1,8 +1,9 @@
 from __future__ import annotations
 
+import json
 import re
 
-from typing import Callable, TYPE_CHECKING
+from typing import Any, Callable, Iterable, TYPE_CHECKING
 
 import torch
 
@@ -229,7 +230,13 @@ class MimoV2Model(TextModel):
 
 
 @ModelBase.register("MiMoV2ForCausalLM")
-class MiMoV2VisionModel(MmprojModel):
+class MiMoV2VisionAudioModel(MmprojModel):
+    has_audio_encoder = True
+
+    _audio_tok_hparams: dict[str, Any] | None = None
+    _rvq_codebook_sizes: list[int] | None = None
+    _code_embd: dict[int, Tensor] | None = None
+
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
         assert self.hparams_vision is not None
@@ -253,10 +260,22 @@ class MiMoV2VisionModel(MmprojModel):
         self.visual_token_window_size = int(hp.get("visual_token_window_size", -1))
         self.use_sink = bool(hp.get("use_sink", False))
 
+    def get_audio_config(self) -> dict[str, Any] | None:
+        if self._audio_tok_hparams is None:
+            path = self.dir_model / "audio_tokenizer" / "config.json"
+            with open(path, "r", encoding="utf-8") as f:
+                cfg = json.load(f)
+            # aliases so MmprojModel.find_aparam() / n_block_keys can resolve them
+            cfg["hidden_size"] = cfg["d_model"]
+            cfg["intermediate_size"] = cfg["encoder_ffn_dim"]
+            cfg["num_attention_heads"] = cfg["encoder_attention_heads"]
+            self._audio_tok_hparams = cfg
+        return self._audio_tok_hparams
+
     def set_gguf_parameters(self):
         super().set_gguf_parameters()
 
-        self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MIMOVL)
+        self.gguf_writer.add_clip_vision_projector_type(gguf.VisionProjectorType.MIMOVL)
         self.gguf_writer.add_vision_use_silu(True)
         self.gguf_writer.add_vision_head_count_kv(self.num_kv_heads)
         self.gguf_writer.add_vision_spatial_merge_size(self.spatial_merge_size)
@@ -266,19 +285,45 @@ class MiMoV2VisionModel(MmprojModel):
         self.gguf_writer.add_vision_min_pixels(int(self.preprocessor_config["min_pixels"]))
         self.gguf_writer.add_vision_max_pixels(int(self.preprocessor_config["max_pixels"]))
 
+        assert self.hparams_audio is not None
+        self.gguf_writer.add_clip_audio_projector_type(gguf.VisionProjectorType.MIMO_AUDIO)
+        self.gguf_writer.add_audio_num_mel_bins(self.hparams_audio["n_mels"])
+        self.gguf_writer.add_audio_attention_layernorm_eps(self.hparams_audio.get("layer_norm_eps", 1e-5))
+
+        assert self._rvq_codebook_sizes is not None
+        self.gguf_writer.add_audio_rvq_num_quantizers(len(self._rvq_codebook_sizes))
+        self.gguf_writer.add_audio_rvq_codebook_size(self._rvq_codebook_sizes)
+
+        n_layer = self.hparams_audio["encoder_layers"]
+        swa_per_block = self.hparams_audio.get("swa_per_block", 1)
+        if self.hparams_audio.get("hybrid_attention") and swa_per_block > 1:
+            wa_pattern = [0 if i % swa_per_block < swa_per_block - 1 else -1 for i in range(n_layer)]
+        else:
+            wa_pattern = [-1] * n_layer
+        self.gguf_writer.add_audio_wa_pattern_mode(wa_pattern)
+        self.gguf_writer.add_audio_window_size(int(self.hparams_audio["encoder_attn_window_size"][0]))
+
+        audio_cfg = self.global_config["audio_config"]
+        self.gguf_writer.add_audio_local_block_count(int(audio_cfg["input_local_layers"]))
+        self.gguf_writer.add_audio_local_group_size(int(audio_cfg["group_size"]))
+
     def tensor_force_quant(self, name, new_name, bid, n_dims):
-        # Sinks must be F32: any sink-style softmax/mask add in ggml requires
-        # F32, and we fold sinks into a host-built F32 mask at encode time.
-        if new_name.endswith(".attn_sinks"):
+        # for audio encoder: keep codebook in F32
+        if new_name in (
+            gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK] + ".weight",
+            gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.A_MM_CODE_EMBD] + ".weight",
+        ):
+            return gguf.GGMLQuantizationType.F32
+        if ("encoder.conv" in name or "encoder.down_sample_layer" in name) and name.endswith(".weight"):
             return gguf.GGMLQuantizationType.F32
         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, _ = item
-        if not name.startswith("visual."):
-            return None
-        return super().filter_tensors(item)
+        if name.startswith("visual.") or name.startswith("speech_embeddings.") or name.startswith("audio_encoder."):
+            return super().filter_tensors(item)
+        return None
 
     def modify_tensors(self, data_torch, name, bid):
         # Conv3D patch embed: split along the temporal axis (kt=2) into two Conv2D
@@ -292,4 +337,64 @@ class MiMoV2VisionModel(MmprojModel):
             yield (embd_name + ".weight.1", data_torch[:, :, 1, ...])
             return
 
+        if m := re.match(r"^speech_embeddings\.(\d+)\.weight$", name):
+            if self._code_embd is None:
+                self._code_embd = {}
+            self._code_embd[int(m.group(1))] = data_torch
+
+            n_channels = int(self.global_config["audio_config"]["audio_channels"])
+            if len(self._code_embd) < n_channels:
+                return
+            merged = torch.stack([self._code_embd.pop(i) for i in range(n_channels)], dim=0)
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MM_CODE_EMBD), merged)
+            return
+
+        if "conv1.bias" in name or "conv2.bias" in name:
+            # transpose conv1/conv2 bias so it broadcasts against [n_frames, C_out, 1]
+            data_torch = data_torch.unsqueeze(-1)
+
+        if name == "audio_encoder.projection.mlp.0.weight":
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 1), data_torch)
+            return
+        if name == "audio_encoder.projection.mlp.2.weight":
+            yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_MMPROJ, 2), data_torch)
+            return
+
         yield from super().modify_tensors(data_torch, name, bid)
+
+    def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]:
+        # note: audio encoder is in its own subdir "audio_tokenizer"
+        from safetensors.torch import load_file
+
+        tok_dir = self.dir_model / "audio_tokenizer"
+        state_dict = load_file(tok_dir / "model.safetensors")
+
+        codebook_re = re.compile(r"^encoder\.quantizer\.vq\.layers\.(\d+)\._codebook\.embed$")
+        codebooks: dict[int, Tensor] = {}
+
+        # EMA/training-only RVQ buffers - not needed for inference (nearest-codebook
+        # lookup only reads "_codebook.embed")
+        skip_suffixes = (
+            "_codebook.cluster_size",
+            "_codebook.embed_avg",
+            "_codebook.inited",
+        )
+        for name, tensor in state_dict.items():
+            if name.endswith(skip_suffixes):
+                continue
+            if m := codebook_re.match(name):
+                codebooks[int(m.group(1))] = tensor
+                continue
+            yield name, tensor
+
+        # gather codebooks and merge into 3D tensor, similar to MoE MLP tensors
+        n_q = len(codebooks)
+        ordered = [codebooks[i] for i in range(n_q)]
+        self._rvq_codebook_sizes = [int(cb.shape[0]) for cb in ordered]
+        max_bins = max(self._rvq_codebook_sizes)
+        dim = ordered[0].shape[1]
+        merged = ordered[0].new_zeros(n_q, max_bins, dim)
+        for i, cb in enumerate(ordered):
+            merged[i, : cb.shape[0], :] = cb
+
+        yield (self.format_tensor_name(gguf.MODEL_TENSOR.A_ENC_RVQ_CODEBOOK), merged)
index 78e3c29a0d38815f6930aded4a8ead06d1030715..f9264425f6e1a7dc3519f359127dedf8740a5a51 100644 (file)
@@ -376,6 +376,12 @@ class Keys:
         CONV_KERNEL_SIZE    = "clip.audio.conv_kernel_size"
         MAX_POS_EMB         = "clip.audio.max_pos_emb"
         FEATURE_LAYERS      = "clip.audio.feature_layer" # Granite Speech Plus
+        RVQ_NUM_QUANTIZERS  = "clip.audio.rvq.num_quantizers"
+        RVQ_CODEBOOK_SIZE   = "clip.audio.rvq.codebook_size"
+        WA_PATTERN_MODE     = "clip.audio.wa_pattern_mode" # per-layer -1 (full) / 0 (windowed)
+        WINDOW_SIZE         = "clip.audio.window_size"
+        LOCAL_BLOCK_COUNT   = "clip.audio.local_block_count" # mimo-v2.5: input_local_transformer layer count
+        LOCAL_GROUP_SIZE    = "clip.audio.local_group_size"  # mimo-v2.5: input_local_transformer grouping size
 
         class Attention:
             HEAD_COUNT      = "clip.audio.attention.head_count"
@@ -945,6 +951,9 @@ class MODEL_TENSOR(IntEnum):
     A_ENC_FFN_SCALE_1     = auto() # gemma3n
     A_ENC_FFN_GATE_1      = auto() # lfm2, gemma3n
     A_ENC_FFN_DOWN_1      = auto() # lfm2, gemma3n
+    A_ENC_DOWNSAMPLE_CONV = auto() # mimo-audio-tokenizer: post-transformer downsample conv
+    A_ENC_DOWNSAMPLE_NORM = auto() # mimo-audio-tokenizer: post-transformer downsample norm
+    A_ENC_RVQ_CODEBOOK    = auto() # mimo-audio-tokenizer: residual vector quantizer codebook, per quantizer index
     A_MMPROJ              = auto()
     A_MMPROJ_FC           = auto()
     A_MM_NORM_PRE         = auto()
@@ -953,6 +962,17 @@ class MODEL_TENSOR(IntEnum):
     A_MM_HARD_EMB_NORM    = auto() # gemma3n
     A_MM_SOFT_EMB_NORM    = auto() # gemma3n
     A_MM_INP_PROJ         = auto() # gemma3n
+    A_MM_CODE_EMBD        = auto() # mimo: text-side RVQ code embedding table ("text codebook"), merged 3D [n_channels, vocab, dim]
+    A_MM_LOCAL_ATTN_Q     = auto() # mimo: input_local_transformer (LLM-side connector)
+    A_MM_LOCAL_ATTN_K     = auto()
+    A_MM_LOCAL_ATTN_V     = auto()
+    A_MM_LOCAL_ATTN_OUT   = auto()
+    A_MM_LOCAL_FFN_GATE   = auto()
+    A_MM_LOCAL_FFN_UP     = auto()
+    A_MM_LOCAL_FFN_DOWN   = auto()
+    A_MM_LOCAL_LN1        = auto()
+    A_MM_LOCAL_LN2        = auto()
+    A_MM_LOCAL_NORM       = auto() # final norm after all input_local_transformer layers
     A_PER_DIM_K_SCALE     = auto() # gemma4
     A_PER_DIM_SCALE       = auto() # gemma4
     # nextn/mtp
@@ -1532,6 +1552,9 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.A_ENC_FFN_UP_1:            "a.blk.{bid}.ffn_up_1",
     MODEL_TENSOR.A_ENC_FFN_GATE_1:          "a.blk.{bid}.ffn_gate_1",
     MODEL_TENSOR.A_ENC_FFN_DOWN_1:          "a.blk.{bid}.ffn_down_1",
+    MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV:     "a.downsample.conv",
+    MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM:     "a.downsample.norm",
+    MODEL_TENSOR.A_ENC_RVQ_CODEBOOK:        "a.rvq.codebook",
     MODEL_TENSOR.A_MMPROJ:                  "mm.a.mlp.{bid}",
     MODEL_TENSOR.A_MMPROJ_FC:               "mm.a.fc",
     MODEL_TENSOR.A_MM_NORM_PRE:             "mm.a.norm_pre",
@@ -1540,6 +1563,17 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
     MODEL_TENSOR.A_MM_SOFT_EMB_NORM:        "mm.a.soft_emb_norm",         # gemma3n
     MODEL_TENSOR.A_MM_EMBEDDING:            "mm.a.embedding",             # gemma3n
     MODEL_TENSOR.A_MM_HARD_EMB_NORM:        "mm.a.hard_emb_norm",         # gemma3n
+    MODEL_TENSOR.A_MM_CODE_EMBD:            "mm.a.code_embd",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_Q:         "mm.a.local_blk.{bid}.attn_q",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_K:         "mm.a.local_blk.{bid}.attn_k",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_V:         "mm.a.local_blk.{bid}.attn_v",
+    MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT:       "mm.a.local_blk.{bid}.attn_out",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_GATE:       "mm.a.local_blk.{bid}.ffn_gate",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_UP:         "mm.a.local_blk.{bid}.ffn_up",
+    MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN:       "mm.a.local_blk.{bid}.ffn_down",
+    MODEL_TENSOR.A_MM_LOCAL_LN1:            "mm.a.local_blk.{bid}.ln1",
+    MODEL_TENSOR.A_MM_LOCAL_LN2:            "mm.a.local_blk.{bid}.ln2",
+    MODEL_TENSOR.A_MM_LOCAL_NORM:           "mm.a.local_norm",
     MODEL_TENSOR.A_PER_DIM_K_SCALE:         "a.blk.{bid}.per_dim_k_scale", # gemma4
     MODEL_TENSOR.A_PER_DIM_SCALE:           "a.blk.{bid}.per_dim_scale",   # gemma4
     # lfm2 audio
@@ -1741,10 +1775,24 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.A_ENC_FFN_UP_1,
         MODEL_TENSOR.A_ENC_FFN_GATE_1,
         MODEL_TENSOR.A_ENC_FFN_DOWN_1,
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV,
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM,
+        MODEL_TENSOR.A_ENC_RVQ_CODEBOOK,
         MODEL_TENSOR.A_MMPROJ,
         MODEL_TENSOR.A_MMPROJ_FC,
         MODEL_TENSOR.A_MM_NORM_PRE,
         MODEL_TENSOR.A_MM_NORM_MID,
+        MODEL_TENSOR.A_MM_CODE_EMBD,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_Q,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_K,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_V,
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_GATE,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_UP,
+        MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN,
+        MODEL_TENSOR.A_MM_LOCAL_LN1,
+        MODEL_TENSOR.A_MM_LOCAL_LN2,
+        MODEL_TENSOR.A_MM_LOCAL_NORM,
         MODEL_TENSOR.A_ENC_NORM_CONV,
         MODEL_TENSOR.A_ENC_LINEAR_POS,
         MODEL_TENSOR.A_ENC_POS_BIAS_U,
@@ -4804,6 +4852,7 @@ class VisionProjectorType:
     MINICPMV4_6    = "minicpmv4_6"
     GRANITE_SPEECH = "granite_speech"  # audio
     MIMOVL         = "mimovl"
+    MIMO_AUDIO     = "mimo_audio"
     GRANITE4_VISION = "granite4_vision"
 
 
index ecf1f17ee8b3b18ee9deb557b68957d6e120d9c2..bd8629aa119e9a157d9e0d38ebca81b6d328c27a 100644 (file)
@@ -1350,6 +1350,24 @@ class GGUFWriter:
     def add_audio_num_mel_bins(self, value: int) -> None:
         self.add_uint32(Keys.ClipAudio.NUM_MEL_BINS, value)
 
+    def add_audio_rvq_num_quantizers(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.RVQ_NUM_QUANTIZERS, value)
+
+    def add_audio_rvq_codebook_size(self, values: Sequence[int]) -> None:
+        self.add_array(Keys.ClipAudio.RVQ_CODEBOOK_SIZE, values)
+
+    def add_audio_wa_pattern_mode(self, modes: Sequence[int]) -> None:
+        self.add_array(Keys.ClipAudio.WA_PATTERN_MODE, modes)
+
+    def add_audio_window_size(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.WINDOW_SIZE, value)
+
+    def add_audio_local_block_count(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.LOCAL_BLOCK_COUNT, value)
+
+    def add_audio_local_group_size(self, value: int) -> None:
+        self.add_uint32(Keys.ClipAudio.LOCAL_GROUP_SIZE, value)
+
     def add_audio_stack_factor(self, value: int) -> None:
         self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value)
 
index 62d7a827e35c0e5eaba99ec2e90a02bfea16f276..8299ac25b432d905239c7219c01b57d60c58825c 100644 (file)
@@ -2095,6 +2095,7 @@ class TensorNameMap:
             "conformer.pre_encode.conv.{bid}", # lfm2
             "model.audio_tower.subsample_conv_projection.conv_{bid}.conv", # gemma3n
             "conformer.subsample_conv_projection.layer{bid}.conv", # gemma4
+            "encoder.conv{bid}", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_CONV1D_NORM: (
@@ -2119,6 +2120,7 @@ class TensorNameMap:
         MODEL_TENSOR.A_POST_NORM: (
             "audio_tower.layer_norm", # ultravox
             "audio_tower.ln_post", # qwen2omni
+            "encoder.layer_norm", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_Q: (
@@ -2127,6 +2129,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.attn.q_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.q_proj", # gemma4
             "encoder.layers.{bid}.attn.to_q", # granite_speech
+            "encoder.layers.{bid}.self_attn.q_proj", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K: (
@@ -2135,6 +2138,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.attn.k_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.k_proj", # gemma4
             "encoder.layers.{bid}.attn.to_k", # granite_speech (split from to_kv)
+            "encoder.layers.{bid}.self_attn.k_proj", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_V: (
@@ -2143,6 +2147,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.attn.v_proj", # gemma3n
             "conformer.layers.{bid}.self_attn.v_proj", # gemma4
             "encoder.layers.{bid}.attn.to_v", # granite_speech (split from to_kv)
+            "encoder.layers.{bid}.self_attn.v_proj", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_ATTN_K_REL: (
@@ -2171,6 +2176,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.norm_self_att", # lfm2
             "conformer.layers.{bid}.attention.pre_attn_norm", # gemma3n
             "encoder.layers.{bid}.attn.pre_norm", # granite_speech
+            "encoder.layers.{bid}.self_attn_layer_norm", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_OUTPUT: (
@@ -2179,6 +2185,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.attention.post", # gemma3n
             "conformer.layers.{bid}.self_attn.post", # gemma4
             "encoder.layers.{bid}.attn.to_out", # granite_speech
+            "encoder.layers.{bid}.self_attn.out_proj", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_OUTPUT_NORM: (
@@ -2186,6 +2193,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.norm_out", # lfm2
             "conformer.layers.{bid}.attention.post_norm", # gemma3n
             "encoder.layers.{bid}.post_norm", # granite_speech
+            "encoder.layers.{bid}.final_layer_norm", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_FFN_NORM: (
@@ -2210,6 +2218,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_1", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_1", # gemma4
             "encoder.layers.{bid}.ff1.up_proj", # granite_speech
+            "encoder.layers.{bid}.fc1", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_FFN_GATE: (),
@@ -2220,6 +2229,7 @@ class TensorNameMap:
             "conformer.layers.{bid}.ffw_layer_start.ffw_layer_2", # gemma3n
             "conformer.layers.{bid}.feed_forward1.ffw_layer_2", # gemma4
             "encoder.layers.{bid}.ff1.down_proj", # granite_speech
+            "encoder.layers.{bid}.fc2", # mimo-audio-tokenizer
         ),
 
         MODEL_TENSOR.A_ENC_FFN_UP_1: (
@@ -2243,6 +2253,19 @@ class TensorNameMap:
             "encoder.layers.{bid}.ff2.pre_norm", # granite_speech
         ),
 
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_CONV: (
+            "encoder.down_sample_layer.0", # mimo-audio-tokenizer
+        ),
+
+        MODEL_TENSOR.A_ENC_DOWNSAMPLE_NORM: (
+            "encoder.down_sample_norm", # mimo-audio-tokenizer
+        ),
+
+        # note: the raw per-quantizer "encoder.quantizer.vq.layers.{i}._codebook.embed"
+        # tensors are merged (padded + stacked, like MoE experts) into this single 3D
+        # tensor in conversion code, so no raw-name mapping is registered here.
+        MODEL_TENSOR.A_ENC_RVQ_CODEBOOK: (),
+
         MODEL_TENSOR.A_ENC_FFN_POST_NORM_1: (
             "conformer.layers.{bid}.ffw_layer_end.post_layer_norm", # gemma3n
             "conformer.layers.{bid}.feed_forward2.post_layer_norm", # gemma4
@@ -2294,6 +2317,42 @@ class TensorNameMap:
             "audio.multi_modal_projector.ln_mid", # ultravox
         ),
 
+        # note: the raw per-channel "speech_embeddings.{i}" tensors are merged
+        # (stacked, like MoE experts) into this single 3D tensor in conversion
+        # code, so no raw-name mapping is registered here.
+        MODEL_TENSOR.A_MM_CODE_EMBD: (),
+
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_Q: (
+            "audio_encoder.input_local_transformer.layers.{bid}.self_attn.q_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_K: (
+            "audio_encoder.input_local_transformer.layers.{bid}.self_attn.k_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_V: (
+            "audio_encoder.input_local_transformer.layers.{bid}.self_attn.v_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_ATTN_OUT: (
+            "audio_encoder.input_local_transformer.layers.{bid}.self_attn.o_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_FFN_GATE: (
+            "audio_encoder.input_local_transformer.layers.{bid}.mlp.gate_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_FFN_UP: (
+            "audio_encoder.input_local_transformer.layers.{bid}.mlp.up_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_FFN_DOWN: (
+            "audio_encoder.input_local_transformer.layers.{bid}.mlp.down_proj", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_LN1: (
+            "audio_encoder.input_local_transformer.layers.{bid}.input_layernorm", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_LN2: (
+            "audio_encoder.input_local_transformer.layers.{bid}.post_attention_layernorm", # mimo-v2.5
+        ),
+        MODEL_TENSOR.A_MM_LOCAL_NORM: (
+            "audio_encoder.input_local_transformer.norm", # mimo-v2.5
+        ),
+
         MODEL_TENSOR.A_ENC_CONV_DW: (
             "conformer.layers.{bid}.conv.depthwise_conv", # lfm2
             "conformer.layers.{bid}.lconv1d.depthwise_conv1d", # gemma3n
index 7c0bac07d09614ddc0d941599e8a5864af3c5ace..92ebc11b99f3ba5ad99adb1a42c646540cd26b72 100644 (file)
@@ -359,6 +359,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
     quantize &= name.find(".patch_embd")    == std::string::npos;
     quantize &= name.find(".patch_merger")  == std::string::npos;
 
+    // audio codebook
+    quantize &= name.find("a.rvq.codebook")  == std::string::npos;
+    quantize &= name.find("mm.a.code_embd")  == std::string::npos;
+
     return quantize;
 }
 
index fd7ddceb0bf070b96a300ecf456c8c75d1bc1831..18a8288ba048f01d18f18306fb6b0988549a11f4 100644 (file)
@@ -51,6 +51,7 @@ add_library(mtmd
             models/qwen3vl.cpp
             models/mimovl.cpp
             models/qwen3a.cpp
+            models/mimo-audio.cpp
             models/step3vl.cpp
             models/siglip.cpp
             models/whisper-enc.cpp
index a95de20a3122fafe5f26471baa6921b6df656e6c..29352abb4c0bc443757a9685e99c20355dda1636 100644 (file)
 
 struct build_vit_opts {
     ggml_tensor * attn_mask = nullptr;
+    // TODO @ngxson : merge attn_mask and attn_mask_layers into one call
+    std::vector<ggml_tensor *> attn_mask_layers; // one per layer
+
+    // hook at layer output embeddings
+    std::function<void(ggml_tensor * cur, int il)> callback_layer_out = nullptr;
+
+    // whether to skip the automatic post-layernorm (model.post_ln_w) applied at the end
+    bool skip_post_ln = false;
 };
 
 struct clip_graph {
index 42374311ce7b60d0ab724e894b233a0c00d822c9..09204113801fd8ebe22cbaedcaf1ec88a23dbaaa 100644 (file)
 #define KEY_A_PROJ_WINDOW_SIZE     "clip.audio.projector.window_size"
 #define KEY_A_PROJ_DOWNSAMPLE_RATE "clip.audio.projector.downsample_rate"
 #define KEY_A_PROJ_HEAD_COUNT      "clip.audio.projector.head_count"
+#define KEY_A_RVQ_NUM_QUANTIZERS   "clip.audio.rvq.num_quantizers"   // mimo-audio-tokenizer
+#define KEY_A_RVQ_CODEBOOK_SIZE    "clip.audio.rvq.codebook_size"    // mimo-audio-tokenizer: per-quantizer bin count
+#define KEY_A_WA_PATTERN_MODE      "clip.audio.wa_pattern_mode"      // mimo-audio-tokenizer, per-layer -1 (full) / 0 (windowed)
+#define KEY_A_ATTN_WINDOW_SIZE     "clip.audio.window_size"          // mimo-audio-tokenizer: sliding-window radius
+#define KEY_A_LOCAL_BLOCK_COUNT    "clip.audio.local_block_count"    // mimo-v2.5: input_local_transformer layer count
+#define KEY_A_LOCAL_GROUP_SIZE     "clip.audio.local_group_size"     // mimo-v2.5: input_local_transformer grouping size
 
 //
 // tensor name constants
 #define TN_MM_NORM_PRE  "mm.a.norm_pre.%s"
 #define TN_MM_NORM_MID  "mm.a.norm_mid.%s"
 
+// mimo-audio-tokenizer
+#define TN_A_DOWNSAMPLE_CONV "a.downsample.conv.%s"
+#define TN_A_DOWNSAMPLE_NORM "a.downsample.norm.%s"
+#define TN_A_RVQ_CODEBOOK    "a.rvq.codebook.%s"
+// mimo-v2.5: text-side RVQ code embedding ("text codebook")
+#define TN_MM_A_CODE_EMBD    "mm.a.code_embd.%s"
+// mimo-v2.5: LLM-side connector (input_local_transformer)
+#define TN_MM_A_LOCAL_ATTN_Q   "mm.a.local_blk.%d.attn_q.%s"
+#define TN_MM_A_LOCAL_ATTN_K   "mm.a.local_blk.%d.attn_k.%s"
+#define TN_MM_A_LOCAL_ATTN_V   "mm.a.local_blk.%d.attn_v.%s"
+#define TN_MM_A_LOCAL_ATTN_OUT "mm.a.local_blk.%d.attn_out.%s"
+#define TN_MM_A_LOCAL_FFN_GATE "mm.a.local_blk.%d.ffn_gate.%s"
+#define TN_MM_A_LOCAL_FFN_UP   "mm.a.local_blk.%d.ffn_up.%s"
+#define TN_MM_A_LOCAL_FFN_DOWN "mm.a.local_blk.%d.ffn_down.%s"
+#define TN_MM_A_LOCAL_LN1      "mm.a.local_blk.%d.ln1.%s"
+#define TN_MM_A_LOCAL_LN2      "mm.a.local_blk.%d.ln2.%s"
+#define TN_MM_A_LOCAL_NORM     "mm.a.local_norm.%s"
+
 // cogvlm
 #define TN_MM_POST_FC_NORM "mm.post_fc_norm.%s"
 #define TN_MM_H_TO_4H      "mm.up.%s"
@@ -374,6 +398,7 @@ enum projector_type {
     PROJECTOR_TYPE_MIMOVL,
     PROJECTOR_TYPE_MINIMAX_M3,
     PROJECTOR_TYPE_GRANITE4_VISION,
+    PROJECTOR_TYPE_MIMO_AUDIO,
     PROJECTOR_TYPE_UNKNOWN,
 };
 
@@ -429,6 +454,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
     { PROJECTOR_TYPE_MIMOVL,            "mimovl"},
     { PROJECTOR_TYPE_MINIMAX_M3,        "minimax_m3"},
     { PROJECTOR_TYPE_GRANITE4_VISION,   "granite4_vision"},
+    { PROJECTOR_TYPE_MIMO_AUDIO,        "mimo_audio"},
 };
 
 static projector_type clip_projector_type_from_string(const std::string & str) {
index 850957d7de1c41d1d9eaa6b338e989a4dded537a..8dc87549766e0877c815cd7094a54a27148ad07e 100644 (file)
@@ -124,6 +124,14 @@ struct clip_hparams {
     int32_t audio_window_len  = -1;
     int32_t audio_hop_len     = -1;
 
+    // mimo-audio-tokenizer: residual vector quantizer
+    int32_t rvq_num_quantizers = 0;
+    std::vector<int32_t> rvq_codebook_size; // per-quantizer bin count (ragged, e.g. 1024/1024/256/128x17)
+
+    // mimo-v2.5: LLM-side connector (input_local_transformer)
+    int32_t audio_local_n_layer = 0;
+    int32_t audio_local_group_size = 0;
+
     // legacy
     bool has_llava_projector = false;
     int minicpmv_version = 0;
@@ -537,6 +545,20 @@ struct clip_model {
     ggml_tensor * mm_norm_pre_b = nullptr;
     ggml_tensor * mm_norm_mid_w = nullptr;
 
+    // mimo-audio-tokenizer: post-transformer downsample + RVQ codebook
+    ggml_tensor * downsample_conv_w = nullptr; // no bias
+    ggml_tensor * downsample_norm_w = nullptr;
+    ggml_tensor * downsample_norm_b = nullptr;
+    ggml_tensor * rvq_codebook = nullptr; // merged 3D [n_q, max_bins, dim]
+
+    // mimo-v2.5: text-side RVQ code embedding ("text codebook")
+    ggml_tensor * mm_a_code_embd = nullptr; // merged 3D [n_channels, vocab, dim]
+
+    // mimo-v2.5: LLM-side connector (input_local_transformer, separate from the
+    // audio_tokenizer's own encoder `layers`)
+    std::vector<clip_layer> mm_a_local_layers;
+    ggml_tensor * mm_a_local_norm_w = nullptr;
+
     // qwen3a
     ggml_tensor * conv2d_1_w = nullptr;
     ggml_tensor * conv2d_1_b = nullptr;
index e0e2107a0be342e2679d5f399b2b4433ccc4d47f..04614b93bd277dfed8309e4e7db759a4ccd7ef6e 100644 (file)
@@ -340,6 +340,11 @@ ggml_tensor * clip_graph::build_vit(
         auto & layer = model.layers[il];
         ggml_tensor * cur = inpL; // inpL = residual, cur = hidden_states
 
+        ggml_tensor * attn_mask = opts.attn_mask;
+        if (opts.attn_mask_layers.size() > (size_t) il) {
+            attn_mask = opts.attn_mask_layers[il];
+        }
+
         // layernorm1
         cur = build_norm(cur, layer.ln_1_w, layer.ln_1_b, norm_t, eps, il);
         cb(cur, "layer_inp_normed", il);
@@ -452,7 +457,7 @@ ggml_tensor * clip_graph::build_vit(
 
             // build_attn returns a flat 2D [n_embd, n_pos*B]
             cur = build_attn(layer.o_w, layer.o_b,
-                Qcur, Kcur, Vcur, opts.attn_mask, kq_scale, il);
+                Qcur, Kcur, Vcur, attn_mask, kq_scale, il);
             cb(cur, "attn_out", il);
         }
 
@@ -471,6 +476,10 @@ ggml_tensor * clip_graph::build_vit(
 
         inpL = cur; // inpL = residual, cur = hidden_states
 
+        if (opts.callback_layer_out) {
+            opts.callback_layer_out(cur, il);
+        }
+
         cb(cur, "ffn_inp", il);
 
         // layernorm2 (pre-ffn norm)
@@ -519,7 +528,7 @@ ggml_tensor * clip_graph::build_vit(
     }
 
     // post-layernorm
-    if (model.post_ln_w) {
+    if (model.post_ln_w && !opts.skip_post_ln) {
         inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, norm_t, eps, -1);
     }
 
@@ -1012,6 +1021,10 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
             {
                 builder = std::make_unique<clip_graph_qwen3a>(ctx, img);
             } break;
+        case PROJECTOR_TYPE_MIMO_AUDIO:
+            {
+                builder = std::make_unique<clip_graph_mimo_audio>(ctx, img);
+            } break;
         case PROJECTOR_TYPE_YOUTUVL:
             {
                 builder = std::make_unique<clip_graph_youtuvl>(ctx, img);
@@ -1575,6 +1588,45 @@ struct clip_model_loader {
                         hparams.audio_window_len   = 400;
                         hparams.audio_hop_len      = 160;
                     } break;
+                case PROJECTOR_TYPE_MIMO_AUDIO:
+                    {
+                        get_u32(KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers, false);
+                        get_arr_int(KEY_A_RVQ_CODEBOOK_SIZE, hparams.rvq_codebook_size, false);
+                        if (hparams.rvq_num_quantizers <= 0) {
+                            throw std::runtime_error(string_format("%s: mimo_audio: missing %s\n", __func__, KEY_A_RVQ_NUM_QUANTIZERS));
+                        }
+                        if ((int) hparams.rvq_codebook_size.size() != hparams.rvq_num_quantizers) {
+                            throw std::runtime_error(string_format(
+                                "%s: mimo_audio: %s length (%zu) must equal %s (%d)\n", __func__,
+                                KEY_A_RVQ_CODEBOOK_SIZE, hparams.rvq_codebook_size.size(),
+                                KEY_A_RVQ_NUM_QUANTIZERS, hparams.rvq_num_quantizers));
+                        }
+                        hparams.ffn_op = FFN_GELU_ERF; // PyTorch F.gelu default (approximate="none")
+                        hparams.rope_theta = 10000.0f;
+
+                        // audio preprocessing params (mel spectrogram)
+                        hparams.audio_sample_rate = 24000;
+                        hparams.audio_n_fft       = 960;
+                        hparams.audio_window_len  = 960;
+                        hparams.audio_hop_len     = 240;
+
+                        get_u32(KEY_A_ATTN_WINDOW_SIZE, hparams.attn_window_size);
+                        std::vector<int> wa_pattern;
+                        get_arr_int(KEY_A_WA_PATTERN_MODE, wa_pattern, true);
+                        if ((int) wa_pattern.size() != hparams.n_layer) {
+                            throw std::runtime_error(string_format(
+                                "%s: mimo_audio: %s length (%zu) must equal n_layer (%d)\n", __func__,
+                                KEY_A_WA_PATTERN_MODE, wa_pattern.size(), hparams.n_layer));
+                        }
+                        hparams.wa_pattern_mode.assign(wa_pattern.begin(), wa_pattern.end());
+
+                        get_u32(KEY_A_LOCAL_BLOCK_COUNT, hparams.audio_local_n_layer);
+                        get_u32(KEY_A_LOCAL_GROUP_SIZE, hparams.audio_local_group_size);
+                        if (hparams.audio_local_group_size <= 0) {
+                            throw std::runtime_error(string_format(
+                                "%s: mimo_audio: %s must be > 0\n", __func__, KEY_A_LOCAL_GROUP_SIZE));
+                        }
+                    } break;
                 case PROJECTOR_TYPE_PADDLEOCR:
                     {
                         hparams.n_merge = 2;
@@ -2444,6 +2496,54 @@ struct clip_model_loader {
                     model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"));
                     model.mm_2_b = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "bias"));
                 } break;
+            case PROJECTOR_TYPE_MIMO_AUDIO:
+                {
+                    model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight"));
+                    model.conv1d_1_b = get_tensor(string_format(TN_CONV1D, 1, "bias"));
+                    model.conv1d_2_w = get_tensor(string_format(TN_CONV1D, 2, "weight"));
+                    model.conv1d_2_b = get_tensor(string_format(TN_CONV1D, 2, "bias"));
+                    model.downsample_conv_w = get_tensor(string_format(TN_A_DOWNSAMPLE_CONV, "weight"));
+                    model.downsample_norm_w = get_tensor(string_format(TN_A_DOWNSAMPLE_NORM, "weight"));
+                    model.downsample_norm_b = get_tensor(string_format(TN_A_DOWNSAMPLE_NORM, "bias"));
+                    model.rvq_codebook   = get_tensor(string_format(TN_A_RVQ_CODEBOOK, "weight"), false);
+                    model.mm_a_code_embd = get_tensor(string_format(TN_MM_A_CODE_EMBD, "weight"), false);
+                    if (!model.rvq_codebook || !model.mm_a_code_embd) {
+                        throw std::runtime_error(string_format("%s: mimo_audio: missing %s or %s\n", __func__,
+                            TN_A_RVQ_CODEBOOK, TN_MM_A_CODE_EMBD));
+                    }
+                    // hparams.rvq_codebook_size comes from GGUF metadata and is independent of the
+                    // tensors' actual shapes - bound it so codebook/code_embd views built from it
+                    // (mimo-audio.cpp) can never read past either tensor's allocated bins/vocab.
+                    for (int32_t bins : hparams.rvq_codebook_size) {
+                        if (bins <= 0 || bins > model.rvq_codebook->ne[1] || bins > model.mm_a_code_embd->ne[1]) {
+                            throw std::runtime_error(string_format(
+                                "%s: mimo_audio: %s entry (%d) out of range for codebook/code_embd tensors\n",
+                                __func__, KEY_A_RVQ_CODEBOOK_SIZE, bins));
+                        }
+                    }
+
+                    // LLM-side connector: input_local_transformer + projection
+                    model.mm_a_local_layers.resize(hparams.audio_local_n_layer);
+                    for (int il = 0; il < hparams.audio_local_n_layer; il++) {
+                        auto & layer = model.mm_a_local_layers[il];
+                        layer.q_w    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_Q,   il, "weight"));
+                        layer.q_b    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_Q,   il, "bias"));
+                        layer.k_w    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_K,   il, "weight"));
+                        layer.k_b    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_K,   il, "bias"));
+                        layer.v_w    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_V,   il, "weight"));
+                        layer.v_b    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_V,   il, "bias"));
+                        layer.o_w    = get_tensor(string_format(TN_MM_A_LOCAL_ATTN_OUT, il, "weight"));
+                        layer.ff_gate_w = get_tensor(string_format(TN_MM_A_LOCAL_FFN_GATE, il, "weight"));
+                        layer.ff_up_w   = get_tensor(string_format(TN_MM_A_LOCAL_FFN_UP,   il, "weight"));
+                        layer.ff_down_w = get_tensor(string_format(TN_MM_A_LOCAL_FFN_DOWN, il, "weight"));
+                        layer.ln_1_w = get_tensor(string_format(TN_MM_A_LOCAL_LN1, il, "weight"));
+                        layer.ln_2_w = get_tensor(string_format(TN_MM_A_LOCAL_LN2, il, "weight"));
+                    }
+                    model.mm_a_local_norm_w = get_tensor(string_format(TN_MM_A_LOCAL_NORM, "weight"));
+
+                    model.mm_1_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 1, "weight"));
+                    model.mm_2_w = get_tensor(string_format(TN_MM_AUDIO_MLP, 2, "weight"));
+                } break;
             case PROJECTOR_TYPE_VOXTRAL:
                 {
                     model.conv1d_1_w = get_tensor(string_format(TN_CONV1D, 1, "weight"));
@@ -3549,6 +3649,15 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) {
             {
                 n_patches = img->nx();  // no downsampling: one token per raw waveform frame
             } break;
+        case PROJECTOR_TYPE_MIMO_AUDIO:
+            {
+                // conv1(s=1) + conv2(s=2) -> RVQ-encoder downsample conv(k=2,s=2)
+                int n = img->nx();
+                n = (n - 1) / 2 + 1;         // conv1 + conv2
+                n = (n - 2) / 2 + 1;         // downsample conv
+                const int group_size = params.audio_local_group_size;
+                n_patches = (n + group_size - 1) / group_size;
+            } break;
         case PROJECTOR_TYPE_GRANITE_SPEECH:
             {
                 const int ws = ctx->model.hparams.audio_proj_window_size;
@@ -4376,6 +4485,58 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32
                     set_input_f32("pos_emb", pos_emb);
                 }
             } break;
+        case PROJECTOR_TYPE_MIMO_AUDIO:
+            {
+                GGML_ASSERT(imgs.entries.size() == 1);
+                const int n_frames = imgs.entries.front().nx();
+                const int n_pos    = (n_frames - 1) / 2 + 1; // matches conv1(s=1)+conv2(s=2) output length
+
+                std::vector<int32_t> positions(n_pos);
+                for (int i = 0; i < n_pos; i++) {
+                    positions[i] = i;
+                }
+                set_input_i32("mimo_audio_positions", positions);
+
+                const int window = hparams.attn_window_size;
+                GGML_ASSERT(window > 0);
+
+                const float neg_inf = std::numeric_limits<float>::lowest();
+                std::vector<float> full_mask((size_t) n_pos * n_pos);
+                std::vector<float> window_mask((size_t) n_pos * n_pos);
+                for (int q = 0; q < n_pos; q++) {
+                    for (int k = 0; k < n_pos; k++) {
+                        const bool causal_ok = k <= q;
+                        full_mask[(size_t) q * n_pos + k]   = causal_ok ? 0.0f : neg_inf;
+                        window_mask[(size_t) q * n_pos + k] = (causal_ok && (q - k) <= window) ? 0.0f : neg_inf;
+                    }
+                }
+                set_input_f32("mimo_audio_full_mask", full_mask);
+                set_input_f32("mimo_audio_window_mask", window_mask);
+
+                // input_local_transformer: block-diagonal mask + in-group positions
+                {
+                    const int n_pos_ds   = (n_pos - 2) / 2 + 1; // matches downsample conv (k=2,s=2,p=0)
+                    const int group_size = hparams.audio_local_group_size;
+                    GGML_ASSERT(group_size > 0);
+                    const int n_groups = (n_pos_ds + group_size - 1) / group_size;
+                    const int n_padded = n_groups * group_size;
+
+                    std::vector<int32_t> local_positions(n_padded);
+                    for (int i = 0; i < n_padded; i++) {
+                        local_positions[i] = i % group_size;
+                    }
+                    set_input_i32("mimo_audio_local_positions", local_positions);
+
+                    std::vector<float> local_mask((size_t) n_padded * n_padded);
+                    for (int q = 0; q < n_padded; q++) {
+                        for (int k = 0; k < n_padded; k++) {
+                            const bool same_group = (q / group_size) == (k / group_size);
+                            local_mask[(size_t) q * n_padded + k] = same_group ? 0.0f : neg_inf;
+                        }
+                    }
+                    set_input_f32("mimo_audio_local_mask", local_mask);
+                }
+            } break;
         case PROJECTOR_TYPE_LFM2A:
             {
                 GGML_ASSERT(imgs.entries.size() == 1);
@@ -4678,6 +4839,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) {
             return ctx->model.qf_proj_blocks.size() * ctx->model.hparams.projection_dim;
         case PROJECTOR_TYPE_GLM4V:
             return ctx->model.mm_ffn_down_w->ne[1];
+        case PROJECTOR_TYPE_MIMO_AUDIO:
+            return ctx->model.mm_2_w->ne[1];
         default:
             GGML_ABORT("Unknown projector type");
     }
diff --git a/tools/mtmd/models/mimo-audio.cpp b/tools/mtmd/models/mimo-audio.cpp
new file mode 100644 (file)
index 0000000..481b36c
--- /dev/null
@@ -0,0 +1,218 @@
+#include "models.h"
+
+ggml_cgraph * clip_graph_mimo_audio::build() {
+    ggml_tensor * inp = build_inp_raw(1); // [n_frames, n_mel, 1]
+
+    ggml_tensor * cur = ggml_conv_1d_ph(ctx0, model.conv1d_1_w, inp, 1, 1);
+    cur = ggml_add(ctx0, cur, model.conv1d_1_b);
+    cur = ggml_gelu_erf(ctx0, cur);
+
+    cur = ggml_conv_1d_ph(ctx0, model.conv1d_2_w, cur, 2, 1);
+    cur = ggml_add(ctx0, cur, model.conv1d_2_b);
+    cur = ggml_gelu_erf(ctx0, cur);
+
+    ggml_tensor * inpL = ggml_cont(ctx0, ggml_transpose(ctx0, cur)); // [n_embd, n_pos]
+    const int64_t n_pos = inpL->ne[1];
+    cb(inpL, "after_conv1d", -1);
+
+    GGML_ASSERT((int) hparams.wa_pattern_mode.size() == n_layer);
+
+    ggml_tensor * inp_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos);
+    ggml_set_name(inp_pos, "mimo_audio_positions");
+    ggml_set_input(inp_pos);
+
+    ggml_tensor * full_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
+    ggml_set_name(full_mask, "mimo_audio_full_mask");
+    ggml_set_input(full_mask);
+
+    ggml_tensor * window_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, n_pos);
+    ggml_set_name(window_mask, "mimo_audio_window_mask");
+    ggml_set_input(window_mask);
+
+    build_vit_opts opts;
+    opts.attn_mask_layers.resize(n_layer);
+    for (int il = 0; il < n_layer; il++) {
+        opts.attn_mask_layers[il] = hparams.wa_pattern_mode[il] == -1 ? full_mask : window_mask;
+    }
+    // the skip connection below must be added before the post-transformer  norm,
+    // so build_vit must not apply that norm itself
+    opts.skip_post_ln = true;
+
+    // encoder_skip_layer_id=3 (1-indexed) -> capture output of layer index 2
+    const int skip_capture_il = 2;
+    GGML_ASSERT(n_layer > skip_capture_il);
+    ggml_tensor * skip_hidden = nullptr;
+    opts.callback_layer_out = [&](ggml_tensor * layer_cur, int il) {
+        if (il == skip_capture_il) {
+            skip_hidden = layer_cur;
+        }
+    };
+
+    auto add_pos = [&](ggml_tensor * x, const clip_layer &) {
+        return ggml_rope_ext(ctx0, x, inp_pos, nullptr, d_head,
+                             GGML_ROPE_TYPE_NEOX, 0, hparams.rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+    };
+
+    inpL = build_vit(inpL, n_pos, NORM_TYPE_NORMAL, hparams.ffn_op, nullptr, add_pos, opts);
+    inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_pos); // build_vit restores a (size-1) batch dim
+
+    GGML_ASSERT(skip_hidden != nullptr);
+    inpL = ggml_add(ctx0, inpL, skip_hidden);
+
+    inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, NORM_TYPE_NORMAL, eps, -1);
+    cb(inpL, "after_transformer", -1);
+
+    // downsample: strided conv (no bias) + gelu + layernorm
+    {
+        ggml_tensor * ds = ggml_cont(ctx0, ggml_transpose(ctx0, inpL)); // [n_pos, n_embd]
+        ds = ggml_conv_1d(ctx0, model.downsample_conv_w, ds, 2, 0, 1);
+        ds = ggml_gelu_erf(ctx0, ds);
+        ds = ggml_cont(ctx0, ggml_transpose(ctx0, ds)); // [n_embd, n_pos/2]
+        ds = build_norm(ds, model.downsample_norm_w, model.downsample_norm_b, NORM_TYPE_NORMAL, eps, -1);
+        inpL = ds;
+    }
+    cb(inpL, "after_downsample", -1);
+
+    // RVQ quantize: codebook ne=[dim, max_bins, n_q]
+    // quantize input vector to codes (type=I32)
+    std::vector<ggml_tensor *> codes;
+    {
+        GGML_ASSERT(model.rvq_codebook != nullptr);
+        const int64_t dim = model.rvq_codebook->ne[0];
+        GGML_ASSERT(dim == inpL->ne[0]);
+        GGML_ASSERT((int64_t) hparams.rvq_codebook_size.size() == model.rvq_codebook->ne[2]);
+
+        ggml_tensor * residual = inpL; // [dim, n_pos_ds]
+
+        for (size_t q = 0; q < hparams.rvq_codebook_size.size(); q++) {
+            const int64_t bins = hparams.rvq_codebook_size[q];
+            ggml_tensor * codebook_q = ggml_view_2d(ctx0, model.rvq_codebook, dim, bins,
+                model.rvq_codebook->nb[1], q * model.rvq_codebook->nb[2]);
+            codebook_q = ggml_cont(ctx0, codebook_q);
+
+            ggml_tensor * codebook_norm = ggml_sum_rows(ctx0, ggml_sqr(ctx0, codebook_q)); // [1, bins]
+            codebook_norm = ggml_cont(ctx0, ggml_transpose(ctx0, codebook_norm));          // [bins, 1]
+
+            ggml_tensor * dot    = ggml_mul_mat(ctx0, codebook_q, residual); // [bins, n_pos_ds]
+            ggml_tensor * scores = ggml_sub(ctx0, ggml_scale(ctx0, dot, 2.0f), codebook_norm);
+
+            ggml_tensor * idx = ggml_argmax(ctx0, scores); // [n_pos_ds]
+            codes.push_back(idx);
+
+            ggml_tensor * quant = ggml_get_rows(ctx0, codebook_q, idx); // [dim, n_pos_ds]
+            residual = ggml_sub(ctx0, residual, quant);
+            cb(idx, "rvq_code", (int) q);
+        }
+    }
+
+    // convert codes to LLM embeddings
+    ggml_tensor * code_embd_sum = nullptr;
+    {
+        GGML_ASSERT(model.mm_a_code_embd != nullptr);
+        const int64_t dim   = model.mm_a_code_embd->ne[0];
+        const int64_t vocab = model.mm_a_code_embd->ne[1];
+        GGML_ASSERT((int64_t) codes.size() == model.mm_a_code_embd->ne[2]);
+        GGML_ASSERT(dim == inpL->ne[0]);
+
+        for (size_t i = 0; i < codes.size(); i++) {
+            ggml_tensor * table_i = ggml_view_2d(ctx0, model.mm_a_code_embd, dim, vocab,
+                model.mm_a_code_embd->nb[1], i * model.mm_a_code_embd->nb[2]);
+            table_i = ggml_cont(ctx0, table_i);
+
+            ggml_tensor * embd_i = ggml_get_rows(ctx0, table_i, codes[i]); // [dim, n_pos_ds]
+            code_embd_sum = code_embd_sum ? ggml_add(ctx0, code_embd_sum, embd_i) : embd_i;
+        }
+        cb(code_embd_sum, "code_embd_sum", -1);
+    }
+
+    // input_local_transformer
+    // groups of `group_size` consecutive downsampled frames are processed together, attending only within their own group.
+    // Implemented as a block-diagonal mask + in-group-repeating positions
+    // (rather than a real batch dim) - same technique as the encoder's masks above, and as gemma4a's / deepseekocr2's chunked attention.
+
+    // note: hand-rolled here instead of build_vit() because this is a second, independent layer stack
+    // (own layer array/count, RMSNorm instead of LN, SiLU FFN, own RoPE theta)
+
+    ggml_tensor * projected;
+    {
+        const int group_size = hparams.audio_local_group_size;
+        GGML_ASSERT(group_size > 0);
+        const int64_t n_pos_ds = code_embd_sum->ne[1];
+        const int64_t n_groups = (n_pos_ds + group_size - 1) / group_size;
+        const int64_t n_padded = n_groups * group_size;
+
+        ggml_tensor * cur_local = code_embd_sum;
+        if (n_padded != n_pos_ds) {
+            cur_local = ggml_pad(ctx0, cur_local, 0, (int) (n_padded - n_pos_ds), 0, 0);
+        }
+
+        ggml_tensor * local_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_padded);
+        ggml_set_name(local_pos, "mimo_audio_local_positions");
+        ggml_set_input(local_pos);
+
+        ggml_tensor * local_mask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_padded, n_padded);
+        ggml_set_name(local_mask, "mimo_audio_local_mask");
+        ggml_set_input(local_mask);
+
+        const float local_rope_theta = 640000.0f; // audio_config.rope_theta (differs from the encoder's)
+        auto apply_local_rope = [&](ggml_tensor * x) {
+            return ggml_rope_ext(ctx0, x, local_pos, nullptr, d_head,
+                                 GGML_ROPE_TYPE_NEOX, 0, local_rope_theta, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f);
+        };
+
+        for (int il = 0; il < hparams.audio_local_n_layer; il++) {
+            auto & layer = model.mm_a_local_layers[il];
+
+            ggml_tensor * attn_in = build_norm(cur_local, layer.ln_1_w, nullptr, NORM_TYPE_RMS, eps, il);
+
+            ggml_tensor * Qcur = build_mm(layer.q_w, attn_in);
+            if (layer.q_b) {
+                Qcur = ggml_add(ctx0, Qcur, layer.q_b);
+            }
+            ggml_tensor * Kcur = build_mm(layer.k_w, attn_in);
+            if (layer.k_b) {
+                Kcur = ggml_add(ctx0, Kcur, layer.k_b);
+            }
+            ggml_tensor * Vcur = build_mm(layer.v_w, attn_in);
+            if (layer.v_b) {
+                Vcur = ggml_add(ctx0, Vcur, layer.v_b);
+            }
+
+            Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_padded);
+            Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head, n_padded);
+            Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head, n_padded);
+
+            Qcur = apply_local_rope(Qcur);
+            Kcur = apply_local_rope(Kcur);
+
+            ggml_tensor * attn_out = build_attn(layer.o_w, nullptr, Qcur, Kcur, Vcur, local_mask, kq_scale, il);
+            cur_local = ggml_add(ctx0, cur_local, attn_out);
+
+            ggml_tensor * ffn_in = build_norm(cur_local, layer.ln_2_w, nullptr, NORM_TYPE_RMS, eps, il);
+            ggml_tensor * ffn_out = build_ffn(ffn_in,
+                layer.ff_up_w, nullptr,
+                layer.ff_gate_w, nullptr,
+                layer.ff_down_w, nullptr,
+                FFN_SILU, il);
+            cur_local = ggml_add(ctx0, cur_local, ffn_out);
+        }
+
+        cur_local = build_norm(cur_local, model.mm_a_local_norm_w, nullptr, NORM_TYPE_RMS, eps, -1);
+        cb(cur_local, "after_local_transformer", -1);
+
+        // flatten each group of `group_size` frames into one (group_size*n_embd)-dim vector
+        // (matching AudioProjection's flattened input)
+        ggml_tensor * grouped = ggml_reshape_2d(ctx0, cur_local, n_embd * group_size, n_groups);
+
+        // AudioProjection: Linear (no bias) -> GELU -> Linear (no bias)
+        projected = build_ffn(grouped,
+            model.mm_1_w, nullptr,
+            nullptr, nullptr,
+            model.mm_2_w, nullptr,
+            FFN_GELU_ERF, -1);
+        cb(projected, "after_projection", -1);
+    }
+
+    ggml_build_forward_expand(gf, projected);
+    return gf;
+}
index 2d7555da41d2f4b4631e28b1e561c881c17a616f..caed438ec513c547fefd522e24a7b474fd9298ec 100644 (file)
@@ -210,6 +210,11 @@ struct clip_graph_qwen3a : clip_graph {
     ggml_cgraph * build() override;
 };
 
+struct clip_graph_mimo_audio : clip_graph {
+    clip_graph_mimo_audio(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {}
+    ggml_cgraph * build() override;
+};
+
 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;
index b72fd067a508250444afabb86d0cb8f593b7766b..ed68951c0151f236e1ea9e53d7bf7107ae5ab11f 100644 (file)
@@ -725,6 +725,72 @@ bool mtmd_audio_preprocessor_qwen3a::preprocess(const float *                 sa
     return true;
 }
 
+//
+// mtmd_audio_preprocessor_mimo_audio
+//
+// Matches torchaudio.transforms.MelSpectrogram(power=1.0, center=True) followed by
+// log(clip(spec, min=1e-7)): HTK mel scale, no Slaney area norm, magnitude (not power)
+// spectrogram, natural log, reflect-padded by n_fft/2 on each side.
+//
+
+void mtmd_audio_preprocessor_mimo_audio::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,
+        0.0f, hparams.audio_sample_rate / 2.0f,
+        /*slaney_area_norm=*/ false,
+        /*scale=*/ 1.0f,
+        /*use_htk=*/ true
+    );
+}
+
+bool mtmd_audio_preprocessor_mimo_audio::preprocess(const float *                 samples,
+                                                    size_t                        n_samples,
+                                                    std::vector<mtmd_audio_mel> & output) {
+    if (n_samples == 0) {
+        return false;
+    }
+
+    GGML_ASSERT(!cache.sin_vals.empty());
+    GGML_ASSERT(!cache.cos_vals.empty());
+    GGML_ASSERT(!cache.filters.data.empty());
+
+    const int pad = hparams.audio_n_fft / 2;
+
+    std::vector<float> padded(n_samples + 2 * pad, 0.0f);
+    for (int i = 0; i < pad; i++) {
+        int src = pad - i;
+        padded[i] = (src < (int)n_samples) ? samples[src] : 0.0f;
+    }
+    std::copy(samples, samples + n_samples, padded.begin() + pad);
+    for (int i = 0; i < pad; i++) {
+        int src = (int)n_samples - 2 - i;
+        padded[n_samples + pad + i] = (src >= 0) ? samples[src] : 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; // reflect padding already applied above
+    params.use_natural_log  = true;
+    params.use_magnitude    = true;
+    params.mel_floor        = 1e-7f;
+    params.norm_per_feature = false;
+
+    mtmd_audio_mel out;
+    bool ok = log_mel_spectrogram(padded.data(), (int)padded.size(), 4, params, cache, out);
+    if (!ok) {
+        return false;
+    }
+
+    output.push_back(std::move(out));
+    return true;
+}
+
 //
 // mtmd_audio_preprocessor_conformer
 //
index ad96bd847cfc3413ed99a0355a3a3020526c00ed..d8ec72b9d54e2dfcd2cc1add8efb4a86f00e2ee8 100644 (file)
@@ -111,6 +111,15 @@ struct mtmd_audio_preprocessor_qwen3a : mtmd_audio_preprocessor {
     mtmd_audio_cache cache;
 };
 
+struct mtmd_audio_preprocessor_mimo_audio : mtmd_audio_preprocessor {
+    mtmd_audio_preprocessor_mimo_audio(const clip_ctx * ctx) : mtmd_audio_preprocessor(ctx) {}
+    void initialize() override;
+    bool preprocess(const float * samples, size_t n_samples, std::vector<mtmd_audio_mel> & output) override;
+
+  private:
+    mtmd_audio_cache cache;
+};
+
 //
 // streaming ISTFT - converts spectrogram frames back to audio one frame at a time
 //
index bb49b211efb35b84f95e1193809e9fdf69c3de30..6e61cf3e520b1238e765d3a15db4112e1cb89d20 100644 (file)
@@ -730,6 +730,12 @@ struct mtmd_context {
                     aud_end = "<audio|>";
                     audio_preproc = std::make_unique<mtmd_audio_preprocessor_gemma4ua>(ctx_a);
                 } break;
+            case PROJECTOR_TYPE_MIMO_AUDIO:
+                {
+                    aud_beg = "<|mimo_audio_start|>";
+                    aud_end = "<|mimo_audio_end|>";
+                    audio_preproc = std::make_unique<mtmd_audio_preprocessor_mimo_audio>(ctx_a);
+                } break;
             default:
                 throw std::runtime_error(string_format("%s: unexpected audio projector type %d\n", __func__, proj));
         }