File size: 6,824 Bytes
e5030dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | """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()
|