wan_animate ¶

def forward(self,
            hidden_states: torch.Tensor,
            encoder_hidden_states: torch.Tensor | list[torch.Tensor],
            timestep: torch.LongTensor,
            encoder_hidden_states_image: torch.Tensor | list[torch.Tensor] | None = None,
            pose_latents: torch.Tensor | None = None,
            face_pixel_values: torch.Tensor | None = None,
            guidance=None,
            **kwargs) -> torch.Tensor:
    """Denoise one step of an animation/replacement segment.

    hidden_states    [B, 36, T+1, H, W]: noise | 4ch mask | conditional
                     latent y, channel-concatenated by the pipeline. Frame
                     0 is the reference latent slot.
    pose_latents     [B, 16, T, H, W] VAE-encoded skeleton video; exactly
                     one frame fewer than hidden_states (no pose on ref).
    face_pixel_values [B, 3, F, 512, 512] raw face crops,
                     F = 4T - 3 pixel frames for T latent frames.
    encoder_hidden_states_image  CLIP features of the reference image
                     (257 tokens) -- required, the I2V cross-attention
                     splits them off positionally.

    Returns the full [B, 16, T+1, H, W] velocity; the decode stage drops
    the reference slot (and any guidance frames).
    """
    if pose_latents is None or face_pixel_values is None:
        raise ValueError(
            "Wan-Animate needs both a pose video and a face video. Got "
            f"pose_latents={'set' if pose_latents is not None else 'None'}, "
            f"face_pixel_values={'set' if face_pixel_values is not None else 'None'}. The "
            "pipeline supplies these from the preprocessed src_pose/src_face inputs.")
    if encoder_hidden_states_image is None or (isinstance(encoder_hidden_states_image, list)
                                               and not encoder_hidden_states_image):
        raise ValueError("Wan-Animate needs CLIP features of the reference image "
                         "(encoder_hidden_states_image): the I2V cross-attention splits off "
                         "the first 257 context tokens as image tokens positionally, so this "
                         "must fail here rather than misread the text tokens.")
    if pose_latents.shape[2] + 1 != hidden_states.shape[2]:
        raise ValueError(
            f"pose_latents has {pose_latents.shape[2]} frames but must have exactly one fewer "
            f"than hidden_states ({hidden_states.shape[2]}): the reference latent frame "
            "carries no pose.")

    orig_dtype = hidden_states.dtype
    if encoder_hidden_states is not None and not isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = encoder_hidden_states[0]
    if isinstance(encoder_hidden_states_image, list):
        encoder_hidden_states_image = encoder_hidden_states_image[0]

    batch_size, _, num_frames, height, width = hidden_states.shape
    p_t, p_h, p_w = self.patch_size
    post_patch_num_frames = num_frames // p_t
    post_patch_height = height // p_h
    post_patch_width = width // p_w

    freqs_cis = self._rotary_embeddings(post_patch_num_frames, post_patch_height, post_patch_width,
                                        hidden_states.device)

    # 1. Patchify, then add pose tokens to every frame but the reference.
    hidden_states = self._patchify_with_pose(hidden_states, pose_latents)
    hidden_states = hidden_states.flatten(2).transpose(1, 2)
    original_seq_len = hidden_states.shape[1]

    # 2. Face tokens: one 4+1 group per latent frame, zero row for the ref.
    motion_vec = self._encode_face(face_pixel_values)
    if motion_vec.shape[1] != post_patch_num_frames:
        raise ValueError(
            f"face/video misalignment: the face video produced {motion_vec.shape[1]} latent-"
            f"frame token groups (incl. the reference pad) but the video has "
            f"{post_patch_num_frames} latent frames. The face video needs 4*T - 3 pixel "
            f"frames for T video latent frames (e.g. 77 for 20).")

    # 3. Condition embeddings (time, text, CLIP image) -- Wan2.1 timestep logic.
    temb, timestep_proj, encoder_hidden_states, encoder_hidden_states_image = self.condition_embedder(
        timestep, encoder_hidden_states, encoder_hidden_states_image, timestep_seq_len=None)
    timestep_proj = timestep_proj.unflatten(1, (6, -1))
    encoder_hidden_states = torch.concat([encoder_hidden_states_image, encoder_hidden_states], dim=1)
    # The blocks add cross-attention output straight onto hidden_states, so
    # a context in another dtype would silently upcast the residual stream.
    assert encoder_hidden_states.dtype == orig_dtype, (
        f"context dtype {encoder_hidden_states.dtype} != latent dtype {orig_dtype}")

    # 4. The tower, with a face cross-attention residual after every
    #    inject-stride-th block.
    for idx, block in enumerate(self.blocks):
        hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis, original_seq_len)
        if idx % self.inject_face_latents_blocks == 0:
            adapter = self.face_adapter[idx // self.inject_face_latents_blocks]
            hidden_states = hidden_states + adapter(hidden_states, motion_vec)

    # 5. Output norm, projection & unpatchify (single-timestep path).
    shift, scale = (self.scale_shift_table + temb.unsqueeze(1)).chunk(2, dim=1)
    hidden_states = self.norm_out(hidden_states, shift, scale)
    hidden_states = self.proj_out(hidden_states)
    hidden_states = hidden_states.reshape(batch_size, post_patch_num_frames, post_patch_height,
                                          post_patch_width, p_t, p_h, p_w, -1)
    hidden_states = hidden_states.permute(0, 7, 1, 4, 2, 5, 3, 6)
    return hidden_states.flatten(6, 7).flatten(4, 5).flatten(2, 3)

def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None:
    """Rebuild the motion encoder's anti-aliasing blur kernels after meta load.

    TransformerLoader builds the model under ``torch.device("meta")`` and
    streams checkpoint weights in; non-persistent buffers are not in the
    checkpoint and stay on meta until this hook (called by fsdp_load)
    recreates them. Kernels stay fp32 -- forward casts them to the
    activation dtype.
    """
    for module in self.modules():
        if isinstance(module, MotionConv2d) and module.blur and module.blur_kernel.is_meta:
            module.blur_kernel = module._build_blur_kernel().to(device)