durgappc/infinitetalk2

1# Copyright 2024-2025 The Alibaba Wan Team Authors. All rights reserved.2import torch3import torch.cuda.amp as amp4import torch.nn as nn5from diffusers.configuration_utils import register_to_config6 7from .model import WanAttentionBlock, WanModel, sinusoidal_embedding_1d8 9 10class VaceWanAttentionBlock(WanAttentionBlock):11 12    def __init__(self,13                 cross_attn_type,14                 dim,15                 ffn_dim,16                 num_heads,17                 window_size=(-1, -1),18                 qk_norm=True,19                 cross_attn_norm=False,20                 eps=1e-6,21                 block_id=0):22        super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size,23                         qk_norm, cross_attn_norm, eps)24        self.block_id = block_id25        if block_id == 0:26            self.before_proj = nn.Linear(self.dim, self.dim)27            nn.init.zeros_(self.before_proj.weight)28            nn.init.zeros_(self.before_proj.bias)29        self.after_proj = nn.Linear(self.dim, self.dim)30        nn.init.zeros_(self.after_proj.weight)31        nn.init.zeros_(self.after_proj.bias)32 33    def forward(self, c, x, **kwargs):34        if self.block_id == 0:35            c = self.before_proj(c) + x36 37        c = super().forward(c, **kwargs)38        c_skip = self.after_proj(c)39        return c, c_skip40 41 42class BaseWanAttentionBlock(WanAttentionBlock):43 44    def __init__(self,45                 cross_attn_type,46                 dim,47                 ffn_dim,48                 num_heads,49                 window_size=(-1, -1),50                 qk_norm=True,51                 cross_attn_norm=False,52                 eps=1e-6,53                 block_id=None):54        super().__init__(cross_attn_type, dim, ffn_dim, num_heads, window_size,55                         qk_norm, cross_attn_norm, eps)56        self.block_id = block_id57 58    def forward(self, x, hints, context_scale=1.0, **kwargs):59        x = super().forward(x, **kwargs)60        if self.block_id is not None:61            x = x + hints[self.block_id] * context_scale62        return x63 64 65class VaceWanModel(WanModel):66 67    @register_to_config68    def __init__(self,69                 vace_layers=None,70                 vace_in_dim=None,71                 model_type='vace',72                 patch_size=(1, 2, 2),73                 text_len=512,74                 in_dim=16,75                 dim=2048,76                 ffn_dim=8192,77                 freq_dim=256,78                 text_dim=4096,79                 out_dim=16,80                 num_heads=16,81                 num_layers=32,82                 window_size=(-1, -1),83                 qk_norm=True,84                 cross_attn_norm=True,85                 eps=1e-6):86        super().__init__(model_type, patch_size, text_len, in_dim, dim, ffn_dim,87                         freq_dim, text_dim, out_dim, num_heads, num_layers,88                         window_size, qk_norm, cross_attn_norm, eps)89 90        self.vace_layers = [i for i in range(0, self.num_layers, 2)91                           ] if vace_layers is None else vace_layers92        self.vace_in_dim = self.in_dim if vace_in_dim is None else vace_in_dim93 94        assert 0 in self.vace_layers95        self.vace_layers_mapping = {96            i: n for n, i in enumerate(self.vace_layers)97        }98 99        # blocks100        self.blocks = nn.ModuleList([101            BaseWanAttentionBlock(102                't2v_cross_attn',103                self.dim,104                self.ffn_dim,105                self.num_heads,106                self.window_size,107                self.qk_norm,108                self.cross_attn_norm,109                self.eps,110                block_id=self.vace_layers_mapping[i]111                if i in self.vace_layers else None)112            for i in range(self.num_layers)113        ])114 115        # vace blocks116        self.vace_blocks = nn.ModuleList([117            VaceWanAttentionBlock(118                't2v_cross_attn',119                self.dim,120                self.ffn_dim,121                self.num_heads,122                self.window_size,123                self.qk_norm,124                self.cross_attn_norm,125                self.eps,126                block_id=i) for i in self.vace_layers127        ])128 129        # vace patch embeddings130        self.vace_patch_embedding = nn.Conv3d(131            self.vace_in_dim,132            self.dim,133            kernel_size=self.patch_size,134            stride=self.patch_size)135 136    def forward_vace(self, x, vace_context, seq_len, kwargs):137        # embeddings138        c = [self.vace_patch_embedding(u.unsqueeze(0)) for u in vace_context]139        c = [u.flatten(2).transpose(1, 2) for u in c]140        c = torch.cat([141            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],142                      dim=1) for u in c143        ])144 145        # arguments146        new_kwargs = dict(x=x)147        new_kwargs.update(kwargs)148 149        hints = []150        for block in self.vace_blocks:151            c, c_skip = block(c, **new_kwargs)152            hints.append(c_skip)153        return hints154 155    def forward(156        self,157        x,158        t,159        vace_context,160        context,161        seq_len,162        vace_context_scale=1.0,163        clip_fea=None,164        y=None,165    ):166        r"""167        Forward pass through the diffusion model168 169        Args:170            x (List[Tensor]):171                List of input video tensors, each with shape [C_in, F, H, W]172            t (Tensor):173                Diffusion timesteps tensor of shape [B]174            context (List[Tensor]):175                List of text embeddings each with shape [L, C]176            seq_len (`int`):177                Maximum sequence length for positional encoding178            clip_fea (Tensor, *optional*):179                CLIP image features for image-to-video mode180            y (List[Tensor], *optional*):181                Conditional video inputs for image-to-video mode, same shape as x182 183        Returns:184            List[Tensor]:185                List of denoised video tensors with original input shapes [C_out, F, H / 8, W / 8]186        """187        # if self.model_type == 'i2v':188        #     assert clip_fea is not None and y is not None189        # params190        device = self.patch_embedding.weight.device191        if self.freqs.device != device:192            self.freqs = self.freqs.to(device)193 194        # if y is not None:195        #     x = [torch.cat([u, v], dim=0) for u, v in zip(x, y)]196 197        # embeddings198        x = [self.patch_embedding(u.unsqueeze(0)) for u in x]199        grid_sizes = torch.stack(200            [torch.tensor(u.shape[2:], dtype=torch.long) for u in x])201        x = [u.flatten(2).transpose(1, 2) for u in x]202        seq_lens = torch.tensor([u.size(1) for u in x], dtype=torch.long)203        assert seq_lens.max() <= seq_len204        x = torch.cat([205            torch.cat([u, u.new_zeros(1, seq_len - u.size(1), u.size(2))],206                      dim=1) for u in x207        ])208 209        # time embeddings210        with amp.autocast(dtype=torch.float32):211            e = self.time_embedding(212                sinusoidal_embedding_1d(self.freq_dim, t).float())213            e0 = self.time_projection(e).unflatten(1, (6, self.dim))214            assert e.dtype == torch.float32 and e0.dtype == torch.float32215 216        # context217        context_lens = None218        context = self.text_embedding(219            torch.stack([220                torch.cat(221                    [u, u.new_zeros(self.text_len - u.size(0), u.size(1))])222                for u in context223            ]))224 225        # if clip_fea is not None:226        #     context_clip = self.img_emb(clip_fea)  # bs x 257 x dim227        #     context = torch.concat([context_clip, context], dim=1)228 229        # arguments230        kwargs = dict(231            e=e0,232            seq_lens=seq_lens,233            grid_sizes=grid_sizes,234            freqs=self.freqs,235            context=context,236            context_lens=context_lens)237 238        hints = self.forward_vace(x, vace_context, seq_len, kwargs)239        kwargs['hints'] = hints240        kwargs['context_scale'] = vace_context_scale241 242        for block in self.blocks:243            x = block(x, **kwargs)244 245        # head246        x = self.head(x, e)247 248        # unpatchify249        x = self.unpatchify(x, grid_sizes)250        return [u.float() for u in x]251