# Copyright 2025 - Wan2.1 T2V-1.3B multi-GPU Gradio demo
# Runs under `torchrun --nproc_per_node=8` with FSDP (DiT + T5) and xDiT USP (ring).
# NOTE: t2v-1.3B has 12 attention heads; ulysses requires num_heads % ulysses_size == 0,
# so the official 8-GPU config for the 1.3B model is ring_size=8, ulysses_size=1.
import os
import random
import sys
import threading
import time
import warnings

warnings.filterwarnings("ignore")

import torch
import torch.distributed as dist
import gradio as gr

sys.path.insert(0, "/opt/Wan2.1")

import wan
from wan.configs import WAN_CONFIGS
from wan.utils.utils import cache_video
from xfuser.core.distributed import (
    init_distributed_environment,
    initialize_model_parallel,
)

RANK = int(os.getenv("RANK", 0))
WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1))
LOCAL_RANK = int(os.getenv("LOCAL_RANK", 0))
CKPT_DIR = os.environ.get("CKPT_DIR", "/opt/Wan2.1-T2V-1.3B")
ULYSSES_SIZE = 1
RING_SIZE = int(os.environ.get("RING_SIZE", str(WORLD_SIZE)))

wan_t2v = None
_dist_lock = threading.Lock()

EXAMPLE_PROMPT = (
    "Two anthropomorphic cats in comfy boxing gear and bright gloves "
    "fight intensely on a spotlighted stage."
)


def init_distributed():
    torch.cuda.set_device(LOCAL_RANK)
    dist.init_process_group(
        backend="nccl",
        init_method="env://",
        rank=RANK,
        world_size=WORLD_SIZE,
    )
    init_distributed_environment(rank=RANK, world_size=WORLD_SIZE)
    initialize_model_parallel(
        sequence_parallel_degree=WORLD_SIZE,
        ring_degree=RING_SIZE,
        ulysses_degree=ULYSSES_SIZE,
    )


def load_model():
    global wan_t2v
    cfg = WAN_CONFIGS["t2v-1.3B"]
    logging.info(f"[rank {RANK}] Creating WanT2V pipeline (FSDP + USP)")
    wan_t2v = wan.WanT2V(
        config=cfg,
        checkpoint_dir=CKPT_DIR,
        device_id=LOCAL_RANK,
        rank=RANK,
        t5_fsdp=True,
        dit_fsdp=True,
        use_usp=True,
    )


def _distributed_generate(kwargs):
    """Broadcast generation kwargs to all ranks, run the distributed pass."""
    obj = [kwargs] if RANK == 0 else [None]
    dist.broadcast_object_list(obj, src=0)
    kwargs = obj[0]
    video = wan_t2v.generate(**kwargs)
    dist.barrier()
    return video


def generate_video(prompt, resolution, sd_steps, guide_scale, shift_scale, seed, n_prompt):
    """Generate a 5-second 480P video from a text prompt on all 8 GPUs."""
    W = int(resolution.split("*")[0])
    H = int(resolution.split("*")[1])
    seed = int(seed)
    if seed < 0:
        seed = random.randint(0, sys.maxsize)

    kwargs = dict(
        input_prompt=prompt,
        size=(W, H),
        shift=float(shift_scale),
        sampling_steps=int(sd_steps),
        guide_scale=float(guide_scale),
        n_prompt=n_prompt,
        seed=seed,
        offload_model=False,
    )

    with _dist_lock:
        video = _distributed_generate(kwargs)

    if RANK == 0:
        save_file = "/tmp/output.mp4"
        cache_video(
            tensor=video[None],
            save_file=save_file,
            fps=16,
            nrow=1,
            normalize=True,
            value_range=(-1, 1),
        )
        return save_file
    return None


def worker_loop():
    """Ranks 1-7: wait for rank 0 to broadcast a generation request."""
    while True:
        obj = [None]
        dist.broadcast_object_list(obj, src=0)
        kwargs = obj[0]
        if kwargs is None:
            time.sleep(1)
            continue
        with _dist_lock:
            wan_t2v.generate(**kwargs)
            dist.barrier()


def build_ui():
    with gr.Blocks(title="Wan2.1 T2V 1.3B - 8x A100") as demo:
        gr.Markdown("""
            <div style="text-align: center; font-size: 32px; font-weight: bold; margin-bottom: 20px;">
                Wan2.1 (T2V-1.3B) - 8x A100 Multi-GPU
            </div>
            <div style="text-align: center; font-size: 16px; font-weight: normal; margin-bottom: 20px;">
                Wan: Open and Advanced Large-Scale Video Generative Models.<br>
                FSDP + xDiT USP (ring=8) inference across 8x A100 80GB.
            </div>
            """)

        with gr.Row():
            with gr.Column():
                prompt = gr.Textbox(
                    label="Prompt",
                    value=EXAMPLE_PROMPT,
                    lines=3,
                    placeholder="Describe the video you want to generate",
                )
                with gr.Accordion("Advanced Options", open=True):
                    resolution = gr.Dropdown(
                        label="Resolution (Width*Height)",
                        choices=[
                            "480*832",
                            "832*480",
                            "624*624",
                            "704*544",
                            "544*704",
                        ],
                        value="832*480",
                    )
                    with gr.Row():
                        sd_steps = gr.Slider(
                            label="Diffusion steps",
                            minimum=1,
                            maximum=100,
                            value=50,
                            step=1,
                        )
                        guide_scale = gr.Slider(
                            label="Guide scale",
                            minimum=0,
                            maximum=20,
                            value=6.0,
                            step=1,
                        )
                    with gr.Row():
                        shift_scale = gr.Slider(
                            label="Shift scale",
                            minimum=0,
                            maximum=20,
                            value=8.0,
                            step=1,
                        )
                        seed = gr.Slider(
                            label="Seed",
                            minimum=-1,
                            maximum=2147483647,
                            step=1,
                            value=-1,
                        )
                    n_prompt = gr.Textbox(
                        label="Negative Prompt",
                        lines=2,
                        value="",
                    )
                run_button = gr.Button("Generate Video", variant="primary")

            with gr.Column():
                result_video = gr.Video(
                    label="Generated Video", interactive=False, height=600
                )

        gr.Examples(
            examples=[
                [EXAMPLE_PROMPT],
                ["A majestic golden eagle soaring above snow-capped mountains at sunrise, cinematic aerial shot"],
                ["A cute corgi puppy running through a field of sunflowers, golden hour lighting"],
                ["A cyberpunk city street in the rain at night, neon lights reflecting on wet asphalt"],
            ],
            inputs=[prompt],
        )

        run_button.click(
            fn=generate_video,
            inputs=[prompt, resolution, sd_steps, guide_scale, shift_scale, seed, n_prompt],
            outputs=[result_video],
            concurrency_limit=1,
        )

    return demo


def main():
    init_distributed()
    load_model()
    dist.barrier()

    if RANK == 0:
        logging.info("[rank 0] Starting Gradio server on port 7860")
        demo = build_ui()
        demo.queue(max_size=16).launch(
            server_name="0.0.0.0", server_port=7860, share=False
        )
    else:
        worker_loop()


if __name__ == "__main__":
    import logging

    logging.basicConfig(
        level=logging.INFO,
        format="[%(asctime)s] %(levelname)s [rank %(process)d] %(message)s",
        stream=sys.stdout,
    )
    main()