| 1 | from typing import Optional, Tuple, Union |
| 2 | |
| 3 | import torch |
| 4 | from torch import nn |
| 5 | |
| 6 | from ltx_core.model.common.normalization import PixelNorm |
| 7 | from ltx_core.model.transformer.timestep_embedding import PixArtAlphaCombinedTimestepSizeEmbeddings |
| 8 | from ltx_core.model.video_vae.convolution import make_conv_nd, make_linear_nd |
| 9 | from ltx_core.model.video_vae.enums import NormLayerType, PaddingModeType |
| 10 | |
| 11 | |
| 12 | class ResnetBlock3D(nn.Module): |
| 13 | r""" |
| 14 | A Resnet block. |
| 15 | Parameters: |
| 16 | in_channels (`int`): The number of channels in the input. |
| 17 | out_channels (`int`, *optional*, default to be `None`): |
| 18 | The number of output channels for the first conv layer. If None, same as `in_channels`. |
| 19 | dropout (`float`, *optional*, defaults to `0.0`): The dropout probability to use. |
| 20 | groups (`int`, *optional*, default to `32`): The number of groups to use for the first normalization layer. |
| 21 | eps (`float`, *optional*, defaults to `1e-6`): The epsilon to use for the normalization. |
| 22 | """ |
| 23 | |
| 24 | def __init__( |
| 25 | self, |
| 26 | dims: Union[int, Tuple[int, int]], |
| 27 | in_channels: int, |
| 28 | out_channels: Optional[int] = None, |
| 29 | dropout: float = 0.0, |
| 30 | groups: int = 32, |
| 31 | eps: float = 1e-6, |
| 32 | norm_layer: NormLayerType = NormLayerType.PIXEL_NORM, |
| 33 | inject_noise: bool = False, |
| 34 | timestep_conditioning: bool = False, |
| 35 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 36 | ): |
| 37 | super().__init__() |
| 38 | self.in_channels = in_channels |
| 39 | out_channels = in_channels if out_channels is None else out_channels |
| 40 | self.out_channels = out_channels |
| 41 | self.inject_noise = inject_noise |
| 42 | |
| 43 | if norm_layer == NormLayerType.GROUP_NORM: |
| 44 | self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True) |
| 45 | elif norm_layer == NormLayerType.PIXEL_NORM: |
| 46 | self.norm1 = PixelNorm() |
| 47 | |
| 48 | self.non_linearity = nn.SiLU() |
| 49 | |
| 50 | self.conv1 = make_conv_nd( |
| 51 | dims, |
| 52 | in_channels, |
| 53 | out_channels, |
| 54 | kernel_size=3, |
| 55 | stride=1, |
| 56 | padding=1, |
| 57 | causal=True, |
| 58 | spatial_padding_mode=spatial_padding_mode, |
| 59 | ) |
| 60 | |
| 61 | if inject_noise: |
| 62 | self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1))) |
| 63 | |
| 64 | if norm_layer == NormLayerType.GROUP_NORM: |
| 65 | self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True) |
| 66 | elif norm_layer == NormLayerType.PIXEL_NORM: |
| 67 | self.norm2 = PixelNorm() |
| 68 | |
| 69 | self.dropout = torch.nn.Dropout(dropout) |
| 70 | |
| 71 | self.conv2 = make_conv_nd( |
| 72 | dims, |
| 73 | out_channels, |
| 74 | out_channels, |
| 75 | kernel_size=3, |
| 76 | stride=1, |
| 77 | padding=1, |
| 78 | causal=True, |
| 79 | spatial_padding_mode=spatial_padding_mode, |
| 80 | ) |
| 81 | |
| 82 | if inject_noise: |
| 83 | self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1))) |
| 84 | |
| 85 | self.conv_shortcut = ( |
| 86 | make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels) |
| 87 | if in_channels != out_channels |
| 88 | else nn.Identity() |
| 89 | ) |
| 90 | |
| 91 | # Using GroupNorm with 1 group is equivalent to LayerNorm but works with (B, C, ...) layout |
| 92 | # avoiding the need for dimension rearrangement used in standard nn.LayerNorm |
| 93 | self.norm3 = ( |
| 94 | nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True) |
| 95 | if in_channels != out_channels |
| 96 | else nn.Identity() |
| 97 | ) |
| 98 | |
| 99 | self.timestep_conditioning = timestep_conditioning |
| 100 | |
| 101 | if timestep_conditioning: |
| 102 | self.scale_shift_table = nn.Parameter(torch.zeros(4, in_channels)) |
| 103 | |
| 104 | def _feed_spatial_noise( |
| 105 | self, |
| 106 | hidden_states: torch.Tensor, |
| 107 | per_channel_scale: torch.Tensor, |
| 108 | generator: Optional[torch.Generator] = None, |
| 109 | ) -> torch.Tensor: |
| 110 | spatial_shape = hidden_states.shape[-2:] |
| 111 | device = hidden_states.device |
| 112 | dtype = hidden_states.dtype |
| 113 | |
| 114 | # similar to the "explicit noise inputs" method in style-gan |
| 115 | spatial_noise = torch.randn(spatial_shape, device=device, dtype=dtype, generator=generator)[None] |
| 116 | scaled_noise = (spatial_noise * per_channel_scale)[None, :, None, ...] |
| 117 | hidden_states = hidden_states + scaled_noise |
| 118 | |
| 119 | return hidden_states |
| 120 | |
| 121 | def forward( |
| 122 | self, |
| 123 | input_tensor: torch.Tensor, |
| 124 | causal: bool = True, |
| 125 | timestep: Optional[torch.Tensor] = None, |
| 126 | generator: Optional[torch.Generator] = None, |
| 127 | ) -> torch.Tensor: |
| 128 | hidden_states = input_tensor |
| 129 | batch_size = hidden_states.shape[0] |
| 130 | |
| 131 | hidden_states = self.norm1(hidden_states) |
| 132 | if self.timestep_conditioning: |
| 133 | if timestep is None: |
| 134 | raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") |
| 135 | ada_values = self.scale_shift_table[None, ..., None, None, None].to( |
| 136 | device=hidden_states.device, dtype=hidden_states.dtype |
| 137 | ) + timestep.reshape( |
| 138 | batch_size, |
| 139 | 4, |
| 140 | -1, |
| 141 | timestep.shape[-3], |
| 142 | timestep.shape[-2], |
| 143 | timestep.shape[-1], |
| 144 | ) |
| 145 | shift1, scale1, shift2, scale2 = ada_values.unbind(dim=1) |
| 146 | |
| 147 | hidden_states = hidden_states * (1 + scale1) + shift1 |
| 148 | |
| 149 | hidden_states = self.non_linearity(hidden_states) |
| 150 | |
| 151 | hidden_states = self.conv1(hidden_states, causal=causal) |
| 152 | |
| 153 | if self.inject_noise: |
| 154 | hidden_states = self._feed_spatial_noise( |
| 155 | hidden_states, |
| 156 | self.per_channel_scale1.to(device=hidden_states.device, dtype=hidden_states.dtype), |
| 157 | generator=generator, |
| 158 | ) |
| 159 | |
| 160 | hidden_states = self.norm2(hidden_states) |
| 161 | |
| 162 | if self.timestep_conditioning: |
| 163 | hidden_states = hidden_states * (1 + scale2) + shift2 |
| 164 | |
| 165 | hidden_states = self.non_linearity(hidden_states) |
| 166 | |
| 167 | hidden_states = self.dropout(hidden_states) |
| 168 | |
| 169 | hidden_states = self.conv2(hidden_states, causal=causal) |
| 170 | |
| 171 | if self.inject_noise: |
| 172 | hidden_states = self._feed_spatial_noise( |
| 173 | hidden_states, |
| 174 | self.per_channel_scale2.to(device=hidden_states.device, dtype=hidden_states.dtype), |
| 175 | generator=generator, |
| 176 | ) |
| 177 | |
| 178 | input_tensor = self.norm3(input_tensor) |
| 179 | |
| 180 | batch_size = input_tensor.shape[0] |
| 181 | |
| 182 | input_tensor = self.conv_shortcut(input_tensor) |
| 183 | |
| 184 | output_tensor = input_tensor + hidden_states |
| 185 | |
| 186 | return output_tensor |
| 187 | |
| 188 | |
| 189 | class UNetMidBlock3D(nn.Module): |
| 190 | """ |
| 191 | A 3D UNet mid-block [`UNetMidBlock3D`] with multiple residual blocks. |
| 192 | Args: |
| 193 | in_channels (`int`): The number of input channels. |
| 194 | dropout (`float`, *optional*, defaults to 0.0): The dropout rate. |
| 195 | num_layers (`int`, *optional*, defaults to 1): The number of residual blocks. |
| 196 | resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks. |
| 197 | resnet_groups (`int`, *optional*, defaults to 32): |
| 198 | The number of groups to use in the group normalization layers of the resnet blocks. |
| 199 | norm_layer (`str`, *optional*, defaults to `group_norm`): |
| 200 | The normalization layer to use. Can be either `group_norm` or `pixel_norm`. |
| 201 | inject_noise (`bool`, *optional*, defaults to `False`): |
| 202 | Whether to inject noise into the hidden states. |
| 203 | timestep_conditioning (`bool`, *optional*, defaults to `False`): |
| 204 | Whether to condition the hidden states on the timestep. |
| 205 | Returns: |
| 206 | `torch.Tensor`: The output of the last residual block, which is a tensor of shape `(batch_size, |
| 207 | in_channels, height, width)`. |
| 208 | """ |
| 209 | |
| 210 | def __init__( |
| 211 | self, |
| 212 | dims: Union[int, Tuple[int, int]], |
| 213 | in_channels: int, |
| 214 | dropout: float = 0.0, |
| 215 | num_layers: int = 1, |
| 216 | resnet_eps: float = 1e-6, |
| 217 | resnet_groups: int = 32, |
| 218 | norm_layer: NormLayerType = NormLayerType.GROUP_NORM, |
| 219 | inject_noise: bool = False, |
| 220 | timestep_conditioning: bool = False, |
| 221 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 222 | ): |
| 223 | super().__init__() |
| 224 | resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) |
| 225 | |
| 226 | self.timestep_conditioning = timestep_conditioning |
| 227 | |
| 228 | if timestep_conditioning: |
| 229 | self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( |
| 230 | embedding_dim=in_channels * 4, size_emb_dim=0 |
| 231 | ) |
| 232 | |
| 233 | self.res_blocks = nn.ModuleList( |
| 234 | [ |
| 235 | ResnetBlock3D( |
| 236 | dims=dims, |
| 237 | in_channels=in_channels, |
| 238 | out_channels=in_channels, |
| 239 | eps=resnet_eps, |
| 240 | groups=resnet_groups, |
| 241 | dropout=dropout, |
| 242 | norm_layer=norm_layer, |
| 243 | inject_noise=inject_noise, |
| 244 | timestep_conditioning=timestep_conditioning, |
| 245 | spatial_padding_mode=spatial_padding_mode, |
| 246 | ) |
| 247 | for _ in range(num_layers) |
| 248 | ] |
| 249 | ) |
| 250 | |
| 251 | def forward( |
| 252 | self, |
| 253 | hidden_states: torch.Tensor, |
| 254 | causal: bool = True, |
| 255 | timestep: Optional[torch.Tensor] = None, |
| 256 | generator: Optional[torch.Generator] = None, |
| 257 | ) -> torch.Tensor: |
| 258 | timestep_embed = None |
| 259 | if self.timestep_conditioning: |
| 260 | if timestep is None: |
| 261 | raise ValueError("'timestep' parameter must be provided when 'timestep_conditioning' is True") |
| 262 | batch_size = hidden_states.shape[0] |
| 263 | timestep_embed = self.time_embedder( |
| 264 | timestep=timestep.flatten(), |
| 265 | hidden_dtype=hidden_states.dtype, |
| 266 | ) |
| 267 | timestep_embed = timestep_embed.view(batch_size, timestep_embed.shape[-1], 1, 1, 1) |
| 268 | |
| 269 | for resnet in self.res_blocks: |
| 270 | hidden_states = resnet( |
| 271 | hidden_states, |
| 272 | causal=causal, |
| 273 | timestep=timestep_embed, |
| 274 | generator=generator, |
| 275 | ) |
| 276 | |
| 277 | return hidden_states |
| 278 |