import gradio as gr
import torch
from diffusers import StableDiffusionXLPipeline, DPMSolverSinglestepScheduler
from PIL import Image

# Load the Stable Diffusion XL model
def load_model():
    pipe = StableDiffusionXLPipeline.from_pretrained(
        "sd-community/sdxl-flash",
        torch_dtype=torch.float16,
        low_cpu_mem_usage=True
    )#.to("cuda")
    pipe.scheduler = DPMSolverSinglestepScheduler.from_config(
        pipe.scheduler.config,
        timestep_spacing="trailing"
    )
    return pipe

# Function to generate images
def generate_image(prompt, num_inference_steps=3, guidance_scale=1):
    pipe = load_model()
    image = pipe(prompt, num_inference_steps=num_inference_steps, guidance_scale=guidance_scale).images[0]
    return image

# Create the Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=[
        gr.Textbox(label="Enter a prompt for the image"),
        # gr.Slider(minimum=1, maximum=50, label="Number of Inference Steps"),
        # gr.Slider(minimum=1, maximum=20, label="Guidance Scale")
    ],
    outputs="image",
    title="Stable Diffusion XL Text-to-Image",
    description="Generate images from text prompts using Stable Diffusion XL."
)

# Launch the Gradio app
iface.launch()