import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Load Stable Diffusion model (we'll run it on CPU)
pipe = StableDiffusionPipeline.from_pretrained("runwayml/stable-diffusion-v1-5")
pipe.to("cpu")

# Inference function to generate an image from text prompt
def generate_image(prompt):
    image = pipe(prompt).images[0]
    return image

# Gradio interface for text-to-image generation
iface = gr.Interface(
    fn=generate_image,
    inputs="text",  # User will input a text prompt
    outputs="image",  # The output will be an image
    title="Stable Diffusion Text-to-Image",
    description="Generate images from text using Stable Diffusion 1.5",
)

# Launch the interface
iface.launch()