| 1 | import torch |
| 2 | import torch.nn.functional as F |
| 3 | |
| 4 | from ltx_core.model.audio_vae.causality_axis import CausalityAxis |
| 5 | |
| 6 | |
| 7 | class CausalConv2d(torch.nn.Module): |
| 8 | """ |
| 9 | A causal 2D convolution. |
| 10 | This layer ensures that the output at time `t` only depends on inputs |
| 11 | at time `t` and earlier. It achieves this by applying asymmetric padding |
| 12 | to the time dimension (width) before the convolution. |
| 13 | """ |
| 14 | |
| 15 | def __init__( |
| 16 | self, |
| 17 | in_channels: int, |
| 18 | out_channels: int, |
| 19 | kernel_size: int | tuple[int, int], |
| 20 | stride: int = 1, |
| 21 | dilation: int | tuple[int, int] = 1, |
| 22 | groups: int = 1, |
| 23 | bias: bool = True, |
| 24 | causality_axis: CausalityAxis = CausalityAxis.HEIGHT, |
| 25 | ) -> None: |
| 26 | super().__init__() |
| 27 | |
| 28 | self.causality_axis = causality_axis |
| 29 | |
| 30 | # Ensure kernel_size and dilation are tuples |
| 31 | kernel_size = torch.nn.modules.utils._pair(kernel_size) |
| 32 | dilation = torch.nn.modules.utils._pair(dilation) |
| 33 | |
| 34 | # Calculate padding dimensions |
| 35 | pad_h = (kernel_size[0] - 1) * dilation[0] |
| 36 | pad_w = (kernel_size[1] - 1) * dilation[1] |
| 37 | |
| 38 | # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom) |
| 39 | match self.causality_axis: |
| 40 | case CausalityAxis.NONE: |
| 41 | self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2) |
| 42 | case CausalityAxis.WIDTH | CausalityAxis.WIDTH_COMPATIBILITY: |
| 43 | self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2) |
| 44 | case CausalityAxis.HEIGHT: |
| 45 | self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0) |
| 46 | case _: |
| 47 | raise ValueError(f"Invalid causality_axis: {causality_axis}") |
| 48 | |
| 49 | # The internal convolution layer uses no padding, as we handle it manually |
| 50 | self.conv = torch.nn.Conv2d( |
| 51 | in_channels, |
| 52 | out_channels, |
| 53 | kernel_size, |
| 54 | stride=stride, |
| 55 | padding=0, |
| 56 | dilation=dilation, |
| 57 | groups=groups, |
| 58 | bias=bias, |
| 59 | ) |
| 60 | |
| 61 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 62 | # Apply causal padding before convolution |
| 63 | x = F.pad(x, self.padding) |
| 64 | return self.conv(x) |
| 65 | |
| 66 | |
| 67 | def make_conv2d( |
| 68 | in_channels: int, |
| 69 | out_channels: int, |
| 70 | kernel_size: int | tuple[int, int], |
| 71 | stride: int = 1, |
| 72 | padding: tuple[int, int, int, int] | None = None, |
| 73 | dilation: int = 1, |
| 74 | groups: int = 1, |
| 75 | bias: bool = True, |
| 76 | causality_axis: CausalityAxis | None = None, |
| 77 | ) -> torch.nn.Module: |
| 78 | """ |
| 79 | Create a 2D convolution layer that can be either causal or non-causal. |
| 80 | Args: |
| 81 | in_channels: Number of input channels |
| 82 | out_channels: Number of output channels |
| 83 | kernel_size: Size of the convolution kernel |
| 84 | stride: Convolution stride |
| 85 | padding: Padding (if None, will be calculated based on causal flag) |
| 86 | dilation: Dilation rate |
| 87 | groups: Number of groups for grouped convolution |
| 88 | bias: Whether to use bias |
| 89 | causality_axis: Dimension along which to apply causality. |
| 90 | Returns: |
| 91 | Either a regular Conv2d or CausalConv2d layer |
| 92 | """ |
| 93 | if causality_axis is not None: |
| 94 | # For causal convolution, padding is handled internally by CausalConv2d |
| 95 | return CausalConv2d(in_channels, out_channels, kernel_size, stride, dilation, groups, bias, causality_axis) |
| 96 | else: |
| 97 | # For non-causal convolution, use symmetric padding if not specified |
| 98 | if padding is None: |
| 99 | padding = kernel_size // 2 if isinstance(kernel_size, int) else tuple(k // 2 for k in kernel_size) |
| 100 | |
| 101 | return torch.nn.Conv2d( |
| 102 | in_channels, |
| 103 | out_channels, |
| 104 | kernel_size, |
| 105 | stride, |
| 106 | padding, |
| 107 | dilation, |
| 108 | groups, |
| 109 | bias, |
| 110 | ) |
| 111 |