import spaces
import gradio as gr
import torch
from diffusers import StableDiffusionPipeline


model_id = "CompVis/stable-diffusion-v1-4"
device = "cuda"


pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe = pipe.to(device)


@spaces.GPU
# Prediction function
def generate_image(prompt):

    image = pipe(prompt).images[0]  
    
    return image
    


demo = gr.Interface(
    fn=generate_image, 
    inputs=gr.Textbox(label="Enter Prompt"),
    outputs=gr.Image(label="Results"), 
    title="generate Images",
    allow_flagging="never"
)


# Launch the interface
demo.launch()