feat(gguf): serve image input on muse-glimmer ggufs from a muse-glimmer mmproj

python/freetoken/mm/processors/muse_glimmer.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 tokenizer's <|image_start|> / <|image_end|>; the config only names the <|patch|> id
-_IMAGE_START_ID = 200080
-_IMAGE_END_ID = 200081
+from freetoken.models.muse_glimmer.config import IMAGE_END_ID, IMAGE_START_ID
 
 
 class MuseGlimmerMMProcessor(MMProcessor):
@@ -62,7 +59,7 @@ def process(self, images: list[Any]) -> list[MMItem]:
     def prompt_replacement(self, item: MMItem) -> PromptReplacement:
         t, h, w = item.grid_thw
         pads = [self.image_token_id] * ((t * h * w) // (self.merge * self.merge))
-        return PromptReplacement.select_token_id([_IMAGE_START_ID, *pads, _IMAGE_END_ID], self.image_token_id)
+        return PromptReplacement.select_token_id([IMAGE_START_ID, *pads, IMAGE_END_ID], self.image_token_id)
 
     def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]:
         merge = self.merge

python/freetoken/models/muse_glimmer/__init__.py

@@ -1,6 +1,6 @@
 from .attention import MuseGlimmerAttention
 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 MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration
 from .vision import MuseGlimmerVisionModel
 from .weight import iter_vision_weights, iter_weights
@@ -15,7 +15,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/muse_glimmer/config.py

@@ -11,6 +11,10 @@
     detect_compressed_tensors_nvfp4,
 )
 
+# the tokenizer's <|image_start|> / <|image_end|>; the config only names the <|patch|> id
+IMAGE_START_ID = 200080
+IMAGE_END_ID = 200081
+
 _SWA_TYPE = "sliding_attention"
 _FULL_TYPE = "full_attention"
 _VISION_LAYER_TYPES = {"window_attention", "full_attention"}
@@ -197,4 +201,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/muse_glimmer/gguf.py

@@ -2,20 +2,26 @@
 
 The converter (conversion/muse_glimmer.py) adds 1 to the four norms of each block, permutes the rows of each q and k head for llama.cpp's interleaved rope, and writes constant q/k norms that carry ``qk_scale_factor``.
 The loader subtracts the 1, undoes the permute, and reads ``qk_scale_factor`` back from the q norm.
+
+A ``muse-glimmer`` mmproj beside the text file adds the vision tower (images only): its metadata and the family constants below become the HF vision keys and the image processor config, and its tensors load under ``vision_tower.``.
+The converter permutes the vision q/k rows the same way (undone at load) and sums the two temporal slabs of the patch embedding into one; the loader splits the sum back over both slabs, which is exact for an image (one frame in both slabs) but not for video.
 """
 
 from __future__ import annotations
 
+import math
 import re
 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
@@ -122,6 +128,12 @@ def kv(key: str):
         output_multiplier=float(kv("logit_scale")),
         vocab_size=int(shim.vocab_size),
         tie_word_embeddings=bool(shim.tie_word_embeddings),
+        # filled by a muse-glimmer mmproj; a text-only process nulls vision_config on its copy
+        vision_config=getattr(shim, "vision_config", None),
+        image_token_id=getattr(shim, "image_token_id", None),
+        projector_hidden_act=getattr(shim, "projector_hidden_act", None),
+        projector_hidden_size=getattr(shim, "projector_hidden_size", None),
+        out_hidden_size=getattr(shim, "out_hidden_size", None),
     )
 
 
@@ -134,19 +146,156 @@ def _unpermute(rows: torch.Tensor, heads: int) -> torch.Tensor:
     return rows.reshape(heads, rows.shape[0] // heads // 2, 2, *rows.shape[1:]).swapaxes(1, 2).reshape(rows.shape)
 
 
+# ----------------------------------------------------------------------------------------------
+# Vision tower (a muse-glimmer mmproj)
+# ----------------------------------------------------------------------------------------------
+
+_PROJECTOR = "muse-glimmer"
+_IMAGE_TOKEN = "<|patch|>"
+# not stored in the mmproj; llama.cpp hardcodes the same values (clip.cpp) and they match the HF release
+_TEMPORAL_PATCH = 2
+_ROPE_THETA = 10000.0
+_MAX_IMAGE_TOKENS = 4096
+# window attention except every 4th block and the last one
+_FULL_EVERY = 4
+
+_TOP = {
+    "v.position_embd.weight": "patch_embedder.position_embedding_table.weight",
+    "v.pre_ln.weight": "ln_pre.weight",
+    "v.pre_ln.bias": "ln_pre.bias",
+    "v.post_ln.weight": "ln_post.weight",
+    "v.post_ln.bias": "ln_post.bias",
+    "mm.0.weight": "adapter.fc1.weight",
+    "mm.1.weight": "adapter.fc2.weight",
+    "mm.2.weight": "projection.weight",
+}
+_BLOCK = {
+    "attn_out": "attn.proj",
+    "ffn_up": "mlp.fc1",
+    "ffn_down": "mlp.fc2",
+    "ln1": "norm1",
+    "ln2": "norm2",
+}
+
+
+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; Muse-Glimmer 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 that moved them would fail per request instead of here
+    for token, want in (("<|image_start|>", IMAGE_START_ID), ("<|image_end|>", IMAGE_END_ID)):
+        if token not in tokens or tokens.index(token) != want:
+            raise ValueError(f"the text GGUF tokenizer does not carry {token} at the family's id {want}")
+    hidden, patch, merge = int(kv("embedding_length")), int(kv("patch_size")), int(kv("spatial_merge_size"))
+    # ggml order [kw, kh, in, out]: one summed slab, not the HF [out, temporal * in * kh * kw]
+    if tuple(index["v.patch_embd.weight"].ne) != (patch, patch, 3, hidden):
+        raise ValueError(f"mmproj {mmproj.path}: v.patch_embd.weight {tuple(index['v.patch_embd.weight'].ne)} is not the summed [p, p, 3, hidden] layout")
+    # ggml order [hidden, rows]: a square table of side x side positions
+    rows = int(index["v.position_embd.weight"].ne[1])
+    side = math.isqrt(rows)
+    if side * side != rows:
+        raise ValueError(f"mmproj {mmproj.path}: the position table has {rows} rows, not a square")
+    layers = int(kv("block_count"))
+    vision_config = {
+        "hidden_act": "gelu",
+        "hidden_size": hidden,
+        "intermediate_size": int(kv("feed_forward_length")),
+        "num_hidden_layers": layers,
+        "num_attention_heads": int(kv("attention.head_count")),
+        "layer_types": ["full_attention" if (i + 1) % _FULL_EVERY == 0 or i == layers - 1 else "window_attention" for i in range(layers)],
+        "patch_size": patch,
+        "patch_temporal": _TEMPORAL_PATCH,
+        "merge_size": merge,
+        "pos_emb_height": side,
+        "pos_emb_width": side,
+        # stored as f32; the decimal is the HF value
+        "layer_norm_eps": float(str(np.float32(kv("attention.layer_norm_epsilon")))),
+        "rope_parameters": {"rope_theta": _ROPE_THETA, "rope_type": "default"},
+    }
+    # ggml order [in, out]: the adapter reads the pixel-shuffled width and writes the projector width
+    adapter = index["mm.0.weight"].ne
+    hf_fields = {
+        "vision_config": vision_config,
+        "image_token_id": tokens.index(_IMAGE_TOKEN),
+        "projector_hidden_act": "gelu",
+        "projector_hidden_size": int(adapter[1]),
+        "out_hidden_size": int(adapter[0]),
+    }
+    preprocessor = {
+        "image_processor_type": "MuseGlimmerImageProcessor",
+        "do_convert_rgb": True,
+        "do_resize": True,
+        "resample": 1,  # LANCZOS
+        "do_rescale": True,
+        "rescale_factor": 1 / 255,
+        "do_normalize": True,
+        "image_mean": [float(v) for v in kv("image_mean")],
+        "image_std": [float(v) for v in kv("image_std")],
+        "patch_size": patch,
+        "temporal_patch_size": _TEMPORAL_PATCH,
+        "merge_size": merge,
+        "max_image_tokens": _MAX_IMAGE_TOKENS,
+    }
+    return hf_fields, preprocessor
+
+
+def iter_mmproj_weights(mmproj: Mmproj) -> Iterator[tuple[str, torch.Tensor]]:
+    """The tower state dict under ``vision_tower.``, named as the HF checkpoint loads."""
+    index = gguf_tensor_index(mmproj.path)
+    heads = int(load_gguf_metadata(mmproj.path)["clip.vision.attention.head_count"])
+    n_blocks = 1 + max(int(n.split(".")[2]) for n in index if n.startswith("v.blk."))
+
+    # the converter stores W0 + W1 of the two temporal slabs, which is all an image needs (its frame fills both); the bf16 hi + lo split carries that f32 sum without rounding
+    summed = mmproj_tensor(index["v.patch_embd.weight"]).float().flatten(1)
+    hi = summed.bfloat16()
+    lo = (summed - hi.float()).bfloat16()
+    yield "vision_tower.patch_embedder.patch_embedding.weight", torch.cat([hi, lo], dim=1)
+    for gguf_name, name in _TOP.items():
+        yield f"vision_tower.{name}", mmproj_tensor(index[gguf_name])
+
+    consumed = {"v.patch_embd.weight", *_TOP}
+    for layer in range(n_blocks):
+        blk, dst = f"v.blk.{layer}", f"vision_tower.layers.{layer}"
+        for part in ("weight", "bias"):
+            q, k, v = (mmproj_tensor(index[f"{blk}.attn_{p}.{part}"]) for p in "qkv")
+            yield f"{dst}.attn.qkv.{part}", torch.cat([_unpermute(q, heads), _unpermute(k, heads), v], dim=0)
+            consumed.update(f"{blk}.attn_{p}.{part}" for p in "qkv")
+            for src, name in _BLOCK.items():
+                yield f"{dst}.{name}.{part}", mmproj_tensor(index[f"{blk}.{src}.{part}"])
+                consumed.add(f"{blk}.{src}.{part}")
+
+    leftovers = sorted(set(index) - consumed)
+    if leftovers:
+        raise ValueError(f"mmproj {mmproj.path}: unexpected tensors {leftovers[:5]}")
+
+
 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 text state dict from the GGUF."""
+    """The state dict from the GGUF, and the vision tower from its mmproj."""
     if not include_non_moe:
         return
     from freetoken.utils import cached_load_hf_config
 
     hf = _hf_config(cached_load_hf_config(model_path))
+    vision = include_vision and hf.vision_config is not None
+    log_mmproj(model_path, vision)
     index = gguf_tensor_index(model_path)
     hidden, heads, kv_heads = hf.hidden_size, hf.num_attention_heads, hf.num_key_value_heads
     q_rows, kv_rows = heads * hf.head_dim, kv_heads * hf.head_dim
@@ -181,5 +330,8 @@ def t(name: str) -> GgufTensorInfo:
         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

@@ -372,13 +372,15 @@ class ModelSpec:
         mm_processor=_MUSE_GLIMMER_PROCESSOR,
         encoders=_MUSE_GLIMMER_ENCODERS,
     ),
-    # llama.cpp muse-glimmer GGUF (text only): the same model classes with every tensor in its packed ggml type
+    # llama.cpp muse-glimmer GGUF: the same model classes with every tensor in its packed ggml type; a muse-glimmer mmproj adds the vision tower (images only)
     "MuseGlimmerGGUFForCausalLM": ModelSpec(
         "freetoken.models.muse_glimmer",
-        "MuseGlimmerForCausalLM",
+        "MuseGlimmerForConditionalGeneration",
         parse_config="parse_gguf_config",
         iter_weights="iter_gguf_weights",
         packed_modules_mapping=_DENSE_PACKED,
+        mm_processor=_MUSE_GLIMMER_PROCESSOR,
+        encoders=_MUSE_GLIMMER_ENCODERS,
     ),
     "MistralForCausalLM": ModelSpec(
         "freetoken.models.mistral",