ALSv/self-forcing

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import torch3import torch.cuda.amp as amp4from xfuser.core.distributed import (get_sequence_parallel_rank,5                                     get_sequence_parallel_world_size,6                                     get_sp_group)7from xfuser.core.long_ctx_attention import xFuserLongContextAttention8 9from ..modules.model import sinusoidal_embedding_1d10 11 12def pad_freqs(original_tensor, target_len):13    seq_len, s1, s2 = original_tensor.shape14    pad_size = target_len - seq_len15    padding_tensor = torch.ones(16        pad_size,17        s1,18        s2,19        dtype=original_tensor.dtype,20        device=original_tensor.device)21    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)22    return padded_tensor23 24 25@amp.autocast(enabled=False)26def rope_apply(x, grid_sizes, freqs):27    """28    x:          [B, L, N, C].29    grid_sizes: [B, 3].30    freqs:      [M, C // 2].31    """32    s, n, c = x.size(1), x.size(2), x.size(3) // 233    # split freqs34    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1)35 36    # loop over samples37    output = []38    for i, (f, h, w) in enumerate(grid_sizes.tolist()):39        seq_len = f * h * w40 41        # precompute multipliers42        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(43            s, n, -1, 2))44        freqs_i = torch.cat([45            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),46            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),47            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)48        ],49            dim=-1).reshape(seq_len, 1, -1)50 51        # apply rotary embedding52        sp_size = get_sequence_parallel_world_size()53        sp_rank = get_sequence_parallel_rank()54        freqs_i = pad_freqs(freqs_i, s * sp_size)55        s_per_rank = s56        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *57                                                       s_per_rank), :, :]58        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)59        x_i = torch.cat([x_i, x[i, s:]])60 61        # append to collection62        output.append(x_i)63    return torch.stack(output).float()64 65 66def usp_dit_forward(67    self,68    x,69    t,70    context,71    seq_len,72    clip_fea=None,73    y=None,74):75    """76    x:              A list of videos each with shape [C, T, H, W].77    t:              [B].78    context:        A list of text embeddings each with shape [L, C].79    """80    if self.model_type == 'i2v':81        assert clip_fea is not None and y is not None82    # params83    device = self.patch_embedding.weight.device84    if self.freqs.device != device:85        self.freqs = self.freqs.to(device)86 87    if y is not None:88        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]89 90    # embeddings91    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]92    grid_sizes = torch.stack(93        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])94    x = [u.flatten(2).transpose(1, 2) for u in x]95    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)96    assert seq_lens.max() <= seq_len97    x = torch.cat([98        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)99        for u in x100    ])101 102    # time embeddings103    with amp.autocast(dtype=torch.float32):104        e = self.time_embedding(105            sinusoidal_embedding_1d(self.freq_dim, t).float())106        e0 = self.time_projection(e).unflatten(1, (6, self.dim))107        assert e.dtype == torch.float32 and e0.dtype == torch.float32108 109    # context110    context_lens = None111    context = self.text_embedding(112        torch.stack([113            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])114            for u in context115        ]))116 117    if clip_fea is not None:118        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim119        context = torch.concat([context_clip, context], dim=1)120 121    # arguments122    kwargs = dict(123        e=e0,124        seq_lens=seq_lens,125        grid_sizes=grid_sizes,126        freqs=self.freqs,127        context=context,128        context_lens=context_lens)129 130    # Context Parallel131    x = torch.chunk(132        x, get_sequence_parallel_world_size(),133        dim=1)[get_sequence_parallel_rank()]134 135    for block in self.blocks:136        x = block(x, **kwargs)137 138    # head139    x = self.head(x, e)140 141    # Context Parallel142    x = get_sp_group().all_gather(x, dim=1)143 144    # unpatchify145    x = self.unpatchify(x, grid_sizes)146    return [u.float() for u in x]147 148 149def usp_attn_forward(self,150                     x,151                     seq_lens,152                     grid_sizes,153                     freqs,154                     dtype=torch.bfloat16):155    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim156    half_dtypes = (torch.float16, torch.bfloat16)157 158    def half(x):159        return x if x.dtype in half_dtypes else x.to(dtype)160 161    # query, key, value function162    def qkv_fn(x):163        q = self.norm_q(self.q(x)).view(b, s, n, d)164        k = self.norm_k(self.k(x)).view(b, s, n, d)165        v = self.v(x).view(b, s, n, d)166        return q, k, v167 168    q, k, v = qkv_fn(x)169    q = rope_apply(q, grid_sizes, freqs)170    k = rope_apply(k, grid_sizes, freqs)171 172    # TODO: We should use unpaded q,k,v for attention.173    # k_lens = seq_lens // get_sequence_parallel_world_size()174    # if k_lens is not None:175    #     q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)176    #     k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)177    #     v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)178 179    x = xFuserLongContextAttention()(180        None,181        query=half(q),182        key=half(k),183        value=half(v),184        window_size=self.window_size)185 186    # TODO: padding after attention.187    # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)188 189    # output190    x = x.flatten(2)191    x = self.o(x)192    return x193