import gradio as gr
from PIL import Image
from diffusers import DiffusionPipeline
import time

# Load model and scheduler
ldm = DiffusionPipeline.from_pretrained("CompVis/ldm-text2im-large-256")

def generate_image(prompt, negative_prompt="Low quality", width=512, height=512):
    # Run pipeline in inference (sample random noise and denoise)
    start_time = time.time()
    images = ldm([prompt], num_inference_steps=50, eta=0.3, guidance_scale=6, negative_prompts=[negative_prompt]).images
    # Resize image to desired width and height
    resized_images = [image.resize((int(width), int(height))) for image in images]
    # Save images
    for idx, image in enumerate(resized_images):
        image.save(f"squirrel-{idx}.png")
    end_time = time.time()
    elapsed_time = round(end_time - start_time, 2)
    return f"Images generated successfully! (Elapsed Time: {elapsed_time} seconds)"

# Define the interface
iface = gr.Interface(
    fn=generate_image,
    inputs=["text", "text", "number", "number"],
    outputs=gr.outputs.Image(label="Generated Image"),
    layout="vertical",
    title="Image Generation",
    description="Generate images based on prompts.",
    article="For more information, visit the documentation: [link](https://docs.gradio.app/)",
    examples=[["A painting of a squirrel eating a burger", "Low quality", 512, 512]]
)

# Launch the interface
iface.launch()