from pathlib import Path

import torch
from diffusers import AutoPipelineForText2Image

model_name = "stable-diffusion-v1-5/stable-diffusion-v1-5"

device = "cuda" if torch.cuda.is_available() else "cpu"

if device == "cuda":
    pipe = AutoPipelineForText2Image.from_pretrained(
        model_name,
        torch_dtype=torch.float16,
        variant="fp16",
    ).to(device)
else:
    pipe = AutoPipelineForText2Image.from_pretrained(
        model_name,
        torch_dtype=torch.float32,
    ).to(device)

prompt = input("Prompt: ").strip()
if not prompt:
    prompt = "a cute robot in a garden, detailed, cinematic, high quality"

result = pipe(
    prompt=prompt,
    num_inference_steps=30,
    guidance_scale=7.5,
)

image = result.images[0]
output_path = Path("generated_image.png")
image.save(output_path)

print(f"Saved image to: {output_path.resolve()}")