| |
| """Generate one image with the OpenVINO INT4 Z-Image-Turbo pipeline (optimum-intel) and print |
| timing lines in the same format as sd-cli, so SmartTest.sh can parse both. |
| |
| usage: ov_bench.py --model-dir <dir> -p <prompt> -o <out.png> [-W 512] [-H 512] [--steps 8] [--seed 42] |
| |
| Printed lines (stderr): |
| [I] load completed in <s> |
| [I] sampling completed, taking <s> s (sum of per-step wall times, from callback_on_step_end) |
| [I] generate_image completed in <s> s (full pipeline call: text encode + denoise + VAE decode) |
| """ |
| import argparse |
| import sys |
| import time |
|
|
| import torch |
| from optimum.intel import OVDiffusionPipeline |
|
|
|
|
| def main() -> int: |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--model-dir", required=True) |
| ap.add_argument("-p", "--prompt", required=True) |
| ap.add_argument("-o", "--output", required=True) |
| ap.add_argument("-W", "--width", type=int, default=512) |
| ap.add_argument("-H", "--height", type=int, default=512) |
| ap.add_argument("--steps", type=int, default=8) |
| ap.add_argument("--seed", type=int, default=42) |
| args = ap.parse_args() |
|
|
| t0 = time.time() |
| pipe = OVDiffusionPipeline.from_pretrained(args.model_dir, compile=True) |
| print(f"[I] load completed in {time.time() - t0:.2f} s", file=sys.stderr) |
|
|
| step_marks = [] |
|
|
| def on_step_end(pipeline, step, timestep, callback_kwargs): |
| step_marks.append(time.time()) |
| return callback_kwargs |
|
|
| gen = torch.Generator().manual_seed(args.seed) |
| t_gen = time.time() |
| image = pipe( |
| prompt=args.prompt, |
| height=args.height, |
| width=args.width, |
| num_inference_steps=args.steps, |
| guidance_scale=0.0, |
| generator=gen, |
| callback_on_step_end=on_step_end, |
| ).images[0] |
| t_end = time.time() |
|
|
| |
| sampling = (step_marks[-1] - t_gen) if step_marks else float("nan") |
| print(f"[I] sampling completed, taking {sampling:.2f} s", file=sys.stderr) |
| print(f"[I] generate_image completed in {t_end - t_gen:.2f} s", file=sys.stderr) |
|
|
| image.save(args.output) |
| print(f"[I] saved {args.output}", file=sys.stderr) |
| return 0 |
|
|
|
|
| if __name__ == "__main__": |
| sys.exit(main()) |
|
|
