ltx2_audio_vae ¶

def __init__(
    self,
    in_channels: int,
    norm_type: NormType = NormType.GROUP,
) -> None:
    super().__init__()
    self.in_channels = in_channels

    self.norm = build_normalization_layer(in_channels, normtype=norm_type)
    self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
    self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
    self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)
    self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0)

def __init__(
    self,
    *,
    ch: int,
    out_ch: int,
    ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
    num_res_blocks: int,
    attn_resolutions: Set[int],
    resolution: int,
    z_channels: int,
    norm_type: NormType = NormType.GROUP,
    causality_axis: CausalityAxis = CausalityAxis.WIDTH,
    dropout: float = 0.0,
    mid_block_add_attention: bool = True,
    sample_rate: int = 16000,
    mel_hop_length: int = 160,
    is_causal: bool = True,
    mel_bins: int | None = None,
) -> None:
    super().__init__()

    # Internal behavioural defaults
    resamp_with_conv = True
    attn_type = AttentionType.VANILLA

    # Per-channel statistics for denormalizing latents
    self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
    self.sample_rate = sample_rate
    self.mel_hop_length = mel_hop_length
    self.is_causal = is_causal
    self.mel_bins = mel_bins
    self.patchifier = AudioPatchifier(
        patch_size=1,
        audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
        sample_rate=sample_rate,
        hop_length=mel_hop_length,
        is_causal=is_causal,
    )

    self.ch = ch
    self.temb_ch = 0
    self.num_resolutions = len(ch_mult)
    self.num_res_blocks = num_res_blocks
    self.resolution = resolution
    self.out_ch = out_ch
    self.give_pre_end = False
    self.tanh_out = False
    self.norm_type = norm_type
    self.z_channels = z_channels
    self.channel_multipliers = ch_mult
    self.attn_resolutions = attn_resolutions
    self.causality_axis = causality_axis
    self.attn_type = attn_type

    base_block_channels = ch * self.channel_multipliers[-1]
    base_resolution = resolution // (2**(self.num_resolutions - 1))
    self.z_shape = (1, z_channels, base_resolution, base_resolution)

    self.conv_in = make_conv2d(z_channels,
                               base_block_channels,
                               kernel_size=3,
                               stride=1,
                               causality_axis=self.causality_axis)
    self.non_linearity = nn.SiLU()

    self.mid = build_mid_block(
        channels=base_block_channels,
        temb_channels=self.temb_ch,
        dropout=dropout,
        norm_type=self.norm_type,
        causality_axis=self.causality_axis,
        attn_type=self.attn_type,
        add_attention=mid_block_add_attention,
    )

    self.up, final_block_channels = build_upsampling_path(
        ch=ch,
        ch_mult=ch_mult,
        num_resolutions=self.num_resolutions,
        num_res_blocks=num_res_blocks,
        resolution=resolution,
        temb_channels=self.temb_ch,
        dropout=dropout,
        norm_type=self.norm_type,
        causality_axis=self.causality_axis,
        attn_type=self.attn_type,
        attn_resolutions=attn_resolutions,
        resamp_with_conv=resamp_with_conv,
        initial_block_channels=base_block_channels,
    )

    self.norm_out = build_normalization_layer(final_block_channels, normtype=self.norm_type)
    self.conv_out = make_conv2d(final_block_channels,
                                out_ch,
                                kernel_size=3,
                                stride=1,
                                causality_axis=self.causality_axis)

def __init__(
    self,
    *,
    ch: int,
    ch_mult: Tuple[int, ...] = (1, 2, 4, 8),
    num_res_blocks: int,
    attn_resolutions: Set[int],
    dropout: float = 0.0,
    resamp_with_conv: bool = True,
    in_channels: int,
    resolution: int,
    z_channels: int,
    double_z: bool = True,
    attn_type: AttentionType = AttentionType.VANILLA,
    mid_block_add_attention: bool = True,
    norm_type: NormType = NormType.GROUP,
    causality_axis: CausalityAxis = CausalityAxis.WIDTH,
    sample_rate: int = 16000,
    mel_hop_length: int = 160,
    n_fft: int = 1024,
    is_causal: bool = True,
    mel_bins: int = 64,
    **_ignore_kwargs,
) -> None:
    super().__init__()

    self.per_channel_statistics = PerChannelStatistics(latent_channels=ch)
    self.sample_rate = sample_rate
    self.mel_hop_length = mel_hop_length
    self.n_fft = n_fft
    self.is_causal = is_causal
    self.mel_bins = mel_bins

    self.patchifier = AudioPatchifier(
        patch_size=1,
        audio_latent_downsample_factor=LATENT_DOWNSAMPLE_FACTOR,
        sample_rate=sample_rate,
        hop_length=mel_hop_length,
        is_causal=is_causal,
    )

    self.ch = ch
    self.temb_ch = 0
    self.num_resolutions = len(ch_mult)
    self.num_res_blocks = num_res_blocks
    self.resolution = resolution
    self.in_channels = in_channels
    self.z_channels = z_channels
    self.double_z = double_z
    self.norm_type = norm_type
    self.causality_axis = causality_axis
    self.attn_type = attn_type

    # Input convolution
    self.conv_in = make_conv2d(
        in_channels,
        self.ch,
        kernel_size=3,
        stride=1,
        causality_axis=self.causality_axis,
    )

    self.non_linearity = nn.SiLU()

    # Downsampling path
    self.down, block_in = build_downsampling_path(
        ch=ch,
        ch_mult=ch_mult,
        num_resolutions=self.num_resolutions,
        num_res_blocks=num_res_blocks,
        resolution=resolution,
        temb_channels=self.temb_ch,
        dropout=dropout,
        norm_type=self.norm_type,
        causality_axis=self.causality_axis,
        attn_type=self.attn_type,
        attn_resolutions=attn_resolutions,
        resamp_with_conv=resamp_with_conv,
    )

    # Mid block
    self.mid = build_mid_block(
        channels=block_in,
        temb_channels=self.temb_ch,
        dropout=dropout,
        norm_type=self.norm_type,
        causality_axis=self.causality_axis,
        attn_type=self.attn_type,
        add_attention=mid_block_add_attention,
    )

    # Output layers
    self.norm_out = build_normalization_layer(block_in, normtype=self.norm_type)
    self.conv_out = make_conv2d(
        block_in,
        2 * z_channels if double_z else z_channels,
        kernel_size=3,
        stride=1,
        causality_axis=self.causality_axis,
    )

def __init__(
    self,
    patch_size: int = 1,
    sample_rate: int = 16000,
    hop_length: int = 160,
    audio_latent_downsample_factor: int = 4,
    is_causal: bool = True,
):
    self.patch_size = patch_size
    self.sample_rate = sample_rate
    self.hop_length = hop_length
    self.audio_latent_downsample_factor = audio_latent_downsample_factor
    self.is_causal = is_causal

def __init__(
    self,
    in_channels: int,
    out_channels: int,
    kernel_size: int | Tuple[int, int],
    stride: int = 1,
    dilation: int | Tuple[int, int] = 1,
    groups: int = 1,
    bias: bool = True,
    causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
    super().__init__()

    self.causality_axis = causality_axis

    # Ensure kernel_size and dilation are tuples
    kernel_size = nn.modules.utils._pair(kernel_size)
    dilation = nn.modules.utils._pair(dilation)

    # Calculate padding dimensions
    pad_h = (kernel_size[0] - 1) * dilation[0]
    pad_w = (kernel_size[1] - 1) * dilation[1]

    # The padding tuple for F.pad is (pad_left, pad_right, pad_top, pad_bottom)
    if self.causality_axis == CausalityAxis.NONE:
        self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h // 2, pad_h - pad_h // 2)
    elif self.causality_axis in (CausalityAxis.WIDTH, CausalityAxis.WIDTH_COMPATIBILITY):
        self.padding = (pad_w, 0, pad_h // 2, pad_h - pad_h // 2)
    elif self.causality_axis == CausalityAxis.HEIGHT:
        self.padding = (pad_w // 2, pad_w - pad_w // 2, pad_h, 0)
    else:
        raise ValueError(f"Invalid causality_axis: {causality_axis}")

    # The internal convolution layer uses no padding, as we handle it manually
    self.conv = nn.Conv2d(
        in_channels,
        out_channels,
        kernel_size,
        stride=stride,
        padding=0,
        dilation=dilation,
        groups=groups,
        bias=bias,
    )

def __init__(
    self,
    in_channels: int,
    with_conv: bool,
    causality_axis: CausalityAxis = CausalityAxis.WIDTH,
) -> None:
    super().__init__()
    self.with_conv = with_conv
    self.causality_axis = causality_axis

    if self.causality_axis != CausalityAxis.NONE and not self.with_conv:
        raise ValueError("causality is only supported when `with_conv=True`.")

    if self.with_conv:
        self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0)

def __init__(self, config: dict[str, Any]):
    super().__init__()
    self.model: AudioDecoder = AudioDecoderConfigurator.from_config(config)

def __init__(self, config: dict[str, Any]):
    super().__init__()
    self.model: AudioEncoder = AudioEncoderConfigurator.from_config(config)

def __init__(self, config: dict[str, Any]):
    super().__init__()
    self.model: Vocoder = VocoderConfigurator.from_config(config)

def __init__(self, latent_channels: int = 128) -> None:
    super().__init__()
    self.register_buffer("std-of-means", torch.empty(latent_channels))
    self.register_buffer("mean-of-means", torch.empty(latent_channels))

def __init__(self, dim: int = 1, eps: float = 1e-8) -> None:
    super().__init__()
    self.dim = dim
    self.eps = eps

def __init__(
        self,
        channels: int,
        kernel_size: int = 3,
        dilation: Tuple[int, int, int] = (1, 3, 5),
):
    super().__init__()
    self.convs1 = nn.ModuleList([
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"),
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"),
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2], padding="same"),
    ])
    self.convs2 = nn.ModuleList([
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"),
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"),
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=1, padding="same"),
    ])

def __init__(
        self,
        channels: int,
        kernel_size: int = 3,
        dilation: Tuple[int, int] = (1, 3),
):
    super().__init__()
    self.convs = nn.ModuleList([
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0], padding="same"),
        nn.Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1], padding="same"),
    ])

def __init__(
    self,
    *,
    in_channels: int,
    out_channels: int | None = None,
    conv_shortcut: bool = False,
    dropout: float = 0.0,
    temb_channels: int = 512,
    norm_type: NormType = NormType.GROUP,
    causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
    super().__init__()
    self.causality_axis = causality_axis

    if self.causality_axis != CausalityAxis.NONE and norm_type == NormType.GROUP:
        raise ValueError("Causal ResnetBlock with GroupNorm is not supported.")

    self.in_channels = in_channels
    out_channels = in_channels if out_channels is None else out_channels
    self.out_channels = out_channels
    self.use_conv_shortcut = conv_shortcut

    self.norm1 = build_normalization_layer(in_channels, normtype=norm_type)
    self.non_linearity = nn.SiLU()
    self.conv1 = make_conv2d(in_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
    if temb_channels > 0:
        self.temb_proj = nn.Linear(temb_channels, out_channels)
    self.norm2 = build_normalization_layer(out_channels, normtype=norm_type)
    self.dropout = nn.Dropout(dropout)
    self.conv2 = make_conv2d(out_channels, out_channels, kernel_size=3, stride=1, causality_axis=causality_axis)
    if self.in_channels != self.out_channels:
        if self.use_conv_shortcut:
            self.conv_shortcut = make_conv2d(in_channels,
                                             out_channels,
                                             kernel_size=3,
                                             stride=1,
                                             causality_axis=causality_axis)
        else:
            self.nin_shortcut = make_conv2d(in_channels,
                                            out_channels,
                                            kernel_size=1,
                                            stride=1,
                                            causality_axis=causality_axis)

def __init__(
    self,
    in_channels: int,
    with_conv: bool,
    causality_axis: CausalityAxis = CausalityAxis.HEIGHT,
) -> None:
    super().__init__()
    self.with_conv = with_conv
    self.causality_axis = causality_axis
    if self.with_conv:
        self.conv = make_conv2d(in_channels, in_channels, kernel_size=3, stride=1, causality_axis=causality_axis)

def __init__(
    self,
    resblock_kernel_sizes: List[int] | None = None,
    upsample_rates: List[int] | None = None,
    upsample_kernel_sizes: List[int] | None = None,
    resblock_dilation_sizes: List[List[int]] | None = None,
    upsample_initial_channel: int = 1024,
    stereo: bool = True,
    resblock: str = "1",
    output_sample_rate: int = 24000,
    activation: str = "snake",
    use_tanh_at_final: bool = True,
    apply_final_activation: bool = True,
    use_bias_at_final: bool = True,
):
    super().__init__()

    # Initialize default values
    if resblock_kernel_sizes is None:
        resblock_kernel_sizes = [3, 7, 11]
    if upsample_rates is None:
        upsample_rates = [6, 5, 2, 2, 2]
    if upsample_kernel_sizes is None:
        upsample_kernel_sizes = [16, 15, 8, 4, 4]
    if resblock_dilation_sizes is None:
        resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]]

    self.output_sample_rate = output_sample_rate
    self.num_kernels = len(resblock_kernel_sizes)
    self.num_upsamples = len(upsample_rates)
    self.use_tanh_at_final = use_tanh_at_final
    self.apply_final_activation = apply_final_activation
    self.is_amp = resblock == "AMP1"
    in_channels = 128 if stereo else 64
    self.conv_pre = nn.Conv1d(in_channels, upsample_initial_channel, 7, 1, padding=3)
    if resblock == "1":
        resblock_class = ResBlock1
    elif resblock == "AMP1":
        resblock_class = AMPBlock1
    else:
        resblock_class = ResBlock2

    self.ups = nn.ModuleList()
    for i, (stride, kernel_size) in enumerate(zip(upsample_rates, upsample_kernel_sizes, strict=True)):
        self.ups.append(
            nn.ConvTranspose1d(
                upsample_initial_channel // (2**i),
                upsample_initial_channel // (2**(i + 1)),
                kernel_size,
                stride,
                padding=(kernel_size - stride) // 2,
            ))

    self.resblocks = nn.ModuleList()
    for i, _ in enumerate(self.ups):
        ch = upsample_initial_channel // (2**(i + 1))
        for kernel_size, dilations in zip(resblock_kernel_sizes, resblock_dilation_sizes, strict=True):
            if self.is_amp:
                self.resblocks.append(resblock_class(ch, kernel_size, dilations, activation=activation))
            else:
                self.resblocks.append(resblock_class(ch, kernel_size, dilations))

    out_channels = 2 if stereo else 1
    final_channels = upsample_initial_channel // (2**self.num_upsamples)
    if self.is_amp:
        self.act_post: nn.Module = Activation1d(SnakeBeta(final_channels))
    else:
        self.act_post = nn.LeakyReLU()
    self.conv_post = nn.Conv1d(
        final_channels,
        out_channels,
        7,
        1,
        padding=3,
        bias=use_bias_at_final,
    )

    self.upsample_factor = math.prod(layer.stride[0] for layer in self.ups)