import gradio as gr
import spaces
from diffusers import DiffusionPipeline

pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo")

@spaces.GPU(duration=250)
def generate_image(prompt):
    # Run the diffusion model to generate an image
    output = pipe(prompt, num_inference_steps=50, guidance_scale=7.5)
    image_tensor = output.images[0]
    image_array = image_tensor.cpu().numpy().transpose(1, 2, 0)
    image_array = (image_array * 255).astype(np.uint8)
    image = Image.fromarray(image_array)
    return image

gr_interface = gr.Interface(
    fn=generate_image,
    inputs="text",
    outputs="image",
    title="Real-time Image Generation with Diffusion",
    description="Enter a prompt to generate an image",
    theme="soft"
)

# Launch the Gradio app
gr_interface.launch()