| """Isolate VAE tile batching from diffusion drift using one saved final latent.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.benchmark import measure |
| from h3_mlx.compare import image_metrics |
| from h3_mlx.runtime import configure_memory, environment |
| from h3_mlx.vae import VAE, to_image |
|
|
|
|
| def 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("--batches", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32]) |
| p.add_argument("--repeats", type=int, default=3) |
| p.add_argument("--output", default="reports/vae") |
| a = p.parse_args() |
| configure_memory() |
| root = Path(a.output) |
| root.mkdir(parents=True, exist_ok=True) |
| data = mx.load(a.trace) |
| latent = data["latent"] |
| mx.eval(latent) |
| del data |
| record = json.loads(Path(a.record).read_text()) |
| width, height = record["size"] |
| model = VAE(a.vae) |
| baseline = model.decode(latent) |
| mx.eval(baseline) |
| ref = to_image(baseline, width, height) |
| trials = [] |
| for batch in a.batches: |
| if batch < 1: |
| p.error("tile batches must be positive") |
| model.tile_batch = batch |
| timer = measure(lambda: model.decode(latent), warmup=1, repeats=a.repeats) |
| pixels = model.decode(latent) |
| image = to_image(pixels, width, height) |
| image.save(root / f"batch-{batch}.png") |
| item = { |
| "batch": batch, |
| **timer, |
| **image_metrics(ref, image), |
| "max_float_error": mx.max(mx.abs(baseline - pixels)).item(), |
| } |
| trials.append(item) |
| print(json.dumps(item), flush=True) |
| model.clear() |
| (root / "trials.json").write_text( |
| json.dumps( |
| { |
| "environment": environment(), |
| "source_trace": a.trace, |
| "scope": "Same final latent; decoder only, loaded weights and one warmup excluded", |
| "trials": trials, |
| }, |
| indent=2, |
| ) |
| + "\n" |
| ) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
