From 1314e3589c7360f2bf15a228e1ed26c8e16ee4c4 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 7 Sep 2026 21:14:46 -0700 Subject: [PATCH 01/11] add a modular test that pins the canonical latents form each family leaves in the state --- .../anima/test_modular_pipeline_anima.py | 2 ++ .../cosmos/test_modular_pipeline_cosmos3.py | 1 + ...test_modular_pipeline_cosmos3_distilled.py | 1 + .../test_modular_pipeline_ernie_image.py | 1 + .../flux2/test_modular_pipeline_flux2.py | 2 ++ .../test_modular_pipeline_flux2_klein.py | 2 ++ .../test_modular_pipeline_flux2_klein_base.py | 2 ++ .../helios/test_modular_pipeline_helios.py | 2 ++ .../test_modular_pipeline_hunyuan_video1_5.py | 1 + .../test_modular_pipeline_minimax_h3.py | 2 ++ .../test_modular_pipeline_minimax_music3.py | 2 ++ .../qwen/test_modular_pipeline_qwenimage.py | 3 +++ ...est_modular_pipeline_stable_diffusion_3.py | 2 ++ ...st_modular_pipeline_stable_diffusion_xl.py | 3 +++ .../modular_pipelines/testing_utils/common.py | 23 +++++++++++++++++++ .../wan/test_modular_pipeline_wan.py | 1 + .../wan/test_modular_pipeline_wan_vace.py | 1 + .../z_image/test_modular_pipeline_z_image.py | 1 + 18 files changed, 52 insertions(+) diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index c3ae0d9f2e0d..93c9941f11fe 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -108,6 +108,7 @@ class AnimaModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = AnimaModularPipeline pipeline_blocks_class = AnimaAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-anima-modular-pipe" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "height", "width", "negative_prompt"]) batch_params = frozenset(["prompt", "negative_prompt"]) expected_workflow_blocks = ANIMA_TEXT2IMAGE_WORKFLOWS @@ -204,6 +205,7 @@ class AnimaImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = AnimaModularPipeline pipeline_blocks_class = AnimaAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-anima-modular-pipe" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "image", "strength", "height", "width", "negative_prompt"]) batch_params = frozenset(["prompt", "negative_prompt"]) expected_workflow_blocks = ANIMA_IMG2IMG_WORKFLOWS diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index b5bf5fb04393..8f8baf7aa41a 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -125,6 +125,7 @@ class Cosmos3OmniModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Cosmos3OmniModularPipeline pipeline_blocks_class = Cosmos3OmniBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-cosmos3-modular-pipe" + expected_latents_shape = (1, 4, 2, 4, 4) params = frozenset(["prompt", "height", "width", "num_frames", "guidance_scale"]) batch_params = frozenset() optional_params = frozenset(["num_inference_steps", "output_type"]) diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py index 97bb2b478989..a7b485fce90d 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -69,6 +69,7 @@ class Cosmos3DistilledModularPipelineTesterConfig(BaseModularPipelineTesterConfi pipeline_class = Cosmos3DistilledModularPipeline pipeline_blocks_class = Cosmos3DistilledBlocks pretrained_model_name_or_path = TINY_DISTILLED_REPO + expected_latents_shape = (1, 4, 2, 4, 4) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset() optional_params = frozenset(["num_inference_steps", "output_type"]) diff --git a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py index 10df99406326..0de66222be6e 100644 --- a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py +++ b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py @@ -42,6 +42,7 @@ class ErnieImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = ErnieImageModularPipeline pipeline_blocks_class = ErnieImageAutoBlocks pretrained_model_name_or_path = "akshan-main/tiny-ernie-image-modular-pipe" + expected_latents_shape = (1, 16, 2, 2) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents"]) diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py index eb2ed35c46b7..6acd1a290f2e 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py @@ -53,6 +53,7 @@ class Flux2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Flux2ModularPipeline pipeline_blocks_class = Flux2AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width", "guidance_scale"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = FLUX2_TEXT2IMAGE_WORKFLOWS @@ -113,6 +114,7 @@ class Flux2ImageConditionedModularPipelineTesterConfig(BaseModularPipelineTester pipeline_class = Flux2ModularPipeline pipeline_blocks_class = Flux2AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) expected_workflow_blocks = FLUX2_IMAGE_CONDITIONED_WORKFLOWS diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py index 4932efa7dd74..93d5c818e056 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py @@ -52,6 +52,7 @@ class Flux2KleinModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Flux2KleinModularPipeline pipeline_blocks_class = Flux2KleinAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) not_params = frozenset(["negative_prompt"]) @@ -111,6 +112,7 @@ class Flux2KleinImageConditionedModularPipelineTesterConfig(BaseModularPipelineT pipeline_class = Flux2KleinModularPipeline pipeline_blocks_class = Flux2KleinAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width", "image"]) batch_params = frozenset(["prompt", "image"]) not_params = frozenset(["negative_prompt"]) diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py index 30460105ad34..90e45a781df8 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py @@ -52,6 +52,7 @@ class Flux2KleinBaseModularPipelineTesterConfig(BaseModularPipelineTesterConfig) pipeline_class = Flux2KleinBaseModularPipeline pipeline_blocks_class = Flux2KleinBaseAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = FLUX2_KLEIN_BASE_WORKFLOWS @@ -110,6 +111,7 @@ class Flux2KleinBaseImageConditionedModularPipelineTesterConfig(BaseModularPipel pipeline_class = Flux2KleinBaseModularPipeline pipeline_blocks_class = Flux2KleinBaseAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular" + expected_latents_shape = (1, 4, 16, 16) params = frozenset(["prompt", "height", "width", "image"]) batch_params = frozenset(["prompt", "image"]) expected_workflow_blocks = FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS diff --git a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py index 5c270c20d9de..f3e7b43335fb 100644 --- a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py +++ b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py @@ -71,6 +71,7 @@ class HeliosModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = HeliosModularPipeline pipeline_blocks_class = HeliosAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-helios-modular-pipe" + expected_latents_shape = (1, 16, 9, 2, 2) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) @@ -147,6 +148,7 @@ class HeliosPyramidModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = HeliosPyramidModularPipeline pipeline_blocks_class = HeliosPyramidAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-helios-pyramid-modular-pipe" + expected_latents_shape = (1, 16, 9, 8, 8) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["pyramid_num_inference_steps_list", "num_videos_per_prompt", "latents"]) diff --git a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py index f6368d91b011..9ee9980f81e0 100644 --- a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py +++ b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py @@ -53,6 +53,7 @@ class HunyuanVideo15ModularPipelineTesterConfig(BaseModularPipelineTesterConfig) pipeline_class = HunyuanVideo15ModularPipeline pipeline_blocks_class = HunyuanVideo15AutoBlocks pretrained_model_name_or_path = "akshan-main/tiny-hunyuanvideo1_5-modular-pipe" + expected_latents_shape = (1, 32, 3, 2, 2) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index 2f653b28e296..3de1ce2a2319 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -159,6 +159,7 @@ class MiniMaxH3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = MiniMaxH3ModularPipeline pipeline_blocks_class = MiniMaxH3Blocks pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-h3-modular-pipe" + expected_latents_shape = (1, 4, 37, 2, 2) params = frozenset(["prompt", "image", "last_image", "height", "width", "num_frames"]) # MiniMax-H3 packs one request into one sequence and rejects a list of prompts, so nothing is batched. batch_params = frozenset() @@ -421,6 +422,7 @@ class MiniMaxH3Ref2VAModularPipelineTesterConfig(BaseModularPipelineTesterConfig pipeline_class = MiniMaxH3ModularPipeline pipeline_blocks_class = MiniMaxH3Blocks pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-h3-modular-pipe" + expected_latents_shape = (1, 4, 37, 2, 2) params = frozenset(["prompt", "references", "height", "width", "num_frames"]) # MiniMax-H3 packs one request into one sequence and rejects a list of prompts, so nothing is batched. batch_params = frozenset() diff --git a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py index 923c4704dd7b..45096f326be1 100644 --- a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py +++ b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py @@ -74,6 +74,8 @@ class MiniMaxMusic3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): batch_params = frozenset() optional_params = frozenset(["num_inference_steps", "output_type"]) output_name = "audios" + latents_output_name = "latent_chunks" + expected_latents_shape = (1, 8, 17) expected_workflow_defaults = { None: { "components": { diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py index 26a1cba8b30a..85aeb06006d1 100644 --- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py +++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py @@ -138,6 +138,7 @@ class QwenImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = QwenImageModularPipeline pipeline_blocks_class = QwenImageAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-modular" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"]) batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) expected_workflow_blocks = QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS @@ -220,6 +221,7 @@ class QwenImageEditModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = QwenImageEditModularPipeline pipeline_blocks_class = QwenImageEditAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-edit-modular" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image", "mask_image"]) batch_params = frozenset(["prompt", "negative_prompt", "image", "mask_image"]) expected_workflow_blocks = QWEN_IMAGE_EDIT_WORKFLOWS @@ -268,6 +270,7 @@ class QwenImageEditPlusModularPipelineTesterConfig(BaseModularPipelineTesterConf pipeline_class = QwenImageEditPlusModularPipeline pipeline_blocks_class = QwenImageEditPlusAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-edit-plus-modular" + expected_latents_shape = (1, 4, 1, 8, 8) # No `mask_image` yet. params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image"]) batch_params = frozenset(["prompt", "negative_prompt", "image"]) diff --git a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py index 03d0c5687794..e19ddec61263 100644 --- a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py +++ b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py @@ -51,6 +51,7 @@ class StableDiffusion3ModularPipelineTesterConfig(BaseModularPipelineTesterConfi pipeline_class = StableDiffusion3ModularPipeline pipeline_blocks_class = StableDiffusion3AutoBlocks pretrained_model_name_or_path = "AlanPonnachan/tiny-sd3-modular" + expected_latents_shape = (1, 4, 32, 32) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = SD3_TEXT2IMAGE_WORKFLOWS @@ -143,6 +144,7 @@ class StableDiffusion3Img2ImgModularPipelineTesterConfig(BaseModularPipelineTest pipeline_class = StableDiffusion3ModularPipeline pipeline_blocks_class = StableDiffusion3AutoBlocks pretrained_model_name_or_path = "AlanPonnachan/tiny-sd3-modular" + expected_latents_shape = (1, 4, 32, 32) params = frozenset(["prompt", "height", "width", "image"]) batch_params = frozenset(["prompt", "image"]) expected_workflow_blocks = SD3_IMAGE2IMAGE_WORKFLOWS diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index 1e1777b90790..b540f6306bfd 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -336,6 +336,7 @@ class SDXLModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = StableDiffusionXLModularPipeline pipeline_blocks_class = StableDiffusionXLAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-sdxl-modular" + expected_latents_shape = (1, 4, 32, 32) params = frozenset( [ "prompt", @@ -464,6 +465,7 @@ class SDXLImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = StableDiffusionXLModularPipeline pipeline_blocks_class = StableDiffusionXLAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-sdxl-modular" + expected_latents_shape = (1, 4, 32, 32) params = frozenset( [ "prompt", @@ -608,6 +610,7 @@ class SDXLInpaintingModularPipelineFastTests( pipeline_class = StableDiffusionXLModularPipeline pipeline_blocks_class = StableDiffusionXLAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-sdxl-modular" + expected_latents_shape = (1, 4, 32, 32) params = frozenset( [ "prompt", diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 47614fd51005..94067a535487 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -53,6 +53,13 @@ class BaseModularPipelineTesterConfig: # Output type for the pipeline (e.g., "images" for image pipelines, "videos" for video pipelines) # Subclasses can override this to change the expected output type output_name = "images" + # State name of the latents the denoise group hands to the decoder. Chunked families that hand over a list of + # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. + latents_output_name = "latents" + # The one canonical form this family keeps latents in after the denoise group, for the geometry that + # `get_dummy_inputs()` produces: the VAE form (e.g. `(1, 4, 32, 32)`) for pipelines that pack/unpack inside the + # denoise group or the transformer, or the channel-packed form for pipelines whose VAE statistics live there. + expected_latents_shape = None # ==================== Required interface ==================== @@ -361,3 +368,19 @@ def test_num_images_per_prompt(self, batch_sizes=[1, 2], num_images_per_prompts= images = pipe(**inputs, num_images_per_prompt=num_images_per_prompt, output=self.output_name) assert images.shape[0] == batch_size * num_images_per_prompt + + def test_latents_output_in_canonical_form(self): + # The denoise group must hand latents back in the family's one canonical form, so any block (a decoder, a + # latent upsampler, a second denoise group) can consume them without geometry inputs. The decoder reads + # the latents without removing them, so they are checked in the state after a full run. + pipe = self.get_pipeline().to(torch_device) + latents = pipe(**self.get_dummy_inputs(), output=self.latents_output_name) + chunks = list(latents) if isinstance(latents, (list, tuple)) else [latents] + shapes = sorted({tuple(chunk.shape) for chunk in chunks}) + + if self.expected_latents_shape is None: + pytest.skip(f"{type(self).__name__}: `expected_latents_shape` not declared; latents shape is {shapes}") + assert shapes == [tuple(self.expected_latents_shape)], ( + f"Latents left in the state have shape {shapes}, expected the canonical " + f"{tuple(self.expected_latents_shape)}" + ) diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py index d35c21455ba9..2f1e0a7de82b 100644 --- a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py @@ -30,6 +30,7 @@ class WanModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = WanModularPipeline pipeline_blocks_class = WanBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-modular-pipe" + expected_latents_shape = (1, 16, 3, 2, 2) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py index 535f83bd5133..3b22389ec632 100644 --- a/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py @@ -33,6 +33,7 @@ class Wan22VaceModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Wan22VaceModularPipeline pipeline_blocks_class = Wan22VaceBlocks pretrained_model_name_or_path = "akshan-main/tiny-wan22-vace-modular-pipe" + expected_latents_shape = (1, 16, 3, 2, 2) params = frozenset(["prompt", "height", "width", "num_frames", "video", "mask", "reference_images"]) batch_params = frozenset() diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py index b67be5bc88e9..0c58866359d4 100644 --- a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py +++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py @@ -53,6 +53,7 @@ class ZImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = ZImageModularPipeline pipeline_blocks_class = ZImageAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-zimage-modular-pipe" + expected_latents_shape = (1, 16, 16, 16) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = ZIMAGE_WORKFLOWS From 900591630fb22edaf4f5985656c7840656280806 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 7 Sep 2026 21:35:10 -0700 Subject: [PATCH 02/11] move unpack out of the flux, krea2 and ltx decoders into the core denoise groups --- .../modular_pipelines/flux/decoders.py | 62 +++++++++++-- .../flux/modular_blocks_flux.py | 6 +- .../flux/modular_blocks_flux_kontext.py | 11 ++- .../modular_pipelines/krea2/decoders.py | 85 +++++++++++++++--- .../krea2/modular_blocks_krea2.py | 3 +- .../krea2/modular_blocks_krea2_turbo.py | 3 +- .../modular_pipelines/ltx/decoders.py | 88 ++++++++++++++++--- .../ltx/modular_blocks_ltx.py | 8 +- .../flux/test_modular_pipeline_flux.py | 23 +++++ .../krea2/test_modular_pipeline_krea2.py | 21 +++++ .../test_modular_pipeline_krea2_turbo.py | 2 + .../ltx/test_modular_pipeline_ltx.py | 21 +++++ .../ltx2/test_modular_pipeline_ltx2.py | 2 + .../ltx2/test_modular_pipeline_ltx25.py | 2 + .../modular_pipelines/testing_utils/common.py | 15 ++-- 15 files changed, 305 insertions(+), 47 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 5fcde5008680..092bead88d56 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/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( diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py index 1f028f555a1b..a725304c6af6 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks_flux.py +++ b/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): diff --git a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py b/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py index c4f8bffffd1e..f937d2555b1a 100644 --- a/src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py +++ b/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): diff --git a/src/diffusers/modular_pipelines/krea2/decoders.py b/src/diffusers/modular_pipelines/krea2/decoders.py index fd308b5ef648..29ea8d0814a5 100644 --- a/src/diffusers/modular_pipelines/krea2/decoders.py +++ b/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 = ( diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py index ae3b2ac4fb52..1a1f11f60656 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py +++ b/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()), ] ) diff --git a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py b/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py index 79fa5406c4e5..e3c8f1c288b3 100644 --- a/src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py +++ b/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()), ] ) diff --git a/src/diffusers/modular_pipelines/ltx/decoders.py b/src/diffusers/modular_pipelines/ltx/decoders.py index 8664dee25bfe..bdf31e917cea 100644 --- a/src/diffusers/modular_pipelines/ltx/decoders.py +++ b/src/diffusers/modular_pipelines/ltx/decoders.py @@ -18,7 +18,7 @@ from ...configuration_utils import FrozenDict from ...models import AutoencoderKLLTXVideo -from ...utils import logging +from ...utils import deprecate, logging from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..modular_pipeline import ModularPipelineBlocks, PipelineState @@ -39,6 +39,65 @@ def _denormalize_latents( return latents +class LTXUnpackLatentsStep(ModularPipelineBlocks): + model_name = "ltx" + + @property + def description(self) -> str: + return ( + "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, F, 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 expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec( + "pachifier", + LTXVideoPachifier, + config=FrozenDict({"patch_size": 1, "patch_size_t": 1}), + default_creation_method="from_config", + ), + ] + + @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.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The denoised latents of shape `[B, C, F, H, W]` (normalized, not packed).", + ) + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + block_state.latents = components.pachifier.unpack_latents( + block_state.latents, latent_num_frames, latent_height, latent_width + ) + self.set_block_state(state, block_state) + return components, state + + class LTXVaeDecoderStep(ModularPipelineBlocks): model_name = "ltx" @@ -69,14 +128,15 @@ def inputs(self) -> list[tuple[str, Any]]: return [ InputParam.template("latents", required=True), InputParam.template("output_type", default="np"), - InputParam.template("height", default=512), - InputParam.template("width", default=704), - InputParam("num_frames", type_hint=int, default=161), InputParam("decode_timestep", default=0.0), InputParam("decode_noise_scale", default=None), InputParam.template("generator"), InputParam.template("batch_size"), InputParam.template("dtype", required=True), + # Only read on the deprecated path that still accepts packed `[B, S, C]` latents. + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, default=161), ] @property @@ -89,16 +149,18 @@ def __call__(self, components, state: PipelineState) -> PipelineState: vae = components.vae latents = block_state.latents + 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, F, H, W]` instead.", + ) + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + latents = components.pachifier.unpack_latents(latents, latent_num_frames, latent_height, latent_width) - height = block_state.height - width = block_state.width - num_frames = block_state.num_frames - - latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1 - latent_height = height // components.vae_spatial_compression_ratio - latent_width = width // components.vae_spatial_compression_ratio - - latents = components.pachifier.unpack_latents(latents, latent_num_frames, latent_height, latent_width) latents = _denormalize_latents(latents, vae.latents_mean, vae.latents_std, vae.config.scaling_factor) latents = latents.to(block_state.dtype) diff --git a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py index 828c79e1c72d..ae9cc7a58d84 100644 --- a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py +++ b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py @@ -21,7 +21,7 @@ LTXSetTimestepsStep, LTXTextInputStep, ) -from .decoders import LTXVaeDecoderStep +from .decoders import LTXUnpackLatentsStep, LTXVaeDecoderStep from .denoise import LTXDenoiseStep, LTXImage2VideoDenoiseStep from .encoders import LTXTextEncoderStep, LTXVaeEncoderStep @@ -81,8 +81,9 @@ class LTXCoreDenoiseStep(SequentialPipelineBlocks): LTXSetTimestepsStep, LTXPrepareLatentsStep, LTXDenoiseStep, + LTXUnpackLatentsStep, ] - block_names = ["input", "set_timesteps", "prepare_latents", "denoise"] + block_names = ["input", "set_timesteps", "prepare_latents", "denoise", "unpack_latents"] @property def description(self): @@ -148,8 +149,9 @@ class LTXImage2VideoCoreDenoiseStep(SequentialPipelineBlocks): LTXPrepareLatentsStep, LTXImage2VideoPrepareLatentsStep, LTXImage2VideoDenoiseStep, + LTXUnpackLatentsStep, ] - block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise"] + block_names = ["input", "set_timesteps", "prepare_latents", "prepare_i2v_latents", "denoise", "unpack_latents"] @property def description(self): diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 502df3f447bf..e42eb21e8d3f 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -17,6 +17,7 @@ import numpy as np import PIL +import pytest import torch from diffusers.image_processor import VaeImageProcessor @@ -47,6 +48,7 @@ ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), ("denoise.denoise", "FluxDenoiseStep"), + ("denoise.unpack_latents", "FluxUnpackLatentsStep"), ("decode", "FluxDecodeStep"), ] } @@ -56,6 +58,7 @@ class FluxModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = FluxModularPipeline pipeline_blocks_class = FluxAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular" + expected_latents_shape = (1, 1, 8, 8) params = frozenset(["prompt", "height", "width", "guidance_scale"]) batch_params = frozenset(["prompt"]) @@ -77,6 +80,21 @@ def get_dummy_inputs(self, seed=0): class TestFluxModularPipelineFast(FluxModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_decode_accepts_packed_latents_with_deprecation(self): + # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. + inputs = self.get_dummy_inputs() + expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) + + blocks = self.pipeline_blocks_class() + blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") + pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + pipe.to(torch_device) + with pytest.warns(FutureWarning, match="packed latents"): + output = pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert torch.equal(output, expected) + def test_float16_inference(self): super().test_float16_inference(9e-2) @@ -109,6 +127,7 @@ class TestFluxModularPipelineAutoOffload(FluxModularPipelineTesterConfig, Modula ("denoise.before_denoise.prepare_img2img_latents", "FluxImg2ImgPrepareLatentsStep"), ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), ("denoise.denoise", "FluxDenoiseStep"), + ("denoise.unpack_latents", "FluxUnpackLatentsStep"), ("decode", "FluxDecodeStep"), ] } @@ -118,6 +137,7 @@ class FluxImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = FluxModularPipeline pipeline_blocks_class = FluxAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular" + expected_latents_shape = (1, 1, 8, 8) params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) @@ -193,6 +213,7 @@ class TestFluxImg2ImgModularPipelineMemory(FluxImg2ImgModularPipelineTesterConfi ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), ("denoise.before_denoise.prepare_rope_inputs", "FluxRoPEInputsStep"), ("denoise.denoise", "FluxKontextDenoiseStep"), + ("denoise.unpack_latents", "FluxUnpackLatentsStep"), ("decode", "FluxDecodeStep"), ], "image_conditioned": [ @@ -206,6 +227,7 @@ class TestFluxImg2ImgModularPipelineMemory(FluxImg2ImgModularPipelineTesterConfi ("denoise.before_denoise.set_timesteps", "FluxSetTimestepsStep"), ("denoise.before_denoise.prepare_rope_inputs", "FluxKontextRoPEInputsStep"), ("denoise.denoise", "FluxKontextDenoiseStep"), + ("denoise.unpack_latents", "FluxUnpackLatentsStep"), ("decode", "FluxDecodeStep"), ], } @@ -215,6 +237,7 @@ class FluxKontextModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = FluxKontextModularPipeline pipeline_blocks_class = FluxKontextAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-kontext-pipe" + expected_latents_shape = (1, 1, 8, 8) params = frozenset(["prompt", "height", "width", "guidance_scale", "image"]) batch_params = frozenset(["prompt", "image"]) diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index 133b3eb1bf66..df10f4d10350 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -14,8 +14,12 @@ # limitations under the License. +import pytest +import torch + from diffusers.modular_pipelines import Krea2AutoBlocks, Krea2ModularPipeline +from ...testing_utils import torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -33,6 +37,7 @@ ("denoise.set_timesteps", "Krea2SetTimestepsStep"), ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), ("denoise.denoise", "Krea2DenoiseStep"), + ("denoise.unpack_latents", "Krea2UnpackLatentsStep"), ("decode", "Krea2DecodeStep"), ], } @@ -42,6 +47,7 @@ class Krea2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Krea2ModularPipeline pipeline_blocks_class = Krea2AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-modular-pipe" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = KREA2_WORKFLOWS @@ -61,6 +67,21 @@ def get_dummy_inputs(self, seed=0): class TestKrea2ModularPipelineFast(Krea2ModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_decode_accepts_packed_latents_with_deprecation(self): + # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. + inputs = self.get_dummy_inputs() + expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) + + blocks = self.pipeline_blocks_class() + blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") + pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + pipe.to(torch_device) + with pytest.warns(FutureWarning, match="packed latents"): + output = pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert torch.equal(output, expected) + def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py index 40d278156ce9..7dbe1444cf0e 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py @@ -33,6 +33,7 @@ ("denoise.set_timesteps", "Krea2TurboSetTimestepsStep"), ("denoise.prepare_position_ids", "Krea2PreparePositionIdsStep"), ("denoise.denoise", "Krea2TurboDenoiseStep"), + ("denoise.unpack_latents", "Krea2UnpackLatentsStep"), ("decode", "Krea2DecodeStep"), ], } @@ -42,6 +43,7 @@ class Krea2TurboModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Krea2TurboModularPipeline pipeline_blocks_class = Krea2TurboAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-turbo-modular-pipe" + expected_latents_shape = (1, 4, 1, 8, 8) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset(["prompt"]) expected_workflow_blocks = KREA2_TURBO_WORKFLOWS diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index 1a58465c23c6..62af9275fcdf 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -14,9 +14,11 @@ # limitations under the License. import pytest +import torch from diffusers.modular_pipelines import LTXAutoBlocks, LTXModularPipeline +from ...testing_utils import torch_device from ..testing_utils import ( BaseModularPipelineTesterConfig, ModularLoadingTesterMixin, @@ -33,6 +35,7 @@ ("denoise.set_timesteps", "LTXSetTimestepsStep"), ("denoise.prepare_latents", "LTXPrepareLatentsStep"), ("denoise.denoise", "LTXDenoiseStep"), + ("denoise.unpack_latents", "LTXUnpackLatentsStep"), ("decode", "LTXVaeDecoderStep"), ], "image2video": [ @@ -43,6 +46,7 @@ ("denoise.prepare_latents", "LTXPrepareLatentsStep"), ("denoise.prepare_i2v_latents", "LTXImage2VideoPrepareLatentsStep"), ("denoise.denoise", "LTXImage2VideoDenoiseStep"), + ("denoise.unpack_latents", "LTXUnpackLatentsStep"), ("decode", "LTXVaeDecoderStep"), ], } @@ -52,6 +56,7 @@ class LTXModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = LTXModularPipeline pipeline_blocks_class = LTXAutoBlocks pretrained_model_name_or_path = "akshan-main/tiny-ltx-modular-pipe" + expected_latents_shape = (1, 8, 5, 16, 16) params = frozenset(["prompt", "height", "width", "num_frames"]) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) @@ -74,6 +79,22 @@ def get_dummy_inputs(self, seed=0): class TestLTXModularPipelineFast(LTXModularPipelineTesterConfig, ModularPipelineTesterMixin): + def test_decode_accepts_packed_latents_with_deprecation(self): + # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. + inputs = self.get_dummy_inputs() + expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) + + blocks = self.pipeline_blocks_class() + for workflow in blocks.sub_blocks["denoise"].sub_blocks.values(): + workflow.sub_blocks.pop("unpack_latents") + pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) + pipe.load_components(dtype=torch.float32) + pipe.to(torch_device) + with pytest.warns(FutureWarning, match="packed latents"): + output = pipe(**self.get_dummy_inputs(), output=self.output_name) + + assert torch.equal(output, expected) + @pytest.mark.skip(reason="num_videos_per_prompt") def test_num_images_per_prompt(self): pass diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py index 2ee330ece20a..23f9a519defd 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py @@ -104,6 +104,8 @@ class LTX2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = LTX2ModularPipeline pipeline_blocks_class = LTX2AutoBlocks pretrained_model_name_or_path = LTX2_REPO_ID + # Still the packed `[B, S, C]` form: #14612 moves ltx2 to the VAE form and updates this. + expected_latents_shape = (1, 768, 4) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) expected_workflow_blocks = LTX2_WORKFLOWS diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index 69080095c627..d36c12d8e676 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py @@ -101,6 +101,8 @@ class LTX25ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = LTX25ModularPipeline pipeline_blocks_class = LTX25AutoBlocks pretrained_model_name_or_path = LTX25_REPO_ID + # Still the packed `[B, S, C]` form: #14612 moves ltx2 to the VAE form and updates this. + expected_latents_shape = (1, 768, 4) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) expected_workflow_blocks = LTX25_WORKFLOWS diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 94067a535487..54fbd1fc223b 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -56,10 +56,15 @@ class BaseModularPipelineTesterConfig: # State name of the latents the denoise group hands to the decoder. Chunked families that hand over a list of # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. latents_output_name = "latents" - # The one canonical form this family keeps latents in after the denoise group, for the geometry that - # `get_dummy_inputs()` produces: the VAE form (e.g. `(1, 4, 32, 32)`) for pipelines that pack/unpack inside the - # denoise group or the transformer, or the channel-packed form for pipelines whose VAE statistics live there. - expected_latents_shape = None + + @property + def expected_latents_shape(self) -> tuple: + raise NotImplementedError( + "You need to set the attribute `expected_latents_shape` in the child test class: the one canonical form " + "this family keeps latents in after the denoise group, for the geometry `get_dummy_inputs()` produces. " + "That is the VAE form (e.g. `(1, 4, 32, 32)`) for pipelines that pack/unpack inside the denoise group or " + "the transformer, or the channel-packed form for pipelines whose VAE statistics live there." + ) # ==================== Required interface ==================== @@ -378,8 +383,6 @@ def test_latents_output_in_canonical_form(self): chunks = list(latents) if isinstance(latents, (list, tuple)) else [latents] shapes = sorted({tuple(chunk.shape) for chunk in chunks}) - if self.expected_latents_shape is None: - pytest.skip(f"{type(self).__name__}: `expected_latents_shape` not declared; latents shape is {shapes}") assert shapes == [tuple(self.expected_latents_shape)], ( f"Latents left in the state have shape {shapes}, expected the canonical " f"{tuple(self.expected_latents_shape)}" From 713dbbeeeb823b388e699f0168fce2e44ed895ed Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 7 Sep 2026 21:47:44 -0700 Subject: [PATCH 03/11] pop the unpack step from a copy of the blocks in the deprecation tests --- tests/modular_pipelines/flux/test_modular_pipeline_flux.py | 4 +++- tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py | 5 ++++- tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py | 5 ++++- 3 files changed, 11 insertions(+), 3 deletions(-) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index e42eb21e8d3f..8eeeb0880540 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy import random import numpy as np @@ -85,7 +86,8 @@ def test_decode_accepts_packed_latents_with_deprecation(self): inputs = self.get_dummy_inputs() expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) - blocks = self.pipeline_blocks_class() + # Blocksets can be composed from shared block instances, so pop from a copy. + blocks = copy.deepcopy(self.pipeline_blocks_class()) blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index df10f4d10350..569f8acc969d 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -14,6 +14,8 @@ # limitations under the License. +import copy + import pytest import torch @@ -72,7 +74,8 @@ def test_decode_accepts_packed_latents_with_deprecation(self): inputs = self.get_dummy_inputs() expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) - blocks = self.pipeline_blocks_class() + # Blocksets can be composed from shared block instances, so pop from a copy. + blocks = copy.deepcopy(self.pipeline_blocks_class()) blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) pipe.load_components(dtype=torch.float32) diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index 62af9275fcdf..c9de9a51f587 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. +import copy + import pytest import torch @@ -84,7 +86,8 @@ def test_decode_accepts_packed_latents_with_deprecation(self): inputs = self.get_dummy_inputs() expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) - blocks = self.pipeline_blocks_class() + # Blocksets can be composed from shared block instances, so pop from a copy. + blocks = copy.deepcopy(self.pipeline_blocks_class()) for workflow in blocks.sub_blocks["denoise"].sub_blocks.values(): workflow.sub_blocks.pop("unpack_latents") pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) From bd54dd83249c569d607ce1c06dad4e561b614525 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 7 Sep 2026 22:01:21 -0700 Subject: [PATCH 04/11] let families that decode inside the denoise loop opt out of the latents-form test --- tests/modular_pipelines/testing_utils/common.py | 5 ++++- .../wan_animate_2/test_modular_pipeline_wan_animate_2.py | 6 ++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 54fbd1fc223b..e42c513d78d4 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -54,7 +54,8 @@ class BaseModularPipelineTesterConfig: # Subclasses can override this to change the expected output type output_name = "images" # State name of the latents the denoise group hands to the decoder. Chunked families that hand over a list of - # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. + # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. Families that decode inside their + # denoise loop and leave no latents in the state set this to `None`. latents_output_name = "latents" @property @@ -378,6 +379,8 @@ def test_latents_output_in_canonical_form(self): # The denoise group must hand latents back in the family's one canonical form, so any block (a decoder, a # latent upsampler, a second denoise group) can consume them without geometry inputs. The decoder reads # the latents without removing them, so they are checked in the state after a full run. + if self.latents_output_name is None: + pytest.skip("This family decodes inside its denoise loop and leaves no latents in the state.") pipe = self.get_pipeline().to(torch_device) latents = pipe(**self.get_dummy_inputs(), output=self.latents_output_name) chunks = list(latents) if isinstance(latents, (list, tuple)) else [latents] diff --git a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py index 057bf11af5d5..1c2af8c0f94c 100644 --- a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py +++ b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py @@ -84,6 +84,9 @@ class WanAnimate2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = WanAnimate2ModularPipeline pipeline_blocks_class = WanAnimate2Blocks pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-modular" + # Each segment is decoded inside the denoise loop, so no latents are left in the state. + latents_output_name = None + expected_latents_shape = None params = frozenset(["prompt", "image", "driving_video"]) batch_params = frozenset() @@ -157,6 +160,9 @@ class WanAnimate2DistilledModularPipelineTesterConfig(WanAnimate2ModularPipeline pipeline_class = WanAnimate2DistilledModularPipeline pipeline_blocks_class = WanAnimate2DistilledBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-distilled-modular" + # Each segment is decoded inside the denoise loop, so no latents are left in the state. + latents_output_name = None + expected_latents_shape = None expected_workflow_defaults = WAN_ANIMATE_2_DISTILLED_DEFAULTS From 7abc2c02b629afa7244e766dd5427b6a1850f227 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 7 Sep 2026 22:05:21 -0700 Subject: [PATCH 05/11] document the canonical latents form in the modular reference --- .ai/references/modular.md | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 8bfb39ebb999..f25043dcda0f 100644 --- a/.ai/references/modular.md +++ b/.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. From 76984eabfaf8d8b67a21b565be32f8a09d0860f0 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 17 Oct 2026 23:58:26 +0200 Subject: [PATCH 06/11] [docs] canonical latents form: name normalize, patchify and pack separately Co-Authored-By: Claude Fable 5.1 --- .ai/references/modular.md | 15 ++++++++++++--- tests/modular_pipelines/testing_utils/common.py | 5 +++-- 2 files changed, 15 insertions(+), 5 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index dda68a4a2f81..4d40ee9e13f7 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -193,10 +193,19 @@ Don't fall back to the standard-pipeline habit of a config flag branching inside ## 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. +Three transformations can sit between a VAE and a transformer: -- 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`. +- 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` + +- 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`. diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index fe15e73df7a7..26f409ab373d 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -63,8 +63,9 @@ def expected_latents_shape(self) -> tuple: raise NotImplementedError( "You need to set the attribute `expected_latents_shape` in the child test class: the one canonical form " "this family keeps latents in after the denoise group, for the geometry `get_dummy_inputs()` produces. " - "That is the VAE form (e.g. `(1, 4, 32, 32)`) for pipelines that pack/unpack inside the denoise group or " - "the transformer, or the channel-packed form for pipelines whose VAE statistics live there." + "That is the form the VAE encoder blocks produce and the decoder step takes: the plain grid (e.g. " + "`(1, 4, 32, 32)`) for most families, or the patchified grid for families whose VAE statistics are " + "defined over patchified channels (e.g. `AutoencoderKLFlux2`)." ) # ==================== Required interface ==================== From 2c2c0cbe15c9cc853f39395477dde97508776d09 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Sun, 18 Oct 2026 20:39:14 -0700 Subject: [PATCH 07/11] require geometry on the unpack steps and test the deprecated decode path directly --- .ai/references/modular.md | 2 + .../modular_pipelines/flux/decoders.py | 9 ++-- .../modular_pipelines/krea2/decoders.py | 41 +++++++++++++------ .../modular_pipelines/ltx/decoders.py | 11 +++-- .../ltx/modular_blocks_ltx.py | 8 ++-- tests/modular_pipelines/echo/testing_utils.py | 1 + .../flux/test_modular_pipeline_flux.py | 18 ++++---- .../flux3/test_modular_pipeline_flux3.py | 2 + .../krea2/test_modular_pipeline_krea2.py | 21 ++++------ .../ltx/test_modular_pipeline_ltx.py | 32 ++++++++------- .../ltx2/test_modular_pipeline_ltx2.py | 8 +++- .../test_modular_pipeline_qwenimage21.py | 1 + ...test_modular_pipeline_qwenimage21_turbo.py | 1 + .../modular_pipelines/testing_utils/common.py | 2 - .../test_modular_pipeline_wan_animate_2.py | 8 ++-- 15 files changed, 92 insertions(+), 73 deletions(-) diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 4d40ee9e13f7..892d59342d7d 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -203,6 +203,8 @@ Each should be applied and undone in mirrored pairs, and the core denoise group 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. diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index ba7980c438c7..5d1cdb70b22c 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -53,8 +53,7 @@ class FluxUnpackLatentsStep(ModularPipelineBlocks): 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." + "VAE takes (still normalized)." ) @property @@ -66,8 +65,8 @@ def inputs(self) -> list[tuple[str, Any]]: 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), + InputParam("height", required=True), + InputParam("width", required=True), ] @property @@ -144,7 +143,7 @@ def __call__(self, components, state: PipelineState) -> tuple[FluxModularPipelin if latents.ndim == 3: deprecate( "packed latents", - "1.0.0", + "0.46.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.", ) diff --git a/src/diffusers/modular_pipelines/krea2/decoders.py b/src/diffusers/modular_pipelines/krea2/decoders.py index 944e84594f57..3adc66b07036 100644 --- a/src/diffusers/modular_pipelines/krea2/decoders.py +++ b/src/diffusers/modular_pipelines/krea2/decoders.py @@ -27,7 +27,6 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name -# auto_docstring def _unpack_latents(latents, height, width, patch_size, vae_scale_factor): batch_size, _, channels = latents.shape p = patch_size @@ -38,15 +37,32 @@ def _unpack_latents(latents, height, width, patch_size, vae_scale_factor): 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). Closes the core denoise group, so the blocks that follow take the same " - "form the VAE produces and need no geometry inputs." + "the VAE takes (still normalized)." ) @property @@ -58,8 +74,8 @@ def inputs(self) -> list[InputParam]: 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), + InputParam.template("height", required=True), + InputParam.template("width", required=True), ] @property @@ -86,10 +102,11 @@ def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> Pi 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`) @@ -97,12 +114,12 @@ class Krea2DecodeStep(ModularPipelineBlocks): 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`): @@ -114,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 @@ -161,7 +178,7 @@ def __call__( if latents.ndim == 3: deprecate( "packed latents", - "1.0.0", + "0.46.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.", ) diff --git a/src/diffusers/modular_pipelines/ltx/decoders.py b/src/diffusers/modular_pipelines/ltx/decoders.py index 071931e216f1..defa9470e0a4 100644 --- a/src/diffusers/modular_pipelines/ltx/decoders.py +++ b/src/diffusers/modular_pipelines/ltx/decoders.py @@ -46,8 +46,7 @@ class LTXUnpackLatentsStep(ModularPipelineBlocks): def description(self) -> str: return ( "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, F, 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." + "the VAE takes (still normalized)." ) @property @@ -70,9 +69,9 @@ def inputs(self) -> list[tuple[str, Any]]: type_hint=torch.Tensor, description="The denoised latents from the denoising step, packed, of shape `[B, S, C]`.", ), - InputParam.template("height", default=512), - InputParam.template("width", default=704), - InputParam("num_frames", type_hint=int, default=161), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + InputParam("num_frames", type_hint=int, required=True), ] @property @@ -152,7 +151,7 @@ def __call__(self, components, state: PipelineState) -> tuple[LTXModularPipeline if latents.ndim == 3: deprecate( "packed latents", - "1.0.0", + "0.46.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, F, H, W]` instead.", ) diff --git a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py index ae9cc7a58d84..0bdaf6ec8816 100644 --- a/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py +++ b/src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py @@ -299,11 +299,11 @@ class LTXAutoCoreDenoiseStep(AutoPipelineBlocks): Timesteps for the denoising process. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - height (`int`, *optional*, defaults to 512): + height (`int`): The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): + width (`int`): The width in pixels of the generated image. - num_frames (`int`, *optional*, defaults to 161): + num_frames (`int`): TODO: Add description. frame_rate (`int`, *optional*, defaults to 25): TODO: Add description. @@ -372,7 +372,7 @@ class LTXAutoBlocks(SequentialPipelineBlocks): Timesteps for the denoising process. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - num_frames (`int`, *optional*, defaults to 161): + num_frames (`int`): TODO: Add description. frame_rate (`int`, *optional*, defaults to 25): TODO: Add description. diff --git a/tests/modular_pipelines/echo/testing_utils.py b/tests/modular_pipelines/echo/testing_utils.py index 08320560a8f4..b85d1d1930c3 100644 --- a/tests/modular_pipelines/echo/testing_utils.py +++ b/tests/modular_pipelines/echo/testing_utils.py @@ -24,6 +24,7 @@ class EchoModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = EchoModularPipeline pipeline_blocks_class = EchoBlocks pretrained_model_name_or_path = "Echo-Team/tiny-echo-modular-pipe" + expected_latents_shape = (1, 4, 3, 16, 16) params = frozenset( [ "prompt", diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 8eeeb0880540..53fa6e3f03dc 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import random import numpy as np @@ -83,19 +82,16 @@ def get_dummy_inputs(self, seed=0): class TestFluxModularPipelineFast(FluxModularPipelineTesterConfig, ModularPipelineTesterMixin): def test_decode_accepts_packed_latents_with_deprecation(self): # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. - inputs = self.get_dummy_inputs() - expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) + decoder = self.pipeline_blocks_class().sub_blocks["decode"].init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components(dtype=torch.float32) + decoder.to(torch_device) - # Blocksets can be composed from shared block instances, so pop from a copy. - blocks = copy.deepcopy(self.pipeline_blocks_class()) - blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") - pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.load_components(dtype=torch.float32) - pipe.to(torch_device) + # The 8x8 dummy canvas is a `[1, 1, 8, 8]` latent, packed in 2x2 patches to 16 tokens of 4 channels. + packed_latents = torch.randn(1, 16, 4, generator=self.get_generator()) with pytest.warns(FutureWarning, match="packed latents"): - output = pipe(**self.get_dummy_inputs(), output=self.output_name) + images = decoder(latents=packed_latents, height=8, width=8, output_type="pt", output="images") - assert torch.equal(output, expected) + assert images.shape == (1, 3, 8, 8) def test_float16_inference(self): super().test_float16_inference(9e-2) diff --git a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py index f024128022e8..f23df103e908 100644 --- a/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py +++ b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py @@ -178,6 +178,8 @@ class Flux3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = Flux3ModularPipeline pipeline_blocks_class = Flux3AutoBlocks pretrained_model_name_or_path = FLUX3_TINY_REPO + # The patchified grid: `AutoencoderKLFlux2` defines its latent statistics over patchified channels. + expected_latents_shape = (1, 16, 16, 16) params = frozenset(["prompt", "height", "width"]) batch_params = frozenset() expected_workflow_blocks = FLUX3_WORKFLOWS diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index 569f8acc969d..d940b410b14c 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -14,8 +14,6 @@ # limitations under the License. -import copy - import pytest import torch @@ -71,19 +69,16 @@ def get_dummy_inputs(self, seed=0): class TestKrea2ModularPipelineFast(Krea2ModularPipelineTesterConfig, ModularPipelineTesterMixin): def test_decode_accepts_packed_latents_with_deprecation(self): # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. - inputs = self.get_dummy_inputs() - expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) - - # Blocksets can be composed from shared block instances, so pop from a copy. - blocks = copy.deepcopy(self.pipeline_blocks_class()) - blocks.sub_blocks["denoise"].sub_blocks.pop("unpack_latents") - pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.load_components(dtype=torch.float32) - pipe.to(torch_device) + decoder = self.pipeline_blocks_class().sub_blocks["decode"].init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components(dtype=torch.float32) + decoder.to(torch_device) + + # The 32x32 dummy canvas is a `[1, 4, 1, 8, 8]` latent, packed in 2x2 patches to 16 tokens of 16 channels. + packed_latents = torch.randn(1, 16, 16, generator=self.get_generator()) with pytest.warns(FutureWarning, match="packed latents"): - output = pipe(**self.get_dummy_inputs(), output=self.output_name) + images = decoder(latents=packed_latents, height=32, width=32, output_type="pt", output="images") - assert torch.equal(output, expected) + assert images.shape == (1, 3, 32, 32) def test_inference_batch_single_identical(self): super().test_inference_batch_single_identical(expected_max_diff=5e-3) diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index c9de9a51f587..951891a4aa36 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -13,7 +13,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -import copy import pytest import torch @@ -83,20 +82,25 @@ def get_dummy_inputs(self, seed=0): class TestLTXModularPipelineFast(LTXModularPipelineTesterConfig, ModularPipelineTesterMixin): def test_decode_accepts_packed_latents_with_deprecation(self): # The old packed `[B, S, C]` form is still accepted by the decode step for a deprecation window. - inputs = self.get_dummy_inputs() - expected = self.get_pipeline().to(torch_device)(**inputs, output=self.output_name) - - # Blocksets can be composed from shared block instances, so pop from a copy. - blocks = copy.deepcopy(self.pipeline_blocks_class()) - for workflow in blocks.sub_blocks["denoise"].sub_blocks.values(): - workflow.sub_blocks.pop("unpack_latents") - pipe = blocks.init_pipeline(self.pretrained_model_name_or_path) - pipe.load_components(dtype=torch.float32) - pipe.to(torch_device) - with pytest.warns(FutureWarning, match="packed latents"): - output = pipe(**self.get_dummy_inputs(), output=self.output_name) + decoder = self.pipeline_blocks_class().sub_blocks["decode"].init_pipeline(self.pretrained_model_name_or_path) + decoder.load_components(dtype=torch.float32) + decoder.to(torch_device) - assert torch.equal(output, expected) + # The 9-frame 32x32 dummy video is a `[1, 8, 5, 16, 16]` latent, packed to 1280 tokens of 8 channels. + packed_latents = torch.randn(1, 1280, 8, generator=self.get_generator()) + with pytest.warns(FutureWarning, match="packed latents"): + videos = decoder( + latents=packed_latents, + height=32, + width=32, + num_frames=9, + batch_size=1, + dtype=torch.float32, + output_type="pt", + output="videos", + ) + + assert tuple(videos.shape[-2:]) == (32, 32) @pytest.mark.skip(reason="num_videos_per_prompt") def test_num_images_per_prompt(self): diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py index 23f9a519defd..9e0f4ac29573 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py @@ -104,8 +104,6 @@ class LTX2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = LTX2ModularPipeline pipeline_blocks_class = LTX2AutoBlocks pretrained_model_name_or_path = LTX2_REPO_ID - # Still the packed `[B, S, C]` form: #14612 moves ltx2 to the VAE form and updates this. - expected_latents_shape = (1, 768, 4) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) expected_workflow_blocks = LTX2_WORKFLOWS @@ -129,6 +127,12 @@ def get_dummy_inputs(self, seed=0): class LTX2ModularPipelineFastTesterMixin(ModularPipelineTesterMixin): """`ModularPipelineTesterMixin` with the two adjustments every LTX2 workflow needs.""" + @pytest.mark.skip( + reason="ltx2 still leaves packed `[B, S, C]` latents in the state; #14612 moves it to the VAE form." + ) + def test_latents_output_in_canonical_form(self): + pass + @pytest.mark.skip(reason="num_videos_per_prompt") def test_num_images_per_prompt(self): pass diff --git a/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py index 52ca4c7d0fbd..5a15c3effb9b 100644 --- a/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py +++ b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py @@ -134,6 +134,7 @@ class QwenImage21ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = QwenImage21ModularPipeline pipeline_blocks_class = QwenImage21AutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage21-modular" + expected_latents_shape = (1, 8, 1, 2, 2) params = frozenset( ["prompt", "negative_prompt", "height", "width", "output_resolution", "attention_kwargs", "image"] ) diff --git a/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py index a6ec0663dcec..d5cd4136c87d 100644 --- a/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py +++ b/tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py @@ -127,6 +127,7 @@ class QwenImage21TurboModularPipelineTesterConfig(BaseModularPipelineTesterConfi pipeline_class = QwenImage21TurboModularPipeline pipeline_blocks_class = QwenImage21TurboAutoBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage21-turbo-modular" + expected_latents_shape = (1, 8, 1, 2, 2) params = frozenset(["prompt", "height", "width", "output_resolution", "attention_kwargs", "image"]) batch_params = frozenset(["prompt"]) not_params = frozenset(["negative_prompt"]) diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 26f409ab373d..d150d59b138d 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -390,8 +390,6 @@ def test_latents_output_in_canonical_form(self): # The denoise group must hand latents back in the family's one canonical form, so any block (a decoder, a # latent upsampler, a second denoise group) can consume them without geometry inputs. The decoder reads # the latents without removing them, so they are checked in the state after a full run. - if self.latents_output_name is None: - pytest.skip("This family decodes inside its denoise loop and leaves no latents in the state.") pipe = self.get_pipeline().to(torch_device) latents = pipe(**self.get_dummy_inputs(), output=self.latents_output_name) chunks = list(latents) if isinstance(latents, (list, tuple)) else [latents] diff --git a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py index 1c2af8c0f94c..e7906dfb6187 100644 --- a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py +++ b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py @@ -85,8 +85,6 @@ class WanAnimate2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_blocks_class = WanAnimate2Blocks pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-modular" # Each segment is decoded inside the denoise loop, so no latents are left in the state. - latents_output_name = None - expected_latents_shape = None params = frozenset(["prompt", "image", "driving_video"]) batch_params = frozenset() @@ -117,6 +115,10 @@ def get_dummy_inputs(self, seed=0): class WanAnimate2ModularPipelineFastTesterMixin(ModularPipelineTesterMixin): """`ModularPipelineTesterMixin` minus the tests that assume a batchable pipeline; shared by both presets.""" + @pytest.mark.skip(reason="Wan-Animate-2 decodes inside its segment loop and leaves no latents in the state.") + def test_latents_output_in_canonical_form(self): + pass + @pytest.mark.skip(reason="Wan-Animate-2 is unbatched: one character image and driving video per call") def test_inference_batch_consistent(self): pass @@ -161,8 +163,6 @@ class WanAnimate2DistilledModularPipelineTesterConfig(WanAnimate2ModularPipeline pipeline_blocks_class = WanAnimate2DistilledBlocks pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-distilled-modular" # Each segment is decoded inside the denoise loop, so no latents are left in the state. - latents_output_name = None - expected_latents_shape = None expected_workflow_defaults = WAN_ANIMATE_2_DISTILLED_DEFAULTS From 80c734ba463083bfca0144a27cbc6d2819b6a599 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 19 Oct 2026 20:01:18 -0700 Subject: [PATCH 08/11] skip the canonical latents test explicitly for ltx 2.5 too --- .../modular_pipelines/ltx2/test_modular_pipeline_ltx25.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index d36c12d8e676..7eb7d4696b79 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py @@ -101,8 +101,6 @@ class LTX25ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_class = LTX25ModularPipeline pipeline_blocks_class = LTX25AutoBlocks pretrained_model_name_or_path = LTX25_REPO_ID - # Still the packed `[B, S, C]` form: #14612 moves ltx2 to the VAE form and updates this. - expected_latents_shape = (1, 768, 4) batch_params = frozenset(["prompt"]) optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) expected_workflow_blocks = LTX25_WORKFLOWS @@ -126,6 +124,12 @@ def get_dummy_inputs(self, seed=0): class LTX25ModularPipelineFastTesterMixin(ModularPipelineTesterMixin): """`ModularPipelineTesterMixin` with the two adjustments every LTX-2.5 workflow needs.""" + @pytest.mark.skip( + reason="LTX-2.5 still leaves packed `[B, S, C]` latents in the state; #14612 moves it to the VAE form." + ) + def test_latents_output_in_canonical_form(self): + pass + @pytest.mark.skip(reason="num_videos_per_prompt") def test_num_images_per_prompt(self): pass From 060e96d5c8c9814fc77ec2650a86189884c031f2 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 19 Oct 2026 20:09:40 -0700 Subject: [PATCH 09/11] remove the packed latents path in 0.47.0, five releases after 0.42 --- src/diffusers/modular_pipelines/flux/decoders.py | 2 +- src/diffusers/modular_pipelines/krea2/decoders.py | 2 +- src/diffusers/modular_pipelines/ltx/decoders.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 5d1cdb70b22c..573adbb7ff2f 100644 --- a/src/diffusers/modular_pipelines/flux/decoders.py +++ b/src/diffusers/modular_pipelines/flux/decoders.py @@ -143,7 +143,7 @@ def __call__(self, components, state: PipelineState) -> tuple[FluxModularPipelin if latents.ndim == 3: deprecate( "packed latents", - "0.46.0", + "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.", ) diff --git a/src/diffusers/modular_pipelines/krea2/decoders.py b/src/diffusers/modular_pipelines/krea2/decoders.py index 3adc66b07036..751b312037d0 100644 --- a/src/diffusers/modular_pipelines/krea2/decoders.py +++ b/src/diffusers/modular_pipelines/krea2/decoders.py @@ -178,7 +178,7 @@ def __call__( if latents.ndim == 3: deprecate( "packed latents", - "0.46.0", + "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.", ) diff --git a/src/diffusers/modular_pipelines/ltx/decoders.py b/src/diffusers/modular_pipelines/ltx/decoders.py index defa9470e0a4..2917eca125d0 100644 --- a/src/diffusers/modular_pipelines/ltx/decoders.py +++ b/src/diffusers/modular_pipelines/ltx/decoders.py @@ -151,7 +151,7 @@ def __call__(self, components, state: PipelineState) -> tuple[LTXModularPipeline if latents.ndim == 3: deprecate( "packed latents", - "0.46.0", + "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, F, H, W]` instead.", ) From 09903d82b425f1716875aeffe36389e7e699b93d Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 19 Oct 2026 20:15:27 -0700 Subject: [PATCH 10/11] update the latents_output_name comment for the explicit skips --- tests/modular_pipelines/testing_utils/common.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index d150d59b138d..545ca0f31ef0 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -54,8 +54,8 @@ class BaseModularPipelineTesterConfig: # Subclasses can override this to change the expected output type output_name = "images" # State name of the latents the denoise group hands to the decoder. Chunked families that hand over a list of - # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. Families that decode inside their - # denoise loop and leave no latents in the state set this to `None`. + # per-chunk latents (e.g. `latent_chunks`) are checked chunk by chunk. A family that leaves no latents in the + # state, or not yet in its canonical form, skips `test_latents_output_in_canonical_form` with a reason. latents_output_name = "latents" @property From 436058a86d44963319e71a93604eb1605c59d1e0 Mon Sep 17 00:00:00 2001 From: Akshan Krithick Date: Mon, 19 Oct 2026 20:18:57 -0700 Subject: [PATCH 11/11] move the packed test latents to the test device --- tests/modular_pipelines/flux/test_modular_pipeline_flux.py | 2 +- tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py | 2 +- tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 53fa6e3f03dc..0d076477ebbe 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -87,7 +87,7 @@ def test_decode_accepts_packed_latents_with_deprecation(self): decoder.to(torch_device) # The 8x8 dummy canvas is a `[1, 1, 8, 8]` latent, packed in 2x2 patches to 16 tokens of 4 channels. - packed_latents = torch.randn(1, 16, 4, generator=self.get_generator()) + packed_latents = torch.randn(1, 16, 4, generator=self.get_generator()).to(torch_device) with pytest.warns(FutureWarning, match="packed latents"): images = decoder(latents=packed_latents, height=8, width=8, output_type="pt", output="images") diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index d940b410b14c..0573823eeaba 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -74,7 +74,7 @@ def test_decode_accepts_packed_latents_with_deprecation(self): decoder.to(torch_device) # The 32x32 dummy canvas is a `[1, 4, 1, 8, 8]` latent, packed in 2x2 patches to 16 tokens of 16 channels. - packed_latents = torch.randn(1, 16, 16, generator=self.get_generator()) + packed_latents = torch.randn(1, 16, 16, generator=self.get_generator()).to(torch_device) with pytest.warns(FutureWarning, match="packed latents"): images = decoder(latents=packed_latents, height=32, width=32, output_type="pt", output="images") diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index 951891a4aa36..3f9df2e52f6c 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -87,7 +87,7 @@ def test_decode_accepts_packed_latents_with_deprecation(self): decoder.to(torch_device) # The 9-frame 32x32 dummy video is a `[1, 8, 5, 16, 16]` latent, packed to 1280 tokens of 8 channels. - packed_latents = torch.randn(1, 1280, 8, generator=self.get_generator()) + packed_latents = torch.randn(1, 1280, 8, generator=self.get_generator()).to(torch_device) with pytest.warns(FutureWarning, match="packed latents"): videos = decoder( latents=packed_latents,