[new-model] Add Wan2.2-Animate-14B character animation/replacement inference

docs/design/inference_schema_parity_inventory.yaml

@@ -481,6 +481,11 @@ surfaces:
       last_image: request.inputs.last_image
       references: request.inputs.references
       video_path: request.inputs.video_path
+      pose_video_path: request.inputs.pose_video_path
+      face_video_path: request.inputs.face_video_path
+      background_video_path: request.inputs.background_video_path
+      mask_video_path: request.inputs.mask_video_path
+      animate_mode: request.inputs.animate_mode
       latents: request.inputs.latents
       audio_latents: request.inputs.audio_latents
       mouse_cond: request.inputs.mouse_cond

docs/inference/support_matrix.md

@@ -78,6 +78,7 @@ column links a runnable script in `examples/inference/basic/` where one exists.
 | wan | `Wan-AI/Wan2.1-T2V-14B-Diffusers`<br>`FastVideo/Wan2.1-VSA-T2V-14B-720P-Diffusers` | T2V | — |
 | wan | `Wan-AI/Wan2.1-I2V-14B-480P-Diffusers` | I2V | — |
 | wan | `Wan-AI/Wan2.1-I2V-14B-720P-Diffusers` | I2V | — |
+| wan | `Wan-AI/Wan2.2-Animate-14B-Diffusers` | I2V | [basic_wan_animate.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan_animate.py) |
 | wan | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | I2V | — |
 | wan | `IRMChen/Wan2.1-Fun-1.3B-Control-Diffusers` | — | [basic_wan2_2_Fun.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_wan2_2_Fun.py) |
 | wan | `FastVideo/FastWan2.1-T2V-1.3B-Diffusers`<br>`FastVideo/FastWan2.1-T2V-14B-480P-Diffusers` | T2V | [basic_dmd.py](https://github.com/hao-ai-lab/FastVideo/blob/main/examples/inference/basic/basic_dmd.py) |
@@ -162,6 +163,7 @@ optimizations: absence means **untested**, not incompatible.
 | Lucy Edit Dev 5B*** | `decart-ai/Lucy-Edit-Dev` | 480P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | Wan2.2 T2V A14B | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
 | Wan2.2 I2V A14B | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480P<br>720P | ❌ | ❌ | ✅ | ⭕ | ⭕ |
+| Wan2.2 Animate 14B | `Wan-AI/Wan2.2-Animate-14B-Diffusers` | 720P | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ |
 | HunyuanVideo | `hunyuanvideo-community/HunyuanVideo` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
 | FastHunyuan | `FastVideo/FastHunyuan-diffusers` | 720px1280p<br>544px960p | ❌ | ✅ | ✅ | ⭕ | ⭕ |
 | Wan2.1 T2V 1.3B | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | 480P | ✅ | ✅ | ✅ | ⭕ | ⭕ |

examples/inference/basic/basic_wan_animate.py

@@ -0,0 +1,52 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Wan2.2-Animate-14B: animate a character image with someone else's performance.
+
+The reference image fixes who the character is; a preprocessed skeleton video
+drives the body and a preprocessed face-crop video drives the expression. The
+driving inputs are the artifacts the official preprocessing pipeline emits
+(https://github.com/Wan-Video/Wan2.2, ``wan/modules/animate/preprocess``):
+``src_pose.mp4``, ``src_face.mp4``, and the aligned reference ``src_ref.png``
+(replace mode adds ``src_bg.mp4`` + ``src_mask.mp4`` and typically the
+relighting LoRA -- see scripts/checkpoint_conversion/wan_animate_relight_lora.py).
+
+CFG is off by default (guidance 1.0): the prompt is non-core for this model.
+"""
+from fastvideo import VideoGenerator
+
+OUTPUT_PATH = "video_samples_wan_animate"
+
+# Outputs of the official preprocessing run on your reference image + driving video.
+REF_IMAGE_PATH = "preprocessed/src_ref.png"
+POSE_VIDEO_PATH = "preprocessed/src_pose.mp4"
+FACE_VIDEO_PATH = "preprocessed/src_face.mp4"
+
+
+def main():
+    generator = VideoGenerator.from_pretrained(
+        "Wan-AI/Wan2.2-Animate-14B-Diffusers",
+        num_gpus=1,
+        use_fsdp_inference=False,
+        dit_cpu_offload=True,
+        vae_cpu_offload=False,
+        text_encoder_cpu_offload=True,
+        pin_cpu_memory=True,
+    )
+
+    generator.generate_video(
+        "视频中的人在做动作",  # the official default prompt; non-core for this model
+        image_path=REF_IMAGE_PATH,
+        pose_video_path=POSE_VIDEO_PATH,
+        face_video_path=FACE_VIDEO_PATH,
+        output_path=OUTPUT_PATH,
+        save_video=True,
+        height=720,
+        width=1280,
+        num_frames=77,
+        fps=30,
+        guidance_scale=1.0,
+        num_inference_steps=20,
+    )
+
+
+if __name__ == "__main__":
+    main()

fastvideo/api/sampling_param.py

@@ -31,6 +31,17 @@ class SamplingParam:
     # Video inputs
     video_path: str | None = None
 
+    # Wan-Animate driving inputs: paths to the *preprocessed* artifacts the
+    # official preprocessing pipeline produces (src_pose.mp4, src_face.mp4;
+    # replace mode adds src_bg.mp4 + src_mask.mp4). `animate_mode` selects
+    # animation (character on the reference image's background) vs replace
+    # (character composited into the background video).
+    pose_video_path: str | None = None
+    face_video_path: str | None = None
+    background_video_path: str | None = None
+    mask_video_path: str | None = None
+    animate_mode: str | None = None
+
     # Optional pre-generated diffusion latents. Used by parity/debug harnesses
     # and advanced callers that need deterministic latent reuse.
     latents: Any | None = None

fastvideo/api/schema.py

@@ -125,6 +125,11 @@ class InputConfig:
     prompt_path: str | None = None
     image_path: str | list[str] | None = None
     video_path: str | list[str] | None = None
+    pose_video_path: str | list[str] | None = None
+    face_video_path: str | list[str] | None = None
+    background_video_path: str | list[str] | None = None
+    mask_video_path: str | list[str] | None = None
+    animate_mode: str | None = None
     pil_image: Any | None = None
     last_image: Any | None = None
     references: list[Any] | None = None

fastvideo/configs/models/dits/wan_animate.py

@@ -0,0 +1,114 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Arch config for Wan2.2-Animate-14B (character animation / replacement).
+
+Every value is transcribed verbatim from
+``Wan-AI/Wan2.2-Animate-14B-Diffusers/transformer/config.json`` -- do not
+"tidy" them. Field names match that file's keys one-for-one because
+``update_model_arch`` overlays any matching key from the checkpoint's
+config.json at load time; a renamed field silently stops receiving its
+checkpoint value.
+
+The tower itself is the Wan-I2V one (this class only switches on the I2V
+knobs ``image_dim``/``added_kv_proj_dim``); why that is the right base is
+explained on ``WanAnimate14BConfig`` in ``configs/pipelines/wan.py``. What
+Animate adds on top:
+
+* ``in_channels = 36 = 16 (noise) + 4 (mask) + 16 (conditional latent)``.
+* ``pose_patch_embedding`` -- a second patchifier whose output is added to
+  the video tokens (skipping the reference latent frame).
+* A LIA-style motion encoder (``motion_*`` fields) turning 512x512 face crops
+  into 20-dim motion codes. The checkpoint stores its conv/linear weights
+  **unit-scale**: the model must apply the StyleGAN2 runtime factor
+  (``1/sqrt(fan_in)``) in forward. Loading these into vanilla
+  ``nn.Conv2d``/``nn.Linear`` forwards succeeds and is silently wrong.
+* A face encoder + ``face_adapter`` cross-attention blocks. The checkpoint
+  indexes the adapters densely (``face_adapter.0 .. .7``) with no record of
+  which transformer block each serves; adapter ``i`` serves block
+  ``i * inject_face_latents_blocks``. That correspondence exists only here,
+  so it is asserted in ``__post_init__`` -- a wrong value loads cleanly and
+  steers the wrong blocks.
+
+Naming contract with the model (``fastvideo/models/dits/wan_animate.py``):
+the Animate-specific modules keep checkpoint-identical parameter names
+(``motion_encoder.*``, ``face_encoder.*``, ``face_adapter.*``) so the loader
+passes them through verbatim; only ``pose_patch_embedding`` needs a mapping
+entry, mirroring how the base ``patch_embedding`` is wrapped in ``.proj``.
+"""
+from dataclasses import dataclass, field
+
+from fastvideo.configs.models.dits.base import DiTArchConfig, DiTConfig
+from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
+from fastvideo.platforms import AttentionBackendEnum
+
+
+def _dit_blocks_only(n: str, m) -> bool:
+    """FSDP/compile unit = one transformer block. The base predicate
+    ("blocks" anywhere in the name) would also match the motion encoder's
+    tiny ``res_blocks.N`` convs and make each its own unit."""
+    return n.startswith("blocks.") and n.split(".")[-1].isdigit()
+
+
+@dataclass
+class WanAnimateArchConfig(WanVideoArchConfig):
+    _fsdp_shard_conditions: list = field(default_factory=lambda: [_dit_blocks_only])
+    # The face adapter's per-frame attention is implemented on these two only;
+    # notably no VSA -- the checkpoint has no gate weights for VSA blocks.
+    _supported_attention_backends: tuple[AttentionBackendEnum,
+                                         ...] = (AttentionBackendEnum.FLASH_ATTN, AttentionBackendEnum.TORCH_SDPA)
+
+    # Appending last is safe despite first-match-wins: the base
+    # `^patch_embedding\.` regex cannot match `pose_patch_embedding.*`.
+    param_names_mapping: dict = field(default_factory=lambda: WanVideoArchConfig().param_names_mapping |
+                                      {r"^pose_patch_embedding\.(.*)$": r"pose_patch_embedding.proj.\1"})
+    # The relighting LoRA was trained on the official (native-naming) model,
+    # whose I2V cross-attention carries k_img/v_img projections the base LoRA
+    # mapping does not cover. Harmless for adapters that do not target them.
+    lora_param_names_mapping: dict = field(
+        default_factory=lambda: WanVideoArchConfig().lora_param_names_mapping | {
+            r"^blocks\.(\d+)\.cross_attn\.k_img\.(.*)$": r"blocks.\1.attn2.add_k_proj.\2",
+            r"^blocks\.(\d+)\.cross_attn\.v_img\.(.*)$": r"blocks.\1.attn2.add_v_proj.\2",
+        })
+
+    # --- I2V skeleton knobs (config.json: in_channels 36 / image_dim 1280) ---
+    in_channels: int = 36
+    out_channels: int = 16
+    image_dim: int | None = 1280
+    added_kv_proj_dim: int | None = 5120
+
+    # --- Animate-specific (names match config.json keys exactly) ---
+    # in_channels = 2 * latent_channels + 4 mask channels; asserted below.
+    latent_channels: int = 16
+    # None -> the model uses the LIA channel table keyed by feature-map size
+    # (512: 32ch ... 4: 512ch). config.json ships null here.
+    motion_encoder_channel_sizes: dict[str, int] | None = None
+    motion_encoder_size: int = 512  # face crops are motion_encoder_size**2 RGB
+    motion_style_dim: int = 512  # appearance vector width
+    motion_dim: int = 20  # the identity-squeezing bottleneck
+    motion_encoder_dim: int = 512  # per-frame motion vector width fed to the face encoder
+    face_encoder_hidden_dim: int = 1024
+    face_encoder_num_heads: int = 4  # 4 face tokens (+1 padding token) per latent frame
+    inject_face_latents_blocks: int = 5  # adapter i serves block i * this
+    # An inference-memory knob (face crops go through the motion encoder this
+    # many frames at a time), but config.json ships it, so it lives here.
+    motion_encoder_batch_size: int = 8
+
+    def __post_init__(self) -> None:
+        super().__post_init__()
+        # Re-checked after update_model_arch overlays the checkpoint's
+        # config.json: a checkpoint this config cannot represent should fail
+        # here, not as a shape error mid-load.
+        assert self.in_channels == 2 * self.latent_channels + 4, (
+            f"in_channels ({self.in_channels}) must be 2 * latent_channels "
+            f"({self.latent_channels}) + 4 mask channels: noise | mask | conditional latent")
+        assert self.inject_face_latents_blocks > 0, "inject_face_latents_blocks must be positive"
+        assert self.num_layers % self.inject_face_latents_blocks == 0, (
+            f"num_layers ({self.num_layers}) must be an exact multiple of the face-adapter "
+            f"injection stride ({self.inject_face_latents_blocks}): the checkpoint stores adapters "
+            f"densely (face_adapter.0..N) and adapter i serves block i * stride")
+
+
+@dataclass
+class WanAnimateConfig(DiTConfig):
+    arch_config: DiTArchConfig = field(default_factory=WanAnimateArchConfig)
+
+    prefix: str = "WanAnimate"

fastvideo/configs/pipelines/wan.py

@@ -6,6 +6,7 @@
 
 from fastvideo.configs.models import DiTConfig, EncoderConfig, VAEConfig
 from fastvideo.configs.models.dits import WanVideoConfig
+from fastvideo.configs.models.dits.wan_animate import WanAnimateConfig as WanAnimateDiTConfig
 from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig
 from fastvideo.configs.models.encoders import (BaseEncoderOutput, CLIPVisionConfig, T5Config,
                                                WAN2_1ControlCLIPVisionConfig)
@@ -323,3 +324,24 @@ class SelfForcingWan2_2_T2V480PConfig(Wan2_2_T2V_A14B_Config):
     def __post_init__(self) -> None:
         self.vae_config.load_encoder = True
         self.vae_config.load_decoder = True
+
+
+# =============================================
+# ========== Wan2.2 Animate (14B) =============
+# =============================================
+@dataclass
+class WanAnimate14BConfig(WanI2V720PConfig):
+    """Configuration for Wan2.2-Animate-14B character animation / replacement.
+
+    Despite the Wan2.2 branding this is a *dense* 14B DiT on the Wan2.1-I2V
+    skeleton (single expert, Wan2.1 16-channel VAE, CLIP image branch), so it
+    subclasses the I2V config -- not the A14B MoE one. The inherited pieces
+    are all correct as-is: CLIPVisionConfig matches the checkpoint's ViT-H
+    image encoder (image_dim 1280), the I2V ``__post_init__`` loads the VAE
+    encoder (needed for the reference image, the pose video and the previous
+    segment's guidance frames), and flow_shift 5.0 matches the official
+    runner (``wan_animate_14B.py``: sample_shift 5.0, 20 steps, CFG 1.0 --
+    steps and guidance land in the preset, not here).
+    """
+
+    dit_config: DiTConfig = field(default_factory=WanAnimateDiTConfig)