"""INT4 single-image inference for Qwen-Image-2.1 OpenVINO INT4.

Original model recommendation (https://huggingface.co/Qwen/Qwen-Image-2.1):
  QwenImage21Pipeline, num_inference_steps=40, true_cfg_scale=1.0 (no guidance),
  width/height per aspect ratio (task fixes 1024x1024).
"""
import torch
from optimum.intel import OVDiffusionPipeline


def generate_image(
    prompt: str,
    model_dir: str = "HelloSun/Qwen-Image-2.1-OpenVINO-INT4",
    num_inference_steps: int = 40,
    height: int = 1024,
    width: int = 1024,
    seed: int = 42,
    output_path: str = None,
):
    print(f"Loading model from {model_dir}...")
    pipeline = OVDiffusionPipeline.from_pretrained(model_dir, compile=True)
    generator = torch.Generator().manual_seed(seed)
    print(f"Generating: {prompt[:80]}...")
    image = pipeline(
        prompt=prompt,
        num_inference_steps=num_inference_steps,
        height=height,
        width=width,
        generator=generator,
    ).images[0]
    if output_path:
        image.save(output_path)
        print(f"Saved to {output_path}")
    return image


if __name__ == "__main__":
    import argparse
    parser = argparse.ArgumentParser(description="Qwen-Image-2.1 INT4 Inference")
    parser.add_argument("--prompt", type=str, required=True)
    parser.add_argument("--model-dir", type=str,
                        default="HelloSun/Qwen-Image-2.1-OpenVINO-INT4")
    parser.add_argument("--steps", type=int, default=40)
    parser.add_argument("--height", type=int, default=1024)
    parser.add_argument("--width", type=int, default=1024)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--output", type=str, default=None)
    args = parser.parse_args()
    generate_image(args.prompt, args.model_dir, args.steps,
                   args.height, args.width, args.seed, args.output)