edge/python/d3_mlx_video.py

"""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}}