causal_wanvideo ¶

def forward(self, 
            q: torch.Tensor,
            k: torch.Tensor,
            v: torch.Tensor,
            freqs_cis: tuple[torch.Tensor, torch.Tensor],
            block_mask: BlockMask,
            kv_cache: dict | None = None,
            current_start: int = 0,
            cache_start: int | None = None,
            frame_seqlen: int = 1560):
    r"""
    Args:
        x(Tensor): Shape [B, L, num_heads, C / num_heads]
        seq_lens(Tensor): Shape [B]
        grid_sizes(Tensor): Shape [B, 3], the second dimension contains (F, H, W)
        freqs(Tensor): Rope freqs, shape [1024, C / num_heads / 2]
        frame_seqlen (int): Number of tokens per latent frame,
            e.g. 1560 for 480x832 resolution.
    """
    if cache_start is None:
        cache_start = current_start

    cos, sin = freqs_cis
    # relativistic defers roping until the cache window is known (and caches raw k)
    relativistic = self.rope_cache_policy == "relativistic" and kv_cache is not None
    if not relativistic:
        roped_query = _apply_rotary_emb(q, cos, sin, is_neox_style=False).type_as(v)
        roped_key = _apply_rotary_emb(k, cos, sin, is_neox_style=False).type_as(v)

    if kv_cache is None:
        # Padding for flex attention
        padded_length = math.ceil(q.shape[1] / 128) * 128 - q.shape[1]
        padded_roped_query = torch.cat(
            [roped_query,
                torch.zeros([q.shape[0], padded_length, q.shape[2], q.shape[3]],
                            device=q.device, dtype=v.dtype)],
            dim=1
        )

        padded_roped_key = torch.cat(
            [roped_key, torch.zeros([k.shape[0], padded_length, k.shape[2], k.shape[3]],
                                    device=k.device, dtype=v.dtype)],
            dim=1
        )

        padded_v = torch.cat(
            [v, torch.zeros([v.shape[0], padded_length, v.shape[2], v.shape[3]],
                            device=v.device, dtype=v.dtype)],
            dim=1
        )

        x = flex_attention(
            query=padded_roped_query.transpose(2, 1),
            key=padded_roped_key.transpose(2, 1),
            value=padded_v.transpose(2, 1),
            block_mask=block_mask
        )[:, :, :-padded_length].transpose(2, 1)
    else:
        current_end = current_start + q.shape[1]
        sink_tokens = self.sink_size * frame_seqlen
        if self.local_attn_size == -1:
            max_attention_size = (GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES * frame_seqlen)
        else:
            max_attention_size = self.local_attn_size * frame_seqlen
        if self.local_attn_size == -1 and current_end > max_attention_size:
            raise ValueError(
                "Causal Wan local_attn_size=-1 keeps the previous "
                f"{GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES}-latent-frame KV "
                "window for compatibility. Set local_attn_size for "
                f"longer rollouts; got current_end={current_end} tokens "
                f"with frame_seqlen={frame_seqlen}.")
        # If we are using local attention and the current KV cache size is larger than the local attention size, we need to truncate the KV cache
        kv_cache_size = kv_cache["k"].shape[1]
        num_new_tokens = q.shape[1]
        stored_key = k if relativistic else roped_key  # raw vs roped in cache
        global_end_index = (
            int(kv_cache["global_end_index"].item())
            if isinstance(kv_cache["global_end_index"], torch.Tensor)
            else int(kv_cache["global_end_index"])
        )
        local_end_index_prev = (
            int(kv_cache["local_end_index"].item())
            if isinstance(kv_cache["local_end_index"], torch.Tensor)
            else int(kv_cache["local_end_index"])
        )
        if self.local_attn_size != -1 and (current_end > global_end_index) and (
                num_new_tokens + local_end_index_prev > kv_cache_size):
            # Calculate the number of new tokens added in this step
            # Shift existing cache content left to discard oldest tokens
            # Clone the source slice to avoid overlapping memory error
            num_evicted_tokens = num_new_tokens + local_end_index_prev - kv_cache_size
            num_rolled_tokens = local_end_index_prev - num_evicted_tokens - sink_tokens
            kv_cache["k"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
                kv_cache["k"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
            kv_cache["v"][:, sink_tokens:sink_tokens + num_rolled_tokens] = \
                kv_cache["v"][:, sink_tokens + num_evicted_tokens:sink_tokens + num_evicted_tokens + num_rolled_tokens].clone()
            # Insert the new keys/values at the end
            local_end_index = local_end_index_prev + current_end - \
                global_end_index - num_evicted_tokens
            local_start_index = local_end_index - num_new_tokens
            kv_cache["k"][:, local_start_index:local_end_index] = stored_key
            kv_cache["v"][:, local_start_index:local_end_index] = v
        else:
            # Assign new keys/values directly up to current_end
            local_end_index = local_end_index_prev + current_end - global_end_index
            local_start_index = local_end_index - num_new_tokens
            kv_cache["k"] = kv_cache["k"].detach()
            kv_cache["v"] = kv_cache["v"].detach()
            # logger.info("kv_cache['k'] is in comp graph: %s", kv_cache["k"].requires_grad or kv_cache["k"].grad_fn is not None)
            kv_cache["k"][:, local_start_index:local_end_index] = stored_key
            kv_cache["v"][:, local_start_index:local_end_index] = v
        key_window = kv_cache["k"][:, max(0, local_end_index - max_attention_size):local_end_index]
        value_window = kv_cache["v"][:, max(0, local_end_index - max_attention_size):local_end_index]
        if relativistic:
            window_len, query_lo, query_hi = relativistic_window_offsets(
                local_end_index, num_new_tokens, max_attention_size)
            roped_query = _apply_rotary_emb(
                q, cos[query_lo:query_hi], sin[query_lo:query_hi],
                is_neox_style=False).type_as(v)
            key_window = _apply_rotary_emb(
                key_window, cos[:window_len], sin[:window_len],
                is_neox_style=False).type_as(v)
        x = self.attn(roped_query, key_window, value_window)
        if isinstance(kv_cache["global_end_index"], torch.Tensor):
            kv_cache["global_end_index"].fill_(current_end)
        else:
            kv_cache["global_end_index"] = current_end
        if isinstance(kv_cache["local_end_index"], torch.Tensor):
            kv_cache["local_end_index"].fill_(local_end_index)
        else:
            kv_cache["local_end_index"] = local_end_index

    return x

def unpatchify(self, x, grid_sizes):
    r"""


    Args:
        x (List[Tensor]):
            List of patchified features, each with shape [L, C_out * prod(patch_size)]
        grid_sizes (Tensor):
            Original spatial-temporal grid dimensions before patching,


    Returns:
        Tensor:
            Reconstructed video tensors with shape [B, C_out, F, H / 8, W / 8]
    """

    c = self.out_channels
    out = []
    for u, v in zip(x, grid_sizes.tolist()):
        u = u[:math.prod(v)].view(*v, *self.patch_size, c)
        u = u.permute(6, 0, 3, 1, 4, 2, 5)
        # u = torch.einsum('fhwpqrc->cfphqwr', u.contiguous())
        u = u.reshape(c, *[i * j for i, j in zip(v, self.patch_size)])
        out.append(u)
    return out