#!/usr/bin/env python3
"""Generate an image with a local SpaceV1.7 Image / Diffusers checkpoint."""
from __future__ import annotations

import argparse
from pathlib import Path


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", default="SpaceV1.7-Image-model", help="local folder or Hugging Face model ID")
    parser.add_argument("--prompt", required=True)
    parser.add_argument("--negative-prompt", default=None)
    parser.add_argument("--output", default="spacev17-image.png")
    parser.add_argument("--steps", type=int, default=20)
    parser.add_argument("--guidance-scale", type=float, default=7.5)
    parser.add_argument("--seed", type=int, default=42)
    parser.add_argument("--device", choices=("cuda", "cpu"), default="cuda")
    args = parser.parse_args()

    import torch
    from diffusers import DiffusionPipeline

    if args.steps < 1:
        parser.error("--steps must be positive")
    if args.device == "cuda" and not torch.cuda.is_available():
        parser.error("CUDA was requested but no CUDA device is available; use --device cpu")
    dtype = torch.float16 if args.device == "cuda" else torch.float32
    pipe = DiffusionPipeline.from_pretrained(args.model, torch_dtype=dtype)
    pipe = pipe.to(args.device)
    generator = torch.Generator(device=args.device).manual_seed(args.seed)
    result = pipe(
        args.prompt,
        negative_prompt=args.negative_prompt,
        num_inference_steps=args.steps,
        guidance_scale=args.guidance_scale,
        generator=generator,
    )
    output = Path(args.output)
    output.parent.mkdir(parents=True, exist_ok=True)
    result.images[0].save(output)
    print(f"Saved image: {output}")
    return 0


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