multimodalart/EchoMimic-zero

1# Adapted from https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention.py2 3from typing import Any, Dict, Optional4 5import torch6from diffusers.models.attention import AdaLayerNorm, Attention, FeedForward7from diffusers.models.embeddings import SinusoidalPositionalEmbedding8from einops import rearrange9from torch import nn10 11 12class BasicTransformerBlock(nn.Module):13    r"""14    A basic Transformer block.15 16    Parameters:17        dim (`int`): The number of channels in the input and output.18        num_attention_heads (`int`): The number of heads to use for multi-head attention.19        attention_head_dim (`int`): The number of channels in each head.20        dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.21        cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.22        activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.23        num_embeds_ada_norm (:24            obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.25        attention_bias (:26            obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.27        only_cross_attention (`bool`, *optional*):28            Whether to use only cross-attention layers. In this case two cross attention layers are used.29        double_self_attention (`bool`, *optional*):30            Whether to use two self-attention layers. In this case no cross attention layers are used.31        upcast_attention (`bool`, *optional*):32            Whether to upcast the attention computation to float32. This is useful for mixed precision training.33        norm_elementwise_affine (`bool`, *optional*, defaults to `True`):34            Whether to use learnable elementwise affine parameters for normalization.35        norm_type (`str`, *optional*, defaults to `"layer_norm"`):36            The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.37        final_dropout (`bool` *optional*, defaults to False):38            Whether to apply a final dropout after the last feed-forward layer.39        attention_type (`str`, *optional*, defaults to `"default"`):40            The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.41        positional_embeddings (`str`, *optional*, defaults to `None`):42            The type of positional embeddings to apply to.43        num_positional_embeddings (`int`, *optional*, defaults to `None`):44            The maximum number of positional embeddings to apply.45    """46 47    def __init__(48        self,49        dim: int,50        num_attention_heads: int,51        attention_head_dim: int,52        dropout=0.0,53        cross_attention_dim: Optional[int] = None,54        activation_fn: str = "geglu",55        num_embeds_ada_norm: Optional[int] = None,56        attention_bias: bool = False,57        only_cross_attention: bool = False,58        double_self_attention: bool = False,59        upcast_attention: bool = False,60        norm_elementwise_affine: bool = True,61        norm_type: str = "layer_norm",  # 'layer_norm', 'ada_norm', 'ada_norm_zero', 'ada_norm_single'62        norm_eps: float = 1e-5,63        final_dropout: bool = False,64        attention_type: str = "default",65        positional_embeddings: Optional[str] = None,66        num_positional_embeddings: Optional[int] = None,67    ):68        super().__init__()69        self.only_cross_attention = only_cross_attention70 71        self.use_ada_layer_norm_zero = (72            num_embeds_ada_norm is not None73        ) and norm_type == "ada_norm_zero"74        self.use_ada_layer_norm = (75            num_embeds_ada_norm is not None76        ) and norm_type == "ada_norm"77        self.use_ada_layer_norm_single = norm_type == "ada_norm_single"78        self.use_layer_norm = norm_type == "layer_norm"79 80        if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:81            raise ValueError(82                f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"83                f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."84            )85 86        if positional_embeddings and (num_positional_embeddings is None):87            raise ValueError(88                "If `positional_embedding` type is defined, `num_positition_embeddings` must also be defined."89            )90 91        if positional_embeddings == "sinusoidal":92            self.pos_embed = SinusoidalPositionalEmbedding(93                dim, max_seq_length=num_positional_embeddings94            )95        else:96            self.pos_embed = None97 98        # Define 3 blocks. Each block has its own normalization layer.99        # 1. Self-Attn100        if self.use_ada_layer_norm:101            self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)102        elif self.use_ada_layer_norm_zero:103            self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)104        else:105            self.norm1 = nn.LayerNorm(106                dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps107            )108 109        self.attn1 = Attention(110            query_dim=dim,111            heads=num_attention_heads,112            dim_head=attention_head_dim,113            dropout=dropout,114            bias=attention_bias,115            upcast_attention=upcast_attention,116        )117 118        # 3. Feed-forward119        if not self.use_ada_layer_norm_single:120            self.norm3 = nn.LayerNorm(121                dim, elementwise_affine=norm_elementwise_affine, eps=norm_eps122            )123 124        self.ff = FeedForward(125            dim,126            dropout=dropout,127            activation_fn=activation_fn,128            final_dropout=final_dropout,129        )130 131        # 4. Fuser132        if attention_type == "gated" or attention_type == "gated-text-image":133            self.fuser = GatedSelfAttentionDense(134                dim, cross_attention_dim, num_attention_heads, attention_head_dim135            )136 137        # 5. Scale-shift for PixArt-Alpha.138        if self.use_ada_layer_norm_single:139            self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)140 141        # let chunk size default to None142        self._chunk_size = None143        self._chunk_dim = 0144 145    def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int = 0):146        # Sets chunk feed-forward147        self._chunk_size = chunk_size148        self._chunk_dim = dim149 150    def forward(151        self,152        hidden_states: torch.FloatTensor,153        attention_mask: Optional[torch.FloatTensor] = None,154        encoder_hidden_states: Optional[torch.FloatTensor] = None,155        encoder_attention_mask: Optional[torch.FloatTensor] = None,156        timestep: Optional[torch.LongTensor] = None,157        cross_attention_kwargs: Dict[str, Any] = None,158        class_labels: Optional[torch.LongTensor] = None,159    ) -> torch.FloatTensor:160        # Notice that normalization is always applied before the real computation in the following blocks.161        # 0. Self-Attention162        batch_size = hidden_states.shape[0]163 164        if self.use_ada_layer_norm:165            norm_hidden_states = self.norm1(hidden_states, timestep)166        elif self.use_ada_layer_norm_zero:167            norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(168                hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype169            )170        elif self.use_layer_norm:171            norm_hidden_states = self.norm1(hidden_states)172        elif self.use_ada_layer_norm_single:173            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (174                self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)175            ).chunk(6, dim=1)176            norm_hidden_states = self.norm1(hidden_states)177            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa178            norm_hidden_states = norm_hidden_states.squeeze(1)179        else:180            raise ValueError("Incorrect norm used")181 182        if self.pos_embed is not None:183            norm_hidden_states = self.pos_embed(norm_hidden_states)184 185        # 1. Retrieve lora scale.186        lora_scale = (187            cross_attention_kwargs.get("scale", 1.0)188            if cross_attention_kwargs is not None189            else 1.0190        )191 192        # 2. Prepare GLIGEN inputs193        cross_attention_kwargs = (194            cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}195        )196        gligen_kwargs = cross_attention_kwargs.pop("gligen", None)197 198        attn_output = self.attn1(199            norm_hidden_states,200            attention_mask=attention_mask,201            **cross_attention_kwargs,202        )203        if self.use_ada_layer_norm_zero:204            attn_output = gate_msa.unsqueeze(1) * attn_output205        elif self.use_ada_layer_norm_single:206            attn_output = gate_msa * attn_output207 208        hidden_states = attn_output + hidden_states209        if hidden_states.ndim == 4:210            hidden_states = hidden_states.squeeze(1)211 212        # 2.5 GLIGEN Control213        if gligen_kwargs is not None:214            hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])215 216 217        # 4. Feed-forward218        if not self.use_ada_layer_norm_single:219            norm_hidden_states = self.norm3(hidden_states)220 221        if self.use_ada_layer_norm_zero:222            norm_hidden_states = (223                norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]224            )225 226        if self.use_ada_layer_norm_single:227            norm_hidden_states = self.norm2(hidden_states)228            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp229 230        ff_output = self.ff(norm_hidden_states, scale=lora_scale)231 232        if self.use_ada_layer_norm_zero:233            ff_output = gate_mlp.unsqueeze(1) * ff_output234        elif self.use_ada_layer_norm_single:235            ff_output = gate_mlp * ff_output236 237        hidden_states = ff_output + hidden_states238        if hidden_states.ndim == 4:239            hidden_states = hidden_states.squeeze(1)240 241        return hidden_states242 243 244class TemporalBasicTransformerBlock(nn.Module):245    def __init__(246        self,247        dim: int,248        num_attention_heads: int,249        attention_head_dim: int,250        dropout=0.0,251        cross_attention_dim: Optional[int] = None,252        activation_fn: str = "geglu",253        num_embeds_ada_norm: Optional[int] = None,254        attention_bias: bool = False,255        only_cross_attention: bool = False,256        upcast_attention: bool = False,257        unet_use_cross_frame_attention=None,258        unet_use_temporal_attention=None,259    ):260        super().__init__()261        self.only_cross_attention = only_cross_attention262        self.use_ada_layer_norm = num_embeds_ada_norm is not None263        self.unet_use_cross_frame_attention = unet_use_cross_frame_attention264        self.unet_use_temporal_attention = unet_use_temporal_attention265 266        # SC-Attn267        self.attn1 = Attention(268            query_dim=dim,269            heads=num_attention_heads,270            dim_head=attention_head_dim,271            dropout=dropout,272            bias=attention_bias,273            upcast_attention=upcast_attention,274        )275        self.norm1 = (276            AdaLayerNorm(dim, num_embeds_ada_norm)277            if self.use_ada_layer_norm278            else nn.LayerNorm(dim)279        )280 281        # Audio Cross-Attn282        if cross_attention_dim is not None:283            self.attn2 = Attention(284                query_dim=dim,285                cross_attention_dim=cross_attention_dim,286                heads=num_attention_heads,287                dim_head=attention_head_dim,288                dropout=dropout,289                bias=attention_bias,290                upcast_attention=upcast_attention,291            )292        else:293            self.attn2 = None294 295        if cross_attention_dim is not None:296            self.norm2 = (297                AdaLayerNorm(dim, num_embeds_ada_norm)298                if self.use_ada_layer_norm299                else nn.LayerNorm(dim)300            )301        else:302            self.norm2 = None303 304        # Feed-forward305        self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)306        self.norm3 = nn.LayerNorm(dim)307        self.use_ada_layer_norm_zero = False308 309        # Temp-Attn310        assert unet_use_temporal_attention is not None311        if unet_use_temporal_attention:312            self.attn_temp = Attention(313                query_dim=dim,314                heads=num_attention_heads,315                dim_head=attention_head_dim,316                dropout=dropout,317                bias=attention_bias,318                upcast_attention=upcast_attention,319            )320            nn.init.zeros_(self.attn_temp.to_out[0].weight.data)321            self.norm_temp = (322                AdaLayerNorm(dim, num_embeds_ada_norm)323                if self.use_ada_layer_norm324                else nn.LayerNorm(dim)325            )326 327    def forward(328        self,329        hidden_states,330        encoder_hidden_states=None,331        audio_cond_fea = None,332        timestep=None,333        attention_mask=None,334        video_length=None,335    ):336        ## implemented in mutual_self_attention.py337        pass338        return hidden_states339