File size: 1,931 Bytes
63e9703 | 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 | """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()
|