Wan-Animate-2 (authored by @kelseyee)

[yiyixuxu]

part of #14412

import time as _clock

_t0 = _clock.time()
import torch

from diffusers import ModularPipeline
from diffusers.utils import export_to_video, load_image, load_video


# to use base checkpoint, you just need to swap the repo id, it will map to the  blocks preset with own sampling defaults (10 steps/ no CFG for distilled, 40 steps / CFG 3.0 for base)
# pipe = ModularPipeline.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Diffusers", revision="refs/pr/2")
pipe = ModularPipeline.from_pretrained("Wan-AI/Wan2.2-Animate-2-14B-Distilled-Diffusers", revision="refs/pr/1")
pipe.load_components(dtype=torch.bfloat16)

# The in-context attention runs on the flex backend; compiling fuses it. Uncompiled flex
# materialises the full score matrix, which does not fit at the default `segment_frame_length`.
from diffusers.hooks import apply_group_offloading

# Stream the transformer's blocks: weights are ~33 GB and the KV cache another ~36 GB at this
# resolution, which does not co-reside on one 80 GB card.
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("assets/demo1/template.mp4", return_fps=True)

_t1 = _clock.time()
videos = pipe(
    image=load_image("assets/demo1/reference.png"),
    driving_video=driving_video,
    driving_video_fps=driving_video_fps,
    prompt="人物外观描述：一只银灰色虎斑纹的小猫，拥有圆润的脸庞、竖立的耳朵和巨大的圆形眼睛。它身穿一套深蓝色的制服套装，包括一件带有金色纽扣的西装外套和一条百褶裙。外套里面搭配着白色衬衫，领口处系着一个红色的蝴蝶结，袖口露出白色的衬衫边缘。背景描述：背景为纯白色，光线均匀明亮，无其他杂物或装饰。",
    height=800,
    width=640,
    generator=torch.Generator("cuda").manual_seed(42),
    output="videos",
)

print(
    f"WALL CLOCK: {(_clock.time() - _t0) / 60:.1f} min  PIPELINE: {(_clock.time() - _t1) / 60:.1f} min  "
    f"PEAK: {torch.cuda.max_memory_allocated() / 1e9:.1f} GB  FRAMES: {len(videos[0])}",
    flush=True,
)
export_to_video(videos[0], "output.mp4", fps=24)

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[yiyixuxu]

@kelseyee
we replaced the cv2/decord-based resizing (resize_by_area / padding_resize) with diffusers' PIL/torch video processor, so cv2 and decord are no longer dependencies -> it cause a very slight difference in output

[kelseyee]

@yiyixuxu Fine on our side. The letterboxing is preprocessing, not part of the trained graph, so a difference at this level is acceptable. No changes needed from us.