move unpack out of the flux, krea2 and ltx decoders into the core denoise groups

.ai/references/modular.md

@@ -191,6 +191,15 @@ Default to taking that option. The only reason not to split is when the variant
 
 Don't fall back to the standard-pipeline habit of a config flag branching inside a shared block (`ConfigSpec(name="is_distilled")` + `if components.config.is_distilled:`). That keeps both variants' behavior bundled in one blockset — and the input surface is the one thing it can never fix: a repo can override components and config values per checkpoint, but never which inputs the blocks declare, so the distilled checkpoint would still accept `negative_prompt` and silently ignore it.
 
+## Key pattern: One canonical latents form
+
+Two transformations can 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 is applied and undone in mirrored pairs, and the core denoise group hands latents back in the same form it received them. So the `latents` a family leaves in the state always has 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.
+
+- Pattern 1 (the common one): `encode -> norm -> pack -> denoise -> unpack -> denorm -> decode`. The pack/unpack lives in the core denoise group, either at block level (a prepare-latents step packs, an unpack step such as `Flux2UnpackLatentsStep` closes the group) or inside the transformer's forward.
+- 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 happens inside the VAE blocks as well: `encode -> pack -> norm -> denoise -> denorm -> unpack -> decode`.
+
+Don't pack inside the denoise group and unpack inside the decoder: that strands packed latents in the state and makes the decoder carry geometry inputs it doesn't otherwise need. `test_latents_output_in_canonical_form` in `tests/modular_pipelines/testing_utils/common.py` pins each family's form through the tester's `expected_latents_shape`.
+
 ## Key pattern: Standalone block reusability
 
 One of the core reason a pipeline is split into blocks at all: each block (text encoder, VAE encoder, prepare-latents, denoise, decoder) must be runnable on its own, and its output must be reusable as the input to a different downstream chain.

src/diffusers/modular_pipelines/flux/decoders.py

@@ -20,7 +20,7 @@
 
 from ...configuration_utils import FrozenDict
 from ...models import AutoencoderKL
-from ...utils import logging
+from ...utils import deprecate, logging
 from ...video_processor import VaeImageProcessor
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
@@ -45,6 +45,50 @@ def _unpack_latents(latents, height, width, vae_scale_factor):
     return latents
 
 
+class FluxUnpackLatentsStep(ModularPipelineBlocks):
+    model_name = "flux"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, H, W]` form the "
+            "VAE takes (still normalized). Closes the core denoise group, so the blocks that follow take the same "
+            "form the VAE encoder produces and need no geometry inputs."
+        )
+
+    @property
+    def inputs(self) -> list[tuple[str, Any]]:
+        return [
+            InputParam(
+                "latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The denoised latents from the denoising step, packed, of shape `[B, S, C]`.",
+            ),
+            InputParam("height", default=1024),
+            InputParam("width", default=1024),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[str]:
+        return [
+            OutputParam(
+                "latents",
+                type_hint=torch.Tensor,
+                description="The denoised latents of shape `[B, C, H, W]` (normalized, not packed).",
+            )
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        block_state.latents = _unpack_latents(
+            block_state.latents, block_state.height, block_state.width, components.vae_scale_factor
+        )
+        self.set_block_state(state, block_state)
+        return components, state
+
+
 class FluxDecodeStep(ModularPipelineBlocks):
     model_name = "flux"
 
@@ -68,14 +112,15 @@ def description(self) -> str:
     def inputs(self) -> list[tuple[str, Any]]:
         return [
             InputParam("output_type", default="pil"),
-            InputParam("height", default=1024),
-            InputParam("width", default=1024),
             InputParam(
                 "latents",
                 required=True,
                 type_hint=torch.Tensor,
-                description="The denoised latents from the denoising step",
+                description="The denoised latents from the denoising step, of shape `[B, C, H, W]`.",
             ),
+            # Only read on the deprecated path that still accepts packed `[B, S, C]` latents.
+            InputParam("height", default=1024),
+            InputParam("width", default=1024),
         ]
 
     @property
@@ -95,7 +140,14 @@ def __call__(self, components, state: PipelineState) -> PipelineState:
 
         if not block_state.output_type == "latent":
             latents = block_state.latents
-            latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor)
+            if latents.ndim == 3:
+                deprecate(
+                    "packed latents",
+                    "1.0.0",
+                    "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group "
+                    "now unpacks them. Pass latents of shape `[B, C, H, W]` instead.",
+                )
+                latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor)
             latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor
             block_state.images = vae.decode(latents, return_dict=False)[0]
             block_state.images = components.image_processor.postprocess(

src/diffusers/modular_pipelines/flux/modular_blocks_flux.py

@@ -22,7 +22,7 @@
     FluxRoPEInputsStep,
     FluxSetTimestepsStep,
 )
-from .decoders import FluxDecodeStep
+from .decoders import FluxDecodeStep, FluxUnpackLatentsStep
 from .denoise import FluxDenoiseStep
 from .encoders import (
     FluxProcessImagesInputStep,
@@ -480,8 +480,8 @@ class FluxCoreDenoiseStep(SequentialPipelineBlocks):
     """
 
     model_name = "flux"
-    block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxDenoiseStep]
-    block_names = ["input", "before_denoise", "denoise"]
+    block_classes = [FluxAutoInputStep, FluxAutoBeforeDenoiseStep, FluxDenoiseStep, FluxUnpackLatentsStep]
+    block_names = ["input", "before_denoise", "denoise", "unpack_latents"]
 
     @property
     def description(self):

src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py

@@ -21,7 +21,7 @@
     FluxRoPEInputsStep,
     FluxSetTimestepsStep,
 )
-from .decoders import FluxDecodeStep
+from .decoders import FluxDecodeStep, FluxUnpackLatentsStep
 from .denoise import FluxKontextDenoiseStep
 from .encoders import (
     FluxKontextProcessImagesInputStep,
@@ -481,8 +481,13 @@ class FluxKontextCoreDenoiseStep(SequentialPipelineBlocks):
     """
 
     model_name = "flux-kontext"
-    block_classes = [FluxKontextAutoInputStep, FluxKontextAutoBeforeDenoiseStep, FluxKontextDenoiseStep]
-    block_names = ["input", "before_denoise", "denoise"]
+    block_classes = [
+        FluxKontextAutoInputStep,
+        FluxKontextAutoBeforeDenoiseStep,
+        FluxKontextDenoiseStep,
+        FluxUnpackLatentsStep,
+    ]
+    block_names = ["input", "before_denoise", "denoise", "unpack_latents"]
 
     @property
     def description(self):

src/diffusers/modular_pipelines/krea2/decoders.py

@@ -18,7 +18,7 @@
 from ...configuration_utils import FrozenDict
 from ...image_processor import VaeImageProcessor
 from ...models import AutoencoderKLQwenImage
-from ...utils import logging
+from ...utils import deprecate, logging
 from ..modular_pipeline import ModularPipelineBlocks, PipelineState
 from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
 from .modular_pipeline import Krea2ModularPipeline
@@ -28,6 +28,64 @@
 
 
 # auto_docstring
+def _unpack_latents(latents, height, width, patch_size, vae_scale_factor):
+    batch_size, _, channels = latents.shape
+    p = patch_size
+    height = p * (int(height) // (vae_scale_factor * p))
+    width = p * (int(width) // (vae_scale_factor * p))
+    latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
+    latents = latents.permute(0, 3, 1, 4, 2, 5)
+    return latents.reshape(batch_size, channels // (p * p), 1, height, width)
+
+
+class Krea2UnpackLatentsStep(ModularPipelineBlocks):
+    model_name = "krea2"
+
+    @property
+    def description(self) -> str:
+        return (
+            "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, 1, H, W]` form "
+            "the VAE takes (still normalized). Closes the core denoise group, so the blocks that follow take the same "
+            "form the VAE produces and need no geometry inputs."
+        )
+
+    @property
+    def inputs(self) -> list[InputParam]:
+        return [
+            InputParam(
+                name="latents",
+                required=True,
+                type_hint=torch.Tensor,
+                description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.",
+            ),
+            InputParam.template("height", default=1024),
+            InputParam.template("width", default=1024),
+        ]
+
+    @property
+    def intermediate_outputs(self) -> list[OutputParam]:
+        return [
+            OutputParam(
+                "latents",
+                type_hint=torch.Tensor,
+                description="The denoised latents of shape `[B, C, 1, H, W]` (normalized, not packed).",
+            )
+        ]
+
+    @torch.no_grad()
+    def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState:
+        block_state = self.get_block_state(state)
+        block_state.latents = _unpack_latents(
+            block_state.latents,
+            block_state.height,
+            block_state.width,
+            components.patch_size,
+            components.vae_scale_factor,
+        )
+        self.set_block_state(state, block_state)
+        return components, state
+
+
 class Krea2DecodeStep(ModularPipelineBlocks):
     """
     Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the VAE's
@@ -77,14 +135,15 @@ def expected_components(self) -> list[ComponentSpec]:
     def inputs(self) -> list[InputParam]:
         return [
             InputParam.template("output_type", default="pil"),
-            InputParam.template("height", default=1024),
-            InputParam.template("width", default=1024),
             InputParam(
                 name="latents",
                 required=True,
                 type_hint=torch.Tensor,
-                description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.",
+                description="The denoised latents of shape `[B, C, 1, H, W]` from the denoising group.",
             ),
+            # Only read on the deprecated path that still accepts packed `[B, S, C]` latents.
+            InputParam.template("height", default=1024),
+            InputParam.template("width", default=1024),
         ]
 
     @property
@@ -96,15 +155,17 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi
         block_state = self.get_block_state(state)
 
         vae = components.vae
-        p = components.patch_size
         latents = block_state.latents
-
-        batch_size, _, channels = latents.shape
-        height = p * (int(block_state.height) // (components.vae_scale_factor * p))
-        width = p * (int(block_state.width) // (components.vae_scale_factor * p))
-        latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p)
-        latents = latents.permute(0, 3, 1, 4, 2, 5)
-        latents = latents.reshape(batch_size, channels // (p * p), 1, height, width)
+        if latents.ndim == 3:
+            deprecate(
+                "packed latents",
+                "1.0.0",
+                "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group "
+                "now unpacks them. Pass latents of shape `[B, C, 1, H, W]` instead.",
+            )
+            latents = _unpack_latents(
+                latents, block_state.height, block_state.width, components.patch_size, components.vae_scale_factor
+            )
 
         latents = latents.to(vae.dtype)
         latents_mean = (

src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py

@@ -22,7 +22,7 @@
     Krea2SetTimestepsStep,
     Krea2TextInputsStep,
 )
-from .decoders import Krea2DecodeStep
+from .decoders import Krea2DecodeStep, Krea2UnpackLatentsStep
 from .denoise import Krea2DenoiseStep
 from .encoders import Krea2TextEncoderStep
 
@@ -37,6 +37,7 @@
         ("set_timesteps", Krea2SetTimestepsStep()),
         ("prepare_position_ids", Krea2PreparePositionIdsStep()),
         ("denoise", Krea2DenoiseStep()),
+        ("unpack_latents", Krea2UnpackLatentsStep()),
     ]
 )
 

src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py

@@ -22,7 +22,7 @@
     Krea2TurboSetTimestepsStep,
     Krea2TurboTextInputsStep,
 )
-from .decoders import Krea2DecodeStep
+from .decoders import Krea2DecodeStep, Krea2UnpackLatentsStep
 from .denoise import Krea2TurboDenoiseStep
 from .encoders import Krea2TurboTextEncoderStep
 
@@ -37,6 +37,7 @@
         ("set_timesteps", Krea2TurboSetTimestepsStep()),
         ("prepare_position_ids", Krea2PreparePositionIdsStep()),
         ("denoise", Krea2TurboDenoiseStep()),
+        ("unpack_latents", Krea2UnpackLatentsStep()),
     ]
 )