import gradio as gr
from diffusers import StableDiffusionPipeline

# Load pre-trained model
model_id = "CompVis/stable-diffusion-v1-4"
pipe = StableDiffusionPipeline.from_pretrained(model_id)

def generate_image(prompt):
    # Generate image from text prompt
    images = pipe(prompt, num_images=1).images
    return images[0]

# Define Gradio UI
iface = gr.Interface(
    fn=generate_image,
    inputs=gr.components.Textbox(lines=2, placeholder="Enter a text prompt here..."),
    outputs=gr.components.Image(),
    title="DreamBooth Image Generation",
    description="Enter a text prompt to generate an image using DreamBooth.",
    examples=[
        ["A cute 2D bear wearing glasses, in a cartoon style"],
        ["An adorable 2D cartoon bear with round glasses"],
        ["A kawaii 2D bear character wearing glasses, with a cute and whimsical style"]
    ]
)

# Launch the UI
iface.launch(share=True)