import gradio as gr
from transformers import StableDiffusionPipeline

def generate_image(prompt):
    model_id = "CompVis/stable-diffusion-v1-4"
    pipe = StableDiffusionPipeline.from_pretrained(model_id)
    pipe = pipe.to("cuda")
    
    image = pipe(prompt).images[0]
    return image

iface = gr.Interface(
    fn=generate_image,
    inputs=gr.inputs.Textbox(lines=2, placeholder="Enter a prompt for the AI to generate an image."),
    outputs="image",
    title="Stable Diffusion Image Generator",
    description="Enter a prompt and the AI will generate an image for you!",
)

if __name__ == "__main__":
    iface.launch()