# 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