feat(gguf): serve image input on minimax-m3 ggufs from a minimax_m3 mmproj

python/freetoken/mm/processors/minimax_m3.py

@@ -11,10 +11,7 @@
 from freetoken.mm import mm_pad_value
 from freetoken.mm.config import MultimodalConfig
 from freetoken.mm.processor import MMProcessor, PromptReplacement, content_hash
-
-# the config names only the image token; the reference processor reads these two from the tokenizer's fixed added tokens
-IMAGE_START_ID = 200029  # ]<]start of image[>[
-IMAGE_END_ID = 200030  # ]<]end of image[>[
+from freetoken.models.minimax_m3.config import IMAGE_END_ID, IMAGE_START_ID
 
 
 def _compression(vc: Any, key: str) -> int:

python/freetoken/models/minimax_m3/__init__.py

@@ -1,5 +1,5 @@
 from .config import VisionConfig, parse_config, parse_vision_config
-from .gguf import iter_gguf_weights, parse_gguf_config, to_gguf
+from .gguf import gguf_mm_config, iter_gguf_weights, iter_mmproj_weights, parse_gguf_config, to_gguf
 from .model import MiniMaxM3ForCausalLM, MiniMaxM3ForConditionalGeneration
 from .vision import MiniMaxM3VisionModel
 from .weight import (
@@ -18,7 +18,9 @@
     "parse_vision_config",
     "iter_vision_weights",
     "iter_weights",
+    "gguf_mm_config",
     "iter_gguf_weights",
+    "iter_mmproj_weights",
     "parse_gguf_config",
     "to_gguf",
 ]

python/freetoken/models/minimax_m3/config.py

@@ -37,6 +37,11 @@
 
 from .args import load_args
 
+# the config names only the image token; the reference processor reads these two from the tokenizer's fixed added tokens
+IMAGE_START_ID = 200029  # ]<]start of image[>[
+IMAGE_END_ID = 200030  # ]<]end of image[>[
+
+
 def _text_config(hf_config: Any) -> Any:
     return getattr(hf_config, "text_config", None) or hf_config
 
@@ -233,4 +238,4 @@ def parse_config(hf_config: Any) -> ModelConfig:
     )
 
 
-__all__ = ["VisionConfig", "parse_config", "parse_vision_config"]
+__all__ = ["IMAGE_END_ID", "IMAGE_START_ID", "VisionConfig", "parse_config", "parse_vision_config"]

python/freetoken/models/minimax_m3/gguf.py

@@ -1,22 +1,30 @@
-"""MiniMax-M3 text from a llama.cpp ``minimax-m3`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
+"""MiniMax-M3 from a llama.cpp ``minimax-m3`` GGUF: config from the GGUF metadata, tensors by name; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``).
 
 The llama.cpp converter (conversion/minimax.py) stacks the w1 / w2 / w3 experts, stores the indexer projections as F32, and adds 1 to every RMSNorm weight so llama.cpp can run plain RMSNorm; the loader subtracts the 1 for this family's (1 + w) norms.
 Both sides run NeoX rope, so q/k carry no rope permute.
+
+A ``minimax_m3`` mmproj beside the text file adds the vision tower under ``vision_tower.``. The converter splits the Conv3d patch embedding into its two temporal halves (``v.patch_embd.weight`` and ``.weight.1``) and reorders the rows of each ViT head's q / k from HF's 3-axis rope layout [t h w | t h w | pad] to [t t | h h | w w | pad]; both are undone at load, every other tower tensor is stored in HF layout.
 """
 
 from __future__ import annotations
 
 import re
 from dataclasses import replace
 from types import SimpleNamespace
-from typing import TYPE_CHECKING, Iterator
+from typing import TYPE_CHECKING, Any, Iterator
 
+import numpy as np
 import torch
 
-from freetoken.models.gguf.reader import GgufTensorInfo, gguf_tensor_index
+from freetoken.models.gguf.mmproj import Mmproj, find_mmproj, log_mmproj, mmproj_tensor
+from freetoken.models.gguf.reader import (
+    GgufTensorInfo,
+    gguf_tensor_index,
+    load_gguf_metadata,
+)
 from freetoken.models.gguf.tensors import float_tensor, fused, matrix
 
-from .config import parse_config
+from .config import IMAGE_END_ID, IMAGE_START_ID, parse_config
 
 if TYPE_CHECKING:
     from freetoken.models.config import ModelConfig
@@ -97,6 +105,10 @@ def kv(key: str):
     return SimpleNamespace(
         model_type="minimax_m3",
         architectures=list(shim.architectures),
+        # filled by a minimax_m3 mmproj (gguf_mm_config); a text-only process nulls vision_config on its copy
+        vision_config=getattr(shim, "vision_config", None),
+        projector_hidden_size=getattr(shim, "projector_hidden_size", None),
+        image_token_id=getattr(shim, "image_token_id", None),
         hidden_act="swigluoai",
         swiglu_alpha=_SWIGLU_ALPHA,
         swiglu_limit=_SWIGLU_LIMIT,
@@ -139,19 +151,159 @@ def parse_gguf_config(shim: GgufConfigShim) -> ModelConfig:
     return replace(parse_config(_hf_config(shim)), expert_quant="gguf")
 
 
+# ----------------------------------------------------------------------------------------------
+# Vision tower (a minimax_m3 mmproj)
+# ----------------------------------------------------------------------------------------------
+
+_PROJECTOR = "minimax_m3"
+_IMAGE_TOKEN = "]<]image[>["
+_WRAPPERS = {"]<]start of image[>[": IMAGE_START_ID, "]<]end of image[>[": IMAGE_END_ID}
+# no mmproj key carries the rope theta; llama.cpp hardcodes the HF value too (clip.cpp)
+_ROPE_THETA = 10000.0
+# MiniMaxM3VLImageProcessor's own size: the HF preprocessor_config.json stores a stale [672, 672] that the class replaces with this
+_IMAGE_SIZE = {"shortest_edge": 3136, "longest_edge": 451584}
+_PATCH = ("v.patch_embd.weight", "v.patch_embd.weight.1")
+_TOWER = "vision_tower"
+
+_TOP = {
+    "v.pre_ln.weight": "vision_model.pre_layrnorm.weight",
+    "v.pre_ln.bias": "vision_model.pre_layrnorm.bias",
+    "mm.1.weight": "multi_modal_projector.linear_1.weight",
+    "mm.1.bias": "multi_modal_projector.linear_1.bias",
+    "mm.2.weight": "multi_modal_projector.linear_2.weight",
+    "mm.2.bias": "multi_modal_projector.linear_2.bias",
+    "mm.merger.fc1.weight": "patch_merge_mlp.linear_1.weight",
+    "mm.merger.fc1.bias": "patch_merge_mlp.linear_1.bias",
+    "mm.merger.fc2.weight": "patch_merge_mlp.linear_2.weight",
+    "mm.merger.fc2.bias": "patch_merge_mlp.linear_2.bias",
+}
+_BLOCK = {
+    "attn_out": "self_attn.out_proj",
+    "ln1": "layer_norm1",
+    "ln2": "layer_norm2",
+    "ffn_up": "mlp.fc1",
+    "ffn_down": "mlp.fc2",
+}
+_QKV = ("attn_q", "attn_k", "attn_v")
+_BLOCK_RE = re.compile(r"^v\.blk\.(\d+)\.(\w+)\.(weight|bias)$")
+
+
+def _f32_value(v: float) -> float:
+    """The shortest decimal of a float32 metadata value: 1e-05, as the HF config writes it, not 9.99999975e-06."""
+    return float(str(np.float32(v)))
+
+
+def gguf_mm_config(text_md: dict[str, Any], mmproj: Mmproj, vocab_size: int, *, model_path: str = "") -> tuple[dict[str, Any], dict[str, Any]]:
+    """The HF config.json keys and the preprocessor_config.json of the tower in ``mmproj``, for the text GGUF whose metadata is ``text_md``."""
+    if mmproj.projector != _PROJECTOR:
+        raise ValueError(f"mmproj {mmproj.path} is a {mmproj.projector!r} tower; MiniMax-M3 needs {_PROJECTOR!r}")
+    md, index = load_gguf_metadata(mmproj.path), gguf_tensor_index(mmproj.path)
+
+    def kv(key: str):
+        return md[f"clip.vision.{key}"]
+
+    text_hidden = int(text_md[f"{_ARCH}.embedding_length"])
+    if int(kv("projection_dim")) != text_hidden:
+        raise ValueError(f"mmproj {mmproj.path} projects to {int(kv('projection_dim'))} but the text model is {text_hidden} wide")
+    tokens = text_md["tokenizer.ggml.tokens"]
+    if _IMAGE_TOKEN not in tokens:
+        raise ValueError(f"the text GGUF tokenizer has no {_IMAGE_TOKEN} token")
+    # the processor hardcodes the wrapper ids; a tokenizer with them elsewhere would fail per request instead of here
+    for token, want in _WRAPPERS.items():
+        got = tokens.index(token) if token in tokens else None
+        if got != want:
+            raise ValueError(f"the text GGUF tokenizer has {token} at id {got}, not the family's {want}")
+    # exactly the tensor names llama.cpp loads, both patch-embedding halves included; other spellings (mm.merge.*) fail there too
+    unexpected = [n for n in index if n not in _PATCH and n not in _TOP and not ((m := _BLOCK_RE.match(n)) and m[2] in (*_BLOCK, *_QKV))]
+    missing = [n for n in (*_PATCH, *_TOP) if n not in index]
+    if unexpected or missing:
+        raise ValueError(f"mmproj {mmproj.path} is not in the minimax_m3 tensor layout (unexpected {unexpected[:3]}, missing {missing[:3]})")
+    merge = int(kv("spatial_merge_size"))
+    vision_config = {
+        # llama.cpp's minimax_m3 graph runs erf-GELU whatever clip.use_gelu says
+        "hidden_act": "gelu",
+        "hidden_size": int(kv("embedding_length")),
+        "num_hidden_layers": int(kv("block_count")),
+        "num_attention_heads": int(kv("attention.head_count")),
+        "intermediate_size": int(kv("feed_forward_length")),
+        # ggml order [kw, kh, in, out]
+        "num_channels": int(index[_PATCH[0]].ne[2]),
+        "patch_size": int(kv("patch_size")),
+        "temporal_patch_size": 2,
+        "spatial_merge_size": merge,
+        "layer_norm_eps": _f32_value(kv("attention.layer_norm_epsilon")),
+        "rope_theta": _ROPE_THETA,
+    }
+    hf_fields = {
+        "vision_config": vision_config,
+        # ggml order [in, out]
+        "projector_hidden_size": int(index["mm.1.weight"].ne[1]),
+        "image_token_id": tokens.index(_IMAGE_TOKEN),
+    }
+    preprocessor = {
+        "image_processor_type": "MiniMaxM3VLImageProcessor",
+        "size": dict(_IMAGE_SIZE),
+        "patch_size": vision_config["patch_size"],
+        "temporal_patch_size": 2,
+        "merge_size": merge,
+        "image_mean": [float(v) for v in kv("image_mean")],
+        "image_std": [float(v) for v in kv("image_std")],
+    }
+    return hf_fields, preprocessor
+
+
+def _hf_qk_rows(head_dim: int) -> torch.Tensor:
+    """The mmproj row of each HF q / k row within one head (the inverse of conversion/minimax.py ``_permute_vit_qk``)."""
+    band = (2 * (head_dim // 2) // 3) // 2
+    stored = []
+    for axis in range(3):
+        stored += [*range(axis * band, (axis + 1) * band), *range((3 + axis) * band, (4 + axis) * band)]
+    stored += range(6 * band, head_dim)
+    return torch.argsort(torch.tensor(stored))
+
+
+def iter_mmproj_weights(mmproj: Mmproj) -> Iterator[tuple[str, torch.Tensor]]:
+    """The tower state dict under ``vision_tower.``, named as the HF checkpoint loads (q / k / v fused)."""
+    md, index = load_gguf_metadata(mmproj.path), gguf_tensor_index(mmproj.path)
+    heads = int(md["clip.vision.attention.head_count"])
+    rows = _hf_qk_rows(int(md["clip.vision.embedding_length"]) // heads)
+
+    def hf_qk(t: torch.Tensor) -> torch.Tensor:
+        return t.reshape(heads, rows.numel(), -1)[:, rows].reshape(t.shape)
+
+    yield f"{_TOWER}.vision_model.embeddings.patch_embedding.weight", torch.stack([mmproj_tensor(index[n]) for n in _PATCH], dim=2)
+    for name, info in index.items():
+        if name in _PATCH:
+            continue
+        if name in _TOP:
+            yield f"{_TOWER}.{_TOP[name]}", mmproj_tensor(info)
+            continue
+        # gguf_mm_config refused any other name
+        layer, part, kind = _BLOCK_RE.match(name).groups()
+        dst = f"{_TOWER}.vision_model.encoder.layers.{layer}"
+        if part in _BLOCK:
+            yield f"{dst}.{_BLOCK[part]}.{kind}", mmproj_tensor(info)
+        elif part == "attn_q":
+            q, k, v = (mmproj_tensor(index[f"v.blk.{layer}.{p}.{kind}"]) for p in _QKV)
+            yield f"{dst}.self_attn.qkv.{kind}", torch.cat([hf_qk(q), hf_qk(k), v])
+
+
 def iter_gguf_weights(
     model_path: str,
     device: torch.device,
     *,
     include_moe_experts: bool,
     include_non_moe: bool,
+    include_vision: bool = True,
 ) -> Iterator[tuple[str, torch.Tensor]]:
-    """The dense text state dict from the GGUF; the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
+    """The dense state dict from the GGUF (and the tower from its mmproj); the routed experts load as GGUF expert banks (``freetoken.models.gguf.experts``)."""
     if not include_non_moe:
         return
     from freetoken.utils import cached_load_hf_config
 
     config = parse_gguf_config(cached_load_hf_config(model_path))
+    vision = include_vision and config.is_multimodal
+    log_mmproj(model_path, vision)
     args = config.m3_args
     index = gguf_tensor_index(model_path)
     hidden = config.hidden_size
@@ -194,5 +346,8 @@ def norm(name: str) -> torch.Tensor:
         yield from fused(f"{mlp}.gate_up_proj", [gate, up], [matrix(gate), matrix(up)], hidden)
         yield from fused(f"{mlp}.down_proj", [down], [matrix(down)], int(down.ne[0]))
 
+    if vision:
+        yield from iter_mmproj_weights(find_mmproj(model_path))
+
 
-__all__ = ["iter_gguf_weights", "parse_gguf_config", "to_gguf"]
+__all__ = ["gguf_mm_config", "iter_gguf_weights", "iter_mmproj_weights", "parse_gguf_config", "to_gguf"]

python/freetoken/models/register.py

@@ -225,13 +225,15 @@ class ModelSpec:
         "MiniMaxM3ForCausalLM",
         packed_modules_mapping=_MINIMAX_M3_PACKED,
     ),
-    # llama.cpp minimax-m3 GGUF (text only): the same model classes with every tensor in its packed ggml type
+    # llama.cpp minimax-m3 GGUF: the same model classes with every tensor in its packed ggml type; a minimax_m3 mmproj adds the vision tower
     "MiniMaxM3GGUFForCausalLM": ModelSpec(
         "freetoken.models.minimax_m3",
-        "MiniMaxM3ForCausalLM",
+        "MiniMaxM3ForConditionalGeneration",
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
         packed_modules_mapping=_MINIMAX_M3_GGUF_PACKED,
+        mm_processor=_MINIMAX_M3_PROCESSOR,
+        encoders=_MINIMAX_M3_ENCODERS,
     ),
     "DeepseekV4ForCausalLM": ModelSpec(
         "freetoken.models.deepseek_v4",