refactor transformer to avoid var mutation

[sayakpaul]

To stop variable mutation from the forward as much as possible? I think it's a good practice. I have verified this with the example provided in huggingface#14413 to ensure that outputs bit-wise match.

Use the script below to confirm that the changes here don't cause any regressions:

Unfold

#!/usr/bin/env python
"""Compare Wan-Animate-2 outputs between two branches of diffusers.

Explicitly checks out both branches — PR 14413's head and the refactor branch —
into their own temporary git worktrees, runs the PR description's code example
against each source tree with identical deterministic inputs (seed 42, a
synthetic reference image and 161-frame driving video), and asserts the decoded
frames are bitwise identical.

Usage:
    python compare_pr14413.py                       # PR 14413 vs yiyixuxu-animate2-refactor
    python compare_pr14413.py <ref-a> <ref-b>       # any two branches/refs/shas

Requires one CUDA GPU (~62 GB peak) and imageio + imageio-ffmpeg. The first run
downloads the distilled checkpoint. Takes ~10-15 min.
"""

import os
import subprocess
import sys
import tempfile
from pathlib import Path

import numpy as np


REF_A = "pull/14413/head"  # PR 14413 (huggingface/diffusers#14413)
REF_B = "yiyixuxu-animate2-refactor"  # the transformer-refactor branch

REPO = Path(
    subprocess.run(
        ["git", "-C", str(Path(__file__).resolve().parent), "rev-parse", "--show-toplevel"],
        check=True,
        capture_output=True,
        text=True,
    ).stdout.strip()
)

PROMPT = (
    "人物外观描述：一只银灰色虎斑纹的小猫，拥有圆润的脸庞、竖立的耳朵和巨大的圆形眼睛。它身穿一套深蓝色的制服套装，"
    "包括一件带有金色纽扣的西装外套和一条百褶裙。外套里面搭配着白色衬衫，领口处系着一个红色的蝴蝶结，袖口露出白色的"
    "衬衫边缘。背景描述：背景为纯白色，光线均匀明亮，无其他杂物或装饰。"
)


def git(*args: str) -> str:
    return subprocess.run(["git", "-C", str(REPO), *args], check=True, capture_output=True, text=True).stdout.strip()


def resolve_commit(ref: str) -> str:
    """Resolve a local branch/tag/sha, or fetch the ref from origin (e.g. pull/14413/head)."""
    local = subprocess.run(
        ["git", "-C", str(REPO), "rev-parse", "--verify", "--quiet", f"{ref}^{{commit}}"],
        capture_output=True,
        text=True,
    )
    if local.returncode == 0:
        return local.stdout.strip()
    git("fetch", "origin", ref)
    return git("rev-parse", "FETCH_HEAD")


def make_inputs(workdir: Path):
    """Deterministic synthetic inputs, generated once and shared by both runs."""
    import imageio.v2 as imageio

    def frame(t, h=800, w=640):
        y, x = np.mgrid[0:h, 0:w].astype(np.float32)
        img = np.stack(
            [
                120 + 60 * np.sin(2 * np.pi * x / w),
                120 + 60 * np.cos(2 * np.pi * y / h),
                np.full((h, w), 140, np.float32),
            ],
            axis=-1,
        )
        cy, cx = h * 0.55 + 40 * np.sin(t / 6), w * 0.5 + 80 * np.sin(t / 9)
        body = (((y - cy) / 180) ** 2 + ((x - cx) / 70) ** 2) < 1
        head = (((y - (cy - 220)) / 60) ** 2 + ((x - cx) / 60) ** 2) < 1
        img[body] = [200, 80, 60]
        img[head] = [230, 190, 160]
        return img.clip(0, 255).astype(np.uint8)

    imageio.imwrite(workdir / "reference.png", frame(0))
    writer = imageio.get_writer(workdir / "driving.mp4", fps=24)
    for t in range(161):
        writer.append_data(frame(t))
    writer.close()
    print("inputs written: reference.png, driving.mp4 (161 frames)")


def generate(tag: str):
    """The PR description's code example; runs in a subprocess whose PYTHONPATH
    points at one branch's `src`, and saves the raw decoded frames."""
    import torch

    from diffusers import ModularPipeline
    from diffusers.hooks import apply_group_offloading
    from diffusers.utils import load_image, load_video

    pipe = ModularPipeline.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers", revision="refs/pr/1")
    pipe.load_components(dtype=torch.bfloat16)

    apply_group_offloading(
        pipe.transformer,
        onload_device=torch.device("cuda"),
        offload_device=torch.device("cpu"),
        offload_type="block_level",
        num_blocks_per_group=4,
        use_stream=True,
    )
    pipe.text_encoder.to("cuda")
    pipe.image_encoder.to("cuda")
    pipe.vae.to("cuda")
    pipe.transformer.compile_repeated_blocks(fullgraph=False)

    driving_video, driving_video_fps = load_video("driving.mp4", return_fps=True)

    videos = pipe(
        image=load_image("reference.png"),
        driving_video=driving_video,
        driving_video_fps=driving_video_fps,
        prompt=PROMPT,
        height=800,
        width=640,
        generator=torch.Generator("cuda").manual_seed(42),
        output="videos",
    )

    frames = np.stack([np.asarray(f) for f in videos[0]])
    np.save(f"frames_{tag}.npy", frames)
    print(f"run {tag}: {frames.shape} {frames.dtype}", flush=True)


def main():
    refs = sys.argv[1:3] if len(sys.argv) > 1 else [REF_A, REF_B]
    workdir = Path(tempfile.mkdtemp(prefix="compare-wan-animate-2-"))
    worktrees = []
    try:
        make_inputs(workdir)

        for tag, ref in zip(("a", "b"), refs):
            commit = resolve_commit(ref)
            worktree = workdir / f"worktree-{tag}"
            git("worktree", "add", "--detach", str(worktree), commit)
            worktrees.append(worktree)
            print(f"[{tag}] {ref} -> {git('log', '--format=%h %s', '-1', commit)}")

            print(f"=== generating with {ref} ===", flush=True)
            subprocess.run(
                [sys.executable, __file__, "--generate", tag],
                check=True,
                cwd=workdir,
                env={**os.environ, "PYTHONPATH": str(worktree / "src")},
            )

        a = np.load(workdir / "frames_a.npy")
        b = np.load(workdir / "frames_b.npy")
        assert a.shape == b.shape, (a.shape, b.shape)
        if np.array_equal(a, b):
            print(f"MATCH: all {a.shape[0]} frames of {a.shape[1]}x{a.shape[2]} are bitwise identical")
        else:
            diff = np.abs(a.astype(np.float64) - b.astype(np.float64))
            print(f"MISMATCH: max abs diff {diff.max()}, mean {diff.mean()}, differing {(diff > 0).mean() * 100:.4f}%")
            sys.exit(1)
    finally:
        for worktree in worktrees:
            subprocess.run(["git", "-C", str(REPO), "worktree", "remove", "--force", str(worktree)])


if __name__ == "__main__":
    if sys.argv[1:2] == ["--generate"]:
        generate(sys.argv[2])
    else:
        main()

[sayakpaul]

Many of these can be accessed through self.config

[sayakpaul]

There are two places where we compute self.freqs. So, this is just a helper to circumvent that. Happy to fold them in the actual caller sites.

[yiyixuxu]

thanks