modeling_prunavae.py

# Copyright 2025 The Lightricks team, The HuggingFace Team, and Pruna AI.
# All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Pruna variant of the LTX-2 video decoder / autoencoder.

This module is intentionally kept structurally identical to upstream
Diffusers' ``autoencoder_kl_ltx2.py``. It exists to support checkpoints
produced by ``PrunaVAED``, which prunes the internal ResNet width of the
decoder's up-blocks while preserving wider skip connections between decoder
stages than stock LTX-2 assumes.

There are exactly three intentional deviations from upstream, each isolated
to a single class so that future rebases against upstream Diffusers can diff
each class independently:

1. ``PrunaLTX2VideoUpBlock3d``
   The ``conv_in`` projection is constructed against the pre-upsampler
   channel width (``out_channels * upscale_factor``) rather than the
   ResNet width (``out_channels``). Upstream implicitly assumes
   ``in_channels == out_channels`` is the only case that needs no
   projection; Pruna's decoder keeps wider skip tensors between stages, so
   that assumption no longer holds.

2. ``PrunaLTX2VideoDecoder3d``
   Up-block input widths are tracked via an explicit ``current_channels``
   accumulator (the true width of the tensor leaving the previous stage)
   instead of being re-derived from ``block_out_channels[i] //
   upsample_factor[i]``. This is the direct consequence of deviation (1):
   once skip widths are no longer implicitly recoverable from
   ``block_out_channels`` alone, the decoder must track them explicitly.
   It also instantiates ``PrunaLTX2VideoUpBlock3d`` in place of the
   upstream ``LTX2VideoUpBlock3d``.

3. ``PrunaAutoencoderKLLTX2Video``
   The constructor is otherwise identical to
   ``AutoencoderKLLTX2Video.__init__``; the only change is that
   ``self.decoder`` is built from ``PrunaLTX2VideoDecoder3d`` instead of
   ``LTX2VideoDecoder3d``. Every other method (``encode``, ``decode``,
   ``forward``, ``tiled_encode``, ``tiled_decode``, etc.) is inherited
   unchanged.

``forward()`` is unchanged in every class below relative to upstream: none
of the three deviations touch execution semantics, only module
construction. This checkpoint topology matches what ``PrunaVAED`` produces
while remaining forward-compatible with the original LTX-2 decoder.
"""

from __future__ import annotations

import torch
import torch.nn as nn

from diffusers.configuration_utils import register_to_config
from diffusers.models.autoencoders.autoencoder_kl_ltx2 import (
    AutoencoderKLLTX2Video,
    LTX2VideoCausalConv3d,
    LTX2VideoMidBlock3d,
    LTX2VideoResnetBlock3d,
    LTX2VideoUpsampler3d,
    PerChannelRMSNorm,
)
from diffusers.models.embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings


# Deliberately NOT imported, since this module replaces them:
#   LTX2VideoDecoder3d, LTX2VideoUpBlock3d


class PrunaLTX2VideoUpBlock3d(nn.Module):
    r"""
    Pruna variant of ``LTX2VideoUpBlock3d``.

    This implementation differs from the upstream Diffusers version in one
    important way:

    The optional ``conv_in`` projection operates on the **pre-upsampler**
    channel width rather than the ResNet width.

    Upstream compares

        in_channels != out_channels

    which assumes the incoming tensor has already been pruned down to the
    block's internal ResNet width before it arrives.

    Pruna preserves wider skip connections between decoder stages and only
    prunes the internal ResNet channels, so we compare against the
    pre-upsampler width instead:

        pre_upsample_channels = out_channels * upscale_factor

    Example
    -------
        incoming tensor : 384 channels
        ResNet width    : 128 channels
        upscale_factor  : 2

    The ResNet therefore expects a 256-channel tensor before the upsampler,
    requiring a 384 -> 256 projection that upstream's narrower comparison
    would never trigger.

    This exactly matches the checkpoint topology produced by ``PrunaVAED``
    while remaining forward-compatible with the original LTX-2 decoder.

    Args:
        in_channels (`int`):
            Number of input channels.
        out_channels (`int`, *optional*):
            Number of output channels. If None, defaults to `in_channels`.
        num_layers (`int`, defaults to `1`):
            Number of resnet layers.
        dropout (`float`, defaults to `0.0`):
            Dropout rate.
        resnet_eps (`float`, defaults to `1e-6`):
            Epsilon value for normalization layers.
        resnet_act_fn (`str`, defaults to `"swish"`):
            Activation function to use.
        spatio_temporal_scale (`bool`, defaults to `True`):
            Whether or not to use an upsampling layer. If not used, output
            dimension would be same as input dimension.
        upscale_factor (`int`, defaults to `1`):
            Channel upscale factor applied by the upsampler.
    """

    _supports_gradient_checkpointing = True

    def __init__(
        self,
        in_channels: int,
        out_channels: int | None = None,
        num_layers: int = 1,
        dropout: float = 0.0,
        resnet_eps: float = 1e-6,
        resnet_act_fn: str = "swish",
        spatio_temporal_scale: bool = True,
        upsample_type: str = "spatiotemporal",
        inject_noise: bool = False,
        timestep_conditioning: bool = False,
        upsample_residual: bool = False,
        upscale_factor: int = 1,
        spatial_padding_mode: str = "zeros",
    ):
        super().__init__()

        out_channels = out_channels or in_channels

        #
        # ------------------------------------------------------------------
        # PRUNA CHANGE (1 of 1 in this class)
        #
        # Width immediately before the upsampler.
        #
        # Stock Diffusers compares:
        #
        #     in_channels != out_channels
        #
        # which assumes the incoming tensor has already been pruned.
        #
        # Pruna preserves wider skip tensors between decoder stages.
        # Therefore we compare against the pre-upsampler width instead.
        # ------------------------------------------------------------------
        #
        pre_upsample_channels = out_channels * upscale_factor

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

        self.conv_in = None
        if in_channels != pre_upsample_channels:
            self.conv_in = LTX2VideoResnetBlock3d(
                in_channels=in_channels,
                out_channels=pre_upsample_channels,
                dropout=dropout,
                eps=resnet_eps,
                non_linearity=resnet_act_fn,
                inject_noise=inject_noise,
                timestep_conditioning=timestep_conditioning,
                spatial_padding_mode=spatial_padding_mode,
            )

        self.upsamplers = None
        if spatio_temporal_scale:
            self.upsamplers = nn.ModuleList()

            if upsample_type == "spatial":
                upsample_stride = (1, 2, 2)
            elif upsample_type == "temporal":
                upsample_stride = (2, 1, 1)
            elif upsample_type == "spatiotemporal":
                upsample_stride = (2, 2, 2)
            else:
                # Upstream leaves this branch implicit; making the failure
                # explicit improves debuggability without changing behavior
                # for any valid configuration.
                raise ValueError(f"Unknown upsample_type: {upsample_type}")

            self.upsamplers.append(
                LTX2VideoUpsampler3d(
                    in_channels=pre_upsample_channels,
                    stride=upsample_stride,
                    residual=upsample_residual,
                    upscale_factor=upscale_factor,
                    spatial_padding_mode=spatial_padding_mode,
                )
            )

        resnets = []
        for _ in range(num_layers):
            resnets.append(
                LTX2VideoResnetBlock3d(
                    in_channels=out_channels,
                    out_channels=out_channels,
                    dropout=dropout,
                    eps=resnet_eps,
                    non_linearity=resnet_act_fn,
                    inject_noise=inject_noise,
                    timestep_conditioning=timestep_conditioning,
                    spatial_padding_mode=spatial_padding_mode,
                )
            )
        self.resnets = nn.ModuleList(resnets)

        self.gradient_checkpointing = False

    # Identical to upstream `LTX2VideoUpBlock3d.forward` -- the Pruna change
    # is purely in module construction (`__init__`), not execution.
    def forward(
        self,
        hidden_states: torch.Tensor,
        temb: torch.Tensor | None = None,
        generator: torch.Generator | None = None,
        causal: bool = True,
    ) -> torch.Tensor:
        if self.conv_in is not None:
            hidden_states = self.conv_in(hidden_states, temb, generator, causal=causal)

        if self.time_embedder is not None:
            temb = self.time_embedder(
                timestep=temb.flatten(),
                resolution=None,
                aspect_ratio=None,
                batch_size=hidden_states.size(0),
                hidden_dtype=hidden_states.dtype,
            )
            temb = temb.view(hidden_states.size(0), -1, 1, 1, 1)

        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                hidden_states = upsampler(hidden_states, causal=causal)

        for resnet in self.resnets:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                hidden_states = self._gradient_checkpointing_func(
                    resnet, hidden_states, temb, generator, causal
                )
            else:
                hidden_states = resnet(hidden_states, temb, generator, causal=causal)

        return hidden_states


class PrunaLTX2VideoDecoder3d(nn.Module):
    r"""
    Pruna variant of ``LTX2VideoDecoder3d``.

    Deliberately **not** a subclass of ``LTX2VideoDecoder3d``: the entire
    `__init__` would be overridden anyway (the up-block construction loop
    must change), so subclassing would only add coupling to the upstream
    constructor's private attributes while saving nothing but the ~50-line
    `forward()` method, which is reproduced verbatim below instead.

    There are exactly two semantic changes relative to upstream:

    1. Up-block input widths are tracked with an explicit
       ``current_channels`` accumulator -- the true width of the tensor
       leaving the previous stage -- rather than being re-derived from
       ``block_out_channels[i] // upsample_factor[i]``. Upstream's
       re-derivation implicitly assumes stage skip widths collapse to the
       ResNet width; Pruna's decoder does not make that assumption.
    2. Each stage instantiates ``PrunaLTX2VideoUpBlock3d`` instead of
       ``LTX2VideoUpBlock3d``.

    Everything else -- ``conv_in``, ``mid_block``, ``norm_out``,
    ``conv_out``, timestep conditioning, and ``forward()`` -- is copied
    unchanged from upstream.

    Args:
        in_channels (`int`, defaults to 128):
            Number of latent channels.
        out_channels (`int`, defaults to 3):
            Number of output channels.
        block_out_channels (`tuple[int, ...]`, defaults to `(256, 512, 1024)`):
            The number of output channels for each block.
        spatio_temporal_scaling (`tuple[bool, ...]`, defaults to `(True, True, True)`):
            Whether a block should contain spatio-temporal upscaling layers or not.
        layers_per_block (`tuple[int, ...]`, defaults to `(5, 5, 5, 5)`):
            The number of layers per block.
        patch_size (`int`, defaults to `4`):
            The size of spatial patches.
        patch_size_t (`int`, defaults to `1`):
            The size of temporal patches.
        resnet_norm_eps (`float`, defaults to `1e-6`):
            Epsilon value for ResNet normalization layers.
        is_causal (`bool`, defaults to `False`):
            Whether this layer behaves causally (future frames depend only on past frames) or not.
        timestep_conditioning (`bool`, defaults to `False`):
            Whether to condition the model on timesteps.
    """

    _supports_gradient_checkpointing = True

    def __init__(
        self,
        in_channels: int = 128,
        out_channels: int = 3,
        block_out_channels: tuple[int, ...] = (256, 512, 1024),
        spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True),
        layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
        upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
        patch_size: int = 4,
        patch_size_t: int = 1,
        resnet_norm_eps: float = 1e-6,
        is_causal: bool = False,
        inject_noise: bool | tuple[bool, ...] = (False, False, False),
        timestep_conditioning: bool = False,
        upsample_residual: bool | tuple[bool, ...] = (True, True, True),
        upsample_factor: tuple[int, ...] = (2, 2, 2),
        spatial_padding_mode: str = "reflect",
    ) -> None:
        super().__init__()
        num_decoder_blocks = len(layers_per_block)
        if isinstance(spatio_temporal_scaling, bool):
            spatio_temporal_scaling = (spatio_temporal_scaling,) * (num_decoder_blocks - 1)
        if isinstance(inject_noise, bool):
            inject_noise = (inject_noise,) * num_decoder_blocks
        if isinstance(upsample_residual, bool):
            upsample_residual = (upsample_residual,) * (num_decoder_blocks - 1)

        self.patch_size = patch_size
        self.patch_size_t = patch_size_t
        self.out_channels = out_channels * patch_size**2
        self.is_causal = is_causal

        block_out_channels = tuple(reversed(block_out_channels))
        spatio_temporal_scaling = tuple(reversed(spatio_temporal_scaling))
        layers_per_block = tuple(reversed(layers_per_block))
        inject_noise = tuple(reversed(inject_noise))
        upsample_residual = tuple(reversed(upsample_residual))
        upsample_factor = tuple(reversed(upsample_factor))
        output_channel = block_out_channels[0]

        self.conv_in = LTX2VideoCausalConv3d(
            in_channels=in_channels,
            out_channels=output_channel,
            kernel_size=3,
            stride=1,
            spatial_padding_mode=spatial_padding_mode,
        )

        self.mid_block = LTX2VideoMidBlock3d(
            in_channels=output_channel,
            num_layers=layers_per_block[0],
            resnet_eps=resnet_norm_eps,
            inject_noise=inject_noise[0],
            timestep_conditioning=timestep_conditioning,
            spatial_padding_mode=spatial_padding_mode,
        )

        # up blocks
        num_block_out_channels = len(block_out_channels)
        self.up_blocks = nn.ModuleList([])

        #
        # ------------------------------------------------------------------
        # PRUNA CHANGE (1 of 2 in this class)
        #
        # Upstream re-derives each stage's input width from
        # `block_out_channels[i] // upsample_factor[i]`, which implicitly
        # assumes the tensor leaving a stage is exactly that stage's ResNet
        # width. Pruna's decoder keeps wider skip connections between
        # stages, so we instead track the *actual* channel width of the
        # tensor as it flows from stage to stage.
        #
        # After `conv_in` + `mid_block`, that width is `output_channel`
        # (== block_out_channels[0]); after each up-block it becomes that
        # block's `resnet_width`.
        # ------------------------------------------------------------------
        #
        current_channels = output_channel

        for i in range(num_block_out_channels):
            resnet_width = block_out_channels[i] // upsample_factor[i]

            #
            # ------------------------------------------------------------------
            # PRUNA CHANGE (2 of 2 in this class)
            #
            # Instantiate the Pruna up-block, which projects from the true
            # incoming skip width (`current_channels`) rather than assuming
            # it already equals the ResNet width.
            # ------------------------------------------------------------------
            #
            up_block = PrunaLTX2VideoUpBlock3d(
                in_channels=current_channels,
                out_channels=resnet_width,
                num_layers=layers_per_block[i + 1],
                resnet_eps=resnet_norm_eps,
                spatio_temporal_scale=spatio_temporal_scaling[i],
                upsample_type=upsample_type[i],
                inject_noise=inject_noise[i + 1],
                timestep_conditioning=timestep_conditioning,
                upsample_residual=upsample_residual[i],
                upscale_factor=upsample_factor[i],
                spatial_padding_mode=spatial_padding_mode,
            )

            self.up_blocks.append(up_block)
            current_channels = resnet_width

        output_channel = current_channels

        # out
        self.norm_out = PerChannelRMSNorm()
        self.conv_act = nn.SiLU()
        self.conv_out = LTX2VideoCausalConv3d(
            in_channels=output_channel,
            out_channels=self.out_channels,
            kernel_size=3,
            stride=1,
            spatial_padding_mode=spatial_padding_mode,
        )

        # timestep embedding
        self.time_embedder = None
        self.scale_shift_table = None
        self.timestep_scale_multiplier = None
        if timestep_conditioning:
            self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0, dtype=torch.float32))
            self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings(output_channel * 2, 0)
            self.scale_shift_table = nn.Parameter(torch.randn(2, output_channel) / output_channel**0.5)

        self.gradient_checkpointing = False

    # Identical to upstream `LTX2VideoDecoder3d.forward` -- both Pruna
    # changes above are construction-time only.
    def forward(
        self,
        hidden_states: torch.Tensor,
        temb: torch.Tensor | None = None,
        causal: bool | None = None,
    ) -> torch.Tensor:
        causal = causal or self.is_causal

        hidden_states = self.conv_in(hidden_states, causal=causal)

        if self.timestep_scale_multiplier is not None:
            temb = temb * self.timestep_scale_multiplier

        if torch.is_grad_enabled() and self.gradient_checkpointing:
            hidden_states = self._gradient_checkpointing_func(self.mid_block, hidden_states, temb, None, causal)

            for up_block in self.up_blocks:
                hidden_states = self._gradient_checkpointing_func(up_block, hidden_states, temb, None, causal)
        else:
            hidden_states = self.mid_block(hidden_states, temb, causal=causal)

            for up_block in self.up_blocks:
                hidden_states = up_block(hidden_states, temb, causal=causal)

        hidden_states = self.norm_out(hidden_states)

        if self.time_embedder is not None:
            temb = self.time_embedder(
                timestep=temb.flatten(),
                resolution=None,
                aspect_ratio=None,
                batch_size=hidden_states.size(0),
                hidden_dtype=hidden_states.dtype,
            )
            temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1))
            temb = temb + self.scale_shift_table[None, ..., None, None, None]
            shift, scale = temb.unbind(dim=1)
            hidden_states = hidden_states * (1 + scale) + shift

        hidden_states = self.conv_act(hidden_states)
        hidden_states = self.conv_out(hidden_states, causal=causal)

        p = self.patch_size
        p_t = self.patch_size_t

        batch_size, num_channels, num_frames, height, width = hidden_states.shape
        hidden_states = hidden_states.reshape(batch_size, -1, p_t, p, p, num_frames, height, width)
        hidden_states = hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3).flatten(6, 7).flatten(4, 5).flatten(2, 3)

        return hidden_states


class PrunaAutoencoderKLLTX2Video(AutoencoderKLLTX2Video):
    r"""
    Pruna variant of ``AutoencoderKLLTX2Video``.

    Differs from upstream in exactly one constructor line: ``self.decoder``
    is built from :class:`PrunaLTX2VideoDecoder3d` instead of
    ``LTX2VideoDecoder3d``. The encoder, buffers, tiling configuration, and
    framewise decoding setup are all copied verbatim from upstream.

    Every other method -- ``encode``, ``decode``, ``forward``,
    ``tiled_encode``, ``tiled_decode``, ``enable_tiling``,
    ``enable_slicing``, etc. -- is inherited unchanged from
    ``AutoencoderKLLTX2Video``, since none of them depend on the decoder's
    internal channel-width bookkeeping.
    """

    _supports_gradient_checkpointing = True

    @register_to_config
    def __init__(
        self,
        in_channels: int = 3,
        out_channels: int = 3,
        latent_channels: int = 128,
        block_out_channels: tuple[int, ...] = (256, 512, 1024, 2048),
        down_block_types: tuple[str, ...] = (
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
            "LTX2VideoDownBlock3D",
        ),
        decoder_block_out_channels: tuple[int, ...] = (256, 512, 1024),
        layers_per_block: tuple[int, ...] = (4, 6, 6, 2, 2),
        decoder_layers_per_block: tuple[int, ...] = (5, 5, 5, 5),
        spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True, True),
        decoder_spatio_temporal_scaling: bool | tuple[bool, ...] = (True, True, True),
        decoder_inject_noise: bool | tuple[bool, ...] = (False, False, False, False),
        downsample_type: tuple[str, ...] = ("spatial", "temporal", "spatiotemporal", "spatiotemporal"),
        upsample_type: tuple[str, ...] = ("spatiotemporal", "spatiotemporal", "spatiotemporal"),
        upsample_residual: bool | tuple[bool, ...] = (True, True, True),
        upsample_factor: tuple[int, ...] = (2, 2, 2),
        timestep_conditioning: bool = False,
        patch_size: int = 4,
        patch_size_t: int = 1,
        resnet_norm_eps: float = 1e-6,
        scaling_factor: float = 1.0,
        encoder_causal: bool = True,
        decoder_causal: bool = True,
        encoder_spatial_padding_mode: str = "zeros",
        decoder_spatial_padding_mode: str = "reflect",
        spatial_compression_ratio: int = None,
        temporal_compression_ratio: int = None,
    ) -> None:
        # Bypass AutoencoderKLLTX2Video.__init__ (and its `self.decoder =
        # LTX2VideoDecoder3d(...)` line) entirely; go straight to
        # nn.Module.__init__ via the mixin chain, exactly as upstream does.
        super(AutoencoderKLLTX2Video, self).__init__()

        num_encoder_blocks = len(layers_per_block)
        num_decoder_blocks = len(decoder_layers_per_block)
        if isinstance(spatio_temporal_scaling, bool):
            spatio_temporal_scaling = (spatio_temporal_scaling,) * (num_encoder_blocks - 1)
        if isinstance(decoder_spatio_temporal_scaling, bool):
            decoder_spatio_temporal_scaling = (decoder_spatio_temporal_scaling,) * (num_decoder_blocks - 1)
        if isinstance(decoder_inject_noise, bool):
            decoder_inject_noise = (decoder_inject_noise,) * num_decoder_blocks
        if isinstance(upsample_residual, bool):
            upsample_residual = (upsample_residual,) * (num_decoder_blocks - 1)

        # Import the encoder + downstream block type lazily from upstream so
        # this file never needs to redefine anything on the encoder side.
        from diffusers.models.autoencoders.autoencoder_kl_ltx2 import LTX2VideoEncoder3d

        self.encoder = LTX2VideoEncoder3d(
            in_channels=in_channels,
            out_channels=latent_channels,
            block_out_channels=block_out_channels,
            down_block_types=down_block_types,
            spatio_temporal_scaling=spatio_temporal_scaling,
            layers_per_block=layers_per_block,
            downsample_type=downsample_type,
            patch_size=patch_size,
            patch_size_t=patch_size_t,
            resnet_norm_eps=resnet_norm_eps,
            is_causal=encoder_causal,
            spatial_padding_mode=encoder_spatial_padding_mode,
        )

        #
        # ------------------------------------------------------------------
        # PRUNA CHANGE (the only change in this class)
        #
        # Instantiate PrunaLTX2VideoDecoder3d instead of LTX2VideoDecoder3d.
        # Every argument passed is identical to upstream.
        # ------------------------------------------------------------------
        #
        self.decoder = PrunaLTX2VideoDecoder3d(
            in_channels=latent_channels,
            out_channels=out_channels,
            block_out_channels=decoder_block_out_channels,
            spatio_temporal_scaling=decoder_spatio_temporal_scaling,
            layers_per_block=decoder_layers_per_block,
            upsample_type=upsample_type,
            patch_size=patch_size,
            patch_size_t=patch_size_t,
            resnet_norm_eps=resnet_norm_eps,
            is_causal=decoder_causal,
            timestep_conditioning=timestep_conditioning,
            inject_noise=decoder_inject_noise,
            upsample_residual=upsample_residual,
            upsample_factor=upsample_factor,
            spatial_padding_mode=decoder_spatial_padding_mode,
        )

        latents_mean = torch.zeros((latent_channels,), requires_grad=False)
        latents_std = torch.ones((latent_channels,), requires_grad=False)
        self.register_buffer("latents_mean", latents_mean, persistent=True)
        self.register_buffer("latents_std", latents_std, persistent=True)

        self.spatial_compression_ratio = (
            patch_size * 2 ** sum(spatio_temporal_scaling)
            if spatial_compression_ratio is None
            else spatial_compression_ratio
        )
        self.temporal_compression_ratio = (
            patch_size_t * 2 ** sum(spatio_temporal_scaling)
            if temporal_compression_ratio is None
            else temporal_compression_ratio
        )

        # When decoding a batch of video latents at a time, one can save memory by slicing across the batch dimension
        # to perform decoding of a single video latent at a time.
        self.use_slicing = False

        # When decoding spatially large video latents, the memory requirement is very high. By breaking the video latent
        # frames spatially into smaller tiles and performing multiple forward passes for decoding, and then blending the
        # intermediate tiles together, the memory requirement can be lowered.
        self.use_tiling = False

        # When decoding temporally long video latents, the memory requirement is very high. By decoding latent frames
        # at a fixed frame batch size (based on `self.num_latent_frames_batch_sizes`), the memory requirement can be lowered.
        self.use_framewise_encoding = False
        self.use_framewise_decoding = False

        # This can be configured based on the amount of GPU memory available.
        # `16` for sample frames and `2` for latent frames are sensible defaults for consumer GPUs.
        # Setting it to higher values results in higher memory usage.
        self.num_sample_frames_batch_size = 16
        self.num_latent_frames_batch_size = 2

        # The minimal tile height and width for spatial tiling to be used
        self.tile_sample_min_height = 512
        self.tile_sample_min_width = 512
        self.tile_sample_min_num_frames = 16

        # The minimal distance between two spatial tiles
        self.tile_sample_stride_height = 448
        self.tile_sample_stride_width = 448
        self.tile_sample_stride_num_frames = 8

    # encode(), decode(), forward(), tiled_encode(), tiled_decode(), and all
    # other methods are inherited unchanged from AutoencoderKLLTX2Video.