| 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.causal_conv_2d import make_conv2d |
| 7 | from ltx_core.model.audio_vae.causality_axis import CausalityAxis |
| 8 | from ltx_core.model.audio_vae.resnet import ResnetBlock |
| 9 | from ltx_core.model.common.normalization import NormType |
| 10 | |
| 11 | |
| 12 | class Upsample(torch.nn.Module): |
| 13 | def __init__( |
| 14 | self, |
| 15 | in_channels: int, |
| 16 | with_conv: bool, |
| 17 | causality_axis: CausalityAxis = CausalityAxis.HEIGHT, |
| 18 | ) -> None: |
| 19 | super().__init__() |
| 20 | self.with_conv = with_conv |
| 21 | self.causality_axis = causality_axis |
| 22 | if self.with_conv: |
| 23 | self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis) |
| 24 | |
| 25 | def forward(self, x: torch.Tensor) -> torch.Tensor: |
| 26 | x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="nearest") |
| 27 | if self.with_conv: |
| 28 | x = self.conv(x) |
| 29 | # Drop FIRST element in the causal axis to undo encoder's padding, while keeping the length 1 + 2 * n. |
| 30 | # For example, if the input is [0, 1, 2], after interpolation, the output is [0, 0, 1, 1, 2, 2]. |
| 31 | # The causal convolution will pad the first element as [-, -, 0, 0, 1, 1, 2, 2], |
| 32 | # So the output elements rely on the following windows: |
| 33 | # 0: [-,-,0] |
| 34 | # 1: [-,0,0] |
| 35 | # 2: [0,0,1] |
| 36 | # 3: [0,1,1] |
| 37 | # 4: [1,1,2] |
| 38 | # 5: [1,2,2] |
| 39 | # Notice that the first and second elements in the output rely only on the first element in the input, |
| 40 | # while all other elements rely on two elements in the input. |
| 41 | # So we can drop the first element to undo the padding (rather than the last element). |
| 42 | # This is a no-op for non-causal convolutions. |
| 43 | match self.causality_axis: |
| 44 | case CausalityAxis.NONE: |
| 45 | pass # x remains unchanged |
| 46 | case CausalityAxis.HEIGHT: |
| 47 | x = x[:, :, 1:, :] |
| 48 | case CausalityAxis.WIDTH: |
| 49 | x = x[:, :, :, 1:] |
| 50 | case CausalityAxis.WIDTH_COMPATIBILITY: |
| 51 | pass # x remains unchanged |
| 52 | case _: |
| 53 | raise ValueError(f"Invalid causality_axis: {self.causality_axis}") |
| 54 | |
| 55 | return x |
| 56 | |
| 57 | |
| 58 | def build_upsampling_path( # noqa: PLR0913 |
| 59 | *, |
| 60 | ch: int, |
| 61 | ch_mult: Tuple[int, ...], |
| 62 | num_resolutions: int, |
| 63 | num_res_blocks: int, |
| 64 | resolution: int, |
| 65 | temb_channels: int, |
| 66 | dropout: float, |
| 67 | norm_type: NormType, |
| 68 | causality_axis: CausalityAxis, |
| 69 | attn_type: AttentionType, |
| 70 | attn_resolutions: Set[int], |
| 71 | resamp_with_conv: bool, |
| 72 | initial_block_channels: int, |
| 73 | ) -> tuple[torch.nn.ModuleList, int]: |
| 74 | """Build the upsampling path with residual blocks, attention, and upsampling layers.""" |
| 75 | up_modules = torch.nn.ModuleList() |
| 76 | block_in = initial_block_channels |
| 77 | curr_res = resolution // (2 ** (num_resolutions - 1)) |
| 78 | |
| 79 | for level in reversed(range(num_resolutions)): |
| 80 | stage = torch.nn.Module() |
| 81 | stage.block = torch.nn.ModuleList() |
| 82 | stage.attn = torch.nn.ModuleList() |
| 83 | block_out = ch * ch_mult[level] |
| 84 | |
| 85 | for _ in range(num_res_blocks + 1): |
| 86 | stage.block.append( |
| 87 | ResnetBlock( |
| 88 | in_channels=block_in, |
| 89 | out_channels=block_out, |
| 90 | temb_channels=temb_channels, |
| 91 | dropout=dropout, |
| 92 | norm_type=norm_type, |
| 93 | causality_axis=causality_axis, |
| 94 | ) |
| 95 | ) |
| 96 | block_in = block_out |
| 97 | if curr_res in attn_resolutions: |
| 98 | stage.attn.append(make_attn(block_in, attn_type=attn_type, norm_type=norm_type)) |
| 99 | |
| 100 | if level != 0: |
| 101 | stage.upsample = Upsample(block_in, resamp_with_conv, causality_axis=causality_axis) |
| 102 | curr_res *= 2 |
| 103 | |
| 104 | up_modules.insert(0, stage) |
| 105 | |
| 106 | return up_modules, block_in |
| 107 |