A transformer block that combines self-attention, cross-attention and MLP layers with AdaLN modulation. Each component (self-attention, cross-attention, MLP) has its own layer normalization and AdaLN modulation. Parameters: x_dim (int): Dimension of input features conte
| 935 | |
| 936 | |
| 937 | class Block(nn.Module): |
| 938 | """ |
| 939 | A transformer block that combines self-attention, cross-attention and MLP layers with AdaLN modulation. |
| 940 | Each component (self-attention, cross-attention, MLP) has its own layer normalization and AdaLN modulation. |
| 941 | |
| 942 | Parameters: |
| 943 | x_dim (int): Dimension of input features |
| 944 | context_dim (int): Dimension of context features for cross-attention |
| 945 | num_heads (int): Number of attention heads |
| 946 | mlp_ratio (float): Multiplier for MLP hidden dimension. Default: 4.0 |
| 947 | use_adaln_lora (bool): Whether to use AdaLN-LoRA modulation. Default: False |
| 948 | adaln_lora_dim (int): Hidden dimension for AdaLN-LoRA layers. Default: 256 |
| 949 | |
| 950 | The block applies the following sequence: |
| 951 | 1. Self-attention with AdaLN modulation |
| 952 | 2. Cross-attention with AdaLN modulation |
| 953 | 3. MLP with AdaLN modulation |
| 954 | |
| 955 | Each component uses skip connections and layer normalization. |
| 956 | """ |
| 957 | |
| 958 | def __init__( |
| 959 | self, |
| 960 | x_dim: int, |
| 961 | context_dim: int, |
| 962 | num_heads: int, |
| 963 | mlp_ratio: float = 4.0, |
| 964 | use_adaln_lora: bool = False, |
| 965 | adaln_lora_dim: int = 256, |
| 966 | self_attention_backend: str = "transformer_engine", |
| 967 | cross_attention_backend: str = "transformer_engine", |
| 968 | ): |
| 969 | super().__init__() |
| 970 | self.x_dim = x_dim |
| 971 | self.layer_norm_self_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) |
| 972 | self.self_attn = Attention( |
| 973 | x_dim, |
| 974 | None, |
| 975 | num_heads, |
| 976 | x_dim // num_heads, |
| 977 | qkv_format="bshd", |
| 978 | backend=self_attention_backend, |
| 979 | ) |
| 980 | |
| 981 | self.layer_norm_cross_attn = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) |
| 982 | self.cross_attn = Attention( |
| 983 | x_dim, context_dim, num_heads, x_dim // num_heads, qkv_format="bshd", backend=cross_attention_backend |
| 984 | ) |
| 985 | |
| 986 | self.layer_norm_mlp = nn.LayerNorm(x_dim, elementwise_affine=False, eps=1e-6) |
| 987 | self.mlp = GPT2FeedForward(x_dim, int(x_dim * mlp_ratio)) |
| 988 | |
| 989 | self.use_adaln_lora = use_adaln_lora |
| 990 | if self.use_adaln_lora: |
| 991 | self.adaln_modulation_self_attn = nn.Sequential( |
| 992 | nn.SiLU(), |
| 993 | nn.Linear(x_dim, adaln_lora_dim, bias=False), |
| 994 | nn.Linear(adaln_lora_dim, 3 * x_dim, bias=False), |