File size: 6,949 Bytes
5e32a69 | 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 | """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") # 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}}
|