# app.py

import uuid
import gc
from pathlib import Path

import numpy as np
import torch
import gradio as gr
import spaces
import os

from diffusers import (
    WanImageToVideoPipeline,
    WanTransformer3DModel,
    FlowMatchEulerDiscreteScheduler,
    AutoencoderKLWan,
)
from diffusers.utils import export_to_video

from copyright_classifier import contains_copyrighted_ip

BASE_MODEL_ID = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
I2V_TRANSFORMER_REPO = "aidealab/AnimeGen-I2V"

OUTPUT_DIR = Path("outputs")
OUTPUT_DIR.mkdir(exist_ok=True)

token = os.getenv("HF_TOKEN")
NG_WORDS=eval(os.getenv("NG_WORD"))
NG_WORDS.extend(eval(os.getenv("NG_WORD_JA")))

def clear_memory():
    gc.collect()

    if torch.cuda.is_available():
        torch.cuda.empty_cache()
        torch.cuda.ipc_collect()


def load_pipeline():
    scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)

    transformer_high = WanTransformer3DModel.from_pretrained(
        I2V_TRANSFORMER_REPO,
        subfolder="transformer",
        torch_dtype=torch.bfloat16,
    )

    transformer_low = WanTransformer3DModel.from_pretrained(
        I2V_TRANSFORMER_REPO,
        subfolder="transformer_2",
        torch_dtype=torch.bfloat16,
    )

    vae = AutoencoderKLWan.from_pretrained(
        BASE_MODEL_ID,
        subfolder="vae",
        torch_dtype=torch.float32,
    )

    pipe = WanImageToVideoPipeline.from_pretrained(
        BASE_MODEL_ID,
        transformer=transformer_high,
        transformer_2=transformer_low,
        scheduler=scheduler,
        vae=vae,
        torch_dtype=torch.bfloat16,
    )

    pipe.load_lora_weights(
        "lightx2v/Wan2.2-Lightning",
        weight_name=(
            "Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/"
            "high_noise_model.safetensors"
        ),
        adapter_name="high",
    )

    pipe.load_lora_weights(
        "lightx2v/Wan2.2-Lightning",
        weight_name=(
            "Wan2.2-I2V-A14B-4steps-lora-rank64-Seko-V1/"
            "low_noise_model.safetensors"
        ),
        adapter_name="low",
        load_into_transformer_2=True,
    )

    pipe.set_adapters(
        ["high", "low"],
        adapter_weights=[1.0, 1.0],
    )

    transformer_high.enable_layerwise_casting(
        storage_dtype=torch.float8_e4m3fn,
        compute_dtype=torch.bfloat16,
    )

    transformer_low.enable_layerwise_casting(
        storage_dtype=torch.float8_e4m3fn,
        compute_dtype=torch.bfloat16,
    )

    pipe.enable_model_cpu_offload()

    return pipe


def resize_first_and_last_images(first_image, last_image, max_area):
    aspect_ratio = first_image.height / first_image.width

    mod_value = pipe.vae_scale_factor_spatial * pipe.transformer.config.patch_size[1]

    height = round(np.sqrt(max_area * aspect_ratio))
    width = round(np.sqrt(max_area / aspect_ratio))

    height = max(mod_value, height // mod_value * mod_value)
    width = max(mod_value, width // mod_value * mod_value)

    first_image = first_image.resize((width, height))
    last_image = last_image.resize((width, height))

    return first_image, last_image, width, height


# 起動時にロード
pipe = load_pipeline()


@spaces.GPU(duration=120)
def generate_video(
    first_image,
    last_image,
    prompt,
    negative_prompt,
):
    if first_image is None:
        raise gr.Error("Please upload the first image.")

    if last_image is None:
        raise gr.Error("Please upload the last image.")

    clear_memory()

    max_area = 832*480

    first_image, last_image, width, height = resize_first_and_last_images(
        first_image,
        last_image,
        max_area,
    )

    num_frames = int(16 * 3 + 1)

    full_prompt = "Japanese anime style, " + prompt.strip()

    # prompt filtering
    print(NG_WORDS)
    prompt_for_check = full_prompt.lower()
    for word in NG_WORDS:
        if word in prompt_for_check:
            print(f"error: {word} .")
            raise Exception()
    
    # LLM filtering
    result = contains_copyrighted_ip(full_prompt)
    if(result):
        print(f"error: {full_prompt} .")
        raise Exception()
        
    frames = pipe(
        image=first_image,
        last_image=last_image,
        prompt=full_prompt,
        negative_prompt=negative_prompt,
        height=height,
        width=width,
        num_frames=num_frames,
        guidance_scale=1.0,
        num_inference_steps=4,
    ).frames[0]

    output_path = OUTPUT_DIR / f"{uuid.uuid4().hex}.mp4"
    export_to_video(frames, str(output_path), fps=16)

    clear_memory()

    return str(output_path)

default_prompt = "Make this character press both hands together in prayer and close their eyes."
default_negative_prompt = "3d, cg, photo, stop, wait"


with gr.Blocks(title="AnimeGen Frame Interpolation") as demo:
    gr.Markdown("# AnimeGen Frame Interpolation")
    gr.Markdown(
        "Generate a video from a first frame and a last frame using AnimeGen I2V."
    )

    with gr.Row():
        with gr.Column(scale=1):
            with gr.Row():
                first_image = gr.Image(
                    label="First image",
                    type="pil",
                )

                last_image = gr.Image(
                    label="Last image",
                    type="pil",
                )

            prompt = gr.Textbox(
                label="Prompt",
                value=default_prompt,
                lines=4,
            )

            negative_prompt = gr.Textbox(
                label="Negative prompt",
                value=default_negative_prompt,
                lines=2,
            )

            generate_button = gr.Button("Generate", variant="primary")

        with gr.Column(scale=1):
            video = gr.Video(label="Output")
            
    generate_button.click(
        fn=generate_video,
        inputs=[
            first_image,
            last_image,
            prompt,
            negative_prompt,
        ],
        outputs=[
            video,
        ],
    )


if __name__ == "__main__":
    demo.queue(max_size=10).launch()