File size: 2,308 Bytes
bc3209b | 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 | #!/usr/bin/env python3
"""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, # Turbo: no CFG, one transformer forward per step
generator=gen,
callback_on_step_end=on_step_end,
).images[0]
t_end = time.time()
# sampling = time spent between pipeline start and the last denoising step
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())
|