AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

move unpack out of the flux, krea2 and ltx decoders into the core denoise groups (#14737)

Browse Source 
* add a modular test that pins the canonical latents form each family leaves in the state

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

* pop the unpack step from a copy of the blocks in the deprecation tests

* let families that decode inside the denoise loop opt out of the latents-form test

* document the canonical latents form in the modular reference

* [docs] canonical latents form: name normalize, patchify and pack separately

Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

* require geometry on the unpack steps and test the deprecated decode path directly

* skip the canonical latents test explicitly for ltx 2.5 too

* remove the packed latents path in 0.47.0, five releases after 0.42

* update the latents_output_name comment for the explicit skips

* move the packed test latents to the test device

---------

Co-authored-by: YiYi Xu <yixu310@gmail.com>
Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com>
This commit is contained in:
Akshan Krithick
2026-10-20 01:16:18 -07:00
committed by GitHub
parent e04f0c5f24
commit 6184688f71
38 changed files with 420 additions and 51 deletions
Download Patch File Download Diff File Expand all files Collapse all files

20
.ai/references/modular.md
View File

@@ -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.

61
src/diffusers/modular_pipelines/flux/decoders.py
View File

@@ -20,7 +20,7 @@ import torch
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 @@ class FluxDecodeStep(ModularPipelineBlocks):
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 @@ class FluxDecodeStep(ModularPipelineBlocks):
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(

6
src/diffusers/modular_pipelines/flux/modular_blocks_flux.py
View File

@@ -22,7 +22,7 @@ from .before_denoise import (
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):

11
src/diffusers/modular_pipelines/flux/modular_blocks_flux_kontext.py
View File

@@ -21,7 +21,7 @@ from .before_denoise import (
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):

114
src/diffusers/modular_pipelines/krea2/decoders.py
View File

@@ -18,7 +18,7 @@ import torch
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 @@ from .modular_pipeline import Krea2ModularPipeline
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 @@ class Krea2DecodeStep(ModularPipelineBlocks):
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 @@ class Krea2DecodeStep(ModularPipelineBlocks):
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 = (

3
src/diffusers/modular_pipelines/krea2/modular_blocks_krea2.py
View File

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

3
src/diffusers/modular_pipelines/krea2/modular_blocks_krea2_turbo.py
View File

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

87
src/diffusers/modular_pipelines/ltx/decoders.py
View File

@@ -18,7 +18,7 @@ import torch
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 @@ class LTXVaeDecoderStep(ModularPipelineBlocks):
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 @@ class LTXVaeDecoderStep(ModularPipelineBlocks):
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)

16
src/diffusers/modular_pipelines/ltx/modular_blocks_ltx.py
View File

@@ -21,7 +21,7 @@ from .before_denoise import (
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.

2
tests/modular_pipelines/anima/test_modular_pipeline_anima.py
View File

@@ -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

1
tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py
View File

@@ -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"])

1
tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py
View File

@@ -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"])

1
tests/modular_pipelines/echo/testing_utils.py
View File

@@ -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",

1
tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py
View File

@@ -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"])

21
tests/modular_pipelines/flux/test_modular_pipeline_flux.py
View File

@@ -17,6 +17,7 @@ import random
import numpy as np
import PIL
import pytest
import torch
from diffusers.image_processor import VaeImageProcessor
@@ -47,6 +48,7 @@ FLUX_TEXT2IMAGE_WORKFLOWS = {
("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 @@ class FluxModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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 @@ FLUX_IMAGE2IMAGE_WORKFLOWS = {
("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 @@ FLUX_KONTEXT_WORKFLOWS = {
("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 @@ FLUX_KONTEXT_WORKFLOWS = {
("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"])

2
tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py
View File

@@ -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

2
tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py
View File

@@ -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"])

2
tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py
View File

@@ -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

2
tests/modular_pipelines/flux3/test_modular_pipeline_flux3.py
View File

@@ -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

2
tests/modular_pipelines/helios/test_modular_pipeline_helios.py
View File

@@ -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"])

1
tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py
View File

@@ -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"])

19
tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py
View File

@@ -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 @@ KREA2_WORKFLOWS = {
("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 @@ class Krea2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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)

2
tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py
View File

@@ -33,6 +33,7 @@ KREA2_TURBO_WORKFLOWS = {
("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

28
tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py
View File

@@ -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 @@ LTX_WORKFLOWS = {
("denoise.set_timesteps", "LTXSetTimestepsStep"),
("denoise.prepare_latents", "LTXPrepareLatentsStep"),
("denoise.denoise", "LTXDenoiseStep"),
("denoise.unpack_latents", "LTXUnpackLatentsStep"),
("decode", "LTXVaeDecoderStep"),
],
"image2video": [
@@ -43,6 +47,7 @@ LTX_WORKFLOWS = {
("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 @@ class LTXModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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

6
tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py
View File

@@ -127,6 +127,12 @@ class LTX2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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

6
tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py
View File

@@ -124,6 +124,12 @@ class LTX25ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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

2
tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py
View File

@@ -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()

2
tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py
View File

@@ -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": {

3
tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py
View File

@@ -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"])

1
tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21.py
View File

@@ -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"]
)

1
tests/modular_pipelines/qwenimage21/test_modular_pipeline_qwenimage21_turbo.py
View File

@@ -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"])

2
tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py
View File

@@ -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

3
tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py
View File

@@ -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",

28
tests/modular_pipelines/testing_utils/common.py
View File

@@ -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 @@ class ModularPipelineTesterMixin(BaseModularPipelineOutputMixin):
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)}"
)

1
tests/modular_pipelines/wan/test_modular_pipeline_wan.py
View File

@@ -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"])

1
tests/modular_pipelines/wan/test_modular_pipeline_wan_vace.py
View File

@@ -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()

6
tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py
View File

@@ -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 @@ class WanAnimate2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
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

1
tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py
View File

@@ -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