wan_s2v ¶

def forward(self, hidden_states: torch.Tensor | list[torch.Tensor],
            encoder_hidden_states: torch.Tensor | list[torch.Tensor], timestep: torch.Tensor,
            ref_latents: torch.Tensor | list[torch.Tensor] | None = None,
            motion_latents: torch.Tensor | list[torch.Tensor] | None = None,
            cond_states: torch.Tensor | list[torch.Tensor] | None = None,
            audio_input: torch.Tensor | None = None,
            motion_frames: tuple[int, int] = (17, 5), add_last_motion: int = 2,
            drop_motion_frames: bool = False, **kwargs) -> list[torch.Tensor]:
    """Denoise one step of an audio-driven video.

    hidden_states  [B, C, T, H, W] or list of [C, T, H, W] noisy video latents
    ref_latents    reference-image latents; required, the model is image-conditioned
    motion_latents previously generated frames; None on the first clip
    cond_states    pose/control latents, or None when unused
    audio_input    [B, 25, C_a, T_a] stacked wav2vec2 hidden states; required

    Returns a batched [B, C, T, H, W] fp32 tensor when ``hidden_states`` came
    in batched (the DenoisingStage path -- CFG arithmetic and scheduler.step
    need a tensor), or a list of per-sample tensors when it came in as a
    list (the reference-implementation path).
    """
    batched_input = isinstance(hidden_states, torch.Tensor)
    hidden_states = self._as_sample_list(hidden_states)
    ref_latents = self._as_sample_list(ref_latents)
    motion_latents = self._as_sample_list(motion_latents)
    cond_states = self._as_sample_list(cond_states)
    if isinstance(encoder_hidden_states, torch.Tensor):
        encoder_hidden_states = list(encoder_hidden_states)
    elif isinstance(encoder_hidden_states, list) and encoder_hidden_states and \
            encoder_hidden_states[0].dim() == 3:
        # DenoisingStage passes a list with one batched [B, L, C] tensor per
        # text encoder (same convention wanvideo.py unwraps with [0]).
        encoder_hidden_states = list(encoder_hidden_states[0])
    if audio_input is None or ref_latents is None:
        raise ValueError(
            "Wan-S2V needs both a reference image and audio. Got "
            f"ref_latents={'set' if ref_latents is not None else 'None'}, "
            f"audio_input={'set' if audio_input is not None else 'None'}. The pipeline supplies "
            "these from batch.image_latent and batch.audio_embeds.")
    if motion_latents is None:
        # First clip: no history yet. The official runner drops motion tokens
        # here (drop_first_motion=True in its config); the zeros only exist so
        # the frame packer has an input to (then discard) -- their values are
        # never attended to.
        drop_motion_frames = True
        motion_latents = [torch.zeros_like(u[:, :1]) for u in hidden_states]
    if cond_states is None and self.cond_encoder is not None:
        # The reference always encodes a cond tensor (zeros when unused), and
        # cond_encoder has a bias -- skipping it entirely would shift every
        # video token relative to the official implementation.
        cond_states = [torch.zeros_like(u) for u in hidden_states]

    # Every conditioning input above is batch-1 by construction (one
    # reference image, one motion history, one audio track) while
    # ``hidden_states`` carries ``num_videos_per_prompt`` samples. Repeat them
    # so each sample is conditioned: the zips below would otherwise truncate
    # to the shorter list and silently collapse the whole batch to 1, while
    # ``context``/``audio_emb`` keep their own batch and then mismatch.
    n_samples = len(hidden_states)
    ref_latents = self._repeat_batch1(ref_latents, n_samples)
    motion_latents = self._repeat_batch1(motion_latents, n_samples)
    if cond_states is not None:
        cond_states = self._repeat_batch1(cond_states, n_samples)
    if audio_input.size(0) == 1 and n_samples > 1:
        audio_input = audio_input.repeat(n_samples, 1, 1, 1)

    add_last_motion = int(self.add_last_motion) * add_last_motion
    audio_emb, audio_emb_global = self._embed_audio(audio_input, motion_frames)
    freqs = self.freqs.to(self.patch_embedding.weight.device)

    # 1. video tokens (+ pose conditioning added in patch space)
    x = [self.patch_embedding(u.unsqueeze(0)) for u in hidden_states]
    if cond_states is not None and self.cond_encoder is not None:
        x = [x_ + self.cond_encoder(c.unsqueeze(0)) for x_, c in zip(x, cond_states, strict=True)]
    original_grid_sizes = torch.stack([torch.tensor(u.shape[2:], dtype=torch.long) for u in x])
    x = [u.flatten(2).transpose(1, 2) for u in x]
    video_len = x[0].size(1)
    grid_sizes = [[torch.zeros_like(original_grid_sizes), original_grid_sizes, original_grid_sizes]]

    # 2. reference-image tokens, parked far away in time (see REF_TIME_INDEX)
    ref = [self.patch_embedding(r.unsqueeze(0)) for r in ref_latents]
    bsz, h, w = len(ref), ref[0].shape[3], ref[0].shape[4]
    grid_sizes.append([
        torch.tensor([REF_TIME_INDEX, 0, 0]).unsqueeze(0).repeat(bsz, 1),
        torch.tensor([REF_TIME_INDEX + 1, h, w]).unsqueeze(0).repeat(bsz, 1),
        torch.tensor([1, h, w]).unsqueeze(0).repeat(bsz, 1),
    ])
    x = [torch.cat([u, r.flatten(2).transpose(1, 2)], dim=1) for u, r in zip(x, ref, strict=True)]

    # 3. token-kind mask: 0 video, 1 reference (2 = motion, tagged on append)
    mask = [torch.zeros([1, u.shape[1]], dtype=torch.long, device=u.device) for u in x]
    for m in mask:
        m[:, video_len:] = 1

    # 4. RoPE for video+ref, then append motion tokens (which carry their own)
    stacked = torch.cat(x)
    rope = rope_precompute(
        stacked.detach().view(stacked.size(0), stacked.size(1), self.num_attention_heads,
                              self.hidden_size // self.num_attention_heads), grid_sizes, freqs)
    x, rope = [u.unsqueeze(0) for u in stacked], [u.unsqueeze(0) for u in rope]
    x, rope, mask = self._inject_motion(x, rope, mask, motion_latents, drop_motion_frames, add_last_motion)

    x = torch.cat(x, dim=0)
    rope = torch.cat(rope, dim=0)
    x = x + self.trainable_cond_mask(torch.cat(mask, dim=0)).to(x.dtype)

    # 5. conditioning embeddings
    e, block_e = self._timestep_embedding(timestep, video_len)
    context = self.text_embedding(
        torch.stack([
            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))]) for u in encoder_hidden_states
        ]))

    # 6. the tower, with audio re-injected at 12 of the 40 blocks
    for idx, block in enumerate(self.blocks):
        x = block(x, block_e, context, rope)
        x = self.audio_injector(x, idx, audio_emb, audio_emb_global, video_len, self.adain_mode)

    # 7. only the video span is denoised output; ref/motion are context
    out = [u.float() for u in self.unpatchify(self.head(x[:, :video_len], e), original_grid_sizes)]
    # Batched callers (DenoisingStage) need a tensor back for CFG arithmetic
    # and scheduler.step; list callers may have per-sample shapes, keep lists.
    return torch.stack(out) if batched_input else out

def materialize_non_persistent_buffers(self, device: torch.device, dtype: torch.dtype | None = None) -> None:
    """Rebuild the RoPE tables after meta-device construction.

    TransformerLoader builds the DiT under ``torch.device("meta")`` and then
    streams checkpoint weights in. Non-persistent buffers are not in the
    checkpoint, so they stay on meta until this hook (called by fsdp_load)
    recreates them with real storage. Complex dtype is deliberate -- these
    are rotation factors, not activations, and must not follow the model dtype.
    """
    head_dim = self.hidden_size // self.num_attention_heads
    if self.freqs.is_meta:
        self.freqs = rope_freqs(head_dim, self.rope_max_seq_len).to(device)
    packer = getattr(self, "frame_packer", None)
    if packer is not None and packer.freqs.is_meta:
        packer.freqs = rope_freqs(packer.inner_dim // packer.num_heads).to(device)