import torch
from diffusers import StableDiffusionPipeline
import imageio
import os
import gradio as gr

# Ensure the output directory exists
OUTPUT_FOLDER = "generated_frames"
os.makedirs(OUTPUT_FOLDER, exist_ok=True)

# Load Stable Diffusion Model
print("Loading Stable Diffusion model...")
pipe = StableDiffusionPipeline.from_pretrained("CompVis/stable-diffusion-v1-4")
pipe.to("cpu")  # Force to run on CPU

def generate_images(prompt, num_images=5):
    """Generates images from a prompt and creates an animated GIF."""
    
    images = []
    for i in range(num_images):
        print(f"Generating image {i+1}/{num_images}...")
        image = pipe(prompt).images[0]  # Generate image
        image_path = os.path.join(OUTPUT_FOLDER, f"frame_{i}.png")
        image.save(image_path)
        images.append(imageio.v3.imread(image_path))  # Use v3 API

    # Create GIF animation
    gif_path = os.path.join(OUTPUT_FOLDER, "animation.gif")
    imageio.mimsave(gif_path, images, duration=0.5)
    
    return gif_path  # Return GIF path for download

# Gradio UI
with gr.Blocks() as demo:
    gr.Markdown("# Stable Diffusion Animation Generator 🎨✨")
    
    with gr.Row():
        prompt_input = gr.Textbox(label="Enter your prompt")
        generate_button = gr.Button("Generate Animation 🎥")
    
    gif_output = gr.File(label="Download your GIF")

    generate_button.click(fn=generate_images, inputs=prompt_input, outputs=gif_output)

# Launch Gradio App
if __name__ == "__main__":
    demo.launch()