huggingface / huggingface/diffusers

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

Open
#14,307 2 comments 1 reaction 0 assignees View on GitHub
Dominant language
Python
Stars
34.5k
Forks
7.3k
Avg merge
3d 3h
Merged PRs (30d)
91

Description

# [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`:

https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py#L600-L611

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

https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py#L624-L632

and `forward` runs `conv_in` then the upsampler:

https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py#L659-L674

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.**

https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py#L928-L947

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](https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py#L954)
and [L650](https://github.com/huggingface/diffusers/blob/e87b2a7ad836f843295dbc60428e29ec4d5aadb0/src/diffusers/models/autoencoders/autoencoder_kl_ltx.py#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

Contributor guide

Open the contributing guide

Research direction

Start in src/diffusers/models/autoencoders/autoencoder_kl_ltx2.py, reading LTX2VideoUpBlock3d.__init__ and LTX2VideoDecoder3d's up-block construction, then trace the existing forward path. Exercise a non-nominal decoder_block_out_channels configuration and verify that AutoencoderKLLTX2Video can load matching conv_in weights and complete forward without shape errors.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.