| """Export complete native VAE tiles to Core ML and inspect real device placement. |
| |
| The source BF16 weights/input values are preserved through FP16 conversion. |
| Internal FP16 rounding differs from the native BF16 decoder; compare images. |
| Run export/benchmark in .venv-coreml, then use v2_decode.py for image evidence. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import hashlib |
| import json |
| import math |
| import statistics |
| import time |
| from pathlib import Path |
|
|
| import coremltools as ct |
| import numpy as np |
| import torch |
| import torch.nn.functional as F |
| from safetensors import safe_open |
|
|
| from experiments.v2_coreml import devices |
|
|
|
|
| class TileDecoder(torch.nn.Module): |
| def __init__(self, path, batch): |
| super().__init__() |
| self.batch = batch |
| for p in sorted(Path(path).glob("*.safetensors")): |
| with safe_open(p, framework="pt", device="cpu") as f: |
| for key in f.keys(): |
| if key.startswith("decoder.") or key.startswith("post_quant_conv."): |
| self.register_buffer( |
| key.replace(".", "__"), f.get_tensor(key).bfloat16().half() |
| ) |
| |
| coords = 2 * ((torch.arange(16, dtype=torch.bfloat16) + 0.5) / 16) - 1 |
| yy, xx = torch.meshgrid(coords, coords, indexing="ij") |
| pos = torch.stack((yy, xx), dim=-1).reshape(-1, 2) |
| pos = torch.cat((pos, torch.zeros(5, 2, dtype=torch.bfloat16))) |
| inv = 1 / (100 ** torch.arange(0, 1, 0.125, dtype=torch.float32)) |
| scaled = (2 * math.pi * pos.float()).bfloat16().float() |
| angles = (scaled[:, :, None] * inv[None, None, :]).reshape(261, 16) |
| angles = torch.cat((angles, angles), dim=-1) |
| self.register_buffer("cosine", angles.cos().bfloat16().half()[None, :, None, :]) |
| self.register_buffer("sine", angles.sin().bfloat16().half()[None, :, None, :]) |
|
|
| def w(self, key): |
| return getattr(self, key.replace(".", "__")) |
|
|
| def linear(self, x, key): |
| return F.linear(x, self.w(key + ".weight"), self.w(key + ".bias")) |
|
|
| def norm(self, x, key=None): |
| xf = x.float() |
| y = xf * torch.rsqrt(torch.mean(xf * xf, dim=-1, keepdim=True) + 1e-5) |
| if key is not None: |
| y = y * self.w(key).float() |
| return y.half() |
|
|
| def rope(self, x): |
| rotary, tail = x[..., :32], x[..., 32:] |
| a, b = rotary.chunk(2, dim=-1) |
| return torch.cat( |
| (rotary * self.cosine + torch.cat((-b, a), dim=-1) * self.sine, tail), |
| dim=-1, |
| ) |
|
|
| def forward(self, raw): |
| x = raw.permute(0, 2, 3, 1) |
| x = F.linear( |
| x, |
| self.w("post_quant_conv.weight").reshape(24, 24), |
| self.w("post_quant_conv.bias"), |
| ) |
| x = self.linear(x, "decoder.x_embedder").reshape(self.batch, 256, 1024) |
| x = torch.cat( |
| ( |
| x, |
| self.w("decoder.register_tokens").expand(self.batch, 4, 1024), |
| torch.zeros(self.batch, 1, 1024, dtype=torch.float16), |
| ), |
| dim=1, |
| ) |
| for i in range(24): |
| p = f"decoder.transformer_blocks.{i}" |
| h = self.norm(x, p + ".norm1.weight") |
| qkv = self.linear(h, p + ".attn.to_qkv").reshape(self.batch, 261, 16, 3, 64) |
| q, k, v = (qkv[:, :, :, j] for j in range(3)) |
| q, k = self.rope(self.norm(q)), self.rope(self.norm(k)) |
| attn = ( |
| F.scaled_dot_product_attention( |
| q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2) |
| ) |
| .transpose(1, 2) |
| .reshape(self.batch, 261, 1024) |
| ) |
| x = x + self.linear(attn, p + ".attn.to_out") * self.w(p + ".scale1") |
| h = self.norm(x, p + ".norm2.weight") |
| gate, up = self.linear(h, p + ".ff.w1").chunk(2, dim=-1) |
| x = x + self.linear(F.silu(gate) * up, p + ".ff.w2") * self.w(p + ".scale2") |
| x = F.layer_norm( |
| x.float(), |
| (1024,), |
| self.w("decoder.norm_out.weight").float(), |
| self.w("decoder.norm_out.bias").float(), |
| 1e-5, |
| ).half() |
| x = self.linear(x, "decoder.proj_out")[:, :256] |
| return ( |
| x.reshape(self.batch, 16, 16, 3, 16, 16) |
| .permute(0, 3, 1, 4, 2, 5) |
| .reshape(self.batch, 3, 256, 256) |
| ) |
|
|
|
|
| def main(a): |
| torch.set_num_threads(a.threads) |
| package = Path(a.package) |
| example = ( |
| np.random.default_rng(731).normal(size=(a.batch, 24, 16, 16)).astype(np.float16) |
| ) |
| if not package.exists(): |
| model = TileDecoder(a.vae, a.batch).eval() |
| with torch.inference_mode(): |
| traced = torch.jit.trace(model, torch.from_numpy(example)) |
| converted = ct.convert( |
| traced, |
| convert_to="mlprogram", |
| inputs=[ct.TensorType(name="raw", shape=example.shape, dtype=np.float16)], |
| outputs=[ct.TensorType(name="pixels", dtype=np.float16)], |
| minimum_deployment_target=ct.target.macOS15, |
| compute_precision=ct.precision.FLOAT16, |
| skip_model_load=True, |
| ) |
| files = sorted(Path(a.vae).glob("*.safetensors")) |
| identity = [ |
| (str(p.resolve()), p.stat().st_size, p.stat().st_mtime_ns) for p in files |
| ] |
| converted.user_defined_metadata.update( |
| { |
| "h3.decoder.format": "h3-coreml-tile-v1", |
| "h3.decoder.source": hashlib.sha256( |
| json.dumps(identity).encode() |
| ).hexdigest(), |
| "h3.decoder.batch": str(a.batch), |
| "h3.decoder.precision": "Source BF16 weights and inputs, FP16 internal arithmetic, BF16 tile assembly", |
| } |
| ) |
| package.parent.mkdir(parents=True, exist_ok=True) |
| converted.save(str(package)) |
| del model, traced, converted |
| gc.collect() |
| report = { |
| "coremltools": ct.__version__, |
| "torch": torch.__version__, |
| "package": str(package), |
| "batch": a.batch, |
| "scope": "Complete 24-layer tile decoder. FP16 internals, source weights rounded to BF16 before conversion. NumPy I/O included; tile assembly and normalization excluded.", |
| "trials": [], |
| } |
| dest = Path(a.output) |
| dest.parent.mkdir(parents=True, exist_ok=True) |
| for unit in (getattr(ct.ComputeUnit, name) for name in a.units): |
| row = {"compute_units": unit.name} |
| try: |
| tick = time.perf_counter() |
| model = ct.models.MLModel( |
| str(package), |
| compute_units=unit, |
| optimization_hints={ |
| "specializationStrategy": ct.SpecializationStrategy.FastPrediction |
| }, |
| ) |
| row["load_compile_s"] = time.perf_counter() - tick |
| plan = ct.models.compute_plan.MLComputePlan.load_from_path( |
| model.get_compiled_model_path(), compute_units=unit |
| ) |
| row["placement"] = devices(plan) |
| tick = time.perf_counter() |
| y = model.predict({"raw": example})["pixels"] |
| row["first_predict_ms"] = 1000 * (time.perf_counter() - tick) |
| samples = [] |
| for _ in range(a.repeats): |
| tick = time.perf_counter() |
| y = model.predict({"raw": example})["pixels"] |
| samples.append(1000 * (time.perf_counter() - tick)) |
| row.update( |
| samples_ms=samples, |
| median_ms=statistics.median(samples), |
| finite=bool(np.isfinite(y).all()), |
| ) |
| del model, plan, y |
| gc.collect() |
| except Exception as exc: |
| row["error"] = f"{type(exc).__name__}: {exc}" |
| report["trials"].append(row) |
| dest.write_text(json.dumps(report, indent=2) + "\n") |
| print( |
| json.dumps({k: v for k, v in row.items() if k != "placement"}), flush=True |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--vae", required=True) |
| p.add_argument("--batch", type=int, default=8) |
| p.add_argument("--threads", type=int, default=10) |
| p.add_argument( |
| "--units", |
| nargs="+", |
| choices=("CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"), |
| default=["CPU_ONLY", "CPU_AND_GPU", "CPU_AND_NE", "ALL"], |
| ) |
| p.add_argument("--repeats", type=int, default=3) |
| p.add_argument("--package", default="cache/coreml-v2/vae-batch8.mlpackage") |
| p.add_argument("--output", default="reports/optimization-v2/coreml-vae.json") |
| main(p.parse_args()) |
|
|
