import gradio as gr
import torch
from diffusers import AnimateDiffPipeline, MotionAdapter, DDIMScheduler
from diffusers.utils import export_to_gif
import random

def generate_gif(image, animation_type, prompt, adapter_strength):
    # Load the motion adapter
    adapter = MotionAdapter.from_pretrained("guoyww/animatediff-motion-adapter-v1-5-2", torch_dtype=torch.float16)

    # Load SD 1.5 based finetuned model
    model_id = "SG161222/Realistic_Vision_V6.0_B1_noVAE"
    pipe = AnimateDiffPipeline.from_pretrained(model_id, motion_adapter=adapter, torch_dtype=torch.float16)

    # Scheduler setup
    scheduler = DDIMScheduler(
        clip_sample=False,
        beta_start=0.00085,
        beta_end=0.012,
        beta_schedule="linear",
        timestep_spacing="trailing",
        steps_offset=1
    )
    pipe.scheduler = scheduler

    # Enable memory savings
    pipe.enable_vae_slicing()
    pipe.enable_model_cpu_offload()

    # Load ip_adapter
    pipe.load_ip_adapter("h94/IP-Adapter", subfolder="models", weight_name="ip-adapter_sd15.bin")

    # Load the selected motion adapter
    pipe.load_lora_weights(f"guoyww/animatediff-motion-lora-{animation_type}", adapter_name=animation_type)

    # Generate a random seed
    seed = random.randint(0, 2**32 - 1)

    # Set adapter weights for the selected adapter
    adapter_weight = [adapter_strength]

    pipe.set_adapters([animation_type], adapter_weights=adapter_weight)

    # Generate GIF
    output = pipe(
        prompt=prompt,
        num_frames=16,
        guidance_scale=7.5,
        num_inference_steps=30,
        ip_adapter_image=image,
        generator=torch.Generator("cpu").manual_seed(seed),
    )
    frames = output.frames[0]

    gif_path = "output_animation.gif"
    export_to_gif(frames, gif_path)
    return gif_path

# Gradio interface
iface = gr.Interface(
    fn=generate_gif,
    inputs=[
        gr.Image(type="pil"), 
        gr.Radio(["zoom-out", "tilt-up", "pan-left"]),
        gr.Textbox(label="Prompt"),
        gr.Slider(minimum=0, maximum=4, step=0.1, value=1, label="IP Adapter Strength")
    ],
    outputs=gr.Image(type="pil", label="Generated GIF"),
    title="AnimateDiff + IP Adapter Demo",
    description="Upload an image, select a motion module type, and adjust the prompt and IP Adapter strength to generate a GIF!"
)

iface.launch(debug=True, share=True)