| 1 | from typing import Optional, Tuple |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings |
| 6 | |
| 7 | # Number of AdaLN modulation parameters per transformer block. |
| 8 | # Base: 2 params (shift + scale) x 3 norms (self-attn, feed-forward, output). |
| 9 | ADALN_NUM_BASE_PARAMS = 6 |
| 10 | # Cross-attention AdaLN adds 3 more (scale, shift, gate) for the CA norm. |
| 11 | ADALN_NUM_CROSS_ATTN_PARAMS = 3 |
| 12 | |
| 13 | |
| 14 | def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: |
| 15 | """Total number of AdaLN parameters per block.""" |
| 16 | return ADALN_NUM_BASE_PARAMS + (ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0) |
| 17 | |
| 18 | |
| 19 | class AdaLayerNormSingle(torch.nn.Module): |
| 20 | r""" |
| 21 | Norm layer adaptive layer norm single (adaLN-single). |
| 22 | As proposed in PixArt-Alpha (see: https://arxiv.org/abs/2310.00426; Section 2.3). |
| 23 | Parameters: |
| 24 | embedding_dim (`int`): The size of each embedding vector. |
| 25 | use_additional_conditions (`bool`): To use additional conditions for normalization or not. |
| 26 | """ |
| 27 | |
| 28 | def __init__(self, embedding_dim: int, embedding_coefficient: int = 6): |
| 29 | super().__init__() |
| 30 | |
| 31 | self.emb = PixArtAlphaCombinedTimestepSizeEmbeddings( |
| 32 | embedding_dim, |
| 33 | size_emb_dim=embedding_dim // 3, |
| 34 | ) |
| 35 | |
| 36 | self.silu = torch.nn.SiLU() |
| 37 | self.linear = torch.nn.Linear(embedding_dim, embedding_coefficient * embedding_dim, bias=True) |
| 38 | |
| 39 | def forward( |
| 40 | self, |
| 41 | timestep: torch.Tensor, |
| 42 | hidden_dtype: Optional[torch.dtype] = None, |
| 43 | ) -> Tuple[torch.Tensor, torch.Tensor]: |
| 44 | embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype) |
| 45 | return self.linear(self.silu(embedded_timestep)), embedded_timestep |
| 46 |