patch_diffusers.py

# Copyright 2026 Pruna AI
# Portions adapted from Hugging Face diffusers:
# Copyright 2026 The HuggingFace Team. Licensed under Apache-2.0.

"""Monkey-patch stock diffusers so pruned LTX-2.3 decoder weights load correctly.

Stock diffusers builds the decoder graph from ``decoder_block_out_channels`` using
the *nominal* width of each up block (divided by ``upsample_factor``). It assumes the
tensor leaving block N always matches that nominal width.

PrunaVAED keeps wider skip connections between blocks and only prunes inside each
block. Example on ``up_blocks.2``:

- tensor arriving from ``up_blocks.1``: **384** ch (after 512→384 projection)
- ResNets inside ``up_blocks.2``: **128** ch (50% pruned vs LTX-2.3 decoder 256)
- ``conv_in`` projection inside the block: **384 → 256** (feeds the upsampler)

Stock diffusers wires ``up_blocks.2`` as 192 → 128 and never creates the 384→256
``conv_in``, so ``from_pretrained`` fails with shape mismatches on PrunaVAED weights.

Fix (two ``__init__`` replacements; forward/decode unchanged)
------------------------------------------------------------
1. ``LTX2VideoDecoder3d``: track ``current_channels`` — the actual tensor width
   between blocks — instead of recomputing from the previous block's nominal width.
2. ``LTX2VideoUpBlock3d``: insert ``conv_in`` when ``in_channels != out_channels *
   upscale_factor`` (pre-upsampler width), not when ``in_channels != out_channels``.

The LTX-2.3 decoder is unaffected (no extra projections). Safe to call before any
``AutoencoderKLLTX2Video.from_pretrained``.
"""

from __future__ import annotations

import torch
import torch.nn as nn
from diffusers.models.autoencoders import autoencoder_kl_ltx2 as ltx2


def patch_pruna_ltx2_decoder() -> None:
    if getattr(ltx2, "_PRUNA_LTX2_DECODER_PATCH", False):
        return

    Resnet = ltx2.LTX2VideoResnetBlock3d
    Upsampler = ltx2.LTX2VideoUpsampler3d
    CausalConv = ltx2.LTX2VideoCausalConv3d
    MidBlock = ltx2.LTX2VideoMidBlock3d
    UpBlock = ltx2.LTX2VideoUpBlock3d
    TimeEmb = ltx2.PixArtAlphaCombinedTimestepSizeEmbeddings
    RMSNorm = ltx2.PerChannelRMSNorm

    def up_init(
        self,
        in_channels,
        out_channels=None,
        num_layers=1,
        dropout=0.0,
        resnet_eps=1e-6,
        resnet_act_fn="swish",
        spatio_temporal_scale=True,
        upsample_type="spatiotemporal",
        inject_noise=False,
        timestep_conditioning=False,
        upsample_residual=False,
        upscale_factor=1,
        spatial_padding_mode="zeros",
    ):
        nn.Module.__init__(self)
        out_channels = out_channels or in_channels
        # Width right before the upsampler (ResNet width × upscale_factor).
        # conv_in projects the incoming skip (e.g. 384) down to this width (e.g. 256).
        pre = out_channels * upscale_factor
        self.time_embedder = TimeEmb(in_channels * 4, 0) if timestep_conditioning else None
        self.conv_in = None
        if in_channels != pre:  # stock diffusers compares in_channels != out_channels
            self.conv_in = Resnet(
                in_channels=in_channels,
                out_channels=pre,
                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:
            stride = {
                "spatial": (1, 2, 2),
                "temporal": (2, 1, 1),
                "spatiotemporal": (2, 2, 2),
            }[upsample_type]
            self.upsamplers = nn.ModuleList(
                [
                    Upsampler(
                        in_channels=pre,
                        stride=stride,
                        residual=upsample_residual,
                        upscale_factor=upscale_factor,
                        spatial_padding_mode=spatial_padding_mode,
                    ),
                ]
            )
        self.resnets = nn.ModuleList(
            [
                Resnet(
                    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,
                )
                for _ in range(num_layers)
            ]
        )
        self.gradient_checkpointing = False

    def decoder_init(
        self,
        in_channels=128,
        out_channels=3,
        block_out_channels=(256, 512, 1024),
        spatio_temporal_scaling=(True, True, True),
        layers_per_block=(5, 5, 5, 5),
        upsample_type=("spatiotemporal", "spatiotemporal", "spatiotemporal"),
        patch_size=4,
        patch_size_t=1,
        resnet_norm_eps=1e-6,
        is_causal=False,
        inject_noise=(False, False, False),
        timestep_conditioning=False,
        upsample_residual=(True, True, True),
        upsample_factor=(2, 2, 2),
        spatial_padding_mode="reflect",
    ):
        nn.Module.__init__(self)
        n = len(layers_per_block)
        if isinstance(spatio_temporal_scaling, bool):
            spatio_temporal_scaling = (spatio_temporal_scaling,) * (n - 1)
        if isinstance(inject_noise, bool):
            inject_noise = (inject_noise,) * n
        if isinstance(upsample_residual, bool):
            upsample_residual = (upsample_residual,) * (n - 1)

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

        ch = 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))
        upsample_type = tuple(upsample_type)
        if len(upsample_type) == len(ch) - 1:
            upsample_type = tuple(reversed(upsample_type))

        width = ch[0]
        self.conv_in = CausalConv(
            in_channels=in_channels,
            out_channels=width,
            kernel_size=3,
            stride=1,
            spatial_padding_mode=spatial_padding_mode,
        )
        self.mid_block = MidBlock(
            in_channels=width,
            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,
        )

        # Stock: input_channel = previous_nominal // factor (loses pruned skip widths).
        # Patched: pass the real tensor width from the previous block as in_channels.
        self.up_blocks = nn.ModuleList()
        current = width
        for i in range(len(ch)):
            resnet_w = ch[i] // upsample_factor[i]
            self.up_blocks.append(
                UpBlock(
                    in_channels=current,
                    out_channels=resnet_w,
                    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,
                )
            )
            current = resnet_w

        self.norm_out = RMSNorm()
        self.conv_act = nn.SiLU()
        self.conv_out = CausalConv(
            in_channels=current,
            out_channels=self.out_channels,
            kernel_size=3,
            stride=1,
            spatial_padding_mode=spatial_padding_mode,
        )
        self.time_embedder = self.scale_shift_table = self.timestep_scale_multiplier = None
        if timestep_conditioning:
            self.timestep_scale_multiplier = nn.Parameter(torch.tensor(1000.0))
            self.time_embedder = TimeEmb(width * 2, 0)
            self.scale_shift_table = nn.Parameter(torch.randn(2, width) / width**0.5)
        self.gradient_checkpointing = False

    UpBlock.__init__ = up_init
    ltx2.LTX2VideoDecoder3d.__init__ = decoder_init
    ltx2._PRUNA_LTX2_DECODER_PATCH = True