| """Revalidate final decoder arithmetic using saved, unchanged DiT latents. |
| |
| Retains original outputs in an ignored development directory. Separate decode |
| measurements are never presented as a freshly measured end-to-end request. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import shutil |
| import time |
| from pathlib import Path |
|
|
| import mlx.core as mx |
|
|
| from h3_mlx.runtime import configure_memory |
| from h3_mlx.vae import VAE, to_image |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--vae", required=True) |
| p.add_argument("--runs", default="outputs/suite") |
| a = p.parse_args() |
| configure_memory() |
| root = Path(a.runs) |
| model = VAE(a.vae) |
| for trace in sorted(root.rglob("*.trace.safetensors")): |
| image = trace.with_name(trace.name.replace(".trace.safetensors", ".png")) |
| record = image.with_suffix(".json") |
| if not record.exists(): |
| continue |
| row = json.loads(record.read_text()) |
| if row.get("decoder_revalidated"): |
| continue |
| backup = Path("outputs/pre-layernorm-fix") / image.relative_to(root) |
| backup.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(image, backup) |
| shutil.copyfile(record, backup.with_suffix(".json")) |
| arrays = mx.load(str(trace)) |
| latent = arrays["latent"] |
| mx.eval(latent) |
| del arrays |
| model.tile_batch = row["vae"]["tile_batch"] |
| start = time.perf_counter() |
| pixels = model.decode(latent) |
| result = to_image(pixels, *row["size"]) |
| elapsed = time.perf_counter() - start |
| result.save(image) |
| original = row["timings"].copy() |
| row["source_request_timings_before_decoder_fix"] = original |
| row["timings"].pop("total_s", None) |
| row["timings"]["decode_s"] = elapsed |
| row["decoder_revalidated"] = { |
| "revision": 2, |
| "method": "FP32 LayerNorm affine before BF16 rounding", |
| "same_saved_latent": True, |
| "decode_s": elapsed, |
| } |
| row["timing_scope"] = ( |
| "Saved DiT stages plus separately timed corrected decoder; no fresh end-to-end total" |
| ) |
| row["png_sha256"] = hashlib.sha256(image.read_bytes()).hexdigest() |
| record.write_text(json.dumps(row, indent=2) + "\n") |
| print(json.dumps({"image": str(image), "decode_s": elapsed}), flush=True) |
| model.clear() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
