import gradio as gr
from diffusers import StableDiffusionPipeline
from PIL import Image
import torch

# Load the Stable Diffusion model
model = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4", torch_dtype=torch.float32)
model = model.to("cuda" if torch.cuda.is_available() else "cpu")

def generate_image(prompt):
    # Generate image from the model
    with torch.no_grad():
        image = model(prompt).images[0]
    # Convert to PIL Image to display in Gradio
    #image = Image.fromarray(image.numpy())
    return image

# Create a Gradio interface
interface = gr.Interface(fn=generate_image, inputs='text', outputs='image')

# Launch the interface
interface.launch()