| 1 | from typing import Set, Tuple |
| 2 | |
| 3 | import torch |
| 4 | |
| 5 | from ltx_core.model.audio_vae.attention import AttentionType, make_attn |
| 6 | from ltx_core.model.audio_vae.causality_axis import CausalityAxis |
| 7 | from ltx_core.model.audio_vae.resnet import ResnetBlock |
| 8 | from ltx_core.model.common.normalization import NormType |
| 9 | |
| 10 | |
| 11 | class Downsample(torch.nn.Module): |
| 12 | """ |
| 13 | A downsampling layer that can use either a strided convolution |
| 14 | or average pooling. Supports standard and causal padding for the |
| 15 | convolutional mode. |
| 16 | """ |
| 17 | |
| 18 | def __init__( |
| 19 | self, |
| 20 | in_channels: int, |
| 21 | with_conv: bool, |
| 22 | causality_axis: CausalityAxis = CausalityAxis.WIDTH, |
| 23 | ) -> None: |
| 24 | super().__init__() |
| 25 | self.with_conv = with_conv |
| 26 | self.causality_axis = causality_axis |
| 27 | |
| 28 | if self.causality_axis != CausalityAxis.NONE and not self.with_conv: |
| 29 | raise ValueError("causality is only supported when `with_conv=True`.") |
| 30 | |
| 31 | if self.with_conv: |
| 32 | # Do time downsampling here |
| 33 | # no asymmetric padding in torch conv, must do it ourselves |
| 34 | self.conv = torch.nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) |
| 35 | |
| 36 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 37 | if self.with_conv: |
| 38 | # Padding tuple is in the order: (left, right, top, bottom). |
| 39 | match self.causality_axis: |
| 40 | case CausalityAxis.NONE: |
| 41 | pad = (0, 1, 0, 1) |
| 42 | case CausalityAxis.WIDTH: |
| 43 | pad = (2, 0, 0, 1) |
| 44 | case CausalityAxis.HEIGHT: |
| 45 | pad = (0, 1, 2, 0) |
| 46 | case CausalityAxis.WIDTH_COMPATIBILITY: |
| 47 | pad = (1, 0, 0, 1) |
| 48 | case _: |
| 49 | raise ValueError(f"Invalid causality_axis: {self.causality_axis}") |
| 50 | |
| 51 | x = torch.nn.functional.pad(x, pad, mode="constant", value=0) |
| 52 | x = self.conv(x) |
| 53 | else: |
| 54 | # This branch is only taken if with_conv=False, which implies causality_axis is NONE. |
| 55 | x = torch.nn.functional.avg_pool2d(x, kernel_size=2, stride=2) |
| 56 | |
| 57 | return x |
| 58 | |
| 59 | |
| 60 | def build_downsampling_path( # noqa: PLR0913 |
| 61 | *, |
| 62 | ch: int, |
| 63 | ch_mult: Tuple[int, ...], |
| 64 | num_resolutions: int, |
| 65 | num_res_blocks: int, |
| 66 | resolution: int, |
| 67 | temb_channels: int, |
| 68 | dropout: float, |
| 69 | norm_type: NormType, |
| 70 | causality_axis: CausalityAxis, |
| 71 | attn_type: AttentionType, |
| 72 | attn_resolutions: Set[int], |
| 73 | resamp_with_conv: bool, |
| 74 | ) -> tuple[torch.nn.ModuleList, int]: |
| 75 | """Build the downsampling path with residual blocks, attention, and downsampling layers.""" |
| 76 | down_modules = torch.nn.ModuleList() |
| 77 | curr_res = resolution |
| 78 | in_ch_mult = (1, *tuple(ch_mult)) |
| 79 | block_in = ch |
| 80 | |
| 81 | for i_level in range(num_resolutions): |
| 82 | block = torch.nn.ModuleList() |
| 83 | attn = torch.nn.ModuleList() |
| 84 | block_in = ch * in_ch_mult[i_level] |
| 85 | block_out = ch * ch_mult[i_level] |
| 86 | |
| 87 | for _ in range(num_res_blocks): |
| 88 | block.append( |
| 89 | ResnetBlock( |
| 90 | in_channels=block_in, |
| 91 | out_channels=block_out, |
| 92 | temb_channels=temb_channels, |
| 93 | dropout=dropout, |
| 94 | norm_type=norm_type, |
| 95 | causality_axis=causality_axis, |
| 96 | ) |
| 97 | ) |
| 98 | block_in = block_out |
| 99 | if curr_res in attn_resolutions: |
| 100 | attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) |
| 101 | |
| 102 | down = torch.nn.Module() |
| 103 | down.block = block |
| 104 | down.attn = attn |
| 105 | if i_level != num_resolutions - 1: |
| 106 | down.downsample = Downsample(block_in, resamp_with_conv, causality_axis=causality_axis) |
| 107 | curr_res = curr_res // 2 |
| 108 | down_modules.append(down) |
| 109 | |
| 110 | return down_modules, block_in |
| 111 |