File size: 8,329 Bytes
78a1f66 | 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 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 | """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}
elif list(cdir.glob("d3_E*_p0.mlpackage")): # large checkpoints: each bucket is a chain of parts
self.text = {}
for p0 in sorted(cdir.glob("d3_E*_p0.mlpackage")):
L = int(p0.stem[4:].split("_")[0])
chain = sorted(cdir.glob(f"d3_E{L}_p*.mlpackage"), key=lambda q: int(q.stem.rsplit("_p", 1)[1]))
self.text[L] = [ct.models.MLModel(str(q), compute_units=compute_units) for q in chain]
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 _text(self, L, inputs):
"""One text bucket: a single graph, or its parts in order (hidden flows from part to part)."""
model = self.text[L]
if not isinstance(model, list):
return model.predict(inputs)["final"]
x = inputs["hidden"]
for part in model[:-1]:
x = part.predict({"hidden": x, "cos": inputs["cos"], "sin": inputs["sin"],
"key_valid": inputs["key_valid"]})["out"]
return model[-1].predict({**inputs, "hidden": x})["final"]
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 in self.text:
self._text(L, {"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, {"hidden": x, "cos": cos.numpy().astype(np.float16), "sin": sin.numpy().astype(np.float16),
"key_valid": valid, "last_onehot": last}).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}}
|