File size: 1,489 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 | """Fresh end-to-end timings after the final decoder arithmetic audit."""
from __future__ import annotations
import argparse
import json
import logging
from pathlib import Path
from h3_mlx.runtime import Pipeline
def main():
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--model", required=True)
p.add_argument("--text-encoder", required=True)
a = p.parse_args()
logging.basicConfig(
level=logging.INFO, format="%(asctime)s %(message)s", datefmt="%H:%M:%S"
)
case = next(
c
for c in json.loads(Path("experiments/cases.json").read_text())
if c["id"] == "product"
)
pipe = Pipeline(a.model, a.text_encoder)
try:
for name, trial, kernel, batch in [
("native", 2, "native", 1),
("native", 3, "native", 1),
("fused", 2, "metal-qk", 8),
]:
pipe.dit_options.update(kernel=kernel)
pipe.vae_options.update(tile_batch=batch)
if pipe.model is not None:
pipe.model.kernel = kernel
if pipe.vae is not None:
pipe.vae.tile_batch = batch
pipe.generate(
case["prompt"],
width=2048,
height=2048,
seed=case["seed"],
output=f"outputs/suite/product/{name}-{trial}.png",
save_trace=True,
)
finally:
pipe.clear()
if __name__ == "__main__":
main()
|