"""Loopy — Seamless Video Loop Generation via Anchored Looping Shift of Positional Embedding.

Faithful ZeroGPU port of https://github.com/WeChatCV/Loopy (htdong/Loopy).

The reference implementation is `generate_2.2_new.py` + `Wan2.2/wan2/text2video_roll.py`
+ `Wan2.2/wan2/modules/model_roll.py` from the official repo. It is Wan2.2-T2V-A14B
(two-expert MoE DiT) with four LoRAs merged in:

  * LightX2V 4-step CFG-distillation LoRAs (high / low noise experts)
  * Loopy's own looping LoRAs (high / low noise experts)

plus the paper's core trick, the *anchored looping shift of positional embedding*:
inside every self-attention block `i` the temporal RoPE frequencies are cyclically
rolled by `shift = 0 if i == 0 else (i - 1) % (F - 1) + 1` (F = number of latent
frames). That makes the temporal positions live on a circle, so the last frame is a
neighbour of the first one and the generated clip loops seamlessly.

This app reproduces that logic on top of 🧨 diffusers' `WanPipeline` so it fits on
ZeroGPU (fp8 dynamic quantization of both experts).
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import spaces  # noqa: E402  (must precede torch)

import gc  # noqa: E402
import random  # noqa: E402
import re  # noqa: E402
import tempfile  # noqa: E402
import time  # noqa: E402

import gradio as gr  # noqa: E402
import numpy as np  # noqa: E402
import torch  # noqa: E402
from diffusers import AutoencoderKLWan, UniPCMultistepScheduler, WanPipeline  # noqa: E402
from diffusers.models.transformers.transformer_wan import (  # noqa: E402
    WanRotaryPosEmbed,
    WanTransformer3DModel,
    WanTransformerBlock,
)
from diffusers.utils import export_to_video  # noqa: E402
from huggingface_hub import hf_hub_download  # noqa: E402
from safetensors.torch import load_file  # noqa: E402
from torchao.quantization import (  # noqa: E402
    Float8DynamicActivationFloat8WeightConfig,
    Int8WeightOnlyConfig,
    quantize_,
)

# --------------------------------------------------------------------------------------
# Constants (all taken from the reference: test.sh + wan2/configs/wan_t2v_A14B.py)
# --------------------------------------------------------------------------------------

BASE_REPO = "linoyts/Wan2.2-T2V-A14B-Diffusers-BF16"  # bf16 mirror of Wan-AI/Wan2.2-T2V-A14B-Diffusers
LOOPY_REPO = "htdong/Loopy"
LIGHTX2V_REPO = "lightx2v/Wan2.2-Distill-Loras"

LIGHTX2V_HIGH = "wan2.2_t2v_A14b_high_noise_lora_rank64_lightx2v_4step_1217.safetensors"
LIGHTX2V_LOW = "wan2.2_t2v_A14b_low_noise_lora_rank64_lightx2v_4step_1217.safetensors"

FIXED_FPS = 16  # wan_shared_cfg.sample_fps
DEFAULT_FRAMES = 53  # test.sh --frame_num 53
DEFAULT_STEPS = 4  # test.sh --sample_steps 4
DEFAULT_SHIFT = 12.0  # t2v_A14B.sample_shift
DEFAULT_GUIDANCE = 1.0  # t2v_A14B.sample_guide_scale = (1.0, 1.0) -> CFG disabled
DEFAULT_REPEATS = 3
MAX_SEED = np.iinfo(np.int32).max

# t2v_A14B / wan_shared_cfg.sample_neg_prompt (only used if guidance > 1)
DEFAULT_NEGATIVE_PROMPT = (
    "色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，"
    "最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，"
    "画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，"
    "杂乱的背景，三条腿，背景人很多，倒着走"
)

RESOLUTIONS = {
    "832 × 480 (landscape)": (832, 480),
    "480 × 832 (portrait)": (480, 832),
}
DEFAULT_RESOLUTION = "832 × 480 (landscape)"


# --------------------------------------------------------------------------------------
# 1. Anchored looping shift of positional embedding
#    (port of rope_apply_loop / WanSelfAttention.forward in wan2/modules/model_roll.py)
# --------------------------------------------------------------------------------------


class _RopeWithGrid(tuple):
    """`(freqs_cos, freqs_sin)` that also remembers the latent grid `(F, H, W)`."""

    grid = None


_orig_rope_forward = WanRotaryPosEmbed.forward


def _rope_forward_with_grid(self, hidden_states: torch.Tensor):
    freqs = _orig_rope_forward(self, hidden_states)
    _, _, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    out = _RopeWithGrid(freqs)
    out.grid = (num_frames // p_t, height // p_h, width // p_w)
    return out


def _roll_rope(rope: _RopeWithGrid, shift: int):
    """Cyclically roll the temporal axis of the RoPE frequencies.

    Equivalent to `torch.roll(freqs_3d, shifts=time_shift, dims=0)` in
    `model_roll.rope_apply_loop`: the height/width components are constant along the
    frame axis, so rolling the whole concatenated frequency vector only affects the
    temporal part.
    """
    if shift == 0:
        return (rope[0], rope[1])
    f, h, w = rope.grid
    rolled = []
    for freqs in (rope[0], rope[1]):
        dim = freqs.shape[-1]
        view = freqs.reshape(f, h, w, dim)
        view = torch.roll(view, shifts=shift, dims=0)
        rolled.append(view.reshape(1, f * h * w, 1, dim))
    return tuple(rolled)


def _loop_shift_for_block(block_idx: int, num_latent_frames: int) -> int:
    if block_idx == 0 or num_latent_frames <= 1:
        return 0
    return (block_idx - 1) % (num_latent_frames - 1) + 1


_orig_block_forward = WanTransformerBlock.forward


def _looping_block_forward(self, hidden_states, encoder_hidden_states, temb, rotary_emb):
    block_idx = getattr(self, "_loopy_block_idx", None)
    grid = getattr(rotary_emb, "grid", None)
    if block_idx is not None and grid is not None:
        rotary_emb = _roll_rope(rotary_emb, _loop_shift_for_block(block_idx, grid[0]))
    return _orig_block_forward(self, hidden_states, encoder_hidden_states, temb, rotary_emb)


def enable_looping(transformer: WanTransformer3DModel) -> None:
    for i, block in enumerate(transformer.blocks):
        block._loopy_block_idx = i


WanRotaryPosEmbed.forward = _rope_forward_with_grid
WanTransformerBlock.forward = _looping_block_forward


# --------------------------------------------------------------------------------------
# 2. LoRA merging (port of WanModel_roll.load_lora / lightx2v_lora_adapter.WanLoraWrapper)
#    Both reference loaders do `W += B @ A` with no alpha/rank rescaling.
# --------------------------------------------------------------------------------------

_ATTN_PROJ = {"q": "to_q", "k": "to_k", "v": "to_v", "o": "to_out.0"}
_LORA_KEY_RE = re.compile(r"^(?P<base>.+?)\.lora_(?P<ab>A|B|down|up)(?:\.default)?\.weight$")


def _native_to_diffusers(name: str):
    """`blocks.7.self_attn.q` -> `blocks.7.attn1.to_q` (same mapping diffusers uses)."""
    parts = name.split(".")
    if len(parts) < 3 or parts[0] != "blocks":
        return None
    idx, rest = parts[1], ".".join(parts[2:])
    head = rest.split(".")
    if head[0] == "self_attn" and len(head) == 2 and head[1] in _ATTN_PROJ:
        return f"blocks.{idx}.attn1.{_ATTN_PROJ[head[1]]}"
    if head[0] == "cross_attn" and len(head) == 2 and head[1] in _ATTN_PROJ:
        return f"blocks.{idx}.attn2.{_ATTN_PROJ[head[1]]}"
    if rest == "ffn.0":
        return f"blocks.{idx}.ffn.net.0.proj"
    if rest == "ffn.2":
        return f"blocks.{idx}.ffn.net.2"
    return None


@torch.no_grad()
def merge_wan_lora(transformer: WanTransformer3DModel, path: str, scale: float = 1.0) -> None:
    state_dict = load_file(path)
    pairs: dict[str, dict[str, torch.Tensor]] = {}
    unmatched: list[str] = []
    for key, value in state_dict.items():
        base = key[len("diffusion_model.") :] if key.startswith("diffusion_model.") else key
        match = _LORA_KEY_RE.match(base)
        if match is None:
            unmatched.append(key)
            continue
        which = "A" if match.group("ab") in ("A", "down") else "B"
        pairs.setdefault(match.group("base"), {})[which] = value

    applied, skipped = 0, 0
    for base, ab in sorted(pairs.items()):
        target = _native_to_diffusers(base)
        if target is None or "A" not in ab or "B" not in ab:
            skipped += 1
            continue
        weight = transformer.get_submodule(target).weight
        down = ab["A"].to(device=weight.device, dtype=torch.float32)
        up = ab["B"].to(device=weight.device, dtype=torch.float32)
        delta = (up @ down) * scale
        if tuple(delta.shape) != tuple(weight.shape):
            raise ValueError(f"shape mismatch for {base} -> {target}: {delta.shape} vs {weight.shape}")
        weight.data.add_(delta.to(weight.dtype))
        applied += 1

    print(
        f"[lora] {os.path.basename(path)}: merged {applied} modules "
        f"(skipped {skipped}, unmatched keys {len(unmatched)})",
        flush=True,
    )
    if unmatched:
        print(f"[lora]   unmatched examples: {unmatched[:4]}", flush=True)
    if applied == 0:
        raise RuntimeError(f"No LoRA weights merged from {path}")
    del state_dict, pairs
    gc.collect()


# --------------------------------------------------------------------------------------
# 3. Pipeline
#
# Ordering matters on ZeroGPU: every `.to("cuda")` / `device_map="cuda"` placement has to
# happen *before* torchao quantization, because the ZeroGPU CUDA hijack cannot register
# `Float8Tensor` subclasses (`aten.empty_like` is unimplemented for them). So: place on
# cuda -> merge LoRAs in place -> quantize last. This mirrors the reference ZeroGPU Wan2.2
# Space (`zerogpu-aoti/wan2-2-fp8da-aoti-faster`).
# --------------------------------------------------------------------------------------

print("[load] high-noise expert (bf16)", flush=True)
transformer_high = WanTransformer3DModel.from_pretrained(
    BASE_REPO, subfolder="transformer", torch_dtype=torch.bfloat16, device_map="cuda"
)
print("[load] low-noise expert (bf16)", flush=True)
transformer_low = WanTransformer3DModel.from_pretrained(
    BASE_REPO, subfolder="transformer_2", torch_dtype=torch.bfloat16, device_map="cuda"
)
vae = AutoencoderKLWan.from_pretrained(BASE_REPO, subfolder="vae", torch_dtype=torch.float32)

pipe = WanPipeline.from_pretrained(
    BASE_REPO,
    transformer=transformer_high,
    transformer_2=transformer_low,
    vae=vae,
    torch_dtype=torch.bfloat16,
).to("cuda")

# LightX2V 4-step distillation + Loopy LoRAs, merged into the matching expert.
merge_wan_lora(pipe.transformer, hf_hub_download(LIGHTX2V_REPO, LIGHTX2V_HIGH))
merge_wan_lora(pipe.transformer, hf_hub_download(LOOPY_REPO, "high_noise.safetensors"))
merge_wan_lora(pipe.transformer_2, hf_hub_download(LIGHTX2V_REPO, LIGHTX2V_LOW))
merge_wan_lora(pipe.transformer_2, hf_hub_download(LOOPY_REPO, "low_noise.safetensors"))

enable_looping(pipe.transformer)
enable_looping(pipe.transformer_2)

print("[load] quantizing (int8 text encoder, fp8-dynamic experts)", flush=True)
quantize_(pipe.text_encoder, Int8WeightOnlyConfig())
quantize_(pipe.transformer, Float8DynamicActivationFloat8WeightConfig())
quantize_(pipe.transformer_2, Float8DynamicActivationFloat8WeightConfig())

pipe.set_progress_bar_config(disable=False)
gc.collect()
print("[load] pipeline ready", flush=True)


# --------------------------------------------------------------------------------------
# 4. Inference
# --------------------------------------------------------------------------------------


def _estimate_duration(
    prompt="",
    resolution=DEFAULT_RESOLUTION,
    num_frames=DEFAULT_FRAMES,
    steps=DEFAULT_STEPS,
    *args,
    **kwargs,
):
    """ZeroGPU budget, fitted on measured runs of this Space:

    53 frames / 832×480 / 4 steps -> 35.0 s ; 65 frames / 480×832 / 8 steps -> 81.0 s
    => ~0.575 s per (step x latent frame) at 832x480 plus ~9 s fixed, +15% margin.
    """
    width, height = RESOLUTIONS.get(resolution, RESOLUTIONS[DEFAULT_RESOLUTION])
    latent_frames = (int(num_frames) - 1) // 4 + 1
    area = (width * height) / (832 * 480)
    return int((9.0 + 0.575 * int(steps) * latent_frames * area) * 1.15) + 1


@spaces.GPU(duration=_estimate_duration)
def generate_loop(
    prompt: str,
    resolution: str = DEFAULT_RESOLUTION,
    num_frames: int = DEFAULT_FRAMES,
    steps: int = DEFAULT_STEPS,
    loop_repeats: int = DEFAULT_REPEATS,
    seed: int = 0,
    randomize_seed: bool = True,
    shift: float = DEFAULT_SHIFT,
    guidance_scale: float = DEFAULT_GUIDANCE,
    guidance_scale_2: float = DEFAULT_GUIDANCE,
    negative_prompt: str = DEFAULT_NEGATIVE_PROMPT,
    progress=gr.Progress(track_tqdm=True),
):
    """Generate a seamlessly looping video from a text prompt with Loopy (Wan2.2-T2V-A14B).

    Args:
        prompt: Text description of the scene. Works best with naturally periodic or
            continuous motion (falling snow, drifting clouds, a walking animal, a
            rotating object). English and Chinese are both supported.
        resolution: Output resolution, "832 × 480 (landscape)" or "480 × 832 (portrait)".
        num_frames: Number of frames in one loop, must be 4n+1 (53 = the paper's setting).
        steps: Denoising steps. 4 is what the LightX2V distillation LoRAs are trained for.
        loop_repeats: How many times the loop is repeated in the exported mp4 file, so the
            seam is visible when you scrub / download the video.
        seed: Random seed.
        randomize_seed: Draw a fresh random seed instead of using `seed`.
        shift: Flow-matching timestep shift (reference value: 12.0).
        guidance_scale: CFG scale for the high-noise expert. The distilled LoRAs expect 1.0
            (no classifier-free guidance).
        guidance_scale_2: CFG scale for the low-noise expert (reference value: 1.0).
        negative_prompt: Only used when a guidance scale is above 1.0.
        progress: Gradio progress tracker.

    Returns:
        A tuple of (path to the generated looping mp4, seed used, timing report).
    """
    if not prompt or not prompt.strip():
        raise gr.Error("Please enter a prompt.")

    width, height = RESOLUTIONS.get(resolution, RESOLUTIONS[DEFAULT_RESOLUTION])
    num_frames = int(num_frames)
    if (num_frames - 1) % 4 != 0:
        num_frames = ((num_frames - 1) // 4) * 4 + 1
    used_seed = random.randint(0, MAX_SEED) if randomize_seed else int(seed)

    pipe.scheduler = UniPCMultistepScheduler.from_config(
        pipe.scheduler.config, flow_shift=float(shift)
    )

    start = time.perf_counter()
    frames = pipe(
        prompt=prompt.strip(),
        negative_prompt=negative_prompt or "",
        height=height,
        width=width,
        num_frames=num_frames,
        num_inference_steps=int(steps),
        guidance_scale=float(guidance_scale),
        guidance_scale_2=float(guidance_scale_2),
        generator=torch.Generator(device="cuda").manual_seed(used_seed),
    ).frames[0]
    elapsed = time.perf_counter() - start

    repeats = max(1, int(loop_repeats))
    with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as handle:
        video_path = handle.name
    export_to_video(list(frames) * repeats, video_path, fps=FIXED_FPS)

    report = (
        f"{num_frames} frames @ {FIXED_FPS} fps ({num_frames / FIXED_FPS:.1f}s loop), "
        f"{width}×{height}, {int(steps)} steps — generated in {elapsed:.1f}s "
        f"(mp4 contains {repeats}× the loop)"
    )
    print(f"[gen] {report}", flush=True)
    return video_path, used_seed, report


# --------------------------------------------------------------------------------------
# 5. UI
# --------------------------------------------------------------------------------------

# Prompts showcased by the authors on https://donghaotian123.github.io/Loopy/
# and in the reference repo's prompt.txt.
EXAMPLE_PROMPTS = [
    [
        "An Arctic fox leaps nimbly through the snow while hunting, its white fur blending "
        "seamlessly with the snowflakes. Realistic style; a scene of winter wildlife."
    ],
    [
        "A hummingbird hovers in front of a trumpet flower, its wings beating into a blur, "
        "while its long beak repeatedly dips into the flower's heart to sip nectar. Realistic style."
    ],
    [
        "Auroras dance across the night sky, with green and purple light bands swirling like "
        "silk, while snow-capped mountain silhouettes stand silently. Wide shot."
    ],
    [
        "A hawksbill turtle slowly flaps its front flippers, swimming through a turquoise coral "
        "sea, with bubbles continuously rising from its shell. Realistic style."
    ],
    [
        "A potter presses and rotates a clay form with both hands, wet clay slowly rising and "
        "shaping between the fingers. Realistic style."
    ],
    [
        "A tranquil lake reflects snow-capped mountains and sunset glow, with ripples spreading "
        "across the water as clouds slowly shift in the sky. Realistic style. Wide shot."
    ],
]

CSS = """
.gradio-container .contain{max-width: 1080px !important; margin: 0 auto !important}
"""

with gr.Blocks() as demo:
    gr.Markdown(
        """
        # 🔁 Loopy — seamless looping video generation

        [Loopy](https://huggingface.co/htdong/Loopy) makes
        [Wan2.2-T2V-A14B](https://huggingface.co/Wan-AI/Wan2.2-T2V-A14B) produce videos that
        **loop seamlessly**: the temporal RoPE frequencies are cyclically rolled per attention
        block (*anchored looping shift of positional embedding*), so the last frame becomes a
        neighbour of the first one. 4-step generation via the
        [LightX2V distillation LoRAs](https://huggingface.co/lightx2v/Wan2.2-Distill-Loras).

        *Prompt tip from the authors: describe motion that is naturally periodic or continuous —
        falling snow, drifting clouds, flowing water, a walking animal, a rotating object — and
        state the style and shot type. English and Chinese both work.*
        """
    )
    with gr.Row():
        with gr.Column(scale=1):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="A hummingbird hovers in front of a trumpet flower… Realistic style.",
                lines=4,
            )
            resolution = gr.Radio(
                label="Resolution",
                choices=list(RESOLUTIONS),
                value=DEFAULT_RESOLUTION,
            )
            run = gr.Button("Generate looping video", variant="primary")
            with gr.Accordion("Advanced settings", open=False):
                num_frames = gr.Slider(
                    label="Frames per loop (4n+1)",
                    minimum=25,
                    maximum=65,
                    step=4,
                    value=DEFAULT_FRAMES,
                    info=f"{DEFAULT_FRAMES} frames ≈ {DEFAULT_FRAMES / FIXED_FPS:.1f}s at {FIXED_FPS} fps (paper setting).",
                )
                steps = gr.Slider(
                    label="Denoising steps", minimum=4, maximum=8, step=1, value=DEFAULT_STEPS
                )
                loop_repeats = gr.Slider(
                    label="Loop repeats in exported mp4",
                    minimum=1,
                    maximum=4,
                    step=1,
                    value=DEFAULT_REPEATS,
                )
                seed = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                randomize_seed = gr.Checkbox(label="Randomize seed", value=True)
                shift = gr.Slider(
                    label="Timestep shift", minimum=1.0, maximum=16.0, step=0.5, value=DEFAULT_SHIFT
                )
                guidance_scale = gr.Slider(
                    label="Guidance scale — high-noise expert",
                    minimum=1.0,
                    maximum=6.0,
                    step=0.5,
                    value=DEFAULT_GUIDANCE,
                    info="The distilled 4-step LoRAs are trained CFG-free; keep at 1.0.",
                )
                guidance_scale_2 = gr.Slider(
                    label="Guidance scale — low-noise expert",
                    minimum=1.0,
                    maximum=6.0,
                    step=0.5,
                    value=DEFAULT_GUIDANCE,
                )
                negative_prompt = gr.Textbox(
                    label="Negative prompt (only used when guidance > 1)",
                    value=DEFAULT_NEGATIVE_PROMPT,
                    lines=3,
                )
        with gr.Column(scale=1):
            video = gr.Video(
                label="Looping video",
                autoplay=True,
                loop=True,
                interactive=False,
            )
            info = gr.Markdown()

    inputs = [
        prompt,
        resolution,
        num_frames,
        steps,
        loop_repeats,
        seed,
        randomize_seed,
        shift,
        guidance_scale,
        guidance_scale_2,
        negative_prompt,
    ]
    outputs = [video, seed, info]

    run.click(fn=generate_loop, inputs=inputs, outputs=outputs)
    prompt.submit(fn=generate_loop, inputs=inputs, outputs=outputs)

    gr.Examples(
        examples=EXAMPLE_PROMPTS,
        inputs=[prompt],
        outputs=outputs,
        fn=generate_loop,
        cache_examples=True,
        cache_mode="lazy",
        label="Prompts from the Loopy paper's showcase",
    )

if __name__ == "__main__":
    # Gradio 6 moved `theme` / `css` from the Blocks constructor to `launch()`.
    demo.queue().launch(theme=gr.themes.Citrus(), css=CSS, mcp_server=True)