Upload folder using huggingface_hub

block.py

@@ -0,0 +1,42 @@
+from diffusers.modular_pipelines import ModularPipelineBlocks, PipelineState
+from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
+
+try:
+    from .modeling_prunavae import PrunaAutoencoderKLLTX2Video
+except ImportError:
+    from modeling_prunavae import PrunaAutoencoderKLLTX2Video
+
+
+class LoadPrunaVAE(ModularPipelineBlocks):
+    model_name = "PrunaVAED"
+
+    @property
+    def description(self) -> str:
+        return "Declares the Pruna LTX-2 VAE as an expected pipeline component."
+
+    @property
+    def expected_components(self) -> list[ComponentSpec]:
+        return [
+            ComponentSpec(
+                "vae",
+                PrunaAutoencoderKLLTX2Video,
+                pretrained_model_name_or_path="PrunaAI/PrunaVAED",
+                subfolder="vae",
+            )
+        ]
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return []
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return []
+
+    def __call__(self, components, state: PipelineState):
+        # Nothing to compute here -- this block's only role is to make
+        # `components.vae` (a PrunaAutoencoderKLLTX2Video) available to
+        # every block downstream in the pipeline. Blocks that actually
+        # need it (e.g. a decode step) read it directly off `components`,
+        # not off `state`.
+        return components, state

modeling_prunavae.py

@@ -0,0 +1,672 @@
+# 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.

modular_config.json

@@ -0,0 +1,7 @@
+{
+  "_class_name": "LoadPrunaVAE",
+  "_diffusers_version": "0.39.0",
+  "auto_map": {
+    "ModularPipelineBlocks": "block.LoadPrunaVAE"
+  }
+}