| """One actual image for a Core ML device/batch, with NumPy bridge diagnostics.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
| import numpy as np |
| from PIL import Image |
|
|
| from h3_mlx.compare import image_metrics, make_comparison |
| from h3_mlx.coreml_vae import CoreMLVAE |
| from h3_mlx.runtime import configure_memory |
| from h3_mlx.vae import to_image |
|
|
|
|
| def main(a): |
| configure_memory(wired_gb=5) |
| data = mx.load(a.trace) |
| latent = data["latent"] |
| mx.eval(latent) |
| del data |
| tick = time.perf_counter() |
| model = CoreMLVAE( |
| a.vae, package=a.package, tile_batch=a.batch, compute_units=a.units |
| ) |
| load_s = time.perf_counter() - tick |
| sample = ( |
| np.random.default_rng(731).normal(size=(a.batch, 24, 16, 16)).astype(np.float16) |
| ) |
| raw = model.model.predict({"raw": sample})["pixels"] |
| audit = { |
| "output_shape": list(raw.shape), |
| "output_strides": list(raw.strides), |
| "output_c_contiguous": bool(raw.flags.c_contiguous), |
| "original_mlx_numpy_import_exact": bool( |
| np.array_equal(np.asarray(mx.array(raw)), raw, equal_nan=True) |
| ), |
| } |
| source = json.loads(Path(a.reference).with_suffix(".json").read_text()) |
| tick = time.perf_counter() |
| pixels = model.decode(latent) |
| mx.eval(pixels) |
| seconds = time.perf_counter() - tick |
| image = to_image(pixels, *source["size"]) |
| root = Path(a.output) |
| root.mkdir(parents=True, exist_ok=True) |
| path = root / "image.png" |
| image.save(path) |
| record = { |
| k: source[k] |
| for k in ( |
| "prompt", |
| "seed", |
| "size", |
| "canvas", |
| "rng", |
| "reference_images", |
| "reference_sha256", |
| "sigmas", |
| "steps", |
| ) |
| if k in source |
| } |
| record.update( |
| validation_scope="Same final latent; decoder only", |
| timing_scope="One complete decoder pass; model loading excluded", |
| timings={"decode_s": seconds}, |
| load_s=load_s, |
| vae_backend="coreml", |
| vae={"batch": a.batch, "compute_units": a.units}, |
| source_trace=a.trace, |
| source_trace_sha256=hashlib.sha256(Path(a.trace).read_bytes()).hexdigest(), |
| output=str(path), |
| png_sha256=hashlib.sha256(path.read_bytes()).hexdigest(), |
| numpy_bridge=audit, |
| **image_metrics(Image.open(a.reference), image), |
| ) |
| path.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n") |
| print(json.dumps(record), flush=True) |
| make_comparison( |
| a.reference, |
| [path], |
| root / "comparison", |
| labels=[ |
| "Native BF16 decoder", |
| f"Core ML {a.units}, batch {a.batch}, contiguous bridge", |
| ], |
| ) |
| model.clear() |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--vae", required=True) |
| p.add_argument("--package", required=True) |
| p.add_argument("--batch", type=int, default=8) |
| p.add_argument("--units", default="CPU_AND_NE") |
| p.add_argument( |
| "--trace", default="outputs/suite/product/native-3.trace.safetensors" |
| ) |
| p.add_argument("--reference", default="outputs/suite/product/native-3.png") |
| p.add_argument("--output", required=True) |
| main(p.parse_args()) |
|
|
