File size: 2,228 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 | """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()
|