import torch
from diffusers import StableDiffusionPipeline
import gradio as gr

# Load the Stable Diffusion model (text-to-image generator)
model_id = "runwayml/stable-diffusion-v1-5"
pipe = StableDiffusionPipeline.from_pretrained(model_id, torch_dtype=torch.float16)
pipe.to("cpu")  # If you have a GPU, use it to speed up generation

# Function to generate images from text prompts
def generate_image(prompt):
    # Generate image based on the prompt
    image = pipe(prompt).images[0]
    return image

# Create Gradio interface
iface = gr.Interface(
    fn=generate_image,  # Function to run
    inputs=gr.Textbox(lines=2, placeholder="Enter your prompt..."),  # Text input for prompts
    outputs=gr.Image(),  # Display the generated image
    title="Text-to-Image Generator",
    description="Enter a text prompt to generate an image using Stable Diffusion."
)

# Launch the Gradio interface
iface.launch()