import argparse
from pathlib import Path
from time import perf_counter

from diffusers import AutoPipelineForText2Image
from diffusers.utils import logging
import torch

MODEL = "stabilityai/sd-turbo"
REVISION = "b261bac6fd2cf515557d5d0707481eafa0485ec2"
PROMPT = "Editorial still life of a green notebook and a wooden pencil on a cream desk, soft window light"

parser = argparse.ArgumentParser(description="Generate a small image with a local diffusion model.")
parser.add_argument("output", type=Path)
parser.add_argument("--device", choices=("cpu", "mps", "cuda"), default="cpu")
parser.add_argument("--prompt", default=PROMPT)
parser.add_argument("--seed", type=int, default=42)
parser.add_argument("--steps", type=int, choices=(1, 2, 3, 4), default=2)
parser.add_argument("--offline", action="store_true")
args = parser.parse_args()

if args.output.suffix.lower() != ".png":
    parser.error("Use a .png output filename.")
if args.output.exists():
    parser.error("Choose an output filename that does not exist.")
if args.device == "mps" and not torch.backends.mps.is_available():
    parser.error("This PyTorch installation does not have a working MPS device.")
if args.device == "cuda" and not torch.cuda.is_available():
    parser.error("This PyTorch installation does not have a working CUDA device.")

logging.set_verbosity_error()
try:
    pipe = AutoPipelineForText2Image.from_pretrained(
        MODEL,
        revision=REVISION,
        variant="fp16",
        torch_dtype=torch.float32 if args.device == "cpu" else torch.float16,
        use_safetensors=True,
        local_files_only=args.offline,
    ).to(args.device)
    pipe.set_progress_bar_config(disable=True)
    generator = torch.Generator(device="cpu").manual_seed(args.seed)
    start = perf_counter()
    image = pipe(
        args.prompt,
        height=512,
        width=512,
        num_inference_steps=args.steps,
        guidance_scale=0.0,
        generator=generator,
    ).images[0]
    elapsed = perf_counter() - start
    with args.output.open("xb") as output:
        image.save(output, format="PNG")
except (OSError, RuntimeError, ValueError) as error:
    raise SystemExit(f"Could not generate the image: {error}") from None

print(f"Saved {image.width} × {image.height} to {args.output} in {elapsed:.2f} s after loading.")