| """Matched image and performance trials. Run from the repository root.""" |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import gc |
| import json |
| import logging |
| import statistics |
| from pathlib import Path |
|
|
| from h3_mlx.runtime import Pipeline |
|
|
| PROFILES = { |
| "native": {}, |
| "fused": {"mlp_chunk": 0, "kernel": "metal-qk", "threads": 256, "vae_batch": 8}, |
| "chunked": {"mlp_chunk": 1024}, |
| "wide-mlp": {"mlp_chunk": 0}, |
| "tuned": {"mlp_chunk": 0, "kernel": "metal", "threads": 256, "vae_batch": 8}, |
| "q8": {"bits": 8, "mlp_chunk": 0, "vae_batch": 8}, |
| "q4": {"bits": 4, "mlp_chunk": 0, "vae_batch": 8}, |
| "q8-g32": {"bits": 8, "group_size": 32, "mlp_chunk": 0, "vae_batch": 8}, |
| "preview": {"mlp_chunk": 0, "kernel": "metal", "threads": 256, "vae_batch": 8}, |
| "stream": {"dit_cache_gb": 0}, |
| } |
|
|
|
|
| def main(): |
| p = argparse.ArgumentParser(description=__doc__) |
| p.add_argument("--model", required=True) |
| p.add_argument("--text-encoder", required=True) |
| p.add_argument( |
| "--profiles", |
| nargs="+", |
| choices=PROFILES, |
| default=["native", "tuned", "q8", "q4", "preview"], |
| ) |
| p.add_argument("--cases", nargs="+", default=["product", "portrait"]) |
| p.add_argument("--case-file", default="experiments/cases.json") |
| p.add_argument("--output", default="outputs/suite") |
| p.add_argument( |
| "--repeat", |
| type=int, |
| default=2, |
| help="First trial includes model loading; later trials reuse resident models", |
| ) |
| p.add_argument("--save-trace", action="store_true") |
| a = p.parse_args() |
| if a.repeat < 1: |
| p.error("repeat must be positive") |
| logging.basicConfig( |
| level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S" |
| ) |
| cases = [c for c in json.loads(Path(a.case_file).read_text()) if c["id"] in a.cases] |
| if {c["id"] for c in cases} != set(a.cases): |
| p.error("Requested case not found in case file") |
| root = Path(a.output) |
| root.mkdir(parents=True, exist_ok=True) |
| ledger = root / "trials.jsonl" |
| for name in a.profiles: |
| logging.info("PROFILE %s", name) |
| pipeline = Pipeline(a.model, a.text_encoder, **PROFILES[name]) |
| try: |
| for c in cases: |
| for trial in range(a.repeat): |
| output = root / c["id"] / f"{name}-{trial}.png" |
| _, row = pipeline.generate( |
| c["prompt"], |
| width=c["width"], |
| height=c["height"], |
| seed=c["seed"], |
| output=output, |
| steps=[0, 2, 4] if name == "preview" else None, |
| save_trace=a.save_trace, |
| images=c.get("images", ()), |
| ) |
| row.update( |
| profile=name, |
| case=c["id"], |
| trial=trial, |
| sample_type="resident" |
| if row["dit_was_resident"] |
| else "load_included", |
| ) |
| with ledger.open("a") as h: |
| h.write(json.dumps(row) + "\n") |
| finally: |
| pipeline.clear() |
| del pipeline |
| gc.collect() |
| rows = [json.loads(line) for line in ledger.read_text().splitlines()] |
| summary = [] |
| for profile in a.profiles: |
| for case in a.cases: |
| group = [r for r in rows if r["profile"] == profile and r["case"] == case] |
| resident = [ |
| r["timings"]["total_s"] for r in group if r["sample_type"] == "resident" |
| ] |
| summary.append( |
| { |
| "profile": profile, |
| "case": case, |
| "samples": len(group), |
| "resident_samples": len(resident), |
| "resident_median_s": statistics.median(resident) |
| if resident |
| else None, |
| "load_included_s": [ |
| r["timings"]["total_s"] |
| for r in group |
| if r["sample_type"] == "load_included" |
| ], |
| } |
| ) |
| (root / "summary.json").write_text(json.dumps(summary, indent=2) + "\n") |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
