"""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("", "", 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") # no weights: prompts only 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): # rows are t-major: each frame pair is h*w consecutive patches 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}}