durgappc/infinitetalk

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import numpy as np3import torch4import torch.nn as nn5import torch.cuda.amp as amp6from xfuser.core.distributed import (7    get_sequence_parallel_rank,8    get_sequence_parallel_world_size,9    get_sp_group,10)11from einops import rearrange12from xfuser.core.long_ctx_attention import xFuserLongContextAttention13import xformers.ops14 15from ..modules.model import sinusoidal_embedding_1d16from ..utils.multitalk_utils import get_attn_map_with_target, split_token_counts_and_frame_ids, normalize_and_scale17from ..modules.attention import SingleStreamAttention, SingleStreamMutiAttention18 19 20def pad_freqs(original_tensor, target_len):21    seq_len, s1, s2 = original_tensor.shape22    pad_size = target_len - seq_len23    padding_tensor = torch.ones(24        pad_size,25        s1,26        s2,27        dtype=original_tensor.dtype,28        device=original_tensor.device)29    padded_tensor = torch.cat([original_tensor, padding_tensor], dim=0)30    return padded_tensor31 32 33@amp.autocast(enabled=False)34def rope_apply(x, grid_sizes, freqs):35    """36    x:          [B, L, N, C].37    grid_sizes: [B, 3].38    freqs:      [M, C // 2].39    """40    s, n, c = x.size(1), x.size(2), x.size(3) // 241    # split freqs42    freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) # [[N, head_dim/2], [N, head_dim/2], [N, head_dim/2]] # T H W 极坐标43 44    # loop over samples45    output = []46    for i, (f, h, w) in enumerate(grid_sizes.tolist()):47        seq_len = f * h * w48 49        # precompute multipliers50        x_i = torch.view_as_complex(x[i, :s].to(torch.float64).reshape(51            s, n, -1, 2)) # [L, N, C/2] # 极坐标52        freqs_i = torch.cat([53            freqs[0][:f].view(f, 1, 1, -1).expand(f, h, w, -1),54            freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1),55            freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1)56        ],57                            dim=-1).reshape(seq_len, 1, -1) # seq_lens, 1,  3 * dim / 2 (T H W)58 59        # apply rotary embedding60        sp_size = get_sequence_parallel_world_size()61        sp_rank = get_sequence_parallel_rank()62        freqs_i = pad_freqs(freqs_i, s * sp_size)63        s_per_rank = s64        freqs_i_rank = freqs_i[(sp_rank * s_per_rank):((sp_rank + 1) *65                                                       s_per_rank), :, :]66        x_i = torch.view_as_real(x_i * freqs_i_rank).flatten(2)67        x_i = torch.cat([x_i, x[i, s:]])68 69        # append to collection70        output.append(x_i)71    return torch.stack(output).float()72 73 74def usp_dit_forward_vace(self, x, vace_context, seq_len, kwargs):75    # embeddings76    c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context]77    c = [u.flatten(2).transpose(1, 2) for u in c]78    c = torch.cat([79        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)80        for u in c81    ])82 83    # arguments84    new_kwargs = dict(x=x)85    new_kwargs.update(kwargs)86 87    # Context Parallel88    c = torch.chunk(89        c, get_sequence_parallel_world_size(),90        dim=1)[get_sequence_parallel_rank()]91 92    hints = []93    for block in self.vace_blocks:94        c, c_skip = block(c, **new_kwargs)95        hints.append(c_skip)96    return hints97 98 99def usp_dit_forward(100    self,101    x,102    t,103    context,104    seq_len,105    vace_context=None,106    vace_context_scale=1.0,107    clip_fea=None,108    y=None,109):110    """111    x:              A list of videos each with shape [C, T, H, W].112    t:              [B].113    context:        A list of text embeddings each with shape [L, C].114    """115    if self.model_type == 'i2v':116        assert clip_fea is not None and y is not None117    # params118    device = self.patch_embedding.weight.device119    if self.freqs.device != device:120        self.freqs = self.freqs.to(device)121 122    if self.model_type != 'vace' and y is not None:123        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]124 125    # embeddings126    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]127    grid_sizes = torch.stack(128        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])129    x = [u.flatten(2).transpose(1, 2) for u in x]130    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)131    assert seq_lens.max() <= seq_len132    x = torch.cat([133        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)134        for u in x135    ])136 137    # time embeddings138    with amp.autocast(dtype=torch.float32):139        e = self.time_embedding(140            sinusoidal_embedding_1d(self.freq_dim, t).float())141        e0 = self.time_projection(e).unflatten(1, (6, self.dim))142        assert e.dtype == torch.float32 and e0.dtype == torch.float32143 144    # context145    context_lens = None146    context = self.text_embedding(147        torch.stack([148            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])149            for u in context150        ]))151 152    if self.model_type != 'vace' and clip_fea is not None:153        context_clip = self.img_emb(clip_fea)  # bs x 257 x dim154        context = torch.concat([context_clip, context], dim=1)155 156    # arguments157    kwargs = dict(158        e=e0,159        seq_lens=seq_lens,160        grid_sizes=grid_sizes,161        freqs=self.freqs,162        context=context,163        context_lens=context_lens)164    165    # Context Parallel166    x = torch.chunk(167        x, get_sequence_parallel_world_size(),168        dim=1)[get_sequence_parallel_rank()]169 170    for block in self.blocks:171        x = block(x, **kwargs)172 173    # head174    x = self.head(x, e)175 176    # Context Parallel177    x = get_sp_group().all_gather(x, dim=1)178 179    # unpatchify180    x = self.unpatchify(x, grid_sizes)181    return [u.float() for u in x]182 183 184def usp_attn_forward(self,185                     x,186                     seq_lens,187                     grid_sizes,188                     freqs,189                     dtype=torch.bfloat16):190    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim191    half_dtypes = (torch.float16, torch.bfloat16)192 193    def half(x):194        return x if x.dtype in half_dtypes else x.to(dtype)195 196    # query, key, value function197    def qkv_fn(x):198        q = self.norm_q(self.q(x)).view(b, s, n, d)199        k = self.norm_k(self.k(x)).view(b, s, n, d)200        v = self.v(x).view(b, s, n, d)201        return q, k, v202 203    q, k, v = qkv_fn(x)204    q = rope_apply(q, grid_sizes, freqs)205    k = rope_apply(k, grid_sizes, freqs)206 207    # TODO: We should use unpaded q,k,v for attention.208    # k_lens = seq_lens // get_sequence_parallel_world_size()209    # if k_lens is not None:210    #     q = torch.cat([u[:l] for u, l in zip(q, k_lens)]).unsqueeze(0)211    #     k = torch.cat([u[:l] for u, l in zip(k, k_lens)]).unsqueeze(0)212    #     v = torch.cat([u[:l] for u, l in zip(v, k_lens)]).unsqueeze(0)213 214    x = xFuserLongContextAttention()(215        None,216        query=half(q),217        key=half(k),218        value=half(v),219        window_size=self.window_size)220 221    # TODO: padding after attention.222    # x = torch.cat([x, x.new_zeros(b, s - x.size(1), n, d)], dim=1)223 224    # output225    x = x.flatten(2)226    x = self.o(x)227    return x228 229 230 231 232def usp_dit_forward_multitalk(233    self,234    x,235    t,236    context,237    seq_len,238    clip_fea=None,239    y=None,240    audio=None,241    ref_target_masks=None,242):243    """244    x:              A list of videos each with shape [C, T, H, W].245    t:              [B].246    context:        A list of text embeddings each with shape [L, C].247    """248    249    assert clip_fea is not None and y is not None250    # params251    device = self.patch_embedding.weight.device252    if self.freqs.device != device:253        self.freqs = self.freqs.to(device)254 255    _, T, H, W = x[0].shape256    N_t = T // self.patch_size[0]257    N_h = H // self.patch_size[1]258    N_w = W // self.patch_size[2]259 260    if y is not None:261        x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]262    x[0] = x[0].to(context[0].dtype)263 264    # embeddings265    x = [self.patch_embedding(u.unsqueeze(0)) for u in x]266    grid_sizes = torch.stack(267        [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])268    x = [u.flatten(2).transpose(1, 2) for u in x]269    seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)270    assert seq_lens.max() <= seq_len271    x = torch.cat([272        torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))], dim=1)273        for u in x274    ])275 276    # time embeddings277    with amp.autocast(dtype=torch.float32):278        e = self.time_embedding(279            sinusoidal_embedding_1d(self.freq_dim, t).float())280        e0 = self.time_projection(e).unflatten(1, (6, self.dim))281        assert e.dtype == torch.float32 and e0.dtype == torch.float32282 283    # context284    context_lens = None285    context = self.text_embedding(286        torch.stack([287            torch.cat([u, u.new_zeros(self.text_len - u.size(0), u.size(1))])288            for u in context289        ]))290 291    if clip_fea is not None:292        context_clip = self.img_emb(clip_fea)  293        context = torch.concat([context_clip, context], dim=1)294 295    # get audio token296    audio_cond = audio.to(device=x.device, dtype=x.dtype)297    first_frame_audio_emb_s = audio_cond[:, :1, ...] 298    latter_frame_audio_emb = audio_cond[:, 1:, ...] 299    latter_frame_audio_emb = rearrange(latter_frame_audio_emb, "b (n_t n) w s c -> b n_t n w s c", n=self.vae_scale) 300    middle_index = self.audio_window // 2301    latter_first_frame_audio_emb = latter_frame_audio_emb[:, :, :1, :middle_index+1, ...] 302    latter_first_frame_audio_emb = rearrange(latter_first_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 303    latter_last_frame_audio_emb = latter_frame_audio_emb[:, :, -1:, middle_index:, ...] 304    latter_last_frame_audio_emb = rearrange(latter_last_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 305    latter_middle_frame_audio_emb = latter_frame_audio_emb[:, :, 1:-1, middle_index:middle_index+1, ...] 306    latter_middle_frame_audio_emb = rearrange(latter_middle_frame_audio_emb, "b n_t n w s c -> b n_t (n w) s c") 307    latter_frame_audio_emb_s = torch.concat([latter_first_frame_audio_emb, latter_middle_frame_audio_emb, latter_last_frame_audio_emb], dim=2) 308    audio_embedding = self.audio_proj(first_frame_audio_emb_s, latter_frame_audio_emb_s) 309    human_num = len(audio_embedding)310    audio_embedding = torch.concat(audio_embedding.split(1), dim=2).to(x.dtype)311 312 313    # convert ref_target_masks to token_ref_target_masks314    if ref_target_masks is not None:315        ref_target_masks = ref_target_masks.unsqueeze(0).to(torch.float32) 316        token_ref_target_masks = nn.functional.interpolate(ref_target_masks, size=(N_h, N_w), mode='nearest') 317        token_ref_target_masks = token_ref_target_masks.squeeze(0) 318        token_ref_target_masks = (token_ref_target_masks > 0)319        token_ref_target_masks = token_ref_target_masks.view(token_ref_target_masks.shape[0], -1) 320        token_ref_target_masks = token_ref_target_masks.to(x.dtype)321    322    if self.enable_teacache:323        modulated_inp = e0 if self.use_ret_steps else e324        if self.cnt%3==0: # cond325            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:326                should_calc_cond = True327                self.accumulated_rel_l1_distance_cond = 0328            else:329                rescale_func = np.poly1d(self.coefficients)330                self.accumulated_rel_l1_distance_cond += rescale_func(((modulated_inp-self.previous_e0_cond).abs().mean() / self.previous_e0_cond.abs().mean()).cpu().item())331                # print("accumulated_rel_l1_distance_even", self.accumulated_rel_l1_distance_even)332                if self.accumulated_rel_l1_distance_cond < self.teacache_thresh:333                    should_calc_cond = False334                else:335                    should_calc_cond = True336                    self.accumulated_rel_l1_distance_cond = 0337            self.previous_e0_cond = modulated_inp.clone()338        elif self.cnt%3==1: # drop_text339            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:340                should_calc_drop_text = True341                self.accumulated_rel_l1_distance_drop_text = 0342            else:343                rescale_func = np.poly1d(self.coefficients)344                self.accumulated_rel_l1_distance_drop_text += rescale_func(((modulated_inp-self.previous_e0_drop_text).abs().mean() / self.previous_e0_drop_text.abs().mean()).cpu().item())345                if self.accumulated_rel_l1_distance_drop_text < self.teacache_thresh:346                    should_calc_drop_text = False347                else:348                    should_calc_drop_text = True349                    self.accumulated_rel_l1_distance_drop_text = 0350            self.previous_e0_drop_text = modulated_inp.clone()351        else: # uncond352            if self.cnt < self.ret_steps or self.cnt >= self.cutoff_steps:353                should_calc_uncond = True354                self.accumulated_rel_l1_distance_uncond = 0355            else:356                rescale_func = np.poly1d(self.coefficients)357                self.accumulated_rel_l1_distance_uncond += rescale_func(((modulated_inp-self.previous_e0_uncond).abs().mean() / self.previous_e0_uncond.abs().mean()).cpu().item())358                if self.accumulated_rel_l1_distance_uncond < self.teacache_thresh:359                    should_calc_uncond = False360                else:361                    should_calc_uncond = True362                    self.accumulated_rel_l1_distance_uncond = 0363            self.previous_e0_uncond = modulated_inp.clone()364 365    # Context Parallel366    x = torch.chunk(367        x, get_sequence_parallel_world_size(),368        dim=1)[get_sequence_parallel_rank()]369 370    # arguments371    kwargs = dict(372        e=e0,373        seq_lens=seq_lens,374        grid_sizes=grid_sizes,375        freqs=self.freqs,376        context=context,377        context_lens=context_lens,378        audio_embedding=audio_embedding,379        ref_target_masks=token_ref_target_masks,380        human_num=human_num,381        )382 383    if self.enable_teacache:384        if self.cnt%3==0:385            if not should_calc_cond:386                x +=  self.previous_residual_cond387            else:388                ori_x = x.clone()389                for block in self.blocks:390                    x = block(x, **kwargs)391                self.previous_residual_cond = x - ori_x392        elif self.cnt%3==1:393            if not should_calc_drop_text:394                x +=  self.previous_residual_drop_text395            else:396                ori_x = x.clone()397                for block in self.blocks:398                    x = block(x, **kwargs)399                self.previous_residual_drop_text = x - ori_x400        else:401            if not should_calc_uncond:402                x +=  self.previous_residual_uncond403            else:404                ori_x = x.clone()405                for block in self.blocks:406                    x = block(x, **kwargs)407                self.previous_residual_uncond = x - ori_x408    else:409        for block in self.blocks:410            x = block(x, **kwargs)411 412    # head413    x = self.head(x, e)414 415    # Context Parallel416    x = get_sp_group().all_gather(x, dim=1)417 418    # unpatchify419    x = self.unpatchify(x, grid_sizes)420    if self.enable_teacache:421        self.cnt += 1422        if self.cnt >= self.num_steps:423            self.cnt = 0424        425    return torch.stack(x).float()426 427 428def usp_attn_forward_multitalk(self,429                     x,430                     seq_lens,431                     grid_sizes,432                     freqs,433                     dtype=torch.bfloat16,434                     ref_target_masks=None):435    b, s, n, d = *x.shape[:2], self.num_heads, self.head_dim436    half_dtypes = (torch.float16, torch.bfloat16)437 438    def half(x):439        return x if x.dtype in half_dtypes else x.to(dtype)440 441    # query, key, value function442    def qkv_fn(x):443        q = self.norm_q(self.q(x)).view(b, s, n, d)444        k = self.norm_k(self.k(x)).view(b, s, n, d)445        v = self.v(x).view(b, s, n, d)446        return q, k, v447 448    q, k, v = qkv_fn(x)449    q = rope_apply(q, grid_sizes, freqs)450    k = rope_apply(k, grid_sizes, freqs)451 452 453    x = xFuserLongContextAttention()(454        None,455        query=half(q),456        key=half(k),457        value=half(v),458        window_size=self.window_size)459 460 461    # output462    x = x.flatten(2)463    x = self.o(x)464 465    with torch.no_grad():466        x_ref_attn_map = get_attn_map_with_target(q.type_as(x), k.type_as(x), grid_sizes[0], 467                                            ref_target_masks=ref_target_masks, enable_sp=True) 468 469    return x, x_ref_attn_map470 471 472 473 474def usp_crossattn_multi_forward_multitalk(self, 475                                        x: torch.Tensor, 476                                        encoder_hidden_states: torch.Tensor,  # 1, 21, 64, C477                                        shape=None, 478                                        x_ref_attn_map=None,479                                        human_num=None) -> torch.Tensor:480        481        N_t, N_h, N_w = shape 482        sp_size = get_sequence_parallel_world_size()483        sp_rank = get_sequence_parallel_rank()484        audio_tokens_per_frame = 32485        visual_seqlen, frame_ids = split_token_counts_and_frame_ids(N_t, N_h * N_w, sp_size, sp_rank)486        encoder_hidden_states = encoder_hidden_states[:, min(frame_ids):max(frame_ids)+1, ...]487        encoder_hidden_states = rearrange(encoder_hidden_states, "B T N C -> B (T N) C")488        N_a = len(frame_ids)489        kv_seq = [audio_tokens_per_frame * human_num] * N_a490 491        if human_num == 1:492            return super(SingleStreamMutiAttention, self).forward(x, encoder_hidden_states, shape, enable_sp=True, kv_seq=kv_seq)493 494 495        # get q for hidden_state496        B, N, C = x.shape497        q = self.q_linear(x) 498        q_shape = (B, N, self.num_heads, self.head_dim) 499        q = q.view(q_shape).permute((0, 2, 1, 3))500 501        if self.qk_norm:502            q = self.q_norm(q)503 504        max_values = x_ref_attn_map.max(1).values[:, None, None] 505        min_values = x_ref_attn_map.min(1).values[:, None, None] 506        max_min_values = torch.cat([max_values, min_values], dim=2)507        max_min_values = get_sp_group().all_gather(max_min_values, dim=1)508 509        human1_max_value, human1_min_value = max_min_values[0, :, 0].max(), max_min_values[0, :, 1].min()510        human2_max_value, human2_min_value = max_min_values[1, :, 0].max(), max_min_values[1, :, 1].min()511 512        human1 = normalize_and_scale(x_ref_attn_map[0], (human1_min_value, human1_max_value), (self.rope_h1[0], self.rope_h1[1]))513        human2 = normalize_and_scale(x_ref_attn_map[1], (human2_min_value, human2_max_value), (self.rope_h2[0], self.rope_h2[1]))514        back   = torch.full((x_ref_attn_map.size(1),), self.rope_bak, dtype=human1.dtype).to(human1.device)515        max_indices = x_ref_attn_map.argmax(dim=0)516        normalized_map = torch.stack([human1, human2, back], dim=1)517        normalized_pos = normalized_map[range(x_ref_attn_map.size(1)), max_indices] # N 518        q = self.rope_1d(q, normalized_pos)519 520        encoder_kv = self.kv_linear(encoder_hidden_states) 521        encoder_kv_shape = (B, encoder_hidden_states.size(1), 2, self.num_heads, self.head_dim)522        encoder_kv = encoder_kv.view(encoder_kv_shape).permute((2, 0, 3, 1, 4)) 523        encoder_k, encoder_v = encoder_kv.unbind(0) # B H N C524 525        if self.qk_norm:526            encoder_k = self.add_k_norm(encoder_k)527 528        # position embedding for condition audio embeddings529        per_frame = torch.zeros(audio_tokens_per_frame * human_num, dtype=encoder_k.dtype).to(encoder_k.device)530        per_frame[:audio_tokens_per_frame] = (self.rope_h1[0] + self.rope_h1[1]) / 2531        per_frame[audio_tokens_per_frame:] = (self.rope_h2[0] + self.rope_h2[1]) / 2532        encoder_pos = torch.concat([per_frame]*N_a, dim=0)533        encoder_k = self.rope_1d(encoder_k, encoder_pos)534 535        # get attn536        q = rearrange(q, "B H M K -> B M H K")537        encoder_k = rearrange(encoder_k, "B H M K -> B M H K")538        encoder_v = rearrange(encoder_v, "B H M K -> B M H K")539        attn_bias = xformers.ops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(visual_seqlen, kv_seq)540        x = xformers.ops.memory_efficient_attention(q, encoder_k, encoder_v, attn_bias=attn_bias, op=None,)541        x = rearrange(x, "B M H K -> B H M K")542 543        # linear transform544        x_output_shape = (B, N, C)545        x = x.transpose(1, 2) 546        x = x.reshape(x_output_shape) 547        x = self.proj(x) 548        x = self.proj_drop(x)549 550        return x