huggingface / huggingface/diffusers

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

Abierto
#14,730 3 comentarios 0 reacciones 0 asignados Ver en GitHub
contributions-welcome
Lenguaje dominante
Python
Estrellas
34.5k
Forks
7.3k
Merge medio
3 d 3 h
PR fusionados (30 d)
91

Descripción

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

Guía de contribución

Abrir la guía de contribución

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Evaluación

Stack tecnológico
python, pytorch
Área
machine-learning, testing-qa
Tipo de issue
Refactorización
Dificultad
4/5
Tiempo estimado
3-5 días
Estado de actividad
Activo
Claridad
Bien especificado
Aptitud para principiantes
68/100

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.