| """Compare native, compiled and Core ML decoders on the exact same final latent.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import hashlib |
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import measure |
| from h3_mlx.compare import image_metrics, make_comparison |
| from h3_mlx.coreml_vae import CoreMLVAE |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.vae import VAE, to_image |
|
|
|
|
| def main(a): |
| configure_memory(wired_gb=5) |
| root = Path(a.output) |
| root.mkdir(parents=True, exist_ok=True) |
| data = mx.load(a.trace) |
| latent = data["latent"] |
| mx.eval(latent) |
| del data |
| source = json.loads(Path(a.record).read_text()) |
| width, height = source["size"] |
| results = [] |
| images = [] |
| labels = [] |
| reference = None |
| variants = [ |
| ("native-b1", VAE, {"tile_batch": 1}), |
| ("native-b4", VAE, {"tile_batch": 4}), |
| ("compiled-b4", VAE, {"tile_batch": 4, "compile_blocks": True}), |
| ("native-b8", VAE, {"tile_batch": 8}), |
| ("compiled-b8", VAE, {"tile_batch": 8, "compile_blocks": True}), |
| ] |
| if a.include_fp32: |
| variants.insert( |
| 0, |
| ( |
| "fp32-b4-compiled", |
| VAE, |
| {"tile_batch": 4, "compile_blocks": True, "dtype": "float32"}, |
| ), |
| ) |
| if a.package: |
| variants.extend( |
| ( |
| "coreml-" + unit, |
| CoreMLVAE, |
| { |
| "tile_batch": a.coreml_batch, |
| "package": a.package, |
| "compute_units": unit, |
| }, |
| ) |
| for unit in a.units |
| ) |
| for name, cls, options in variants: |
| model = cls(a.vae, **options) |
| timer = measure( |
| lambda model=model: model.decode(latent), warmup=1, repeats=a.repeats |
| ) |
| pixels = model.decode(latent) |
| image = to_image(pixels, width, height) |
| if reference is None: |
| reference = image |
| path = root / f"{name}.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( |
| vae=options, |
| vae_backend="coreml" if cls is CoreMLVAE else "mlx", |
| environment=environment(), |
| validation_scope="Same final latent; decoder only", |
| timing_scope="Warm decoder only; model loading excluded", |
| timings={"decode_s": timer["median_ms"] / 1000}, |
| 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(), |
| ) |
| path.with_suffix(".json").write_text(json.dumps(record, indent=2) + "\n") |
| row = {"variant": name, **timer, **image_metrics(reference, image)} |
| results.append(row) |
| images.append(path) |
| labels.append(name) |
| print(json.dumps(row), flush=True) |
| (root / "trials.json").write_text( |
| json.dumps( |
| { |
| "source_trace": a.trace, |
| "scope": "Same final latent, decoder-only hot latency", |
| "trials": results, |
| }, |
| indent=2, |
| ) |
| + "\n" |
| ) |
| model.clear() |
| del model |
| gc.collect() |
| mx.clear_cache() |
| make_comparison(images[0], images[1:], root / "comparison", labels=labels) |
|
|
|
|
| if __name__ == "__main__": |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--vae", required=True) |
| p.add_argument("--trace", required=True) |
| p.add_argument("--record", required=True) |
| p.add_argument("--package") |
| p.add_argument( |
| "--include-fp32", |
| action="store_true", |
| help="Use the original FP32 decoder as the same-latent reference", |
| ) |
| p.add_argument("--coreml-batch", type=int, default=8) |
| p.add_argument("--units", nargs="+", default=["CPU_AND_GPU"]) |
| p.add_argument("--repeats", type=int, default=3) |
| p.add_argument("--output", default="comparisons/optimization-v2/decoder-product") |
| main(p.parse_args()) |
|
|
