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, negative_prompt, num_inference_steps, guidance_scale):
    # Run the diffusion model to generate an image
    output = pipe(prompt, negative_prompt, num_inference_steps=50, guidance_scale=7.5)
    return output.images[0]

num_inference_steps=gr.Slider(10, 100, value=50, label="Choose Number of Inference Steps")
guidance_scale=gr.Slider(1, 10, value=7.5, label="Choose Guidance Scale")
prompt = gr.Textbox(label = "Prompt", info = "Describe the subject, the background and the style of image; 77 token limit", placeholder = "Describe what you want to see", lines = 2)
negative_prompt = gr.Textbox(label = "Negative prompt", placeholder = "Describe what you do NOT want to see", value = "Ugly, malformed, noise, blur, watermark")

gr_interface = gr.Interface(
    fn=generate_image,
    inputs=[prompt, negative_prompt],
    additional_inputs=[num_inference_steps, guidance_scale],
    outputs="image",
    examples=[["Astronaut riding a horse on the moon", "Bad quality, watermark"], ["Jungle landscape, photo", "Bad quality, watermark"], ["A woman near gold car", "Bad quality, unrealistic"]],
    title="Real-time Image Generation with Diffusion",
    description="Enter a prompt to generate an image",
    theme="soft"
)

# Launch the Gradio app
gr_interface.launch()