generate_image.py

from diffusers import StableDiffusionXLPipeline
import torch
from PIL import Image
import os

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float16,
    variant="fp16"
).to("cuda" if torch.cuda.is_available() else "cpu")
pipe.safety_checker = None

def generate_scene_image(prompt, idx):
    image = pipe(prompt).images[0]
    out_path = f"assets/images/scene_{idx}.png"
    image.save(out_path)
    return out_path