File size: 7,362 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 124 125 126 127 128 129 | """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" # 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}}
|