"""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("", "") 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" # one package, functions E{L} sharing the weights 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") # one vision call per image; every question reuses the image tokens 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}}