import gradio as gr
from diffusers import StableDiffusionPipeline
import torch

# Load Stable Diffusion pipeline (make sure to use the right model name)
model_name = "CompVis/stable-diffusion-v-1-4-original"
pipe = StableDiffusionPipeline.from_pretrained(model_name, torch_dtype=torch.float16)
pipe.to("cuda")  # Move the model to GPU (if available)

def generate_image(description):
    # Generate an image based on the description
    image = pipe(description).images[0]
    return image

# Set up the Gradio interface
iface = gr.Interface(fn=generate_image, 
                     inputs=gr.Textbox(lines=2, placeholder="Enter the description here...", label="Description"),
                     outputs=gr.Image(type="pil"),
                     title="Traffic and People Walking Image Generator",
                     description="Provide a description of the traffic and people walking scene to generate an image.",
                     live=True)

# Launch the app
iface.launch()