metaencoder_mlx.py

"""MLX inference for MetaEncoder with its official prompt and last-token pooling."""
import argparse
import json
from pathlib import Path
from types import MethodType

import mlx.core as mx
import numpy as np
from mlx_vlm.utils import load_model
from PIL import Image
from transformers.models.muse_glimmer.processing_muse_glimmer import MuseGlimmerProcessor

from upstream_metaencoder import (
    DEFAULT_SYSTEM_INSTRUCTION, VIDEO_FPS, _as_items, _user_text,
)


def _reference_vision_rotary(vision, positions):
    inv_freq = 1.0 / (vision.rope_theta ** (
        mx.arange(0, vision.rope_spatial_dim, 2, dtype=mx.float32) / vision.rope_spatial_dim))
    inv_freq = inv_freq.astype(mx.bfloat16).astype(mx.float32)
    width = positions[:, 0].astype(mx.float32)[:, None] * inv_freq[None, :]
    height = positions[:, 1].astype(mx.float32)[:, None] * inv_freq[None, :]
    frequencies = mx.concatenate([width, height, width, height], axis=-1)
    return mx.cos(frequencies).astype(mx.bfloat16), mx.sin(frequencies).astype(mx.bfloat16)


class MetaEncoderMLX:
    def __init__(self, path):
        self.model = load_model(Path(path), lazy=False)
        self.model.eval()
        # Use the original Transformers processor for exact image/video token routing.
        self.processor = MuseGlimmerProcessor.from_pretrained(path)
        self.processor.tokenizer.padding_side = "left"
        # The reference MetaEncoder rounds rotary frequencies to BF16 at load time.
        for layer in self.model.language_model.model.layers:
            rope = layer.self_attn.rope
            rope._frequencies = rope._frequencies.astype(mx.bfloat16).astype(mx.float32)
        self.model.vision_tower._rotary = MethodType(_reference_vision_rotary, self.model.vision_tower)

    @property
    def embedding_dim(self):
        return self.model.config.text_config.hidden_size

    def encode(self, inputs, instruction=None, max_length=8192):
        items = _as_items(inputs, instruction)
        vectors = []
        # Individual forwards avoid padding entering the decoder's attention masks.
        for item in items:
            content = [{"type": "image", "image": image} for image in item["images"] or []]
            if item["video"]:
                content.append({"type": "video", "video": item["video"]})
            text = _user_text(item["instruction"], item["text"])
            if text or not content:
                content.append({"type": "text", "text": text or " "})
            conversation = [
                {"role": "system", "content": [{"type": "text", "text": DEFAULT_SYSTEM_INSTRUCTION}]},
                {"role": "user", "content": content},
            ]
            prompt = self.processor.apply_chat_template(conversation, add_generation_prompt=False, tokenize=False)
            has_media = bool(item["images"] or item["video"])
            kwargs = dict(text=[prompt], images=item["images"], videos=[item["video"]] if item["video"] else None,
                          return_tensors="np", add_special_tokens=False, padding=False,
                          truncation=not has_media)
            if not has_media:
                kwargs["max_length"] = max_length
            if item["video"]:
                from transformers.video_utils import VideoMetadata
                frames = item["video"]
                kwargs.update(do_sample_frames=False, video_metadata=[VideoMetadata(
                    total_num_frames=len(frames), fps=VIDEO_FPS,
                    duration=len(frames)/VIDEO_FPS, frames_indices=list(range(len(frames))))])
            processed = self.processor(**kwargs)
            for media, pixel_key in ((item["images"], "pixel_values"), (item["video"], "pixel_values_videos")):
                if media and pixel_key not in processed:
                    raise ValueError(f"Processor dropped media: missing {pixel_key}")
            tensors = {key: mx.array(np.asarray(value)) for key, value in processed.items()
                       if key in {"input_ids", "pixel_values", "pixel_values_videos", "image_grid_thw", "video_grid_thw"}}
            ids = tensors["input_ids"].astype(mx.int32)
            language = self.model.language_model.model
            embeddings = language.embed_norm(language.embed_tokens(ids))
            for pixel_key, grid_key, token_id in (
                ("pixel_values", "image_grid_thw", self.model.config.image_token_id),
                ("pixel_values_videos", "video_grid_thw", self.model.config.video_token_id),
            ):
                if pixel_key not in tensors:
                    continue
                features = self.model._encode_image(tensors[pixel_key], tensors[grid_key]).astype(embeddings.dtype)
                positions = np.flatnonzero(np.asarray(ids[0]) == token_id)
                if len(positions) != features.shape[0]:
                    raise ValueError(f"Media token count {len(positions)} differs from feature count {features.shape[0]}")
                embeddings[0, mx.array(positions.astype(np.int32))] = features
            hidden = language(ids, inputs_embeds=embeddings)
            pooled = hidden[:, -1, :].astype(mx.float32)
            pooled = pooled / mx.maximum(mx.linalg.norm(pooled, axis=-1, keepdims=True), 1e-12)
            mx.eval(pooled)
            vectors.append(pooled)
        return mx.concatenate(vectors, axis=0) if vectors else mx.zeros((0, self.embedding_dim))

    def score(self, task, candidates):
        task_vector = self.encode(task)
        candidate_vectors = candidates if isinstance(candidates, mx.array) else self.encode(candidates)
        if candidate_vectors.ndim != 2 or candidate_vectors.shape[1] != self.embedding_dim:
            raise ValueError(f"Expected candidate embeddings of shape (n, {self.embedding_dim})")
        scores = (task_vector @ candidate_vectors.T)[0]
        mx.eval(scores)
        return scores

    def match(self, task, candidates, top_k=None):
        scores = np.asarray(self.score(task, candidates))
        order = np.argsort(-scores, kind="stable")
        if top_k is not None:
            if top_k < 0:
                raise ValueError("top_k must be nonnegative")
            order = order[:top_k]
        return [{"index": int(index), "score": float(scores[index]),
                 "candidate": None if isinstance(candidates, mx.array) else candidates[int(index)]}
                for index in order]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument("--task", required=True)
    parser.add_argument("--candidates", nargs="+", required=True)
    args = parser.parse_args()
    model = MetaEncoderMLX(args.model)
    print(json.dumps(model.match(args.task, args.candidates), indent=2))


if __name__ == "__main__":
    main()