| """d3 video requests on Core ML (Python host). |
| |
| Upstream's own request preparation runs on a weightless (meta-device) d3, so frame sampling, the video processor, |
| the timestamped placeholders and the 3-D positions are exactly upstream's. Core ML runs the compute: the fp32 vision |
| tower once per video (attention within each frame pair, as upstream: one block per temporal group), then the text |
| graph per question with the video tokens spliced in; FP32 readout on the host. |
| |
| m = D3CoreMLVideo("<d3 checkpoint>", "<coreml dir>", max_pixels=100_352) |
| m.system_one(state=..., questions={...}, videos=[frames]) # frames: uint8 [n, h, w, 3] at 2 fps""" |
| import json |
| import sys |
| import time |
| from pathlib import Path |
|
|
| import coremltools as ct |
| import numpy as np |
| import torch |
|
|
| HERE = Path(__file__).parent |
|
|
|
|
| class D3CoreMLVideo: |
| def __init__(self, root, coreml_dir, max_pixels=None, vision_buckets=None, text_buckets=None, |
| compute_units=ct.ComputeUnit.CPU_AND_GPU): |
| """`vision_buckets` / `text_buckets`: load only these sizes (each loaded model costs memory).""" |
| root, cdir = Path(root), Path(coreml_dir) |
| sys.path.insert(0, str(root)) |
| sys.path.insert(0, str(HERE.parent / "vision")) |
| sys.path.insert(0, str(HERE.parent)) |
| import d3_runtime as up |
| if max_pixels: |
| up.VIDEO_MAX_PIXELS = max_pixels |
| from d3_graph import load_text_config |
| from qwen35_export import rope_cos_sin |
| from safetensors import safe_open |
| from safetensors.torch import load_file |
| from vision_export import VisionConfig |
|
|
| self.up, self.rope_cos_sin = up, rope_cos_sin |
| self.prep = up.D3.from_pretrained(str(root), device="meta", verify="none") |
| self.processor, self.temperature = self.prep.processor, self.prep.temperature |
| self.backbone = self.prep.backbone |
| self.vcfg = VisionConfig(json.loads((root / "config.json").read_text())["vision_config"]) |
| self.tcfg = load_text_config(str(root)) |
| for path in sorted(root.glob("model*.safetensors")): |
| with safe_open(str(path), "pt") as f: |
| if "visual.pos_embed.weight" in f.keys(): |
| self.pos_table = f.get_tensor("visual.pos_embed.weight").float() |
| self.readout = load_file(str(root / "readout.safetensors"))["weight"].float().numpy() |
| H = self.tcfg.hidden_size |
| self.embed = np.memmap(cdir / "embeddings.f16", dtype=np.float16, mode="r").reshape(-1, H) |
| self.video_token = self.backbone.config.video_token_id |
| self.vision = {int(p.name[8:]): ct.models.MLModel(str(p / "VisionTower_fp32.mlpackage"), compute_units=compute_units) |
| for p in sorted(cdir.glob("Vision_P*")) if not vision_buckets or int(p.name[8:]) in vision_buckets} |
| self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units) |
| for p in sorted(cdir.glob("d3_E*.mlpackage")) if not text_buckets or int(p.stem[4:]) in text_buckets} |
| self.model_name = root.name + "-coreml" |
| self.timing = {} |
|
|
| def warm(self, frames, state, questions): |
| for _ in range(2): |
| self.system_one(state=state, questions=questions, videos=[frames]) |
|
|
| def _video_tokens(self, pixel_values, grid): |
| """Vision tower for one video: attention within each temporal group (frame pair), padding masked.""" |
| from vision_export import prepare_inputs |
| n = pixel_values.shape[0] |
| N = next(b for b in sorted(self.vision) if b >= n) |
| patches, pe, cs, sn, mask, n = prepare_inputs(self.vcfg, self.pos_table, pixel_values, grid, N) |
| t, h, w = (int(x) for x in grid.reshape(-1)) |
| block = np.full((N, N), -1e4, dtype=np.float32) |
| for g in range(t): |
| s = g * h * w |
| block[s:s + h * w, s:s + h * w] = 0 |
| out = self.vision[N].predict({"patches": patches.numpy(), "pos_embeds": pe.numpy(), "cos": cs.numpy(), |
| "sin": sn.numpy(), "mask": block[None, None]})["tokens"] |
| return out[: n // 4] |
|
|
| def system_one(self, *, state, questions, videos): |
| up, proc = self.up, self.processor |
| t0 = time.perf_counter() |
| prepared = self.prep.prepare(state, questions, videos=videos) |
| media = prepared.media |
| grids = media["video_grid_thw"] |
| tokens, start = [], 0 |
| for grid in grids: |
| n = int(grid.prod()) |
| tokens.append(self._video_tokens(media["pixel_values_videos"][start:start + n], grid[None])) |
| start += n |
| tokens = np.concatenate(tokens) |
| t1 = time.perf_counter() |
| answers, used = {}, 0 |
| for key in prepared.runnable: |
| question = prepared.questions[key] |
| 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)) |
| row = ids[0].numpy() |
| hidden = self.embed[row].astype(np.float32) |
| hidden[np.where(row == self.video_token)[0]] = tokens |
| 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] |
| n = len(row) |
| L = next((b for b in sorted(self.text) if b >= n), None) |
| if L is None: |
| raise ValueError(f"{n} tokens exceed the largest text bucket {max(self.text)}") |
| cos, sin = self.rope_cos_sin(self.tcfg, torch.cat([pos, pos[:, -1:].expand(3, L - n)], dim=1)) |
| x = np.zeros((1, L, hidden.shape[1]), dtype=np.float16); x[0, :n] = hidden |
| valid = np.zeros(L, dtype=np.float16); valid[:n] = 1 |
| last = np.zeros((1, L), dtype=np.float16); last[0, n - 1] = 1 |
| h = self.text[L].predict({"hidden": x, "cos": cos.numpy().astype(np.float16), |
| "sin": sin.numpy().astype(np.float16), "key_valid": valid, |
| "last_onehot": last})["final"].reshape(-1) |
| logits = self.readout[:len(question.keys)].astype(np.float64) @ h.astype(np.float64) / self.temperature |
| p = np.exp(logits - logits.max()) |
| answers[key] = up.product_answer(question, (p / p.sum()).tolist()) |
| used += n |
| self.timing = {"prepare+vision": round(t1 - t0, 3), "text": round(time.perf_counter() - t1, 3)} |
| return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}} |
|
|
