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