import sys
from diffusers import StableDiffusionXLPipeline
import torch

# Command-line args
prompt = sys.argv[1] if len(sys.argv) > 1 else "a futuristic city at sunset"
filename = sys.argv[2] if len(sys.argv) > 2 else "output.png"

# Load model
pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16,
    variant="fp16"
).to("mps")  # Use Apple Metal backend

pipe.set_progress_bar_config(disable=False)

# Generate image
image = pipe(prompt=prompt, guidance_scale=0.0, num_inference_steps=4).images[0]
image.save(filename)
print(f"Saved image to {filename}")