File size: 6,643 Bytes
78a1f66 | 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 | """d3 text and video requests on MLX (Apple GPU): upstream prompt and answer code; mlx-vlm's Qwen3.5 vision tower and language
model; all questions of a request in ONE batched forward (right-padded rows, key-padding mask on the noncausal
full-attention layers, causal Gated DeltaNet); FP32 readout on the host.
m = D3MLXVideo("<d3 checkpoint>", "<mlx dir: renamed weights + config>", max_pixels=100_352)
m.system_one(state=..., questions={...}, videos=[frames]) # or videos=None for a text request"""
import sys, time
from pathlib import Path
import mlx.core as mx
import numpy as np
import torch
_KEY_MASK = None
_LAYER_TYPES = None
def _patch(layer_types):
"""Two fixes to mlx-vlm's Qwen3.5 for d3. (1) Layer kinds come from the checkpoint's `layer_types`: mlx-vlm
assumes full attention every `full_attention_interval` layers, but d3-edge has it at layers 2, 5, 8-11, so half
its layers were built as the wrong kind and ran with random weights. (2) d3's full-attention layers are noncausal
(`noncausal_full_attention`): every real token is visible, only padding keys are masked.
Imported late: importing mlx_vlm's qwen3_5 registers its own Qwen3-VL processor, which upstream d3's request
preparation must not get."""
global _LAYER_TYPES
_LAYER_TYPES = list(layer_types)
import mlx_vlm.models.qwen3_5.language as qlang
if not getattr(qlang, "_d3_patched", False):
init = qlang.Qwen3_5DecoderLayer.__init__
def layer_init(self, args, layer_idx):
if _LAYER_TYPES is not None:
kind = _LAYER_TYPES[layer_idx]
# the interval the stock __init__ consults; make it agree with this layer's kind
args.full_attention_interval = 1 if kind == "full_attention" else 10 ** 6
init(self, args, layer_idx)
causal = qlang._create_qwen3_5_attention_mask
def noncausal(h, cache):
return _KEY_MASK if _KEY_MASK is not None else causal(h, cache)
qlang.Qwen3_5DecoderLayer.__init__ = layer_init
qlang._create_qwen3_5_attention_mask = noncausal
qlang._d3_patched = True
class D3MLXVideo:
def __init__(self, root, mlx_dir, max_pixels=None):
root = Path(root)
sys.path.insert(0, str(root))
import d3_runtime as up
if max_pixels:
up.VIDEO_MAX_PIXELS = max_pixels
from safetensors.torch import load_file
self.up = up
self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none") # prompts and positions only
self.processor, self.temperature, self.backbone = self.prep.processor, self.prep.temperature, self.prep.backbone
self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy()
self.video_token = self.backbone.config.video_token_id
import json
_patch(json.loads((root / "config.json").read_text())["text_config"]["layer_types"])
from mlx_vlm.utils import load_model
self.model = load_model(Path(mlx_dir), strict=True)
self.lm = self.model.language_model.model
self.model_name = root.name + "-mlx"
self.timing = {}
def system_one(self, *, state, questions, videos=None):
global _KEY_MASK
up, proc = self.up, self.processor
t0 = time.perf_counter()
prepared = self.prep.prepare(state, questions, videos=videos)
vis = None
if videos:
media = prepared.media
grids = media["video_grid_thw"]
pv = mx.array(media["pixel_values_videos"].float().numpy())
vis, _ = self.model.vision_tower(pv, mx.array(grids.numpy()))
mx.eval(vis) # MLX is lazy: time the vision tower here
t1 = time.perf_counter()
rows, keys = [], []
for key in prepared.runnable:
if videos:
expanded = proc.get_text_with_replacements([prepared.texts[key]], [], prepared.video_texts)[0][0]
enc = proc.tokenizer([expanded], return_token_type_ids=False, return_tensors="pt")
ids = enc["input_ids"]
types = torch.as_tensor(proc.create_mm_token_type_ids(ids))
with torch.no_grad():
pos = self.backbone.compute_3d_position_ids(
input_ids=ids, inputs_embeds=None, image_grid_thw=None, video_grid_thw=grids,
attention_mask=enc["attention_mask"], past_key_values=None, mm_token_type_ids=types)[:, 0]
rows.append((ids[0].numpy(), pos.numpy()))
else: # text: upstream's token ids, positions 0..n-1 on all three axes
row = np.array(prepared.sequences[key], np.int32)
rows.append((row, np.tile(np.arange(len(row), dtype=np.int32)[None], (3, 1))))
keys.append(key)
B, L = len(rows), max(len(r[0]) for r in rows)
ids = np.zeros((B, L), np.int32); pos = np.zeros((3, B, L), np.int32); valid = np.zeros((B, L), bool)
lengths = []
for b, (r, p) in enumerate(rows):
n = len(r); lengths.append(n)
ids[b, :n] = r; ids[b, n:] = r[-1]
pos[:, b, :n] = p; pos[:, b, n:] = p[:, -1:]
valid[b, :n] = True
emb = self.lm.embed_tokens(mx.array(ids))
if vis is not None:
is_video = mx.array(ids == self.video_token)
vis = vis.astype(emb.dtype)
per_row = int(is_video[0].sum().item())
# every row carries the same video tokens, in the same order
idx = mx.cumsum(is_video.astype(mx.int32), axis=1) - 1
emb = mx.where(is_video[..., None], vis[mx.clip(idx, 0, per_row - 1)], emb)
_KEY_MASK = mx.array(valid)[:, None, None, :]
try:
h = self.lm(None, inputs_embeds=emb, position_ids=mx.array(pos))
finally:
_KEY_MASK = None
last = mx.stack([h[b, lengths[b] - 1] for b in range(B)]).astype(mx.float32)
hs = np.array(last)
t2 = time.perf_counter()
answers, used = {}, 0
for b, key in enumerate(keys):
question = prepared.questions[key]
logits = self.readout[:len(question.keys)].astype(np.float64) @ hs[b].astype(np.float64) / self.temperature
p = np.exp(logits - logits.max())
answers[key] = up.product_answer(question, (p / p.sum()).tolist())
used += lengths[b]
self.timing = {"prepare+vision": round(t1 - t0, 3), "text": round(t2 - t1, 3)}
return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}}
|