diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 9411d5ee6c37..892d59342d7d 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -191,6 +191,26 @@ Default to taking that option. The only reason not to split is when the variant Don't fall back to the standard-pipeline habit of a config flag branching inside a shared block (`ConfigSpec(name="is_distilled")` + `if components.config.is_distilled:`). That keeps both variants' behavior bundled in one blockset — and the input surface is the one thing it can never fix: a repo can override components and config values per checkpoint, but never which inputs the blocks declare, so the distilled checkpoint would still accept `negative_prompt` and silently ignore it. +## Key pattern: One canonical latents form + +Three transformations can sit between a VAE and a transformer: + +- normalize / denormalize: apply or undo the VAE's latent statistics. +- patchify / unpatchify: fold a `p x p` spatial patch into channels, `[B, C, H, W]` <-> `[B, C*p*p, H/p, W/p]`. The result is still a grid. +- pack / unpack: collapse the grid into a token sequence, `[B, C, F, H, W]` <-> `[B, S, D]`. + +Each should be applied and undone in mirrored pairs, and the core denoise group hands latents back in the same form it received them. The `latents` output of the denoise group should always be in one form: unpacked and still normalized, the same form its VAE encoder blocks produce. Whatever runs next (a decoder, a latent upsampler, a second denoise group) consumes that form, and no block outside the group needs `height`/`width`/`num_frames` just to interpret tokens. + +The pattern: `encode -> norm -> pack -> denoise -> unpack -> denorm -> decode` + +In `flux/`: `FluxVaeEncoderStep` encodes and normalizes, `FluxPrepareLatentsStep` packs, `FluxUnpackLatentsStep` closes `FluxCoreDenoiseStep`, and `FluxDecodeStep` denormalizes and decodes. + +- Normalize / denormalize belong to the VAE blocks. The encoder step normalizes right after `vae.encode`, the decoder step denormalizes right before `vae.decode`. Denormalized latents never sit in the state. +- Pack / unpack belong to the denoise group, either at block level (a prepare-latents step packs, an unpack step such as `FluxUnpackLatentsStep` closes the group) or inside the transformer's forward. Packing is a transformer layout detail; nothing outside the group should see tokens. +- Patchify / unpatchify can sit on either side, as long as it is mirrored: usually it is fused with pack in the denoise group (Flux, Krea2, LTX, QwenImage); when the VAE statistics are defined over patchified channels (`AutoencoderKLFlux2`: Flux2, FLUX.3, Ernie-Image) it lives in the VAE blocks and the canonical form is the patchified grid. + +Don't pack inside the denoise group and unpack inside the decoder: that strands packed latents in the state and makes the decoder carry geometry inputs it doesn't otherwise need. `test_latents_output_in_canonical_form` in `tests/modular_pipelines/testing_utils/common.py` pins each family's form through the tester's `expected_latents_shape`. + ## Key pattern: Standalone block reusability One of the core reason a pipeline is split into blocks at all: each block (text encoder, VAE encoder, prepare-latents, denoise, decoder) must be runnable on its own, and its output must be reusable as the input to a different downstream chain. diff --git a/src/diffusers/modular_pipelines/flux/decoders.py b/src/diffusers/modular_pipelines/flux/decoders.py index 796bff258a6b..573adbb7ff2f 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 @@ -46,6 +46,49 @@ def _unpack_latents(latents, height, width, vae_scale_factor): return latents +class FluxUnpackLatentsStep(ModularPipelineBlocks): + model_name = "flux" + + @property + def description(self) -> str: + return ( + "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, H, W]` form the " + "VAE takes (still normalized)." + ) + + @property + def inputs(self) -> list[tuple[str, Any]]: + return [ + InputParam( + "latents", + required=True, + type_hint=torch.Tensor, + description="The denoised latents from the denoising step, packed, of shape `[B, S, C]`.", + ), + InputParam("height", required=True), + InputParam("width", required=True), + ] + + @property + def intermediate_outputs(self) -> list[str]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The denoised latents of shape `[B, C, H, W]` (normalized, not packed).", + ) + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.latents = _unpack_latents( + block_state.latents, block_state.height, block_state.width, components.vae_scale_factor + ) + self.set_block_state(state, block_state) + return components, state + + class FluxDecodeStep(ModularPipelineBlocks): model_name = "flux" @@ -69,14 +112,15 @@ def description(self) -> str: def inputs(self) -> list[tuple[str, Any]]: return [ InputParam("output_type", default="pil"), - InputParam("height", default=1024), - InputParam("width", default=1024), InputParam( "latents", required=True, type_hint=torch.Tensor, - description="The denoised latents from the denoising step", + description="The denoised latents from the denoising step, of shape `[B, C, H, W]`.", ), + # Only read on the deprecated path that still accepts packed `[B, S, C]` latents. + InputParam("height", default=1024), + InputParam("width", default=1024), ] @property @@ -96,7 +140,14 @@ def __call__(self, components, state: PipelineState) -> tuple[FluxModularPipelin if not block_state.output_type == "latent": latents = block_state.latents - latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor) + if latents.ndim == 3: + deprecate( + "packed latents", + "0.47.0", + "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group " + "now unpacks them. Pass latents of shape `[B, C, H, W]` instead.", + ) + latents = _unpack_latents(latents, block_state.height, block_state.width, components.vae_scale_factor) latents = (latents / vae.config.scaling_factor) + vae.config.shift_factor block_state.images = vae.decode(latents, return_dict=False)[0] block_state.images = components.image_processor.postprocess( 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 2c3fe4c0bf20..751b312037d0 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 @@ -27,11 +27,86 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def _unpack_latents(latents, height, width, patch_size, vae_scale_factor): + batch_size, _, channels = latents.shape + p = patch_size + height = p * (int(height) // (vae_scale_factor * p)) + width = p * (int(width) // (vae_scale_factor * p)) + latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p) + latents = latents.permute(0, 3, 1, 4, 2, 5) + return latents.reshape(batch_size, channels // (p * p), 1, height, width) + + +# auto_docstring +class Krea2UnpackLatentsStep(ModularPipelineBlocks): + """ + Unpacks the denoised latents from the transformer's token layout back into the `[B, C, 1, H, W]` form the VAE takes + (still normalized). + + Inputs: + latents (`Tensor`): + The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop. + height (`int`): + The height in pixels of the generated image. + width (`int`): + The width in pixels of the generated image. + + Outputs: + latents (`Tensor`): + The denoised latents of shape `[B, C, 1, H, W]` (normalized, not packed). + """ + + model_name = "krea2" + + @property + def description(self) -> str: + return ( + "Unpacks the denoised latents from the transformer's token layout back into the `[B, C, 1, H, W]` form " + "the VAE takes (still normalized)." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + name="latents", + required=True, + type_hint=torch.Tensor, + description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.", + ), + InputParam.template("height", required=True), + InputParam.template("width", required=True), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="The denoised latents of shape `[B, C, 1, H, W]` (normalized, not packed).", + ) + ] + + @torch.no_grad() + def __call__(self, components: Krea2ModularPipeline, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + block_state.latents = _unpack_latents( + block_state.latents, + block_state.height, + block_state.width, + components.patch_size, + components.vae_scale_factor, + ) + self.set_block_state(state, block_state) + return components, state + + # auto_docstring class Krea2DecodeStep(ModularPipelineBlocks): """ - Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the VAE's - per-channel statistics, and decodes them through the Qwen-Image VAE into images. + Step that de-normalizes the denoised latents with the VAE's per-channel statistics and decodes them through the + Qwen-Image VAE into images. Components: vae (`AutoencoderKLQwenImage`) image_processor (`VaeImageProcessor`) @@ -39,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`): @@ -56,8 +131,8 @@ class Krea2DecodeStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Step that unpacks the denoised packed latents back to the spatial grid, de-normalizes them with the " - "VAE's per-channel statistics, and decodes them through the Qwen-Image VAE into images." + "Step that de-normalizes the denoised latents with the VAE's per-channel statistics and decodes them " + "through the Qwen-Image VAE into images." ) @property @@ -77,14 +152,15 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam.template("output_type", default="pil"), - InputParam.template("height", default=1024), - InputParam.template("width", default=1024), InputParam( name="latents", required=True, type_hint=torch.Tensor, - description="The denoised packed latents (B, image_seq_len, in_channels) from the denoising loop.", + description="The denoised latents of shape `[B, C, 1, H, W]` from the denoising group.", ), + # Only read on the deprecated path that still accepts packed `[B, S, C]` latents. + InputParam.template("height", default=1024), + InputParam.template("width", default=1024), ] @property @@ -98,15 +174,17 @@ def __call__( block_state = self.get_block_state(state) vae = components.vae - p = components.patch_size latents = block_state.latents - - batch_size, _, channels = latents.shape - height = p * (int(block_state.height) // (components.vae_scale_factor * p)) - width = p * (int(block_state.width) // (components.vae_scale_factor * p)) - latents = latents.view(batch_size, height // p, width // p, channels // (p * p), p, p) - latents = latents.permute(0, 3, 1, 4, 2, 5) - latents = latents.reshape(batch_size, channels // (p * p), 1, height, width) + if latents.ndim == 3: + deprecate( + "packed latents", + "0.47.0", + "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group " + "now unpacks them. Pass latents of shape `[B, C, 1, H, W]` instead.", + ) + latents = _unpack_latents( + latents, block_state.height, block_state.width, components.patch_size, components.vae_scale_factor + ) latents = latents.to(vae.dtype) latents_mean = ( 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 9d5238069bcf..2917eca125d0 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,64 @@ 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)." + ) + + @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", required=True), + InputParam.template("width", required=True), + InputParam("num_frames", type_hint=int, required=True), + ] + + @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 +127,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 +148,18 @@ def __call__(self, components, state: PipelineState) -> tuple[LTXModularPipeline vae = components.vae latents = block_state.latents + if latents.ndim == 3: + deprecate( + "packed latents", + "0.47.0", + "Passing packed latents of shape `[B, S, C]` to the decode step is deprecated; the denoise group " + "now unpacks them. Pass latents of shape `[B, C, 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..0bdaf6ec8816 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): @@ -297,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. @@ -370,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/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index ff8a52839df1..625dbd552136 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -110,6 +110,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 @@ -255,6 +256,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 c3bcb267758f..5a2723b84625 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -132,6 +132,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 86bc6ad4d6f2..cf3a38a66b7c 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -71,6 +71,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/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/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/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 502df3f447bf..0d076477ebbe 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,19 @@ 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. + 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 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()).to(torch_device) + with pytest.warns(FutureWarning, match="packed latents"): + images = decoder(latents=packed_latents, height=8, width=8, output_type="pt", output="images") + + assert images.shape == (1, 3, 8, 8) + def test_float16_inference(self): super().test_float16_inference(9e-2) @@ -109,6 +125,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 +135,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 +211,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 +225,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 +235,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/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/flux3/test_modular_pipeline_flux3.py b/tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py index e5b6c07b117c..3d4daa665dc7 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/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/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index 133b3eb1bf66..0573823eeaba 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,19 @@ 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. + 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()).to(torch_device) + with pytest.warns(FutureWarning, match="packed latents"): + images = decoder(latents=packed_latents, height=32, width=32, output_type="pt", output="images") + + 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/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..3f9df2e52f6c 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -13,10 +13,13 @@ # See the License for the specific language governing permissions and # 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 +36,7 @@ ("denoise.set_timesteps", "LTXSetTimestepsStep"), ("denoise.prepare_latents", "LTXPrepareLatentsStep"), ("denoise.denoise", "LTXDenoiseStep"), + ("denoise.unpack_latents", "LTXUnpackLatentsStep"), ("decode", "LTXVaeDecoderStep"), ], "image2video": [ @@ -43,6 +47,7 @@ ("denoise.prepare_latents", "LTXPrepareLatentsStep"), ("denoise.prepare_i2v_latents", "LTXImage2VideoPrepareLatentsStep"), ("denoise.denoise", "LTXImage2VideoDenoiseStep"), + ("denoise.unpack_latents", "LTXUnpackLatentsStep"), ("decode", "LTXVaeDecoderStep"), ], } @@ -52,6 +57,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 +80,28 @@ 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. + 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 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()).to(torch_device) + 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): 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..9e0f4ac29573 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py @@ -127,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/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index 69080095c627..7eb7d4696b79 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py @@ -124,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 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 a5eb5095366e..03239ef82e3d 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 @@ -169,6 +169,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() @@ -608,6 +609,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/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/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 26d7391a286b..545ca0f31ef0 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -53,6 +53,20 @@ 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. 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 + 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 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 ==================== @@ -371,3 +385,17 @@ 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}) + + 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/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..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 @@ -84,6 +84,7 @@ 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. params = frozenset(["prompt", "image", "driving_video"]) batch_params = frozenset() @@ -114,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 @@ -157,6 +162,7 @@ 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. expected_workflow_defaults = WAN_ANIMATE_2_DISTILLED_DEFAULTS 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