wan/distributed/xdit_context_parallel.py

# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.
from time import time
import torch
import torch.cuda.amp as amp
from xfuser.core.distributed import (get_sequence_parallel_rank,
                                     get_sequence_parallel_world_size,
                                     get_sp_group)
from xfuser.core.long_ctx_attention import xFuserLongContextAttention

from ..modules.model import sinusoidal_embedding_1d
from typing import List, Union, Optional, Tuple
import torch.nn.functional as F
import torch


def pad_freqs(original_tensor, target_len):
    seq_len, s1, s2 = original_tensor.shape
    pad_size = target_len - seq_len
    padding_tensor = torch.ones(
        pad_size,
        s1,
        s2,
        dtype=original_tensor.dtype,
        device=original_tensor.device)
    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)
    return padded_tensor


@amp.autocast(enabled=False)
def rope_apply(x, grid_sizes, freqs):
    """
    x:          [B, L, N, C].
    grid_sizes: [B, 3].
    freqs:      [M, C // 2].
    """
    s, n, c = x.size(1), x.size(2), x.size(3) // 2
    # split freqs
    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)

    # loop over samples
    output = []
    for i, (f, h, w) in enumerate(grid_sizes.tolist()):
        seq_len = f * h * w

        # precompute multipliers
        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(
            s, n, -1, 2))
        freqs_i = torch.cat([
            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),
            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),
            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)
        ],
                            dim=-1).reshape(seq_len, 1, -1)

        # apply rotary embedding
        sp_size = get_sequence_parallel_world_size()
        sp_rank = get_sequence_parallel_rank()
        freqs_i = pad_freqs(freqs_i, s * sp_size)
        s_per_rank = s
        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *
                                                       s_per_rank), :, :]
        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)
        x_i = torch.cat([x_i, x[i, s:]])

        # append to collection
        output.append(x_i)
    return torch.stack(output).float()


@torch.no_grad() # Usually don't need gradients for mask generation
def generate_attention_mask(
    attention_map: torch.Tensor,
    grid_sizes: torch.Tensor,
    target_x_shape: Tuple[int, int, int, int], # Target shape: (C, T, H, W)
    batch_index: int = 0,
    target_word_indices: Union[List[int], slice] = None,
    head_index: Optional[int] = None, # Process single head or average
    word_aggregation_method: str = 'mean', # How to combine scores for multiple words
    upsample_mode_spatial: str = 'nearest', # 'nearest', 'bilinear'
    upsample_mode_temporal: str = 'nearest', # 'nearest', 'linear'
    output_dtype: torch.dtype = torch.float32 # or torch.bool for soft mask before threshold
) -> torch.Tensor:
    """
    Generates a binary mask from an attention map based on attention towards target words.

    The mask identifies regions in the video (x) that attend strongly to the specified
    context words, exceeding a given threshold. The mask has the same dimensions as x.

    Args:
        attention_map (torch.Tensor): Attention weights [B, Head_num, Lx, Lctx].
                                      Lx = flattened video tokens (patches),
                                      Lctx = context tokens (words).
        target_word_indices (Union[List[int], slice]): Indices or slice for the target
                                                     word(s) in the Lctx dimension.
        grid_sizes (torch.Tensor): Patch grid dimensions [B, 3] -> (F, H_patch, W_patch)
                                   for each batch item, corresponding to Lx.
                                   F, H_patch, W_patch should be integers.
        target_x_shape (Tuple[int, int, int, int]): The desired output shape [C, T, H, W],
                                                    matching the original video tensor x.
        threshold (float): Value between 0 and 1. Attention scores >= threshold become 1 (True),
                           otherwise 0 (False).
        batch_index (int, optional): Batch item to process. Defaults to 0.
        head_index (Optional[int], optional): Specific head to use. If None, average
                                              attention across all heads. Defaults to None.
        word_aggregation_method (str, optional): How to aggregate scores if multiple
                                                 target_word_indices are given ('mean',
                                                 'sum', 'max'). Defaults to 'mean'.
        upsample_mode_spatial (str, optional): PyTorch interpolate mode for H, W dimensions.
                                           Defaults to 'nearest'.
        upsample_mode_temporal (str, optional): PyTorch interpolate mode for T dimension.
                                            Defaults to 'nearest'.
        output_dtype (torch.dtype, optional): Data type of the output mask.
                                              Defaults to torch.bool.

    Returns:
        torch.Tensor: A binary mask tensor of shape target_x_shape [C, T, H, W].

    Raises:
        TypeError: If inputs are not torch.Tensors.
        ValueError: If tensor dimensions or indices are invalid, or if
                    aggregation/upsample modes are unknown.
        IndexError: If batch_index or head_index are out of bounds.
    """
    # --- Input Validation ---
    if not isinstance(attention_map, torch.Tensor):
        raise TypeError("attention_map must be a torch.Tensor")
    if not isinstance(grid_sizes, torch.Tensor):
        raise TypeError("grid_sizes must be a torch.Tensor")
    if attention_map.dim() != 4:
        raise ValueError(f"attention_map must be [B, H, Lx, Lctx], got {attention_map.dim()} dims")
    if grid_sizes.dim() != 2 or grid_sizes.shape[1] != 3:
        raise ValueError(f"grid_sizes must be [B, 3], got {grid_sizes.shape}")
    if len(target_x_shape) != 4:
         raise ValueError(f"target_x_shape must be [C, T, H, W], got length {len(target_x_shape)}")

    B, H, Lx, Lctx = attention_map.shape
    C_out, T_out, H_out, W_out = target_x_shape

    if not 0 <= batch_index < B:
        raise IndexError(f"batch_index {batch_index} out of range for batch size {B}")
    if head_index is not None and not 0 <= head_index < H:
         raise IndexError(f"head_index {head_index} out of range for head count {H}")
    if word_aggregation_method not in ['mean', 'sum', 'max']:
        raise ValueError(f"Unknown word_aggregation_method: {word_aggregation_method}")
    if upsample_mode_spatial not in ['nearest', 'bilinear']:
        raise ValueError(f"Unknown upsample_mode_spatial: {upsample_mode_spatial}")
    if upsample_mode_temporal not in ['nearest', 'linear']:
         raise ValueError(f"Unknown upsample_mode_temporal: {upsample_mode_temporal}")


    # --- Select Head(s) ---
    if head_index is None:
        # Average across heads. Shape -> [Lx, Lctx]
        attn_map_processed = attention_map[batch_index].mean(dim=0)
    else:
        # Select specific head. Shape -> [Lx, Lctx]
        attn_map_processed = attention_map[batch_index, head_index]

    # --- Select and Aggregate Word Attention ---
    # Ensure target_word_indices are valid before slicing
    if isinstance(target_word_indices, slice):
        _slice_indices = range(*target_word_indices.indices(Lctx))
        if not _slice_indices: # Empty slice
             num_words = 0
        elif _slice_indices.start >= Lctx or _slice_indices.stop < -Lctx : # Basic out of bounds check
             num_words = len(_slice_indices) # Proceed cautiously or add stricter check
        else:
             num_words = len(_slice_indices)
        word_indices_str = f"slice({_slice_indices.start}:{_slice_indices.stop}:{_slice_indices.step})"
        word_attn_scores = attn_map_processed[:, target_word_indices] # Shape -> [Lx, num_words]
    elif isinstance(target_word_indices, list):
         # Check indices are within bounds
         valid_indices = [idx for idx in target_word_indices if -Lctx <= idx < Lctx]
         if not valid_indices:
             num_words = 0
             word_attn_scores = torch.empty((Lx, 0), device=attention_map.device, dtype=attention_map.dtype) # Handle empty case
         else:
            word_attn_scores = attn_map_processed[:, valid_indices] # Shape -> [Lx, num_words]
            num_words = len(valid_indices)
         word_indices_str = str(valid_indices) # Report used indices
    else:
        raise TypeError(f"target_word_indices must be list or slice, got {type(target_word_indices)}")

    if num_words > 1:
        if word_aggregation_method == 'mean':
            aggregated_scores = word_attn_scores.mean(dim=-1)
        elif word_aggregation_method == 'sum':
            aggregated_scores = word_attn_scores.sum(dim=-1)
        elif word_aggregation_method == 'max':
            aggregated_scores = word_attn_scores.max(dim=-1).values
        # aggregated_scores shape -> [Lx]
    elif num_words == 1:
         aggregated_scores = word_attn_scores.squeeze(-1) # Shape -> [Lx]
    else: # No valid words selected
         return torch.zeros(target_x_shape, dtype=output_dtype, device=attention_map.device)

    # --- Reshape to Video Patch Grid ---
    # Ensure grid sizes are integers
    f_patch, h_patch, w_patch = map(int, grid_sizes[batch_index].tolist())
    actual_num_tokens = f_patch * h_patch * w_patch

    if actual_num_tokens == 0:
         return torch.zeros(target_x_shape, dtype=output_dtype, device=attention_map.device)

    # Handle mismatch between expected tokens (from grid) and actual attention length (Lx)
    if actual_num_tokens > Lx:
         # Pad aggregated_scores to actual_num_tokens size
         padding_size = actual_num_tokens - aggregated_scores.numel()
         scores_padded = F.pad(aggregated_scores, (0, padding_size), "constant", 0)
         scores_unpadded = scores_padded # Use the padded version for reshaping
         # This scenario is less common than Lx > actual_num_tokens
    elif actual_num_tokens < Lx:
        scores_unpadded = aggregated_scores[:actual_num_tokens]
    else:
        scores_unpadded = aggregated_scores # Shape [actual_num_tokens]

    try:
        # Reshape to [F_patch, H_patch, W_patch]
        attention_patch_grid = scores_unpadded.reshape(f_patch, h_patch, w_patch)
    except RuntimeError as e:
        raise e

    # --- Upsample to Original Video Resolution ---
    # Add batch and channel dims for interpolation: [1, 1, F_patch, H_patch, W_patch]
    # Note: Assuming attention is channel-agnostic here.
    grid_for_upsample = attention_patch_grid.unsqueeze(0).unsqueeze(0).float() # Interpolate needs float

 
     # --- SIMPLIFIED LOGIC: Always use 3D interpolation ---
    target_size_3d = (T_out, H_out, W_out)

    # Determine the 3D interpolation mode.
    # Default to 'nearest' unless temporal dimension changes AND 'linear' is requested.
    if upsample_mode_temporal == 'linear' and f_patch != T_out:
        upsample_mode_3d = 'trilinear'
        align_corners_3d = False # align_corners usually False for non-nearest modes
    else:
        # Use 'nearest' if T isn't changing, or if temporal mode is 'nearest'.
        # 'nearest' is generally safer and handles spatial modes implicitly.
        upsample_mode_3d = 'nearest'
        align_corners_3d = None # align_corners=None for nearest

    upsampled_scores_grid = F.interpolate(grid_for_upsample,
                                       size=target_size_3d,
                                       mode=upsample_mode_3d,
                                       align_corners=align_corners_3d)
    # Expected shape: [1, 1, T_out, H_out, W_out] == [1, 1, 21, 60, 104]

    # --- END SIMPLIFIED LOGIC ---
 
    # Remove batch and channel dims: [T_out, H_out, W_out]
    upsampled_scores = upsampled_scores_grid.squeeze(0).squeeze(0)

    # --- Thresholding ---
    binary_mask_thw = (upsampled_scores / torch.max(upsampled_scores)) # Shape [T_out, H_out, W_out]

    # --- Expand Channel Dimension ---
    # Repeat the mask across the channel dimension C_out
    # Input shape: [T_out, H_out, W_out]
    # After unsqueeze(0): [1, T_out, H_out, W_out]
    # Target shape:      [C_out, T_out, H_out, W_out]
    # This expand operation is valid as explained above.
    final_mask = binary_mask_thw.unsqueeze(0).expand(C_out, T_out, H_out, W_out)

    return final_mask.to(dtype=output_dtype)


def usp_dit_forward( 
    self,
    x,
    t,
    context,
    seq_len,
    clip_fea=None,
    y=None,
    words_indices=None,
    block_id=-1,
    type=None,
    timestep=None
):
    """
    x:              A list of videos each with shape [C, T, H, W].
    t:              [B].
    context:        A list of text embeddings each with shape [L, C].
    """
    if self.model_type == 'i2v':
        assert clip_fea is not None and y is not None
    # params
    device = self.patch_embedding.weight.device
    if self.freqs.device != device:
        self.freqs = self.freqs.to(device)

    if y is not None:
        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]

    # embeddings
    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]
    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]
    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)
    assert seq_lens.max() <= seq_len
    x = torch.cat([
        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)
        for u in x
    ])

    # time embeddings
    with amp.autocast(dtype=torch.float32):
        e = self.time_embedding(
            sinusoidal_embedding_1d(self.freq_dim, t).float())
        e0 = self.time_projection(e).unflatten(1, (6, self.dim))
        assert e.dtype == torch.float32 and e0.dtype == torch.float32

    # context
    context_lens = None
    context = self.text_embedding(
        torch.stack([
            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])
            for u in context
        ]))

    if clip_fea is not None:
        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim
        context = torch.concat([context_clip, context], dim=1)

    # arguments
    kwargs = dict(
        e=e0,
        seq_lens=seq_lens,
        grid_sizes=grid_sizes,
        freqs=self.freqs,
        context=context,
        context_lens=context_lens,
        collect_attn_map=False)

    # Context Parallel
    x = torch.chunk(
        x, get_sequence_parallel_world_size(),
        dim=1)[get_sequence_parallel_rank()]
    
    save_block_id = block_id 
    attn_map = None
    binary_mask = None
    for i, block in enumerate(self.blocks):
        kwargs["collect_attn_map"] = False
        if i == save_block_id:
            kwargs["collect_attn_map"] = True
            x, attn_map = block(x, **kwargs)
        else:
            x = block(x, **kwargs)

    # head
    x = self.head(x, e)
    # Context Parallel
    x = get_sp_group().all_gather(x, dim=1)

    # unpatchify
    x = self.unpatchify(x, grid_sizes)

    if save_block_id != -1 and words_indices is not None:
        attention_map = get_sp_group().all_gather(attn_map, dim=2)
        binary_mask = generate_attention_mask(
                        attention_map=attention_map, # [1, 12, 32760, 512] batchsize, head_num, l_x, l_context
                        target_word_indices=words_indices,
                        grid_sizes=grid_sizes, # Make sure grid_sizes covers the full batch
                        target_x_shape=x[0].shape, # channel, frames, h, W
                        batch_index=0, # Process the first item in the batch
                        head_index=None, # Average over heads
                        word_aggregation_method='mean'
                        )

    return [u.float() for u in x], binary_mask



def usp_attn_forward(self,
                     x,
                     seq_lens,
                     grid_sizes,
                     freqs,
                     dtype=torch.bfloat16):
    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim
    half_dtypes = (torch.float16, torch.bfloat16)

    def half(x):
        return x if x.dtype in half_dtypes else x.to(dtype)

    # query, key, value function
    def qkv_fn(x):
        q = self.norm_q(self.q(x)).view(b, s, n, d)
        k = self.norm_k(self.k(x)).view(b, s, n, d)
        v = self.v(x).view(b, s, n, d)
        return q, k, v
    q, k, v = qkv_fn(x)
    q = rope_apply(q, grid_sizes, freqs)
    k = rope_apply(k, grid_sizes, freqs)

    # TODO: We should use unpaded q,k,v for attention.
    # k_lens = seq_lens // get_sequence_parallel_world_size()
    # if k_lens is not None:
    #     q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)
    #     k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)
    #     v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)

    x = xFuserLongContextAttention()(
        None,
        query=half(q),
        key=half(k),
        value=half(v),
        window_size=self.window_size)

    # TODO: padding after attention.
    # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)

    # output
    x = x.flatten(2)
    x = self.o(x)
    return x