File size: 4,391 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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 | """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()
|