feat(muse_glimmer): serve image input on Muse-Glimmer-30B

docs/models.md

@@ -30,6 +30,7 @@ These families accept image input by default; pass `--text-model-only` to skip t
 | Gemma-4 26B-A4B, 31B (`gemma4`: ViT tower, streamed under `--mm-encoder-weights host`) | one of the soft-token budgets 70 / 140 / 280 / 560 / 1120, every image scaled to its budget as far as the aspect ratio allows | the maximum picks the largest budget within it, below 70 is refused at start-up; the minimum has no effect | `{"max_soft_tokens": 1120}` |
 | Gemma-4 12B (`gemma4_unified`: linear patch embedder, resident under either placement) | same budgets, one 48x48 super-patch per soft token | same as the tower releases | same |
 | GLM-5.3-Flash (`glm5_next`: ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, dynamic resolution on a canvas zero-padded to a 28-multiple | token counts, passed through as the processor's `min_image_tokens` / `max_image_tokens`; checkpoint defaults 16 to 8000 tokens | `{"max_image_tokens": 2048}` |
+| Muse-Glimmer-30B (windowed ViT tower, streamed under `--mm-encoder-weights host`) | one token per 28x28 pixels of the resized image, aspect ratio kept under a token cap; checkpoint default 4096 tokens | the maximum is the cap (`max_image_tokens`); the minimum has no effect | `{"max_image_tokens": 1024}` |
 
 ## MoE strategies
 

python/freetoken/mm/processors/muse_glimmer.py

@@ -0,0 +1,79 @@
+"""Muse-Glimmer image processor: patch-grid ViT with a 2x2 pixel shuffle, a token-maximum budget and 1-D rope; the replacement adds the start/end wrapper the checkpoint processor renders."""
+
+from __future__ import annotations
+
+import struct
+from typing import Any
+
+import torch
+
+from freetoken.message import MMItem
+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
+
+
+class MuseGlimmerMMProcessor(MMProcessor):
+    def __init__(self, hf_config: Any, model_path: str, mm: MultimodalConfig) -> None:
+        super().__init__(model_path, mm)
+        vc = hf_config.vision_config
+        self.image_token_id = hf_config.image_token_id
+        # the template renders one <|patch|> per image
+        self.placeholder = [self.image_token_id]
+        self.merge = vc.merge_size
+        self.patch_dim = vc.patch_temporal * 3 * vc.patch_size**2
+
+    def get_mm_processor_kwargs(self, mm: MultimodalConfig) -> dict[str, Any]:
+        kwargs: dict[str, Any] = {"return_tensors": "pt"}
+        # the image processor's only budget is a token maximum; a minimum has no knob
+        if mm.image_max_tokens is not None:
+            kwargs["max_image_tokens"] = mm.image_max_tokens
+        return {**kwargs, **mm.processor_kwargs}
+
+    def process(self, images: list[Any]) -> list[MMItem]:
+        processor = self._image_processor()
+        kwargs = self.get_mm_processor_kwargs(self.mm)
+        items: list[MMItem] = []
+        for pil in images:
+            out = processor(images=pil, **kwargs)
+            pixel_values = out["pixel_values"].to(torch.bfloat16)
+            grid = out["image_grid_thw"][0].tolist()
+            t, h, w = int(grid[0]), int(grid[1]), int(grid[2])
+            if t != 1:
+                raise ValueError("video/temporal input is not supported")
+            # hash the bf16 wire buffer: half the bytes, and the exact tensor the cache and radix key agree on
+            item_hash = content_hash(pixel_values, struct.pack("<3i", t, h, w))
+            items.append(
+                MMItem(
+                    modality="image",
+                    hash=item_hash,
+                    pad_value=mm_pad_value(item_hash),
+                    offsets=[],
+                    feature=pixel_values,
+                    model_specific_data={"grid_thw": [t, h, w]},
+                )
+            )
+        return items
+
+    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)
+
+    def dummy_items(self, dtype: torch.dtype, device: torch.device) -> list[MMItem]:
+        merge = self.merge
+        return [MMItem(
+            modality="image",
+            hash=0,
+            pad_value=0,
+            offsets=[[0, 1]],
+            feature=torch.zeros(merge * merge, self.patch_dim, dtype=dtype, device=device),
+            model_specific_data={"grid_thw": [1, merge, merge]},
+        )]
+
+
+__all__ = ["MuseGlimmerMMProcessor"]

python/freetoken/models/muse_glimmer/__init__.py

@@ -1,11 +1,16 @@
 from .attention import MuseGlimmerAttention
-from .config import parse_config
-from .model import MuseGlimmerForCausalLM
+from .config import VisionConfig, parse_config, parse_vision_config
+from .model import MuseGlimmerForCausalLM, MuseGlimmerForConditionalGeneration
+from .vision import MuseGlimmerVisionModel
 from .weight import iter_weights
 
 __all__ = [
     "MuseGlimmerAttention",
     "MuseGlimmerForCausalLM",
+    "MuseGlimmerForConditionalGeneration",
+    "MuseGlimmerVisionModel",
+    "VisionConfig",
     "parse_config",
+    "parse_vision_config",
     "iter_weights",
 ]

python/freetoken/models/muse_glimmer/config.py

@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+from dataclasses import dataclass
 from typing import Any
 
 from freetoken.models.config import (
@@ -12,13 +13,76 @@
 
 _SWA_TYPE = "sliding_attention"
 _FULL_TYPE = "full_attention"
+_VISION_LAYER_TYPES = {"window_attention", "full_attention"}
+
+
+@dataclass(frozen=True)
+class VisionConfig:
+    hidden_size: int
+    intermediate_size: int
+    num_layers: int
+    num_heads: int
+    layer_types: tuple[str, ...]
+    patch_size: int
+    temporal_patch_size: int
+    merge_size: int
+    pos_emb_side: int
+    layer_norm_eps: float
+    rope_theta: float
+    projector_hidden_size: int
+    text_hidden_size: int
+    text_rms_norm_eps: float
+
+    @property
+    def patch_dim(self) -> int:
+        return self.temporal_patch_size * 3 * self.patch_size**2
+
+    @property
+    def out_hidden_size(self) -> int:
+        return self.hidden_size * self.merge_size**2
+
+
+def parse_vision_config(hf_config: Any) -> VisionConfig | None:
+    """None when the config carries no vision section, which is how a text-only engine asks for no tower."""
+    vc = getattr(hf_config, "vision_config", None)
+    if vc is None:
+        return None
+    if vc.hidden_act != "gelu" or hf_config.projector_hidden_act != "gelu":
+        raise NotImplementedError(
+            f"muse_glimmer vision activations {vc.hidden_act!r} / {hf_config.projector_hidden_act!r}; only gelu is implemented"
+        )
+    if vc.pos_emb_height != vc.pos_emb_width:
+        raise NotImplementedError("muse_glimmer vision tower: the position table must be square")
+    layer_types = tuple(vc.layer_types)
+    assert set(layer_types) <= _VISION_LAYER_TYPES, f"unknown vision layer types in {set(layer_types)}"
+    rope_params = getattr(vc, "rope_parameters", None) or {}
+    text = _text_config(hf_config)
+    config = VisionConfig(
+        hidden_size=vc.hidden_size,
+        intermediate_size=vc.intermediate_size,
+        num_layers=vc.num_hidden_layers,
+        num_heads=vc.num_attention_heads,
+        layer_types=layer_types,
+        patch_size=vc.patch_size,
+        temporal_patch_size=vc.patch_temporal,
+        merge_size=vc.merge_size,
+        pos_emb_side=vc.pos_emb_height,
+        layer_norm_eps=vc.layer_norm_eps,
+        rope_theta=float(rope_params.get("rope_theta", 10000.0)),
+        projector_hidden_size=hf_config.projector_hidden_size,
+        text_hidden_size=text.hidden_size,
+        text_rms_norm_eps=text.rms_norm_eps,
+    )
+    assert hf_config.out_hidden_size == config.out_hidden_size, (
+        f"adapter input {hf_config.out_hidden_size} != pixel-shuffled tower width {config.out_hidden_size}"
+    )
+    return config
 
 
 def _text_config(hf_config: Any) -> Any:
     """Muse Glimmer ships as a multimodal wrapper (MuseGlimmerForConditionalGeneration):
     the text tower lives in ``text_config`` and the weights carry a ``language_model.``
-    prefix. Served text-only -- the ~1.8B ViT perception encoder is never built, so
-    ``ModelConfig.vision_config`` stays None and the loader drops the vision tensors."""
+    prefix."""
     text = getattr(hf_config, "text_config", None)
     return text if text is not None else hf_config
 
@@ -108,7 +172,7 @@ def parse_config(hf_config: Any) -> ModelConfig:
         attn_sm_scale=attn_sm_scale,
         final_logit_softcapping=getattr(text, "final_logit_softcapping", None),
         output_multiplier=getattr(text, "output_multiplier", None),
-        vision_config=None,  # served text-only
+        vision_config=parse_vision_config(hf_config),
         image_token_id=getattr(hf_config, "image_token_id", None),
         attn_quant=quant,
         dense_quant=quant,
@@ -133,4 +197,4 @@ def parse_config(hf_config: Any) -> ModelConfig:
     )
 
 
-__all__ = ["parse_config"]
+__all__ = ["VisionConfig", "parse_config", "parse_vision_config"]

python/freetoken/models/muse_glimmer/model.py

@@ -17,12 +17,14 @@
     VocabParallelEmbedding,
     silu_and_mul,
 )
-from freetoken.models.blocks import BaseLLMModel
+from freetoken.models.blocks import BaseLLMModel, embed_input_ids
 from freetoken.utils import nvtx_annotate
 
 from .attention import MuseGlimmerAttention
+from .vision import MuseGlimmerVisionModel
 
 if TYPE_CHECKING:
+    from freetoken.message import MMItem
     from freetoken.models.config import ModelConfig
 
 
@@ -83,6 +85,21 @@ def forward(self, x: torch.Tensor) -> torch.Tensor:
         return residual + h
 
 
+class _NormedEmbedding:
+    """embed_tokens followed by the weightless norm as one lookup for embed_input_ids, so image rows (already normed by the tower) scatter in after it like in the reference."""
+
+    def __init__(self, embed_tokens: VocabParallelEmbedding, embed_norm: GemmaRMSNorm):
+        self.embed_tokens = embed_tokens
+        self.embed_norm = embed_norm
+
+    @property
+    def num_embeddings(self) -> int:
+        return self.embed_tokens.num_embeddings
+
+    def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.embed_norm.forward(self.embed_tokens.forward(input_ids))
+
+
 class MuseGlimmerModel(BaseOP):
     def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_tokens = VocabParallelEmbedding(
@@ -95,6 +112,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.embed_norm = GemmaRMSNorm(
             config.hidden_size, eps=config.rms_norm_eps, with_scale=False
         )
+        self._normed_embedding = _NormedEmbedding(self.embed_tokens, self.embed_norm)
         self.layers = OPList(
             [
                 MuseGlimmerDecoderLayer(config, layer_id, prefix=f"{prefix}.layers.{layer_id}")
@@ -106,7 +124,7 @@ def __init__(self, config: ModelConfig, *, prefix: str = "model"):
         self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
     def forward(self, input_ids: torch.Tensor) -> torch.Tensor:
-        x = self.embed_norm.forward(self.embed_tokens.forward(input_ids))
+        x = embed_input_ids(self._normed_embedding, input_ids, get_global_ctx().batch)
         for layer in self.layers.op_list:
             x = layer.forward(x)
         return self.norm.forward(x)
@@ -139,4 +157,17 @@ def forward(self) -> torch.Tensor:
         return logits
 
 
-__all__ = ["MuseGlimmerForCausalLM"]
+class MuseGlimmerForConditionalGeneration(MuseGlimmerForCausalLM):
+    def __init__(self, config: ModelConfig):
+        super().__init__(config)
+        if config.is_multimodal:
+            self.vision_tower = MuseGlimmerVisionModel(config.vision_config)
+
+    def place_encoder_weights(self, mode: str) -> None:
+        self.vision_tower.place_weights(mode)
+
+    def encode(self, item: MMItem) -> torch.Tensor:
+        return self.vision_tower.forward(item.feature, [item.grid_thw])
+
+
+__all__ = ["MuseGlimmerForCausalLM", "MuseGlimmerForConditionalGeneration"]