import gradio as gr
import numpy as np
import random
import torch
from diffusers import DiffusionPipeline
import imageio
from PIL import Image

device = "cpu"

pipe = DiffusionPipeline.from_pretrained(
    "stabilityai/sdxl-turbo",
    torch_dtype=torch.float32
).to(device)

def generate_animation(prompt):
    image = pipe(
        prompt=prompt,
        num_inference_steps=2,
        guidance_scale=0.0
    ).images[0]


    frames = []
    for i in range(10):
        scale = 1 + (i * 0.03)
        w, h = image.size
        resized = image.resize((int(w*scale), int(h*scale)))
        
        # crop center
        left = (resized.width - w)//2
        top = (resized.height - h)//2
        frame = resized.crop((left, top, left+w, top+h))
        
        frames.append(frame)

    gif_path = "animation.gif"
    imageio.mimsave(gif_path, frames, fps=5)

    return gif_path


with gr.Blocks() as demo:
    gr.Markdown("# 🎬 CPU AI Animation (Lightweight)")

    prompt = gr.Textbox(label="Enter Prompt")

    output = gr.Image(type="filepath")

    btn = gr.Button("Generate Animation")

    btn.click(generate_animation, inputs=prompt, outputs=output)

demo.launch()