| """Build matched galleries, raw record exports and charts from completed trials.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import json |
| import shutil |
| from pathlib import Path |
|
|
| from h3_mlx.compare import make_comparison |
|
|
| LABELS = { |
| "native": "BF16 native / six passes", |
| "tuned": "BF16 / fused AdaLN / tile batch 8", |
| "fused": "BF16 / fused AdaLN + QK-RoPE / tile batch 8", |
| "q8": "DiT Q8 / six passes", |
| "q4": "DiT Q4 / six passes", |
| "preview": "BF16 preview / steps 0,2,4", |
| "q8-g32": "DiT Q8 / group 32 / six passes", |
| "chunked": "BF16 / MLP chunks 1024", |
| } |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--runs", default="outputs/suite") |
| p.add_argument("--output", default="comparisons") |
| p.add_argument( |
| "--cuda-reference-root", |
| help="Optional saved single-bf16-stream CUDA run; quality=lossless", |
| ) |
| a = p.parse_args() |
| runs, output = Path(a.runs), Path(a.output) |
| records = [] |
| for case in sorted(d for d in runs.iterdir() if d.is_dir()): |
| files = sorted(case.glob("*.png")) |
| native = sorted(case.glob("native-*.png")) |
| if not native: |
| continue |
| reference = native[-1] |
| candidates = [] |
| labels = [LABELS["native"]] |
| for name, label in LABELS.items(): |
| if name == "native": |
| continue |
| selected = sorted(case.glob(f"{name}-*.png")) |
| if selected: |
| candidates.append(selected[-1]) |
| labels.append(label) |
| if candidates: |
| crop = (1024, 384, 1536, 896) if case.name == "portrait" else None |
| make_comparison( |
| reference, candidates, output / case.name, labels=labels, crop=crop |
| ) |
| for file in files: |
| record = json.loads(file.with_suffix(".json").read_text()) |
| record["case"] = case.name |
| record["profile"] = file.stem.rsplit("-", 1)[0] |
| record["trial"] = int(file.stem.rsplit("-", 1)[1]) |
| records.append(record) |
| target = ( |
| Path("reports/inference") / case.name / file.with_suffix(".json").name |
| ) |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(file.with_suffix(".json"), target) |
| if a.cuda_reference_root: |
| source = Path(a.cuda_reference_root) |
| image = source / f"{case.name}.png" |
| if image.exists(): |
| rows = [ |
| json.loads(line) |
| for line in (source / "results.jsonl").read_text().splitlines() |
| ] |
| original = next( |
| r for r in rows if r["case"] == case.name and r["trial"] == 1 |
| ) |
| if hashlib.sha256(image.read_bytes()).hexdigest() != original["sha256"]: |
| raise ValueError(f"CUDA source hash mismatch: {image}") |
| native_record = json.loads(reference.with_suffix(".json").read_text()) |
| if tuple(map(int, original["size"].split("x"))) != tuple( |
| native_record["size"] |
| ): |
| continue |
| target = output / "cuda-reference" / f"{case.name}.png" |
| target.parent.mkdir(parents=True, exist_ok=True) |
| shutil.copyfile(image, target) |
| r = { |
| k: native_record[k] |
| for k in ( |
| "prompt", |
| "seed", |
| "size", |
| "canvas", |
| "rng", |
| "sigmas", |
| "steps", |
| ) |
| } |
| |
| r.update( |
| output=str(target), |
| timings={"total_s": original["wall_seconds"]}, |
| environment={ |
| "device": "RTX 5090", |
| "backend": "SGLang CUDA", |
| "quality": original["quality"], |
| }, |
| source_record=original, |
| reference_images=[], |
| reference_sha256=[], |
| ) |
| r["dit"] = {"precision": "BF16"} |
| r["vae"] = {"quality": original["quality"]} |
| target.with_suffix(".json").write_text(json.dumps(r, indent=2) + "\n") |
| make_comparison( |
| target, |
| [reference], |
| output / f"{case.name}-cuda", |
| labels=[ |
| "CUDA BF16 / upstream lossless quality", |
| "MLX BF16 / native tile recipe", |
| ], |
| ) |
| Path("reports").mkdir(exist_ok=True) |
| Path("reports/image-trials.json").write_text(json.dumps(records, indent=2) + "\n") |
| plot(records, output) |
|
|
|
|
| def plot(records, output): |
| import matplotlib |
|
|
| matplotlib.use("Agg") |
| import matplotlib.pyplot as plt |
| from PIL import Image |
|
|
| from h3_mlx.compare import image_metrics |
|
|
| product = [ |
| r for r in records if r["case"] == "product" and r.get("dit_was_resident") |
| ] |
| if not product: |
| return |
| order = [n for n in LABELS if any(r["profile"] == n for r in product)] |
| selected = [ |
| max((r for r in product if r["profile"] == n), key=lambda r: r["trial"]) |
| for n in order |
| ] |
| baseline = max( |
| (r for r in records if r["case"] == "product" and r["profile"] == "native"), |
| key=lambda r: r["trial"], |
| ) |
| ref = Image.open(baseline["output"]) |
| scores = [image_metrics(ref, Image.open(r["output"])) for r in selected] |
| summary = [] |
| for row, score in zip(selected, scores, strict=True): |
| summary.append( |
| { |
| "profile": row["profile"], |
| "passes": len(row["steps"]), |
| "denoise_s": sum( |
| v for k, v in row["timings"].items() if k.startswith("step_") |
| ), |
| "fresh_end_to_end_s": row["timings"].get("total_s"), |
| "peak_mlx_gib": row["memory"]["peak_gb"], |
| **score, |
| } |
| ) |
| Path("reports/tradeoffs.json").write_text(json.dumps(summary, indent=2) + "\n") |
| table = [ |
| "| Profile | Passes | DiT only | Fresh end-to-end | Recorded peak MLX | RGB SSIM | RGB MAE / 255 |", |
| "|---|---:|---:|---:|---:|---:|---:|", |
| ] |
| for row in summary: |
| e2e = ( |
| "—" |
| if row["fresh_end_to_end_s"] is None |
| else f"{row['fresh_end_to_end_s']:.2f} s" |
| ) |
| table.append( |
| f"| {row['profile']} | {row['passes']} | {row['denoise_s']:.2f} s | {e2e} | {row['peak_mlx_gib']:.2f} GiB | {row['ssim_rgb']:.4f} | {row['mae_8bit']:.3f} |" |
| ) |
| result_path = Path("reports/RESULTS.md") |
| if result_path.exists(): |
| content = result_path.read_text() |
| begin, end = "<!-- PRODUCT_TABLE_START -->", "<!-- PRODUCT_TABLE_END -->" |
| if begin in content and end in content: |
| before, rest = content.split(begin, 1) |
| _, after = rest.split(end, 1) |
| result_path.write_text( |
| before + begin + "\n\n" + "\n".join(table) + "\n\n" + end + after |
| ) |
| fig, axes = plt.subplots(1, 3, figsize=(14, 5), layout="constrained") |
| names = [r["profile"] for r in selected] |
| values = [ |
| [ |
| sum(v for k, v in r["timings"].items() if k.startswith("step_")) |
| for r in selected |
| ], |
| [r["memory"]["peak_gb"] for r in selected], |
| [m["ssim_rgb"] for m in scores], |
| ] |
| colors = [ |
| "#456dba" if n in ("native", "tuned", "fused") else "#d28c43" for n in names |
| ] |
| for ax, vs, title in zip( |
| axes, |
| values, |
| [ |
| "DiT denoising only: seconds", |
| "Recorded peak MLX allocation: GiB", |
| "RGB SSIM to MLX native", |
| ], |
| strict=True, |
| ): |
| bars = ax.bar(names, vs, color=colors) |
| ax.set_title(title) |
| ax.tick_params(axis="x", rotation=45) |
| ax.spines[["top", "right"]].set_visible(False) |
| ax.bar_label(bars, fmt="%.3g", padding=3, fontsize=9) |
| ax.set_ylim(0, max(vs) * 1.2) |
| fig.suptitle( |
| "M5 Pro 64 GB · 2048×2048 product · recorded DiT stages, corrected decoder for every image", |
| fontsize=12, |
| ) |
| fig.savefig(Path("reports") / "tradeoffs.png", dpi=160) |
| plt.close(fig) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
