#!/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())