import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

def generate_image(prompt):
    device = "cuda" if torch.cuda.is_available() else "cpu"
    model_id = "CompVis/stable-diffusion-v1-4"
    pipe = StableDiffusionPipeline.from_pretrained(model_id).to(device)
    
    with torch.no_grad():
        image = pipe(prompt).images[0]
    
    return image

# Define the Gradio interface
interface = gr.Interface(
    fn=generate_image, 
    inputs="text", 
    outputs="image",
    title="Stable Diffusion Image Generator",
    description="Generate images from text prompts using Stable Diffusion."
)

# Launch the app
if __name__ == "__main__":
    interface.launch()