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

# Load model and set the device to CPU
pipe = StableDiffusionXLPipeline.from_pretrained("sd-community/sdxl-flash", torch_dtype=torch.float32).to("cpu")

# Ensure sampler uses "trailing" timesteps
pipe.scheduler = DPMSolverSinglestepScheduler.from_config(pipe.scheduler.config, timestep_spacing="trailing")

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

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.Textbox(label="Enter Prompt", placeholder="Enter a description for the image"),
    outputs=gr.Image(label="Generated Image")
)

# Launch the Gradio interface
iface.launch()