[LTX-2 VAE] LTX2VideoUpBlock3d.conv_in is wired to the wrong widths

[mehdisi-mohammed-PrunaAI]

[LTX-2 VAE] LTX2VideoUpBlock3d.conv_in is wired to the wrong widths, making any non-nominal decoder_block_out_channels unloadable

Describe the bug

LTX2VideoUpBlock3d has an optional conv_in projection for the case where the tensor
entering the block is not as wide as the block itself. That branch is never taken by the
released LTX-2 checkpoints, and it is broken in two independent ways as soon as it is
taken. As a result, AutoencoderKLLTX2Video can only be instantiated with decoder widths
that follow one exact pattern; any other decoder_block_out_channels yields a model that
either fails in from_pretrained with a shape mismatch or cannot run forward at all.

Both defects are in __init__ only — forward is fine.

1. conv_in projects onto out_channels, but the upsampler that consumes its output
expects out_channels * upscale_factor.

conv_in is built as in_channels -> out_channels:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py

Lines 600 to 611 in e87b2a7

	self.conv_in = None
	if in_channels != out_channels:
	self.conv_in = LTX2VideoResnetBlock3d(
	in_channels=in_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,
	)

while the upsampler right after it is built with in_channels=out_channels * upscale_factor:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py

Lines 624 to 632 in e87b2a7

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

and forward runs conv_in then the upsampler:

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py

Lines 659 to 674 in e87b2a7

	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)

So whenever conv_in exists and upscale_factor > 1, the block cannot run. The guard
should compare against the pre-upsampler width, i.e. in_channels != out_channels * upscale_factor,
and conv_in should project onto that same width.

2. The in_channels that LTX2VideoDecoder3d passes to each up block is not the width of
the tensor that actually arrives there.

diffusers/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py

Lines 928 to 947 in e87b2a7

	# up blocks
	num_block_out_channels = len(block_out_channels)
	self.up_blocks = nn.ModuleList([])
	for i in range(num_block_out_channels):
	input_channel = output_channel // upsample_factor[i]
	output_channel = block_out_channels[i] // upsample_factor[i]
	up_block = LTX2VideoUpBlock3d(
	in_channels=input_channel,
	out_channels=output_channel,
	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,
	)

The tensor leaving up block i-1 is output_channel wide (the previous iteration's value).
Dividing it by the current block's upsample_factor[i] produces a number that corresponds
to no tensor in the graph, so both the in_channels != out_channels decision and the
conv_in input width are computed from a fictitious value.

Why this is invisible on the released checkpoints. For every published LTX-2 video VAE
(diffusers/LTX-2.3-Diffusers, diffusers/LTX-2.3-Distilled-Diffusers) we have
decoder_block_out_channels = [256, 512, 512, 1024] and upsample_factor = [2, 2, 1, 2],
for which input_channel == output_channel in all four up blocks. conv_in is therefore
never instantiated and the whole path is dead code. It is inherited from
LTXVideoUpBlock3d in autoencoder_kl_ltx.py (same two lines,
L954
and L650),
where it happens to be harmless because LTX-1 configs use upsample_factor = (1, 1, 1, 1),
which makes out_channels * upscale_factor == out_channels.

Practical impact. This blocks loading any LTX-2 VAE variant whose decoder has
non-uniform channel widths — e.g. a channel-pruned or distilled decoder — even though the
architecture is otherwise exactly LTX-2 and the checkpoint contains precisely the
up_blocks.N.conv_in.* weights that LTX2VideoUpBlock3d already knows how to build:

ValueError: Cannot load  because decoder.up_blocks.2.conv_in.conv1.conv.bias
expected shape torch.Size([128]), but got torch.Size([256]).

System Info

• diffusers 0.39.0
• torch 2.12.1+cu130, CUDA 13.0
• Python 3.10.12, Linux 6.8.0 x86_64, glibc 2.35
• huggingface_hub 1.22.0, transformers 5.13.0, accelerate 1.13.0, safetensors 0.8.0
• GPU: NVIDIA H100 80GB HBM3 (both repros above run on CPU)

Who can help?

@dg845 @a-r-r-o-w @sayakpaul @DN6 @llcnt