"""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()