| """d3 image requests on Core ML (Python host): upstream prompt, processor and answer code; Core ML vision tower |
| (fp32) + text graph (fp16, host-spliced embeddings, host M-RoPE); FP32 readout on the host. |
| |
| m = D3CoreMLImage("<d3 checkpoint dir>", "<coreml dir: Vision_P*, d3_E*.mlpackage, embeddings.f16>") |
| m.system_one(state=..., images=[...], questions={...}) # same answers dict as upstream |
| |
| Positions come from transformers' Qwen3_5Model.compute_3d_position_ids on a weightless (meta) model.""" |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import coremltools as ct |
| import numpy as np |
| import torch |
|
|
| HERE = Path(__file__).parent |
|
|
|
|
| class D3CoreMLImage: |
| def __init__(self, root, coreml_dir, compute_units=ct.ComputeUnit.CPU_AND_GPU): |
| 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 |
| from qwen35_export import rope_cos_sin |
| from transformers import AutoConfig |
| from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5Model |
| from vision_export import VisionConfig |
| from d3_graph import load_text_config |
|
|
| self.up, self.rope_cos_sin = up, rope_cos_sin |
| config = json.loads((root / "decision_config.json").read_text()) |
| self.codes, self.temperature = config["codes"], float(config.get("temperature", 1.0)) |
| self.processor = up.load_processor(root) |
| self.vcfg = VisionConfig(json.loads((root / "config.json").read_text())["vision_config"]) |
| self.tcfg = load_text_config(str(root)) |
| with torch.device("meta"): |
| self.positions = Qwen3_5Model(AutoConfig.from_pretrained(str(root))) |
| from safetensors import safe_open |
| self.pos_table = None |
| 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() |
| from safetensors.torch import load_file |
| 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.image_token = AutoConfig.from_pretrained(str(root)).image_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*"))} |
| multi = cdir / "d3_text.mlpackage" |
| if multi.exists(): |
| names = json.loads((cdir / "d3_text.functions.json").read_text()) |
| self.text = {int(n[1:]): ct.models.MLModel(str(multi), compute_units=compute_units, function_name=n) |
| for n in names} |
| else: |
| self.text = {int(p.stem[4:]): ct.models.MLModel(str(p), compute_units=compute_units) |
| for p in sorted(cdir.glob("d3_E*.mlpackage"))} |
| self.model_name = root.name + "-coreml" |
|
|
| def warm(self): |
| """Run every vision and text bucket once: the first call per bucket pays the GPU setup (seconds).""" |
| for N, model in self.vision.items(): |
| model.predict({"patches": np.zeros((N, self.vcfg.patch_dim), np.float32), |
| "pos_embeds": np.zeros((N, self.vcfg.hidden), np.float32), |
| "cos": np.zeros((N, self.vcfg.head_dim), np.float32), |
| "sin": np.zeros((N, self.vcfg.head_dim), np.float32), |
| "mask": np.zeros((1, 1, N, N), np.float32)}) |
| H, R = self.tcfg.hidden_size, self.tcfg.rotary_dim |
| for L, model in self.text.items(): |
| model.predict({"hidden": np.zeros((1, L, H), np.float16), "cos": np.ones((L, R), np.float16), |
| "sin": np.zeros((L, R), np.float16), "key_valid": np.ones(L, np.float16), |
| "last_onehot": np.eye(1, L, L - 1, dtype=np.float16)}) |
|
|
| def _tokens(self, pv, grid): |
| from vision_export import prepare_inputs |
| N = next(b for b in sorted(self.vision) if b >= pv.shape[0]) |
| patches, pe, cs, sn, mask, n = prepare_inputs(self.vcfg, self.pos_table, pv, grid, N) |
| out = self.vision[N].predict({"patches": patches.numpy(), "pos_embeds": pe.numpy(), "cos": cs.numpy(), |
| "sin": sn.numpy(), "mask": mask.numpy()})["tokens"] |
| return out[: n // 4] |
|
|
| def _probs(self, text, images, count, image_tokens): |
| enc = self.processor(text=[text], images=images, return_tensors="pt") |
| ids = enc["input_ids"][0].numpy() |
| hidden = self.embed[ids].astype(np.float32) |
| rows = np.where(ids == self.image_token)[0] |
| hidden[rows] = np.concatenate(image_tokens) |
| with torch.no_grad(): |
| pos = self.positions.compute_3d_position_ids( |
| input_ids=enc["input_ids"], inputs_embeds=None, image_grid_thw=enc["image_grid_thw"], |
| attention_mask=enc["attention_mask"], mm_token_type_ids=enc["mm_token_type_ids"])[:, 0] |
| n = len(ids) |
| 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 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[:count].astype(np.float64) @ h.astype(np.float64) / self.temperature |
| p = np.exp(logits - logits.max()) |
| return (p / p.sum()).tolist(), n |
|
|
| def system_one(self, *, state, questions, images): |
| up = self.up |
| decoded = [up.load_image(v) for v in images] |
| enc = self.processor.image_processor(images=decoded, return_tensors="pt") |
| |
| tokens, start = [], 0 |
| for k, grid in enumerate(enc["image_grid_thw"]): |
| n = int(grid.prod()) |
| tokens.append(self._tokens(enc["pixel_values"][start:start + n], grid[None])) |
| start += n |
| answers, used = {}, 0 |
| for key, q in questions.items(): |
| question = up.normalize_question(q) |
| text = self.processor.apply_chat_template( |
| up.image_messages("d3", state, question.rendered, self.codes, len(decoded)), |
| tokenize=False, add_generation_prompt=True, enable_thinking=False) |
| probs, n = self._probs(text, decoded, len(question.keys), tokens) |
| used += n |
| answers[key] = up.product_answer(question, probs) |
| return {"model": self.model_name, "answers": answers, "usage": {"input_tokens": used, "output_tokens": 0}} |
|
|
