| 1 | from typing import Tuple, Union |
| 2 | |
| 3 | import torch |
| 4 | from einops import rearrange |
| 5 | from torch import nn |
| 6 | from torch.nn import functional as F |
| 7 | |
| 8 | from ltx_core.model.video_vae.enums import PaddingModeType |
| 9 | |
| 10 | |
| 11 | def make_conv_nd( # noqa: PLR0913 |
| 12 | dims: Union[int, Tuple[int, int]], |
| 13 | in_channels: int, |
| 14 | out_channels: int, |
| 15 | kernel_size: int, |
| 16 | stride: int = 1, |
| 17 | padding: int = 0, |
| 18 | dilation: int = 1, |
| 19 | groups: int = 1, |
| 20 | bias: bool = True, |
| 21 | causal: bool = False, |
| 22 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 23 | temporal_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 24 | ) -> nn.Module: |
| 25 | if not (spatial_padding_mode == temporal_padding_mode or causal): |
| 26 | raise NotImplementedError("spatial and temporal padding modes must be equal") |
| 27 | if dims == 2: |
| 28 | return nn.Conv2d( |
| 29 | in_channels=in_channels, |
| 30 | out_channels=out_channels, |
| 31 | kernel_size=kernel_size, |
| 32 | stride=stride, |
| 33 | padding=padding, |
| 34 | dilation=dilation, |
| 35 | groups=groups, |
| 36 | bias=bias, |
| 37 | padding_mode=spatial_padding_mode.value, |
| 38 | ) |
| 39 | elif dims == 3: |
| 40 | if causal: |
| 41 | return CausalConv3d( |
| 42 | in_channels=in_channels, |
| 43 | out_channels=out_channels, |
| 44 | kernel_size=kernel_size, |
| 45 | stride=stride, |
| 46 | dilation=dilation, |
| 47 | groups=groups, |
| 48 | bias=bias, |
| 49 | spatial_padding_mode=spatial_padding_mode, |
| 50 | ) |
| 51 | return nn.Conv3d( |
| 52 | in_channels=in_channels, |
| 53 | out_channels=out_channels, |
| 54 | kernel_size=kernel_size, |
| 55 | stride=stride, |
| 56 | padding=padding, |
| 57 | dilation=dilation, |
| 58 | groups=groups, |
| 59 | bias=bias, |
| 60 | padding_mode=spatial_padding_mode.value, |
| 61 | ) |
| 62 | elif dims == (2, 1): |
| 63 | return DualConv3d( |
| 64 | in_channels=in_channels, |
| 65 | out_channels=out_channels, |
| 66 | kernel_size=kernel_size, |
| 67 | stride=stride, |
| 68 | padding=padding, |
| 69 | bias=bias, |
| 70 | padding_mode=spatial_padding_mode.value, |
| 71 | ) |
| 72 | else: |
| 73 | raise ValueError(f"unsupported dimensions: {dims}") |
| 74 | |
| 75 | |
| 76 | def make_linear_nd( |
| 77 | dims: int, |
| 78 | in_channels: int, |
| 79 | out_channels: int, |
| 80 | bias: bool = True, |
| 81 | ) -> nn.Module: |
| 82 | if dims == 2: |
| 83 | return nn.Conv2d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) |
| 84 | elif dims in (3, (2, 1)): |
| 85 | return nn.Conv3d(in_channels=in_channels, out_channels=out_channels, kernel_size=1, bias=bias) |
| 86 | else: |
| 87 | raise ValueError(f"unsupported dimensions: {dims}") |
| 88 | |
| 89 | |
| 90 | class DualConv3d(nn.Module): |
| 91 | def __init__( |
| 92 | self, |
| 93 | in_channels: int, |
| 94 | out_channels: int, |
| 95 | kernel_size: int, |
| 96 | stride: Union[int, Tuple[int, int, int]] = 1, |
| 97 | padding: Union[int, Tuple[int, int, int]] = 0, |
| 98 | dilation: Union[int, Tuple[int, int, int]] = 1, |
| 99 | groups: int = 1, |
| 100 | bias: bool = True, |
| 101 | padding_mode: str = "zeros", |
| 102 | ) -> None: |
| 103 | super(DualConv3d, self).__init__() |
| 104 | |
| 105 | self.in_channels = in_channels |
| 106 | self.out_channels = out_channels |
| 107 | self.padding_mode = padding_mode |
| 108 | # Ensure kernel_size, stride, padding, and dilation are tuples of length 3 |
| 109 | if isinstance(kernel_size, int): |
| 110 | kernel_size = (kernel_size, kernel_size, kernel_size) |
| 111 | if kernel_size == (1, 1, 1): |
| 112 | raise ValueError("kernel_size must be greater than 1. Use make_linear_nd instead.") |
| 113 | if isinstance(stride, int): |
| 114 | stride = (stride, stride, stride) |
| 115 | if isinstance(padding, int): |
| 116 | padding = (padding, padding, padding) |
| 117 | if isinstance(dilation, int): |
| 118 | dilation = (dilation, dilation, dilation) |
| 119 | |
| 120 | # Set parameters for convolutions |
| 121 | self.groups = groups |
| 122 | self.bias = bias |
| 123 | |
| 124 | # Define the size of the channels after the first convolution |
| 125 | intermediate_channels = out_channels if in_channels < out_channels else in_channels |
| 126 | |
| 127 | # Define parameters for the first convolution |
| 128 | self.weight1 = nn.Parameter( |
| 129 | torch.Tensor( |
| 130 | intermediate_channels, |
| 131 | in_channels // groups, |
| 132 | 1, |
| 133 | kernel_size[1], |
| 134 | kernel_size[2], |
| 135 | ) |
| 136 | ) |
| 137 | self.stride1 = (1, stride[1], stride[2]) |
| 138 | self.padding1 = (0, padding[1], padding[2]) |
| 139 | self.dilation1 = (1, dilation[1], dilation[2]) |
| 140 | if bias: |
| 141 | self.bias1 = nn.Parameter(torch.Tensor(intermediate_channels)) |
| 142 | else: |
| 143 | self.register_parameter("bias1", None) |
| 144 | |
| 145 | # Define parameters for the second convolution |
| 146 | self.weight2 = nn.Parameter(torch.Tensor(out_channels, intermediate_channels // groups, kernel_size[0], 1, 1)) |
| 147 | self.stride2 = (stride[0], 1, 1) |
| 148 | self.padding2 = (padding[0], 0, 0) |
| 149 | self.dilation2 = (dilation[0], 1, 1) |
| 150 | if bias: |
| 151 | self.bias2 = nn.Parameter(torch.Tensor(out_channels)) |
| 152 | else: |
| 153 | self.register_parameter("bias2", None) |
| 154 | |
| 155 | # Initialize weights and biases |
| 156 | self.reset_parameters() |
| 157 | |
| 158 | def reset_parameters(self) -> None: |
| 159 | nn.init.kaiming_uniform_(self.weight1, a=torch.sqrt(5)) |
| 160 | nn.init.kaiming_uniform_(self.weight2, a=torch.sqrt(5)) |
| 161 | if self.bias: |
| 162 | fan_in1, _ = nn.init._calculate_fan_in_and_fan_out(self.weight1) |
| 163 | bound1 = 1 / torch.sqrt(fan_in1) |
| 164 | nn.init.uniform_(self.bias1, -bound1, bound1) |
| 165 | fan_in2, _ = nn.init._calculate_fan_in_and_fan_out(self.weight2) |
| 166 | bound2 = 1 / torch.sqrt(fan_in2) |
| 167 | nn.init.uniform_(self.bias2, -bound2, bound2) |
| 168 | |
| 169 | def forward( |
| 170 | self, |
| 171 | x: torch.Tensor, |
| 172 | use_conv3d: bool = False, |
| 173 | skip_time_conv: bool = False, |
| 174 | ) -> torch.Tensor: |
| 175 | if use_conv3d: |
| 176 | return self.forward_with_3d(x=x, skip_time_conv=skip_time_conv) |
| 177 | else: |
| 178 | return self.forward_with_2d(x=x, skip_time_conv=skip_time_conv) |
| 179 | |
| 180 | def forward_with_3d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: |
| 181 | # First convolution |
| 182 | x = F.conv3d( |
| 183 | x, |
| 184 | self.weight1, |
| 185 | self.bias1, |
| 186 | self.stride1, |
| 187 | self.padding1, |
| 188 | self.dilation1, |
| 189 | self.groups, |
| 190 | padding_mode=self.padding_mode, |
| 191 | ) |
| 192 | |
| 193 | if skip_time_conv: |
| 194 | return x |
| 195 | |
| 196 | # Second convolution |
| 197 | x = F.conv3d( |
| 198 | x, |
| 199 | self.weight2, |
| 200 | self.bias2, |
| 201 | self.stride2, |
| 202 | self.padding2, |
| 203 | self.dilation2, |
| 204 | self.groups, |
| 205 | padding_mode=self.padding_mode, |
| 206 | ) |
| 207 | |
| 208 | return x |
| 209 | |
| 210 | def forward_with_2d(self, x: torch.Tensor, skip_time_conv: bool = False) -> torch.Tensor: |
| 211 | b, _, _, h, w = x.shape |
| 212 | |
| 213 | # First 2D convolution |
| 214 | x = rearrange(x, "b c d h w -> (b d) c h w") |
| 215 | # Squeeze the depth dimension out of weight1 since it's 1 |
| 216 | weight1 = self.weight1.squeeze(2) |
| 217 | # Select stride, padding, and dilation for the 2D convolution |
| 218 | stride1 = (self.stride1[1], self.stride1[2]) |
| 219 | padding1 = (self.padding1[1], self.padding1[2]) |
| 220 | dilation1 = (self.dilation1[1], self.dilation1[2]) |
| 221 | x = F.conv2d( |
| 222 | x, |
| 223 | weight1, |
| 224 | self.bias1, |
| 225 | stride1, |
| 226 | padding1, |
| 227 | dilation1, |
| 228 | self.groups, |
| 229 | padding_mode=self.padding_mode, |
| 230 | ) |
| 231 | |
| 232 | _, _, h, w = x.shape |
| 233 | |
| 234 | if skip_time_conv: |
| 235 | x = rearrange(x, "(b d) c h w -> b c d h w", b=b) |
| 236 | return x |
| 237 | |
| 238 | # Second convolution which is essentially treated as a 1D convolution across the 'd' dimension |
| 239 | x = rearrange(x, "(b d) c h w -> (b h w) c d", b=b) |
| 240 | |
| 241 | # Reshape weight2 to match the expected dimensions for conv1d |
| 242 | weight2 = self.weight2.squeeze(-1).squeeze(-1) |
| 243 | # Use only the relevant dimension for stride, padding, and dilation for the 1D convolution |
| 244 | stride2 = self.stride2[0] |
| 245 | padding2 = self.padding2[0] |
| 246 | dilation2 = self.dilation2[0] |
| 247 | x = F.conv1d( |
| 248 | x, |
| 249 | weight2, |
| 250 | self.bias2, |
| 251 | stride2, |
| 252 | padding2, |
| 253 | dilation2, |
| 254 | self.groups, |
| 255 | padding_mode=self.padding_mode, |
| 256 | ) |
| 257 | x = rearrange(x, "(b h w) c d -> b c d h w", b=b, h=h, w=w) |
| 258 | |
| 259 | return x |
| 260 | |
| 261 | @property |
| 262 | def weight(self) -> torch.Tensor: |
| 263 | return self.weight2 |
| 264 | |
| 265 | |
| 266 | class CausalConv3d(nn.Module): |
| 267 | def __init__( |
| 268 | self, |
| 269 | in_channels: int, |
| 270 | out_channels: int, |
| 271 | kernel_size: int = 3, |
| 272 | stride: Union[int, Tuple[int]] = 1, |
| 273 | dilation: int = 1, |
| 274 | groups: int = 1, |
| 275 | bias: bool = True, |
| 276 | spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS, |
| 277 | ) -> None: |
| 278 | super().__init__() |
| 279 | |
| 280 | self.in_channels = in_channels |
| 281 | self.out_channels = out_channels |
| 282 | |
| 283 | kernel_size = (kernel_size, kernel_size, kernel_size) |
| 284 | self.time_kernel_size = kernel_size[0] |
| 285 | |
| 286 | dilation = (dilation, 1, 1) |
| 287 | |
| 288 | height_pad = kernel_size[1] // 2 |
| 289 | width_pad = kernel_size[2] // 2 |
| 290 | padding = (0, height_pad, width_pad) |
| 291 | |
| 292 | self.conv = nn.Conv3d( |
| 293 | in_channels, |
| 294 | out_channels, |
| 295 | kernel_size, |
| 296 | stride=stride, |
| 297 | dilation=dilation, |
| 298 | padding=padding, |
| 299 | padding_mode=spatial_padding_mode.value, |
| 300 | groups=groups, |
| 301 | bias=bias, |
| 302 | ) |
| 303 | |
| 304 | def forward(self, x: torch.Tensor, causal: bool = True) -> torch.Tensor: |
| 305 | if causal: |
| 306 | first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, self.time_kernel_size - 1, 1, 1)) |
| 307 | x = torch.concatenate((first_frame_pad, x), dim=2) |
| 308 | else: |
| 309 | first_frame_pad = x[:, :, :1, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) |
| 310 | last_frame_pad = x[:, :, -1:, :, :].repeat((1, 1, (self.time_kernel_size - 1) // 2, 1, 1)) |
| 311 | x = torch.concatenate((first_frame_pad, x, last_frame_pad), dim=2) |
| 312 | x = self.conv(x) |
| 313 | return x |
| 314 | |
| 315 | @property |
| 316 | def weight(self) -> torch.Tensor: |
| 317 | return self.conv.weight |
| 318 |