| """Run a recorded sample through the public inference pipeline.""" |
|
|
| import argparse |
| import json |
| from pathlib import Path |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--model", required=True) |
| parser.add_argument("--text-encoder", required=True) |
| parser.add_argument( |
| "--case", |
| choices=("product", "portrait", "holdout-detail", "edit"), |
| default="product", |
| ) |
| parser.add_argument( |
| "--image", help="Original edit reference, required for the edit case" |
| ) |
| parser.add_argument("--cache-dir", default="cache") |
| parser.add_argument("--output", default="outputs/sample.png") |
| args = parser.parse_args() |
| if args.case == "edit" and args.image is None: |
| parser.error( |
| "The edit sample requires --image pointing to the original reference" |
| ) |
| catalog = json.loads(Path(__file__).with_name("catalog.json").read_text()) |
| case = next(e for e in catalog["examples"] if e["case"] == args.case) |
|
|
| from h3_mlx.runtime import Pipeline |
|
|
| pipe = Pipeline( |
| args.model, |
| args.text_encoder, |
| cache_dir=args.cache_dir, |
| compile_blocks=True, |
| eval_every=0, |
| prefetch_blocks=4, |
| vae_compile=True, |
| vae_batch=4, |
| kernel="metal-qk-exact", |
| tune_matmul=True, |
| ) |
| try: |
| _, record = pipe.generate( |
| case["prompt"], |
| seed=case["seed"], |
| width=case["requested_size"][0], |
| height=case["requested_size"][1], |
| images=[args.image] if args.case == "edit" else [], |
| output=args.output, |
| save_trace=True, |
| ) |
| print( |
| json.dumps( |
| {"output": record["output"], "png_sha256": record["png_sha256"]}, |
| indent=2, |
| ) |
| ) |
| finally: |
| pipe.clear() |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
