File size: 3,387 Bytes
3da88e3 | 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 | """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())
|