#!/usr/bin/env python3
"""Generate one image with the INT4 OpenVINO build of stable-diffusion-3.5-medium.

Inference parameters follow the ORIGINAL model card, not the FLUX.2 reference:

    stabilityai/stable-diffusion-3.5-medium
      1024x1024, num_inference_steps=40, guidance_scale=4.5

(FLUX.2-klein-4B-OpenVINO-INT4 uses 4 steps / guidance 1.0 because FLUX.2 is a
4-step distilled model. SD3.5-medium is a full flow-matching model, so the step
count is 10x higher. The scheduler shift of 3.0 comes from the repo's
scheduler/scheduler_config.json and is picked up automatically.)

Usage
-----
    python inference_int4.py --prompt "Astronaut in a jungle" --seed 43 --output out.png
    python inference_int4.py --prompt "..." --steps 28        # faster, lower quality
    python inference_int4.py --prompt "..." --resize 512      # cheap preview
"""

from __future__ import annotations

import argparse
import os
import time
from pathlib import Path

REPO_ID = "HelloSun/SD3.5-medium-OpenVINO-INT4"

# Values from the original model card.
DEFAULT_STEPS = 40
DEFAULT_GUIDANCE = 4.5
DEFAULT_SIZE = 1024


def load_pipeline(repo_id: str, device: str, compile_model: bool):
    import torch
    from optimum.intel import OVDiffusionPipeline

    print(f"loading {repo_id} ...", flush=True)
    t0 = time.perf_counter()
    pipe = OVDiffusionPipeline.from_pretrained(repo_id, device=device, compile=compile_model)
    print(f"load+compile {time.perf_counter() - t0:.2f}s", flush=True)
    return pipe, torch


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--negative-prompt", default="")
    parser.add_argument("--model", default=REPO_ID, help="local path or Hub id")
    parser.add_argument("--output", type=Path, default=Path("out.png"))
    parser.add_argument("--steps", type=int, default=DEFAULT_STEPS)
    parser.add_argument("--guidance", type=float, default=DEFAULT_GUIDANCE)
    parser.add_argument("--width", type=int, default=DEFAULT_SIZE)
    parser.add_argument("--height", type=int, default=DEFAULT_SIZE)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--stepsize", type=int, default=8, help="scheduler shift; 8 is common for SD3 speed/quality")
    parser.add_argument("--resize", type=int, default=None, help="also save a downscaled copy at this size")
    parser.add_argument("--device", default="CPU")
    parser.add_argument("--no-compile", action="store_true", help="skip OV graph compilation")
    parser.add_argument("--num-images", type=int, default=1)
    args = parser.parse_args()

    pipe, torch = load_pipeline(args.model, args.device, not args.no_compile)

    args.output.parent.mkdir(parents=True, exist_ok=True)
    started = time.perf_counter()
    image = pipe(
        prompt=args.prompt,
        negative_prompt=args.negative_prompt,
        num_inference_steps=args.steps,
        guidance_scale=args.guidance,
        width=args.width,
        height=args.height,
        num_images_per_prompt=args.num_images,
        generator=torch.Generator().manual_seed(args.seed),
    ).images[0]
    elapsed = time.perf_counter() - started

    image.save(args.output)
    print(f"saved {args.output}  {image.size[0]}x{image.size[1]}  {elapsed:.2f}s "
          f"({elapsed / max(args.steps, 1):.2f}s/step)", flush=True)

    if args.resize:
        preview = args.output.with_name(f"{args.output.stem}_{args.resize}.png")
        try:
            from PIL import Image as _Image

            _resample = _Image.Resampling.LANCZOS  # Pillow >= 9.1
        except (ImportError, AttributeError):
            _resample = image.LANCZOS
        image.resize((args.resize, args.resize), _resample).save(preview)
        print(f"saved {preview}", flush=True)

    peak_mb = 0.0
    try:
        # Kernel high-water mark: exact and dependency-free on Linux.
        for line in Path("/proc/self/status").read_text().splitlines():
            if line.startswith("VmHWM:"):
                peak_mb = int(line.split()[1]) / 1024
                break
        else:
            import psutil

            peak_mb = psutil.Process(os.getpid()).memory_info().rss / 1024**2
    except (OSError, ValueError, IndexError, ImportError):
        pass
    if peak_mb:
        print(f"peak RSS: {peak_mb:.0f} MB", flush=True)
    return 0


if __name__ == "__main__":
    raise SystemExit(main())