ltx2vae ¶

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

    self.in_channels = in_channels
    self.out_channels = out_channels

    kernel_size = (kernel_size, kernel_size, kernel_size)
    self.time_kernel_size = kernel_size[0]

    dilation = (dilation, 1, 1)

    height_pad = kernel_size[1] // 2
    width_pad = kernel_size[2] // 2
    padding = (0, height_pad, width_pad)

    self.conv = nn.Conv3d(
        in_channels,
        out_channels,
        kernel_size,
        stride=stride,
        dilation=dilation,
        padding=padding,
        padding_mode=spatial_padding_mode.value,
        groups=groups,
        bias=bias,
    )

def __init__(
    self,
    dims: int,
    in_channels: int,
    stride: Tuple[int, int, int],
    residual: bool = False,
    out_channels_reduction_factor: int = 1,
    spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
    super().__init__()
    self.stride = stride
    self.out_channels = math.prod(stride) * in_channels // out_channels_reduction_factor
    self.conv = make_conv_nd(
        dims=dims,
        in_channels=in_channels,
        out_channels=self.out_channels,
        kernel_size=3,
        stride=1,
        causal=True,
        spatial_padding_mode=spatial_padding_mode,
    )
    self.residual = residual
    self.out_channels_reduction_factor = out_channels_reduction_factor

def __init__(self, config: dict[str, Any]):
    super().__init__()
    self.config = config
    self.encoder = VideoEncoderConfigurator.from_config(config)
    self.decoder = VideoDecoderConfigurator.from_config(config)
    self._use_tiling: bool = False
    self._use_channels_last_3d: bool = False

    if _is_env_enabled("FASTVIDEO_LTX2_VAE_CHANNELS_LAST_3D", default="1"):
        self.enable_channels_last_3d()

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

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

def __init__(self, latent_channels: int = 128):
    super().__init__()
    self.register_buffer("std-of-means", torch.ones(latent_channels))
    self.register_buffer("mean-of-means", torch.zeros(latent_channels))
    self.register_buffer("mean-of-stds", torch.ones(latent_channels))
    self.register_buffer("mean-of-stds_over_std-of-means", torch.ones(latent_channels))
    self.register_buffer("channel", torch.arange(latent_channels, dtype=torch.float32))

def __init__(self, embedding_dim: int, size_emb_dim: int = 0):
    super().__init__()
    self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0)
    self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim)

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

def __init__(
    self,
    dims: int,
    in_channels: int,
    out_channels: int | None = None,
    dropout: float = 0.0,
    groups: int = 32,
    eps: float = 1e-6,
    norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
    inject_noise: bool = False,
    timestep_conditioning: bool = False,
    spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
    super().__init__()
    self.in_channels = in_channels
    out_channels = in_channels if out_channels is None else out_channels
    self.out_channels = out_channels
    self.inject_noise = inject_noise

    if norm_layer == NormLayerType.GROUP_NORM:
        self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
    elif norm_layer == NormLayerType.PIXEL_NORM:
        self.norm1 = PixelNorm()

    self.non_linearity = nn.SiLU()

    self.conv1 = make_conv_nd(
        dims,
        in_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        padding=1,
        causal=True,
        spatial_padding_mode=spatial_padding_mode,
    )

    if inject_noise:
        self.per_channel_scale1 = nn.Parameter(torch.zeros((in_channels, 1, 1)))

    if norm_layer == NormLayerType.GROUP_NORM:
        self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
    elif norm_layer == NormLayerType.PIXEL_NORM:
        self.norm2 = PixelNorm()

    self.dropout = nn.Dropout(dropout)

    self.conv2 = make_conv_nd(
        dims,
        out_channels,
        out_channels,
        kernel_size=3,
        stride=1,
        padding=1,
        causal=True,
        spatial_padding_mode=spatial_padding_mode,
    )

    if inject_noise:
        self.per_channel_scale2 = nn.Parameter(torch.zeros((in_channels, 1, 1)))

    self.conv_shortcut = (
        make_linear_nd(dims=dims, in_channels=in_channels, out_channels=out_channels)
        if in_channels != out_channels
        else nn.Identity()
    )

    self.norm3 = (
        nn.GroupNorm(num_groups=1, num_channels=in_channels, eps=eps, affine=True)
        if in_channels != out_channels
        else nn.Identity()
    )

    self.timestep_conditioning = timestep_conditioning

    if timestep_conditioning:
        self.scale_shift_table = nn.Parameter(torch.randn(4, in_channels) / in_channels**0.5)

def __init__(
    self,
    dims: int,
    in_channels: int,
    out_channels: int,
    stride: Tuple[int, int, int],
    spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
    super().__init__()
    self.stride = stride
    self.group_size = in_channels * math.prod(stride) // out_channels
    self.conv = make_conv_nd(
        dims=dims,
        in_channels=in_channels,
        out_channels=out_channels // math.prod(stride),
        kernel_size=3,
        stride=1,
        causal=True,
        spatial_padding_mode=spatial_padding_mode,
    )

def __init__(self, in_channels: int, time_embed_dim: int):
    super().__init__()
    self.linear_1 = nn.Linear(in_channels, time_embed_dim)
    self.act = nn.SiLU()
    self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)

def __init__(self, num_channels: int, flip_sin_to_cos: bool = True, downscale_freq_shift: float = 0):
    super().__init__()
    self.num_channels = num_channels
    self.flip_sin_to_cos = flip_sin_to_cos
    self.downscale_freq_shift = downscale_freq_shift

def __init__(
    self,
    dims: int,
    in_channels: int,
    dropout: float = 0.0,
    num_layers: int = 1,
    resnet_eps: float = 1e-6,
    resnet_groups: int = 32,
    norm_layer: NormLayerType = NormLayerType.GROUP_NORM,
    inject_noise: bool = False,
    timestep_conditioning: bool = False,
    spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
    attention_head_dim: int | None = None,  # unused, for compatibility
):
    super().__init__()
    resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

    self.timestep_conditioning = timestep_conditioning

    if timestep_conditioning:
        self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
            embedding_dim=in_channels * 4, size_emb_dim=0
        )

    self.res_blocks = nn.ModuleList(
        [
            ResnetBlock3D(
                dims=dims,
                in_channels=in_channels,
                out_channels=in_channels,
                eps=resnet_eps,
                groups=resnet_groups,
                dropout=dropout,
                norm_layer=norm_layer,
                inject_noise=inject_noise,
                timestep_conditioning=timestep_conditioning,
                spatial_padding_mode=spatial_padding_mode,
            )
            for _ in range(num_layers)
        ]
    )

def __init__(
    self,
    convolution_dimensions: int = 3,
    in_channels: int = 128,
    out_channels: int = 3,
    decoder_blocks: list[tuple[str, int | dict]] = [],
    patch_size: int = 4,
    norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
    causal: bool = False,
    timestep_conditioning: bool = False,
    decoder_spatial_padding_mode: PaddingModeType = PaddingModeType.REFLECT,
):
    super().__init__()

    self.patch_size = patch_size
    out_channels = out_channels * patch_size**2
    self.causal = causal
    self.timestep_conditioning = timestep_conditioning
    self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS

    self.per_channel_statistics = PerChannelStatistics(latent_channels=in_channels)

    self.decode_noise_scale = 0.025
    self.decode_timestep = 0.05

    feature_channels = in_channels
    for block_name, block_params in list(reversed(decoder_blocks)):
        block_config = block_params if isinstance(block_params, dict) else {}
        if block_name == "res_x_y":
            feature_channels = feature_channels * block_config.get("multiplier", 2)
        if block_name == "compress_all":
            feature_channels = feature_channels * block_config.get("multiplier", 1)
        if block_name == "compress_space":
            feature_channels = feature_channels * block_config.get("multiplier", 1)
        if block_name == "compress_time":
            feature_channels = feature_channels * block_config.get("multiplier", 1)

    self.conv_in = make_conv_nd(
        dims=convolution_dimensions,
        in_channels=in_channels,
        out_channels=feature_channels,
        kernel_size=3,
        stride=1,
        padding=1,
        causal=True,
        spatial_padding_mode=decoder_spatial_padding_mode,
    )

    self.up_blocks = nn.ModuleList([])

    for block_name, block_params in list(reversed(decoder_blocks)):
        block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params

        block, feature_channels = _make_decoder_block(
            block_name=block_name,
            block_config=block_config,
            in_channels=feature_channels,
            convolution_dimensions=convolution_dimensions,
            norm_layer=norm_layer,
            timestep_conditioning=timestep_conditioning,
            norm_num_groups=self._norm_num_groups,
            spatial_padding_mode=decoder_spatial_padding_mode,
        )

        self.up_blocks.append(block)

    if norm_layer == NormLayerType.GROUP_NORM:
        self.conv_norm_out = nn.GroupNorm(
            num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6
        )
    elif norm_layer == NormLayerType.PIXEL_NORM:
        self.conv_norm_out = PixelNorm()

    self.conv_act = nn.SiLU()
    self.conv_out = make_conv_nd(
        dims=convolution_dimensions,
        in_channels=feature_channels,
        out_channels=out_channels,
        kernel_size=3,
        padding=1,
        causal=True,
        spatial_padding_mode=decoder_spatial_padding_mode,
    )

    if timestep_conditioning:
        self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
        self.last_time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(
            embedding_dim=feature_channels * 2, size_emb_dim=0
        )
        self.last_scale_shift_table = nn.Parameter(torch.empty(2, feature_channels))

def __init__(
    self,
    convolution_dimensions: int = 3,
    in_channels: int = 3,
    out_channels: int = 128,
    encoder_blocks: list[tuple[str, int]] | list[tuple[str, dict[str, Any]]] = [],
    patch_size: int = 4,
    norm_layer: NormLayerType = NormLayerType.PIXEL_NORM,
    latent_log_var: LogVarianceType = LogVarianceType.UNIFORM,
    encoder_spatial_padding_mode: PaddingModeType = PaddingModeType.ZEROS,
):
    super().__init__()

    self.patch_size = patch_size
    self.norm_layer = norm_layer
    self.latent_channels = out_channels
    self.latent_log_var = latent_log_var
    self._norm_num_groups = self._DEFAULT_NORM_NUM_GROUPS

    self.per_channel_statistics = PerChannelStatistics(latent_channels=out_channels)

    in_channels = in_channels * patch_size**2
    feature_channels = out_channels

    self.conv_in = make_conv_nd(
        dims=convolution_dimensions,
        in_channels=in_channels,
        out_channels=feature_channels,
        kernel_size=3,
        stride=1,
        padding=1,
        causal=True,
        spatial_padding_mode=encoder_spatial_padding_mode,
    )

    self.down_blocks = nn.ModuleList([])

    for block_name, block_params in encoder_blocks:
        block_config = {"num_layers": block_params} if isinstance(block_params, int) else block_params

        block, feature_channels = _make_encoder_block(
            block_name=block_name,
            block_config=block_config,
            in_channels=feature_channels,
            convolution_dimensions=convolution_dimensions,
            norm_layer=norm_layer,
            norm_num_groups=self._norm_num_groups,
            spatial_padding_mode=encoder_spatial_padding_mode,
        )

        self.down_blocks.append(block)

    if norm_layer == NormLayerType.GROUP_NORM:
        self.conv_norm_out = nn.GroupNorm(
            num_channels=feature_channels, num_groups=self._norm_num_groups, eps=1e-6
        )
    elif norm_layer == NormLayerType.PIXEL_NORM:
        self.conv_norm_out = PixelNorm()

    self.conv_act = nn.SiLU()

    conv_out_channels = out_channels
    if latent_log_var == LogVarianceType.PER_CHANNEL:
        conv_out_channels *= 2
    elif latent_log_var in {LogVarianceType.UNIFORM, LogVarianceType.CONSTANT}:
        conv_out_channels += 1
    elif latent_log_var != LogVarianceType.NONE:
        raise ValueError(f"Invalid latent_log_var: {latent_log_var}")

    self.conv_out = make_conv_nd(
        dims=convolution_dimensions,
        in_channels=feature_channels,
        out_channels=conv_out_channels,
        kernel_size=3,
        padding=1,
        causal=True,
        spatial_padding_mode=encoder_spatial_padding_mode,
    )