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

.ai/references/modular.md

@@ -191,6 +191,26 @@ 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
+
+Three transformations can sit between a VAE and a transformer:
+
+- normalize / denormalize: apply or undo the VAE's latent statistics.
+- patchify / unpatchify: fold a `p x p` spatial patch into channels, `[B, C, H, W]` <-> `[B, C*p*p, H/p, W/p]`. The result is still a grid.
+- pack / unpack: collapse the grid into a token sequence, `[B, C, F, H, W]` <-> `[B, S, D]`.
+
+Each should be applied and undone in mirrored pairs, and the core denoise group hands latents back in the same form it received them. The `latents` output of the denoise group should always be in one form: unpacked and still normalized, the same form its VAE encoder blocks produce. Whatever runs next (a decoder, a latent upsampler, a second denoise group) consumes that form, and no block outside the group needs `height`/`width`/`num_frames` just to interpret tokens.
+
+The pattern: `encode -> norm -> pack -> denoise -> unpack -> denorm -> decode`
+
+In `flux/`: `FluxVaeEncoderStep` encodes and normalizes, `FluxPrepareLatentsStep` packs, `FluxUnpackLatentsStep` closes `FluxCoreDenoiseStep`, and `FluxDecodeStep` denormalizes and decodes.
+
+- Normalize / denormalize belong to the VAE blocks. The encoder step normalizes right after `vae.encode`, the decoder step denormalizes right before `vae.decode`. Denormalized latents never sit in the state.
+- Pack / unpack belong to the denoise group, either at block level (a prepare-latents step packs, an unpack step such as `FluxUnpackLatentsStep` closes the group) or inside the transformer's forward. Packing is a transformer layout detail; nothing outside the group should see tokens.
+- Patchify / unpatchify can sit on either side, as long as it is mirrored: usually it is fused with pack in the denoise group (Flux, Krea2, LTX, QwenImage); when the VAE statistics are defined over patchified channels (`AutoencoderKLFlux2`: Flux2, FLUX.3, Ernie-Image) it lives in the VAE blocks and the canonical form is the patchified grid.
+
+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
@@ -46,6 +46,49 @@ 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)."
+        )
+
+    @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", required=True),
+            InputParam("width", required=True),
+        ]
+
+    @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"
 
@@ -69,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
@@ -96,7 +140,14 @@ def __call__(self, components, state: PipelineState) -> tuple[FluxModularPipelin
 
         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",
+                    "0.47.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
@@ -27,24 +27,99 @@
 logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
 
 
+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)
+
+
+# auto_docstring
+class Krea2UnpackLatentsStep(ModularPipelineBlocks):
+    """
+    Unpacks the denoised latents from the transformer's token layout back into the `[B, C, 1, H, W]` form the VAE takes
+    (still normalized).
+
+      Inputs:
+          latents (`Tensor`):
+              The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.
+          height (`int`):
+              The height in pixels of the generated image.
+          width (`int`):
+              The width in pixels of the generated image.
+
+      Outputs:
+          latents (`Tensor`):
+              The denoised latents of shape `[B, C, 1, H, W]` (normalized, not packed).
+    """
+
+    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)."
+        )
+
+    @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", required=True),
+            InputParam.template("width", required=True),
+        ]
+
+    @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
+
+
 # auto_docstring
 class Krea2DecodeStep(ModularPipelineBlocks):
     """
-    Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the VAE's
-    per-channel statistics, and decodes them through the Qwen-Image VAE into images.
+    Step that de-normalizes the denoised latents with the VAE's per-channel statistics and decodes them through the
+    Qwen-Image VAE into images.
 
       Components:
           vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`)
 
       Inputs:
           output_type (`str`, *optional*, defaults to pil):
               Output format: 'pil', 'np', 'pt'.
+          latents (`Tensor`):
+              The denoised latents of shape `[B, C, 1, H, W]` from the denoising group.
           height (`int`, *optional*, defaults to 1024):
               The height in pixels of the generated image.
           width (`int`, *optional*, defaults to 1024):
               The width in pixels of the generated image.
-          latents (`Tensor`):
-              The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.
 
       Outputs:
           images (`list`):
@@ -56,8 +131,8 @@ class Krea2DecodeStep(ModularPipelineBlocks):
     @property
     def description(self) -> str:
         return (
-            "Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the "
-            "VAE's per-channel statistics, and decodes them through the Qwen-Image VAE into images."
+            "Step that de-normalizes the denoised latents with the VAE's per-channel statistics and decodes them "
+            "through the Qwen-Image VAE into images."
         )
 
     @property
@@ -77,14 +152,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
@@ -98,15 +174,17 @@ def __call__(
         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",
+                "0.47.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()),
     ]
 )