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

[yiyixuxu]

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 [modular] LTX-2.5: two-stage generation as one pipeline #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):

# 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 [modular] LTX-2.5: two-stage generation as one pipeline #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.

[yiyixuxu]

cc @akshan-main if you're interested

[yiyixuxu]

Contributions are welcome, but we prefer existing/experienced contributors for this issue!
let me know if anyone is interested

[akshan-main]

@yiyixuxu happy to take this one!