huggingface / huggingface/diffusers

[modular] move unpack out of the decoder blocks for flux, krea2 and ltx

オープン
#14,730 コメント 3 件 リアクション 0 件 担当者 0 名 GitHub で見る
contributions-welcome
主要言語
Python
スター
34.5k
フォーク
7.3k
平均マージ
3日 3時間
マージ済み PR(30日)
91

説明

**The principle:** there are two possible transformations that sit between a VAE and a transformer — *normalize/denormalize* (the VAE's latent statistics) and *pack/unpack* (`[B, C, F, H, W]` <-> a token sequence `[B, S, D]`). Each must be applied and undone in mirrored pairs, and the core denoise group must hand `latents` back in the same form it received them — so the "`latents`" output is always in one consistent form that any block (a decoder, a latent upsampler, a second denoise group) can consume, and no block outside the group needs `height`/`width`/`num_frames` just to interpret tokens.

There are two acceptable patterns; both follow this principle.

**Pattern 1 (the most common one):** normalize/denormalize on unpacked dimension

```
encode → norm → pack → denoise → unpack → denorm → decode
```

The pack/unpack lives in core-denoise blocks: either at block level (the prepare-latents step packs, a dedicated after-denoise step unpacks) or inside the transformer's `forward` (the model patchifies/unpatchifies internally and blocks never pack at all).

**Pattern 2 (packed-space statistics):** when the VAE's statistics are defined over the *packed* channels, norm/denorm can only run on packed tensors, so pack/unpack has to happen inside the VAE blocks as well:

```
encode → pack → norm → denoise → denorm → unpack → decode
```

**Where the codebase stands** (all 19 modular families):

- Pattern 1, model level (the transformer patchifies internally, blocks never pack): `anima`, `cosmos`, `helios`, `hunyuan_video1_5`, `minimax_music3`, `stable_diffusion_3`, `stable_diffusion_xl`, `wan`, `wan_animate_2`, `z_image` — fine.
- Pattern 1, block level: `qwenimage`, `flux2`, `minimax_h3` — fine. `ltx2` is being moved to this pattern in #14612 (adds `LTX2UnpackLatentsStep` at the end of the core denoise group). `ideogram4` also follows it despite packed-space statistics, by tiling the stats onto the unpacked channels in its decoder (norm/denorm is elementwise and pack is a permutation, so they commute if you permute the stats along).
- Pattern 2: `ernie_image` (decode side: denorm on the packed latents, then unpack, then `vae.decode`; it is t2i-only so there is no encode side) — fine.
- **Outliers — where the fix is needed:** `flux`, `krea2` and `ltx` pack inside the denoise group but unpack inside the decoder block. That asymmetry strands packed `[B, seq_len, dim]` latents in the pipeline state after denoising, where no other block can consume them, and makes the decoders carry geometry inputs they don't otherwise need. (A leftover from porting the standard pipelines, where unpack → denorm → `vae.decode` is just the tail of `__call__`.)

**What to do:**

1. Start with a test in `ModularPipelineTesterMixin` (`tests/modular_pipelines/testing_utils/common.py`) that enforces the principle for every pipeline: each family's tester declares the one canonical form its pipelines keep `latents` in (the VAE form for pattern 1, the channel-packed form for pattern 2), and the test asserts the state actually carries that shape after denoising. A sketch of the idea (not actual code — untested, adapt as needed):

```python
# each tester declares the canonical latents form its family keeps in the state,
# for the geometry that get_dummy_inputs produces
expected_latents_shape = (1, 4, 32, 32)

# on ModularPipelineTesterMixin
def test_latents_output_in_vae_form(self):
# run a partial pipeline without the decode step, and check the latents it leaves in the state
# (existing tests already run sub-blocks this way, see e.g. test_modular_pipeline_stable_diffusion_xl.py)
blocks = self.pipeline_blocks_class()
blocks.sub_blocks.pop("decode")
pipe = blocks.init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components()
latents = pipe(**self.get_dummy_inputs(), output="latents")
assert latents.shape == self.expected_latents_shape
```

Adding the test first surfaces everything that doesn't follow the principle — then look into each failure case to understand why, and whether there is any other pattern we have not identified here.

2. For each pipeline the test catches, add an unpack step at the end of the core denoise group (following `Flux2UnpackLatentsStep`, or `LTX2UnpackLatentsStep` from #14612, as the reference) and remove the unpack + related geometry inputs from the decoder block. Denormalization stays in the decoder, right before `vae.decode` (it just operates on the unpacked latents now).

3. Deprecate the old packed form instead of hard-switching. This is a behavior change: reading `latents` back from the state (e.g. `pipe(..., output="latents")`) currently gives *packed* latents for these three pipelines; after the move it gives unpacked latents. To keep existing workflows running (packed latents that were saved, or fed into a decoder-only pipeline), the decoder block should keep accepting the old packed form for a deprecation window: dispatch on `ndim`, and when latents come in packed, unpack them with a deprecation warning telling users to pass VAE-form latents instead.

**How to verify:** the new mixin test passes across `tests/modular_pipelines/`. Beyond that, this is a pure relocation of a reshape, so image/video outputs must be numerically identical and the rest of the existing tests under `tests/modular_pipelines/flux`, `krea2` and `ltx` should pass. Please also run the slow tests to show the image/video outputs are visually the same.

コントリビューションガイド

コントリビューションガイドを開く

調査の方向性

Start in tests/modular_pipelines/testing_utils/common.py and compare the existing partial-pipeline tests, including test_modular_pipeline_stable_diffusion_xl.py. Add the canonical latent-form check, then trace failures for flux, krea2, and ltx and use the existing Flux2UnpackLatentsStep or LTX2UnpackLatentsStep as references. Done means the mixin and relevant modular tests pass, packed inputs remain supported with deprecation warnings, and image/video outputs remain numerically or visually unchanged.

索引モデルが issue の本文から書いたものです。

評価

技術スタック
python, pytorch
領域
machine-learning, testing-qa
issue の種類
リファクタリング
難易度
4/5
見積もり時間
3〜5日
活発さ
活発
明瞭さ
明確に書かれている
初心者へのやさしさ
68/100

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。