"""image_gen.py — Local image generation on Mac via Stable Diffusion XL Turbo.

SDXL Turbo is a distilled SD model that produces a 512x512 image in
~4 denoising steps. On Apple Silicon (MPS) that's about 5-10 seconds
per image on an M2 Pro after the first run.

First run downloads ~6 GB from HuggingFace. After that it works
fully offline, like the rest of the Mac Stack.

Standalone:
    python image_gen.py "a Lanna temple at sunset" out.png

As a library (used by agent_pro.py):
    from image_gen import generate
    generate("a Lanna temple at sunset", "out.png")
"""
import sys
from pathlib import Path


def _load_pipe():
    """Build the diffusers pipeline. Picks MPS on Apple Silicon, CUDA on Linux/Win GPU, else CPU."""
    import torch
    from diffusers import AutoPipelineForText2Image

    if torch.backends.mps.is_available():
        device, dtype, variant = "mps", torch.float16, "fp16"
    elif torch.cuda.is_available():
        device, dtype, variant = "cuda", torch.float16, "fp16"
    else:
        device, dtype, variant = "cpu", torch.float32, None

    pipe = AutoPipelineForText2Image.from_pretrained(
        "stabilityai/sdxl-turbo",
        torch_dtype=dtype,
        variant=variant,
    )
    return pipe.to(device)


_PIPE = None  # cached so a second generate() in the same process is fast


def generate(prompt: str, out_path: str) -> str:
    """Generate an image from prompt and save to out_path. Returns the path."""
    global _PIPE
    if _PIPE is None:
        _PIPE = _load_pipe()
    image = _PIPE(
        prompt=prompt,
        num_inference_steps=4,   # SDXL Turbo only needs ~4 steps
        guidance_scale=0.0,      # Turbo recommends 0.0 (no classifier-free guidance)
    ).images[0]
    Path(out_path).parent.mkdir(parents=True, exist_ok=True)
    image.save(out_path)
    return out_path


if __name__ == "__main__":
    if len(sys.argv) < 3:
        print("usage: python image_gen.py 'prompt' out.png")
        sys.exit(1)
    saved = generate(sys.argv[1], sys.argv[2])
    print(f"Saved to {saved}")