| """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() |
|
|
