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

split up tests in modular pipelines (#14444)

Browse Source 
* notes from .ai

* ltx2

* cover the remaining tests

* fixes
This commit is contained in:
Sayak Paul
2026-08-14 11:54:01 +02:00
committed by GitHub
parent 8222302f03
commit 102c24934f
36 changed files with 2848 additions and 1784 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
.ai/testing.md
View File

@@ -34,9 +34,15 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers
### Modular pipelines
- Location: `tests/modular_pipelines/<model>/test_modular_pipeline_<model>.py` (one test class per blockset / pipeline variant).
- Subclass `ModularPipelineTesterMixin` (from `..test_modular_pipelines_common`) — it runs the pipeline end-to-end (call signature, batch consistency, float16, device placement) against a tiny checkpoint.
- Set `pipeline_class`, `pipeline_blocks_class`, `pretrained_model_name_or_path`, `params` / `batch_params`, and implement `get_dummy_inputs(seed=0)`. Set `expected_workflow_blocks` to pin the block name → class ordering per workflow (only for blocksets with a `_workflow_map` — with a single workflow the list would just restate the class definition), and `expected_workflow_defaults` to pin each workflow's components, pipeline configs, and inputs — required ones by name, optional ones with their defaults. A pipeline without workflows pins its full blockset under the `None` key. An optional `component_configs` entry pins config values of `from_config` components against their creating spec (e.g. the guider scale that tells a base and a distilled preset apart); pretrained components take their config from the repo, so there is nothing block-level to pin.
- Location: `tests/modular_pipelines/<model>/test_modular_pipeline_<model>.py` (one config class + set of test classes per blockset / pipeline variant).
- **Define one config class**, `<Pipeline>ModularPipelineTesterConfig`, subclassing `BaseModularPipelineTesterConfig` (from `..testing_utils`). Set `pipeline_class`, `pipeline_blocks_class`, `pretrained_model_name_or_path`, `params` / `batch_params`, and implement `get_dummy_inputs(seed=0)`. Set `expected_workflow_blocks` to pin the block name → class ordering per workflow. The config holds the whole testing contract and performs no assertions.
- **Then one test class per concern**, each composing the config with a tester mixin from `..testing_utils`. Keep them separate — pytest reads class-level markers off the whole MRO, so folding a marked mixin (`@is_memory`, ...) into the same class as the others would tag every test in it:
- `ModularPipelineTesterMixin` — call signature, batch consistency, float16, device placement, NaN-free output. Put pipeline-specific tests as methods on this class.
- `ModularLoadingTesterMixin` — `save_pretrained`/`from_pretrained` round-trips, `modular_model_index.json` contents, `load_components`/`unload_components`.
- `ModularWorkflowTesterMixin` — everything driven by the blocks class's `_workflow_map`; skips itself when there is none.
- `ModularMemoryTesterMixin` — auto CPU offload, group offload, device memory reclaimed on unload.
- `ModularGuiderTesterMixin` — only for pipelines with a `guider` component.
- `ModularAutoOffloadTesterMixin` — opt-in, for pipelines with several offloadable model components; asserts on the offload *decisions* under simulated memory pressure.
- `pretrained_model_name_or_path` is a tiny repo with real components (tiny transformer, real scheduler / VAE / tokenizer configs). Develop against a personal repo; tiny repos ultimately live under `hf-internal-testing/` — not merge-blocking, a maintainer moves it before or after merge.
- **The tiny repo must mirror the real checkpoint's shape** — same index file type, same pipeline-level config keys, a scheduler configured like the real one. A fixture that doesn't look like the published repos tests a loading/config path no user will ever hit, while the path users *do* hit stays uncovered. If the model ships variants with different configs (base/distilled, different schedules), make one tiny repo and test class per variant — see the flux2 klein base/distilled split.
- **Bespoke tests go on the tester class as methods**, not as module-level functions — the mixin is pytest-style, so fixtures (`tmp_path`, `pytest.raises`, parametrize) all work in methods.

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

@@ -26,8 +26,16 @@ from diffusers import (
CosmosTransformer3DModel,
)
from ...testing_utils import enable_full_determinism, require_peft_backend
from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin
from ...testing_utils import enable_full_determinism, is_lora, require_peft_backend
from ..testing_utils import (
BaseModularPipelineOutputMixin,
BaseModularPipelineTesterConfig,
ModularGuiderTesterMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
enable_full_determinism()
@@ -96,7 +104,7 @@ class AnimaTextConditionerFastTests(unittest.TestCase):
self.assertTrue(torch.allclose(output[:, 4:], torch.zeros_like(output[:, 4:]), atol=1e-5))
class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
class AnimaModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = AnimaModularPipeline
pipeline_blocks_class = AnimaAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-anima-modular-pipe"
@@ -117,6 +125,8 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest
"output_type": "pt",
}
class TestAnimaModularPipelineFast(AnimaModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_empty_negative_prompt(self):
pipe = self.get_pipeline()
@@ -130,6 +140,8 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-4)
class TestAnimaModularPipelineLoading(AnimaModularPipelineTesterConfig, ModularLoadingTesterMixin):
def test_save_load_components(self):
pipe = self.get_pipeline()
@@ -141,6 +153,9 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest
assert isinstance(pipe.text_conditioner, AnimaTextConditioner)
assert isinstance(pipe.transformer, CosmosTransformer3DModel)
@is_lora
class TestAnimaModularPipelineLoRA(AnimaModularPipelineTesterConfig, BaseModularPipelineOutputMixin):
def test_lora_state_dict_conversion(self):
state_dict = {
"diffusion_model.blocks.0.self_attn.q_proj.lora_A.weight": torch.randn(2, 32),
@@ -173,7 +188,19 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest
assert "dummy" in pipe.text_conditioner.peft_config
class TestAnimaImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
class TestAnimaModularPipelineWorkflow(AnimaModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestAnimaModularPipelineMemory(AnimaModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestAnimaModularPipelineGuider(AnimaModularPipelineTesterConfig, ModularGuiderTesterMixin):
pass
class AnimaImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = AnimaModularPipeline
pipeline_blocks_class = AnimaAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-anima-modular-pipe"
@@ -196,6 +223,8 @@ class TestAnimaImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "pt",
}
class TestAnimaImg2ImgModularPipelineFast(AnimaImg2ImgModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_basic(self):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
@@ -233,3 +262,15 @@ class TestAnimaImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-4)
class TestAnimaImg2ImgModularPipelineLoading(AnimaImg2ImgModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestAnimaImg2ImgModularPipelineWorkflow(AnimaImg2ImgModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestAnimaImg2ImgModularPipelineMemory(AnimaImg2ImgModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -34,7 +34,13 @@ from diffusers.modular_pipelines.cosmos.before_denoise import (
from diffusers.modular_pipelines.cosmos.encoders import Cosmos3TextEncoderStep
from ...testing_utils import torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
TEXT_VISION_WORKFLOW = [
@@ -114,11 +120,10 @@ COSMOS3_OMNI_WORKFLOWS = {
}
class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin):
class Cosmos3OmniModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Cosmos3OmniModularPipeline
pipeline_blocks_class = Cosmos3OmniBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-cosmos3-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames", "guidance_scale"])
batch_params = frozenset()
optional_params = frozenset(["num_inference_steps", "output_type"])
@@ -143,6 +148,8 @@ class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "latent",
}
class TestCosmos3OmniModularPipelineFast(Cosmos3OmniModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="Cosmos3 does not support batched prompts.")
def test_inference_batch_consistent(self):
pass
@@ -159,20 +166,6 @@ class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin):
def test_float16_inference(self):
pass
def test_save_from_pretrained(self, tmp_path):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path))
loaded_pipe.load_components(dtype=torch.float32)
loaded_pipe.disable_safety_checker()
loaded_pipe.to(torch_device)
base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name)
loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name)
assert torch.abs(base_output - loaded_output).max() < 1e-3
def test_vae_encoder_is_standalone_and_validates_conditioning_inputs(self):
pipe = self.get_pipeline()
vae_encoder = pipe.blocks.sub_blocks["vae_encoder"]
@@ -373,3 +366,27 @@ class TestCosmos3OmniModularPipelineFast(ModularPipelineTesterMixin):
torch.testing.assert_close(timesteps_pipe.scheduler.sigmas[:-1], expected_sigmas)
assert native_timesteps.tolist() == [99, 74, 49, 24]
assert not torch.equal(native_timesteps, default_timesteps)
class TestCosmos3OmniModularPipelineLoading(Cosmos3OmniModularPipelineTesterConfig, ModularLoadingTesterMixin):
def test_save_from_pretrained(self, tmp_path):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path))
loaded_pipe.load_components(dtype=torch.float32)
loaded_pipe.disable_safety_checker()
loaded_pipe.to(torch_device)
base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name)
loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name)
assert torch.abs(base_output - loaded_output).max() < 1e-3
class TestCosmos3OmniModularPipelineWorkflow(Cosmos3OmniModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestCosmos3OmniModularPipelineMemory(Cosmos3OmniModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -21,7 +21,13 @@ from diffusers import ModularPipeline
from diffusers.modular_pipelines import Cosmos3DistilledBlocks, Cosmos3DistilledModularPipeline
from ...testing_utils import torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
TINY_DISTILLED_REPO = "hf-internal-testing/tiny-cosmos3-distilled-modular-pipe"
@@ -59,11 +65,10 @@ COSMOS3_DISTILLED_WORKFLOWS = {
}
class TestCosmos3DistilledModularPipelineFast(ModularPipelineTesterMixin):
class Cosmos3DistilledModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Cosmos3DistilledModularPipeline
pipeline_blocks_class = Cosmos3DistilledBlocks
pretrained_model_name_or_path = TINY_DISTILLED_REPO
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset()
optional_params = frozenset(["num_inference_steps", "output_type"])
@@ -86,20 +91,8 @@ class TestCosmos3DistilledModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "latent",
}
def test_save_from_pretrained(self, tmp_path):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path))
loaded_pipe.load_components(torch_dtype=torch.float32)
loaded_pipe.disable_safety_checker()
loaded_pipe.to(torch_device)
base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name)
loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name)
assert torch.abs(base_output - loaded_output).max() < 1e-3
class TestCosmos3DistilledModularPipelineFast(Cosmos3DistilledModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="Cosmos3 does not support batched prompts.")
def test_inference_batch_consistent(self):
pass
@@ -153,3 +146,31 @@ class TestCosmos3DistilledModularPipelineFast(ModularPipelineTesterMixin):
with pytest.raises(ValueError, match="`guidance_scale` must be 1.0"):
pipe(**inputs, output=self.output_name)
class TestCosmos3DistilledModularPipelineLoading(
Cosmos3DistilledModularPipelineTesterConfig, ModularLoadingTesterMixin
):
def test_save_from_pretrained(self, tmp_path):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
loaded_pipe = ModularPipeline.from_pretrained(str(tmp_path))
loaded_pipe.load_components(torch_dtype=torch.float32)
loaded_pipe.disable_safety_checker()
loaded_pipe.to(torch_device)
base_output = base_pipe(**self.get_dummy_inputs(), output=self.output_name)
loaded_output = loaded_pipe(**self.get_dummy_inputs(), output=self.output_name)
assert torch.abs(base_output - loaded_output).max() < 1e-3
class TestCosmos3DistilledModularPipelineWorkflow(
Cosmos3DistilledModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestCosmos3DistilledModularPipelineMemory(Cosmos3DistilledModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -17,7 +17,13 @@ import pytest
from diffusers.modular_pipelines import ErnieImageAutoBlocks, ErnieImageModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
ERNIE_IMAGE_WORKFLOWS = {
@@ -32,11 +38,10 @@ ERNIE_IMAGE_WORKFLOWS = {
}
class TestErnieImageModularPipelineFast(ModularPipelineTesterMixin):
class ErnieImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = ErnieImageModularPipeline
pipeline_blocks_class = ErnieImageAutoBlocks
pretrained_model_name_or_path = "akshan-main/tiny-ernie-image-modular-pipe"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents"])
@@ -53,6 +58,20 @@ class TestErnieImageModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "pt",
}
class TestErnieImageModularPipelineFast(ErnieImageModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="PE generation is non-deterministic on CPU")
def test_float16_inference(self):
pass
class TestErnieImageModularPipelineLoading(ErnieImageModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestErnieImageModularPipelineWorkflow(ErnieImageModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestErnieImageModularPipelineMemory(ErnieImageModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -29,8 +29,14 @@ from diffusers.modular_pipelines import (
)
from ...testing_utils import floats_tensor, torch_device
from ..test_components_manager import ModularPipelineOffloadTesterMixin
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularAutoOffloadTesterMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
FLUX_TEXT2IMAGE_WORKFLOWS = {
@@ -46,7 +52,7 @@ FLUX_TEXT2IMAGE_WORKFLOWS = {
}
class TestFluxModularPipelineFast(ModularPipelineOffloadTesterMixin, ModularPipelineTesterMixin):
class FluxModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = FluxModularPipeline
pipeline_blocks_class = FluxAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular"
@@ -69,10 +75,28 @@ class TestFluxModularPipelineFast(ModularPipelineOffloadTesterMixin, ModularPipe
}
return inputs
class TestFluxModularPipelineFast(FluxModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFluxModularPipelineLoading(FluxModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFluxModularPipelineWorkflow(FluxModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFluxModularPipelineMemory(FluxModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestFluxModularPipelineAutoOffload(FluxModularPipelineTesterConfig, ModularAutoOffloadTesterMixin):
pass
FLUX_IMAGE2IMAGE_WORKFLOWS = {
"image2image": [
("text_encoder", "FluxTextEncoderStep"),
@@ -90,7 +114,7 @@ FLUX_IMAGE2IMAGE_WORKFLOWS = {
}
class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
class FluxImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = FluxModularPipeline
pipeline_blocks_class = FluxAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-modular"
@@ -129,32 +153,38 @@ class TestFluxImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
return inputs
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
pipes.append(base_pipe)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
pipes.append(pipe)
image_slices = []
for pipe in pipes:
inputs = self.get_dummy_inputs()
image = pipe(**inputs, output="images")
image_slices.append(image[0, -3:, -3:, -1].flatten())
assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
class TestFluxImg2ImgModularPipelineFast(FluxImg2ImgModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(8e-2)
class TestFluxImg2ImgModularPipelineLoading(FluxImg2ImgModularPipelineTesterConfig, ModularLoadingTesterMixin):
def test_save_from_pretrained(self, tmp_path, base_pipe_output):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path)
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
# Re-apply the `vae_scale_factor` override `get_pipeline` makes (see the comment there).
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
image = pipe(**self.get_dummy_inputs(), output=self.output_name)
expected_slice = base_pipe_output[0, -3:, -3:, -1].flatten()
image_slice = image[0, -3:, -3:, -1].flatten()
assert torch.abs(expected_slice - image_slice).max() < 1e-3
class TestFluxImg2ImgModularPipelineWorkflow(FluxImg2ImgModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFluxImg2ImgModularPipelineMemory(FluxImg2ImgModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
FLUX_KONTEXT_WORKFLOWS = {
"text2image": [
("text_encoder", "FluxTextEncoderStep"),
@@ -181,7 +211,7 @@ FLUX_KONTEXT_WORKFLOWS = {
}
class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
class FluxKontextModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = FluxKontextModularPipeline
pipeline_blocks_class = FluxKontextAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux-kontext-pipe"
@@ -210,27 +240,32 @@ class TestFluxKontextModularPipelineFast(ModularPipelineTesterMixin):
return inputs
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
pipes.append(base_pipe)
class TestFluxKontextModularPipelineFast(FluxKontextModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFluxKontextModularPipelineLoading(FluxKontextModularPipelineTesterConfig, ModularLoadingTesterMixin):
def test_save_from_pretrained(self, tmp_path, base_pipe_output):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path).to(torch_device)
pipe = ModularPipeline.from_pretrained(tmp_path)
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
pipe.image_processor = VaeImageProcessor(vae_scale_factor=2)
pipes.append(pipe)
image = pipe(**self.get_dummy_inputs(), output=self.output_name)
image_slices = []
for pipe in pipes:
inputs = self.get_dummy_inputs()
image = pipe(**inputs, output="images")
expected_slice = base_pipe_output[0, -3:, -3:, -1].flatten()
image_slice = image[0, -3:, -3:, -1].flatten()
assert torch.abs(expected_slice - image_slice).max() < 1e-3
image_slices.append(image[0, -3:, -3:, -1].flatten())
assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3
class TestFluxKontextModularPipelineWorkflow(FluxKontextModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFluxKontextModularPipelineMemory(FluxKontextModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -25,7 +25,13 @@ from diffusers.modular_pipelines import (
)
from ...testing_utils import floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
FLUX2_TEXT2IMAGE_WORKFLOWS = {
@@ -43,11 +49,10 @@ FLUX2_TEXT2IMAGE_WORKFLOWS = {
}
class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
class Flux2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2ModularPipeline
pipeline_blocks_class = Flux2AutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-modular"
params = frozenset(["prompt", "height", "width", "guidance_scale"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = FLUX2_TEXT2IMAGE_WORKFLOWS
@@ -68,10 +73,24 @@ class TestFlux2ModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestFlux2ModularPipelineFast(Flux2ModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFlux2ModularPipelineLoading(Flux2ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFlux2ModularPipelineWorkflow(Flux2ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFlux2ModularPipelineMemory(Flux2ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
FLUX2_IMAGE_CONDITIONED_WORKFLOWS = {
"image_conditioned": [
("text_encoder", "Flux2TextEncoderStep"),
@@ -90,11 +109,10 @@ FLUX2_IMAGE_CONDITIONED_WORKFLOWS = {
}
class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
class Flux2ImageConditionedModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2ModularPipeline
pipeline_blocks_class = Flux2AutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-modular"
params = frozenset(["prompt", "height", "width", "guidance_scale", "image"])
batch_params = frozenset(["prompt", "image"])
expected_workflow_blocks = FLUX2_IMAGE_CONDITIONED_WORKFLOWS
@@ -120,9 +138,31 @@ class TestFlux2ImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
return inputs
class TestFlux2ImageConditionedModularPipelineFast(
Flux2ImageConditionedModularPipelineTesterConfig, ModularPipelineTesterMixin
):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
@pytest.mark.skip(reason="batched inference is currently not supported")
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=0.0001):
return
class TestFlux2ImageConditionedModularPipelineLoading(
Flux2ImageConditionedModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestFlux2ImageConditionedModularPipelineWorkflow(
Flux2ImageConditionedModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestFlux2ImageConditionedModularPipelineMemory(
Flux2ImageConditionedModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass

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

@@ -25,7 +25,13 @@ from diffusers.modular_pipelines import (
)
from ...testing_utils import floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
FLUX2_KLEIN_WORKFLOWS = {
@@ -42,11 +48,10 @@ FLUX2_KLEIN_WORKFLOWS = {
}
class TestFlux2KleinModularPipelineFast(ModularPipelineTesterMixin):
class Flux2KleinModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2KleinModularPipeline
pipeline_blocks_class = Flux2KleinAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
not_params = frozenset(["negative_prompt"])
@@ -67,10 +72,24 @@ class TestFlux2KleinModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestFlux2KleinModularPipelineFast(Flux2KleinModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFlux2KleinModularPipelineLoading(Flux2KleinModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFlux2KleinModularPipelineWorkflow(Flux2KleinModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFlux2KleinModularPipelineMemory(Flux2KleinModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS = {
"image_conditioned": [
("text_encoder", "Flux2KleinTextEncoderStep"),
@@ -88,11 +107,10 @@ FLUX2_KLEIN_IMAGE_CONDITIONED_WORKFLOWS = {
}
class TestFlux2KleinImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
class Flux2KleinImageConditionedModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2KleinModularPipeline
pipeline_blocks_class = Flux2KleinAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-modular"
params = frozenset(["prompt", "height", "width", "image"])
batch_params = frozenset(["prompt", "image"])
not_params = frozenset(["negative_prompt"])
@@ -118,9 +136,31 @@ class TestFlux2KleinImageConditionedModularPipelineFast(ModularPipelineTesterMix
return inputs
class TestFlux2KleinImageConditionedModularPipelineFast(
Flux2KleinImageConditionedModularPipelineTesterConfig, ModularPipelineTesterMixin
):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
@pytest.mark.skip(reason="batched inference is currently not supported")
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=0.0001):
return
class TestFlux2KleinImageConditionedModularPipelineLoading(
Flux2KleinImageConditionedModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestFlux2KleinImageConditionedModularPipelineWorkflow(
Flux2KleinImageConditionedModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestFlux2KleinImageConditionedModularPipelineMemory(
Flux2KleinImageConditionedModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass

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

@@ -25,7 +25,13 @@ from diffusers.modular_pipelines import (
)
from ...testing_utils import floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
FLUX2_KLEIN_BASE_WORKFLOWS = {
@@ -42,11 +48,10 @@ FLUX2_KLEIN_BASE_WORKFLOWS = {
}
class TestFlux2KleinBaseModularPipelineFast(ModularPipelineTesterMixin):
class Flux2KleinBaseModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2KleinBaseModularPipeline
pipeline_blocks_class = Flux2KleinBaseAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = FLUX2_KLEIN_BASE_WORKFLOWS
@@ -66,10 +71,24 @@ class TestFlux2KleinBaseModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestFlux2KleinBaseModularPipelineFast(Flux2KleinBaseModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestFlux2KleinBaseModularPipelineLoading(Flux2KleinBaseModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestFlux2KleinBaseModularPipelineWorkflow(Flux2KleinBaseModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestFlux2KleinBaseModularPipelineMemory(Flux2KleinBaseModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS = {
"image_conditioned": [
("text_encoder", "Flux2KleinBaseTextEncoderStep"),
@@ -87,11 +106,10 @@ FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS = {
}
class TestFlux2KleinBaseImageConditionedModularPipelineFast(ModularPipelineTesterMixin):
class Flux2KleinBaseImageConditionedModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Flux2KleinBaseModularPipeline
pipeline_blocks_class = Flux2KleinBaseAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-flux2-klein-base-modular"
params = frozenset(["prompt", "height", "width", "image"])
batch_params = frozenset(["prompt", "image"])
expected_workflow_blocks = FLUX2_KLEIN_BASE_IMAGE_CONDITIONED_WORKFLOWS
@@ -116,9 +134,31 @@ class TestFlux2KleinBaseImageConditionedModularPipelineFast(ModularPipelineTeste
return inputs
class TestFlux2KleinBaseImageConditionedModularPipelineFast(
Flux2KleinBaseImageConditionedModularPipelineTesterConfig, ModularPipelineTesterMixin
):
def test_float16_inference(self):
super().test_float16_inference(9e-2)
@pytest.mark.skip(reason="batched inference is currently not supported")
def test_inference_batch_single_identical(self, batch_size=2, expected_max_diff=0.0001):
return
class TestFlux2KleinBaseImageConditionedModularPipelineLoading(
Flux2KleinBaseImageConditionedModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestFlux2KleinBaseImageConditionedModularPipelineWorkflow(
Flux2KleinBaseImageConditionedModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestFlux2KleinBaseImageConditionedModularPipelineMemory(
Flux2KleinBaseImageConditionedModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass

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

@@ -22,7 +22,13 @@ from diffusers.modular_pipelines import (
HeliosPyramidModularPipeline,
)
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
HELIOS_WORKFLOWS = {
@@ -61,11 +67,10 @@ HELIOS_WORKFLOWS = {
}
class TestHeliosModularPipelineFast(ModularPipelineTesterMixin):
class HeliosModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = HeliosModularPipeline
pipeline_blocks_class = HeliosAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-helios-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
@@ -86,11 +91,25 @@ class TestHeliosModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestHeliosModularPipelineFast(HeliosModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
class TestHeliosModularPipelineLoading(HeliosModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestHeliosModularPipelineWorkflow(HeliosModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestHeliosModularPipelineMemory(HeliosModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
HELIOS_PYRAMID_WORKFLOWS = {
"text2video": [
("text_encoder", "HeliosTextEncoderStep"),
@@ -124,11 +143,10 @@ HELIOS_PYRAMID_WORKFLOWS = {
}
class TestHeliosPyramidModularPipelineFast(ModularPipelineTesterMixin):
class HeliosPyramidModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = HeliosPyramidModularPipeline
pipeline_blocks_class = HeliosPyramidAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-helios-pyramid-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["pyramid_num_inference_steps_list", "num_videos_per_prompt", "latents"])
@@ -149,18 +167,28 @@ class TestHeliosPyramidModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestHeliosPyramidModularPipelineFast(HeliosPyramidModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
# Pyramid pipeline injects noise at each stage, so batch vs single can differ more
super().test_inference_batch_single_identical(expected_max_diff=5e-1)
@pytest.mark.skip(reason="Pyramid multi-stage noise makes offload comparison unreliable with tiny models")
def test_components_auto_cpu_offload_inference_consistent(self):
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
class TestHeliosPyramidModularPipelineLoading(HeliosPyramidModularPipelineTesterConfig, ModularLoadingTesterMixin):
@pytest.mark.skip(reason="Pyramid multi-stage noise makes save/load comparison unreliable with tiny models")
def test_save_from_pretrained(self):
pass
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
class TestHeliosPyramidModularPipelineWorkflow(HeliosPyramidModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestHeliosPyramidModularPipelineMemory(HeliosPyramidModularPipelineTesterConfig, ModularMemoryTesterMixin):
@pytest.mark.skip(reason="Pyramid multi-stage noise makes offload comparison unreliable with tiny models")
def test_components_auto_cpu_offload_inference_consistent(self):
pass

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

@@ -17,7 +17,13 @@ import pytest
from diffusers.modular_pipelines import HunyuanVideo15AutoBlocks, HunyuanVideo15ModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
HUNYUANVIDEO15_WORKFLOWS = {
@@ -43,11 +49,10 @@ HUNYUANVIDEO15_WORKFLOWS = {
}
class TestHunyuanVideo15ModularPipelineFast(ModularPipelineTesterMixin):
class HunyuanVideo15ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = HunyuanVideo15ModularPipeline
pipeline_blocks_class = HunyuanVideo15AutoBlocks
pretrained_model_name_or_path = "akshan-main/tiny-hunyuanvideo1_5-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
@@ -67,6 +72,8 @@ class TestHunyuanVideo15ModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestHunyuanVideo15ModularPipelineFast(HunyuanVideo15ModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
@@ -81,3 +88,15 @@ class TestHunyuanVideo15ModularPipelineFast(ModularPipelineTesterMixin):
def test_float16_inference(self):
super().test_float16_inference(expected_max_diff=0.1)
class TestHunyuanVideo15ModularPipelineLoading(HunyuanVideo15ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestHunyuanVideo15ModularPipelineWorkflow(HunyuanVideo15ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestHunyuanVideo15ModularPipelineMemory(HunyuanVideo15ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -16,7 +16,13 @@
from diffusers.modular_pipelines import Krea2AutoBlocks, Krea2ModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
KREA2_WORKFLOWS = {
@@ -32,11 +38,10 @@ KREA2_WORKFLOWS = {
}
class TestKrea2ModularPipelineFast(ModularPipelineTesterMixin):
class Krea2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Krea2ModularPipeline
pipeline_blocks_class = Krea2AutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-modular-pipe"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = KREA2_WORKFLOWS
@@ -54,5 +59,19 @@ class TestKrea2ModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestKrea2ModularPipelineFast(Krea2ModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-3)
class TestKrea2ModularPipelineLoading(Krea2ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestKrea2ModularPipelineWorkflow(Krea2ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestKrea2ModularPipelineMemory(Krea2ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -16,7 +16,13 @@
from diffusers.modular_pipelines import Krea2TurboAutoBlocks, Krea2TurboModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
KREA2_TURBO_WORKFLOWS = {
@@ -32,11 +38,10 @@ KREA2_TURBO_WORKFLOWS = {
}
class TestKrea2TurboModularPipelineFast(ModularPipelineTesterMixin):
class Krea2TurboModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = Krea2TurboModularPipeline
pipeline_blocks_class = Krea2TurboAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-krea2-turbo-modular-pipe"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = KREA2_TURBO_WORKFLOWS
@@ -54,5 +59,19 @@ class TestKrea2TurboModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestKrea2TurboModularPipelineFast(Krea2TurboModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-3)
class TestKrea2TurboModularPipelineLoading(Krea2TurboModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestKrea2TurboModularPipelineWorkflow(Krea2TurboModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestKrea2TurboModularPipelineMemory(Krea2TurboModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -17,7 +17,13 @@ import pytest
from diffusers.modular_pipelines import LTXAutoBlocks, LTXModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
LTX_WORKFLOWS = {
@@ -42,11 +48,10 @@ LTX_WORKFLOWS = {
}
class TestLTXModularPipelineFast(ModularPipelineTesterMixin):
class LTXModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = LTXModularPipeline
pipeline_blocks_class = LTXAutoBlocks
pretrained_model_name_or_path = "akshan-main/tiny-ltx-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
@@ -67,6 +72,20 @@ class TestLTXModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestLTXModularPipelineFast(LTXModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
class TestLTXModularPipelineLoading(LTXModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTXModularPipelineWorkflow(LTXModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestLTXModularPipelineMemory(LTXModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -22,19 +22,88 @@ from diffusers.modular_pipelines import LTX2AutoBlocks, LTX2ModularPipeline
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
LTX2_REPO_ID = "hf-internal-testing/tiny-ltx2-modular-pipe"
LTX2_WORKFLOWS = {
"text2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2DenoiseStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"image2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("vae_encoder", "LTX2VaeEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_i2v_latents", "LTX2Image2VideoPrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2Image2VideoDenoiseStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"condition": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2ConditionPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"in_context": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("reference_encoder", "LTX2ReferenceEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2InContextPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
class LTX2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""Shared configuration for every LTX2 workflow; a variant config adds its own `params` and dummy inputs."""
class LTX2ModularTests(ModularPipelineTesterMixin):
pipeline_class = LTX2ModularPipeline
pipeline_blocks_class = LTX2AutoBlocks
pretrained_model_name_or_path = LTX2_REPO_ID
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
expected_workflow_blocks = LTX2_WORKFLOWS
output_name = "videos"
def get_dummy_inputs(self, seed=0):
@@ -51,6 +120,10 @@ class LTX2ModularTests(ModularPipelineTesterMixin):
"output_type": "pt",
}
class LTX2ModularPipelineFastTesterMixin(ModularPipelineTesterMixin):
"""`ModularPipelineTesterMixin` with the two adjustments every LTX2 workflow needs."""
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
@@ -59,24 +132,13 @@ class LTX2ModularTests(ModularPipelineTesterMixin):
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
class TestLTX2ModularPipelineFast(LTX2ModularTests):
class LTX2Text2VideoModularPipelineTesterConfig(LTX2ModularPipelineTesterConfig):
params = frozenset(["prompt", "height", "width", "num_frames"])
expected_workflow_blocks = {
"text2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2DenoiseStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
class TestLTX2Text2VideoModularPipelineFast(
LTX2Text2VideoModularPipelineTesterConfig, LTX2ModularPipelineFastTesterMixin
):
def test_audio_output(self):
pipe = self.get_pipeline().to("cpu")
@@ -103,25 +165,21 @@ class TestLTX2ModularPipelineFast(LTX2ModularTests):
assert 0 < num_frames <= round(2.0 * inputs["frame_rate"])
class TestLTX2ModularImage2VideoPipelineFast(LTX2ModularTests):
class TestLTX2Text2VideoModularPipelineLoading(LTX2Text2VideoModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX2Text2VideoModularPipelineMemory(LTX2Text2VideoModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
# The four workflows share `LTX2AutoBlocks` and the same repo, so one workflow test class covers all of them.
class TestLTX2ModularPipelineWorkflow(LTX2Text2VideoModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class LTX2Image2VideoModularPipelineTesterConfig(LTX2ModularPipelineTesterConfig):
params = frozenset(["prompt", "image", "height", "width", "num_frames"])
expected_workflow_blocks = {
"image2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("vae_encoder", "LTX2VaeEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_i2v_latents", "LTX2Image2VideoPrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2Image2VideoDenoiseStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
@@ -133,25 +191,22 @@ class TestLTX2ModularImage2VideoPipelineFast(LTX2ModularTests):
return inputs
class TestLTX2ModularConditionPipelineFast(LTX2ModularTests):
class TestLTX2Image2VideoModularPipelineFast(
LTX2Image2VideoModularPipelineTesterConfig, LTX2ModularPipelineFastTesterMixin
):
pass
class TestLTX2Image2VideoModularPipelineLoading(LTX2Image2VideoModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX2Image2VideoModularPipelineMemory(LTX2Image2VideoModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class LTX2ConditionModularPipelineTesterConfig(LTX2ModularPipelineTesterConfig):
params = frozenset(["prompt", "conditions", "height", "width", "num_frames"])
expected_workflow_blocks = {
"condition": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2ConditionPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
@@ -161,28 +216,39 @@ class TestLTX2ModularConditionPipelineFast(LTX2ModularTests):
return inputs
class TestLTX2ModularInContextPipelineFast(LTX2ModularTests):
class TestLTX2ConditionModularPipelineFast(
LTX2ConditionModularPipelineTesterConfig, LTX2ModularPipelineFastTesterMixin
):
pass
class TestLTX2ConditionModularPipelineLoading(LTX2ConditionModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX2ConditionModularPipelineMemory(LTX2ConditionModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class LTX2InContextModularPipelineTesterConfig(LTX2ModularPipelineTesterConfig):
params = frozenset(["prompt", "reference_conditions", "height", "width", "num_frames"])
expected_workflow_blocks = {
"in_context": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("reference_encoder", "LTX2ReferenceEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2InContextPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2VaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
video = torch.rand((1, 5, 3, 32, 32), generator=torch.Generator("cpu").manual_seed(seed))
inputs["reference_conditions"] = LTX2ReferenceCondition(frames=video, strength=1.0)
return inputs
class TestLTX2InContextModularPipelineFast(
LTX2InContextModularPipelineTesterConfig, LTX2ModularPipelineFastTesterMixin
):
pass
class TestLTX2InContextModularPipelineLoading(LTX2InContextModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX2InContextModularPipelineMemory(LTX2InContextModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -22,19 +22,88 @@ from diffusers.modular_pipelines import LTX2AutoBlocks, LTX25AutoBlocks, LTX25Mo
from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition
from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
LTX25_REPO_ID = "hf-internal-testing/tiny-ltx2-5-modular-pipe"
LTX25_WORKFLOWS = {
"text2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2DenoiseStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"image2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("vae_encoder", "LTX2VaeEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_i2v_latents", "LTX2Image2VideoPrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2Image2VideoDenoiseStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"condition": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2ConditionPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
"in_context": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("reference_encoder", "LTX2ReferenceEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2InContextPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
class LTX25ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""Shared configuration for every LTX-2.5 workflow; a variant config adds its own `params` and dummy inputs."""
class LTX25ModularTests(ModularPipelineTesterMixin):
pipeline_class = LTX25ModularPipeline
pipeline_blocks_class = LTX25AutoBlocks
pretrained_model_name_or_path = LTX25_REPO_ID
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
expected_workflow_blocks = LTX25_WORKFLOWS
output_name = "videos"
def get_dummy_inputs(self, seed=0):
@@ -51,6 +120,10 @@ class LTX25ModularTests(ModularPipelineTesterMixin):
"output_type": "pt",
}
class LTX25ModularPipelineFastTesterMixin(ModularPipelineTesterMixin):
"""`ModularPipelineTesterMixin` with the two adjustments every LTX-2.5 workflow needs."""
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
@@ -59,24 +132,13 @@ class LTX25ModularTests(ModularPipelineTesterMixin):
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
class TestLTX25ModularPipelineFast(LTX25ModularTests):
class LTX25Text2VideoModularPipelineTesterConfig(LTX25ModularPipelineTesterConfig):
params = frozenset(["prompt", "height", "width", "num_frames"])
expected_workflow_blocks = {
"text2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2DenoiseStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
class TestLTX25Text2VideoModularPipelineFast(
LTX25Text2VideoModularPipelineTesterConfig, LTX25ModularPipelineFastTesterMixin
):
def test_diffusion_decoder_output(self):
pipe = self.get_pipeline().to("cpu")
@@ -121,25 +183,21 @@ class TestLTX25ModularPipelineFast(LTX25ModularTests):
assert 0 < num_frames <= round(2.0 * inputs["frame_rate"])
class TestLTX25ModularImage2VideoPipelineFast(LTX25ModularTests):
class TestLTX25Text2VideoModularPipelineLoading(LTX25Text2VideoModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX25Text2VideoModularPipelineMemory(LTX25Text2VideoModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
# The four workflows share `LTX25AutoBlocks` and the same repo, so one workflow test class covers all of them.
class TestLTX25ModularPipelineWorkflow(LTX25Text2VideoModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class LTX25Image2VideoModularPipelineTesterConfig(LTX25ModularPipelineTesterConfig):
params = frozenset(["prompt", "image", "height", "width", "num_frames"])
expected_workflow_blocks = {
"image2video": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("vae_encoder", "LTX2VaeEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.set_timesteps", "LTX2SetTimestepsStep"),
("denoise.prepare_latents", "LTX2PrepareLatentsStep"),
("denoise.prepare_i2v_latents", "LTX2Image2VideoPrepareLatentsStep"),
("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2PrepareCoordsStep"),
("denoise.denoise", "LTX2Image2VideoDenoiseStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
@@ -151,25 +209,24 @@ class TestLTX25ModularImage2VideoPipelineFast(LTX25ModularTests):
return inputs
class TestLTX25ModularConditionPipelineFast(LTX25ModularTests):
class TestLTX25Image2VideoModularPipelineFast(
LTX25Image2VideoModularPipelineTesterConfig, LTX25ModularPipelineFastTesterMixin
):
pass
class TestLTX25Image2VideoModularPipelineLoading(
LTX25Image2VideoModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestLTX25Image2VideoModularPipelineMemory(LTX25Image2VideoModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class LTX25ConditionModularPipelineTesterConfig(LTX25ModularPipelineTesterConfig):
params = frozenset(["prompt", "conditions", "height", "width", "num_frames"])
expected_workflow_blocks = {
"condition": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("duration", "LTX2DurationStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2ConditionPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
@@ -179,28 +236,39 @@ class TestLTX25ModularConditionPipelineFast(LTX25ModularTests):
return inputs
class TestLTX25ModularInContextPipelineFast(LTX25ModularTests):
class TestLTX25ConditionModularPipelineFast(
LTX25ConditionModularPipelineTesterConfig, LTX25ModularPipelineFastTesterMixin
):
pass
class TestLTX25ConditionModularPipelineLoading(LTX25ConditionModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX25ConditionModularPipelineMemory(LTX25ConditionModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class LTX25InContextModularPipelineTesterConfig(LTX25ModularPipelineTesterConfig):
params = frozenset(["prompt", "reference_conditions", "height", "width", "num_frames"])
expected_workflow_blocks = {
"in_context": [
("text_encoder.text_encoder", "LTX2TextEncoderStep"),
("text_encoder.connectors", "LTX2TextConnectorStep"),
("condition_encoder", "LTX2ConditionEncoderStep"),
("reference_encoder", "LTX2ReferenceEncoderStep"),
("denoise.input", "LTX2TextInputStep"),
("denoise.prepare_latents", "LTX2InContextPrepareLatentsStep"),
("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"),
("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"),
("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"),
("denoise.denoise", "LTX2ConditionDenoiseStep"),
("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"),
("decode.video_decode", "LTX2DiffusionVaeDecoderStep"),
("decode.audio_decode", "LTX2AudioDecoderStep"),
],
}
def get_dummy_inputs(self, seed=0):
inputs = super().get_dummy_inputs(seed)
video = torch.rand((1, 5, 3, 32, 32), generator=torch.Generator("cpu").manual_seed(seed))
inputs["reference_conditions"] = LTX2ReferenceCondition(frames=video, strength=1.0)
return inputs
class TestLTX25InContextModularPipelineFast(
LTX25InContextModularPipelineTesterConfig, LTX25ModularPipelineFastTesterMixin
):
pass
class TestLTX25InContextModularPipelineLoading(LTX25InContextModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestLTX25InContextModularPipelineMemory(LTX25InContextModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -36,7 +36,13 @@ from diffusers.modular_pipelines.minimax_h3.encoders import (
)
from diffusers.modular_pipelines.minimax_h3.modular_pipeline import MINIMAX_H3_FPS
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
# The blocks every workflow of [`MiniMaxH3Blocks`] runs, in order. A keyframe adds the canvas block and the one
@@ -147,13 +153,12 @@ MINIMAX_H3_REF2VA_WORKFLOW_DEFAULTS = {
}
class TestMiniMaxH3ModularPipelineFast(ModularPipelineTesterMixin):
class MiniMaxH3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""The `t2va` and `fl2va` requests of [`MiniMaxH3Blocks`]: a prompt, optionally with keyframes."""
pipeline_class = MiniMaxH3ModularPipeline
pipeline_blocks_class = MiniMaxH3Blocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-h3-modular-pipe"
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()
@@ -164,6 +169,22 @@ class TestMiniMaxH3ModularPipelineFast(ModularPipelineTesterMixin):
expected_workflow_blocks = MINIMAX_H3_WORKFLOWS
expected_workflow_defaults = MINIMAX_H3_WORKFLOW_DEFAULTS
def get_dummy_inputs(self, seed=0):
return {
"prompt": "a robot dancing",
"generator": self.get_generator(seed),
"num_inference_steps": 2,
# MiniMax-H3 generates 5 to 15 seconds at a fixed 24 fps, so 124 frames (`17 * 7 + 5`, the next length
# the video VAE can decode) is the shortest admissible request; a 32 pixel canvas is one `(1, 2, 2)`
# patch row per latent frame, which is what makes it affordable on CPU.
"height": 32,
"width": 32,
"num_frames": 124,
"output_type": "pt",
}
class TestMiniMaxH3ModularPipelineFast(MiniMaxH3ModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="MiniMax-H3 packs one request into one sequence, so a batch of prompts is not a thing.")
def test_inference_batch_consistent(self):
pass
@@ -189,20 +210,6 @@ class TestMiniMaxH3ModularPipelineFast(ModularPipelineTesterMixin):
with pytest.raises(ValueError, match="rounded up to 362"):
pipe(**inputs)
def get_dummy_inputs(self, seed=0):
return {
"prompt": "a robot dancing",
"generator": self.get_generator(seed),
"num_inference_steps": 2,
# MiniMax-H3 generates 5 to 15 seconds at a fixed 24 fps, so 124 frames (`17 * 7 + 5`, the next length
# the video VAE can decode) is the shortest admissible request; a 32 pixel canvas is one `(1, 2, 2)`
# patch row per latent frame, which is what makes it affordable on CPU.
"height": 32,
"width": 32,
"num_frames": 124,
"output_type": "pt",
}
def test_video_and_audio_outputs(self):
r"""One call denoises both modalities out of the one packed sequence, and returns both."""
pipe = self.get_pipeline()
@@ -396,13 +403,24 @@ class TestMiniMaxH3ModularPipelineFast(ModularPipelineTesterMixin):
pipe(**inputs)
class TestMiniMaxH3Ref2VAModularPipelineFast(ModularPipelineTesterMixin):
class TestMiniMaxH3ModularPipelineLoading(MiniMaxH3ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestMiniMaxH3ModularPipelineWorkflow(MiniMaxH3ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestMiniMaxH3ModularPipelineMemory(MiniMaxH3ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class MiniMaxH3Ref2VAModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""The `ref2va` requests of [`MiniMaxH3Blocks`]: a prompt and an ordered list of references."""
pipeline_class = MiniMaxH3ModularPipeline
pipeline_blocks_class = MiniMaxH3Blocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-h3-modular-pipe"
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()
@@ -431,6 +449,23 @@ class TestMiniMaxH3Ref2VAModularPipelineFast(ModularPipelineTesterMixin):
pipeline.set_progress_bar_config(disable=None)
return pipeline
def get_dummy_inputs(self, seed=0):
return {
"prompt": "a robot dancing",
"references": [MiniMaxH3ImageReference(image=Image.new("RGB", (48, 80)))],
"generator": self.get_generator(seed),
"num_inference_steps": 2,
# MiniMax-H3 generates 5 to 15 seconds at a fixed 24 fps, so 124 frames (`17 * 7 + 5`, the next length
# the video VAE can decode) is the shortest admissible request; a 32 pixel canvas is one `(1, 2, 2)`
# patch row per latent frame, which is what makes it affordable on CPU.
"height": 32,
"width": 32,
"num_frames": 124,
"output_type": "pt",
}
class TestMiniMaxH3Ref2VAModularPipelineFast(MiniMaxH3Ref2VAModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="MiniMax-H3 packs one request into one sequence, so a batch of prompts is not a thing.")
def test_inference_batch_consistent(self):
pass
@@ -459,21 +494,6 @@ class TestMiniMaxH3Ref2VAModularPipelineFast(ModularPipelineTesterMixin):
with pytest.raises(ValueError, match="rounded up to 362"):
pipe(**inputs)
def get_dummy_inputs(self, seed=0):
return {
"prompt": "a robot dancing",
"references": [MiniMaxH3ImageReference(image=Image.new("RGB", (48, 80)))],
"generator": self.get_generator(seed),
"num_inference_steps": 2,
# MiniMax-H3 generates 5 to 15 seconds at a fixed 24 fps, so 124 frames (`17 * 7 + 5`, the next length
# the video VAE can decode) is the shortest admissible request; a 32 pixel canvas is one `(1, 2, 2)`
# patch row per latent frame, which is what makes it affordable on CPU.
"height": 32,
"width": 32,
"num_frames": 124,
"output_type": "pt",
}
def test_video_and_audio_outputs(self):
r"""A reference conditions the request without binding the generated geometry."""
pipe = self.get_pipeline()
@@ -742,6 +762,20 @@ class TestMiniMaxH3Ref2VAModularPipelineFast(ModularPipelineTesterMixin):
pipe(**inputs)
class TestMiniMaxH3Ref2VAModularPipelineLoading(MiniMaxH3Ref2VAModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestMiniMaxH3Ref2VAModularPipelineWorkflow(
MiniMaxH3Ref2VAModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestMiniMaxH3Ref2VAModularPipelineMemory(MiniMaxH3Ref2VAModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestMiniMaxH3Reference:
"""
The three reference dataclasses and the passes that normalize them, neither of which needs a checkpoint.

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

@@ -21,15 +21,30 @@ from diffusers import (
MiniMaxMusic3Blocks,
MiniMaxMusic3ConditionEncoder,
MiniMaxMusic3ModularPipeline,
ModularPipeline,
)
from ...testing_utils import enable_full_determinism, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
enable_full_determinism()
_SAMPLING_SKIP = (
"MiniMax Music 3's autoregressive stage samples discrete audio tokens; any numeric perturbation "
"(dtype, offload-induced kernel changes) legitimately flips sampled tokens, so value-equality "
"comparisons cannot hold. Shape and determinism-at-fixed-seed are covered by the other tests."
)
_UNBATCHED_SKIP = "MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings."
class MiniMaxMusic3ConditionEncoderFastTests(unittest.TestCase):
def test_condition_encoder_output_shape(self):
condition_encoder = MiniMaxMusic3ConditionEncoder(
@@ -50,41 +65,7 @@ class MiniMaxMusic3ConditionEncoderFastTests(unittest.TestCase):
self.assertEqual(output.shape, (1, expected_length, 16))
class TestMiniMaxMusic3ModularPipelineFast(ModularPipelineTesterMixin):
_SAMPLING_SKIP = (
"MiniMax Music 3's autoregressive stage samples discrete audio tokens; any numeric perturbation "
"(dtype, offload-induced kernel changes) legitimately flips sampled tokens, so value-equality "
"comparisons cannot hold. Shape and determinism-at-fixed-seed are covered by the other tests."
)
def test_float16_inference(self):
pytest.skip(self._SAMPLING_SKIP)
def test_components_auto_cpu_offload_inference_consistent(self):
pytest.skip(self._SAMPLING_SKIP)
def test_unload_components_auto_cpu_offload(self):
pytest.skip(self._SAMPLING_SKIP)
def test_num_images_per_prompt(self):
pytest.skip("The pipeline generates a single waveform per call; there is no num_images_per_prompt input.")
def test_save_from_pretrained(self, tmp_path):
# the common implementation indexes 4-D image outputs; compare the audio waveform directly
pipe = self.get_pipeline().to(torch_device)
pipe.save_pretrained(str(tmp_path))
from diffusers import ModularPipeline
reloaded = ModularPipeline.from_pretrained(str(tmp_path))
reloaded.load_components()
reloaded.to(torch_device)
inputs = self.get_dummy_inputs()
audio = pipe(**inputs, output="audios")
inputs = self.get_dummy_inputs()
audio_reloaded = reloaded(**inputs, output="audios")
assert audio.shape == audio_reloaded.shape
assert torch.allclose(audio, audio_reloaded, atol=1e-4)
class MiniMaxMusic3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = MiniMaxMusic3ModularPipeline
pipeline_blocks_class = MiniMaxMusic3Blocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-minimax-music3-modular-pipe"
@@ -126,17 +107,59 @@ class TestMiniMaxMusic3ModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "pt",
}
def test_inference_batch_consistent(self):
pytest.skip("MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings.")
class TestMiniMaxMusic3ModularPipelineFast(MiniMaxMusic3ModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason=_SAMPLING_SKIP)
def test_float16_inference(self):
pass
@pytest.mark.skip(reason="The pipeline generates a single waveform per call; there is no num_images_per_prompt.")
def test_num_images_per_prompt(self):
pass
@pytest.mark.skip(reason=_UNBATCHED_SKIP)
def test_inference_batch_consistent(self):
pass
@pytest.mark.skip(reason=_UNBATCHED_SKIP)
def test_inference_batch_single_identical(self):
pytest.skip("MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings.")
pass
def test_output_is_stereo_waveform(self):
pipe = self.get_pipeline()
audio = pipe(**self.get_dummy_inputs(), output="audios")
audio = pipe(**self.get_dummy_inputs(), output=self.output_name)
assert audio.shape[0] == 1
assert audio.shape[1] == 2
assert audio.abs().max() <= 1.0
class TestMiniMaxMusic3ModularPipelineLoading(MiniMaxMusic3ModularPipelineTesterConfig, ModularLoadingTesterMixin):
def test_save_from_pretrained(self, tmp_path, base_pipe_output):
# the common implementation indexes 4-D image outputs; compare the audio waveform directly
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(str(tmp_path))
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
audio = pipe(**self.get_dummy_inputs(), output=self.output_name)
assert audio.shape == base_pipe_output.shape
assert torch.allclose(audio, base_pipe_output, atol=1e-4)
class TestMiniMaxMusic3ModularPipelineWorkflow(MiniMaxMusic3ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestMiniMaxMusic3ModularPipelineMemory(MiniMaxMusic3ModularPipelineTesterConfig, ModularMemoryTesterMixin):
@pytest.mark.skip(reason=_SAMPLING_SKIP)
def test_components_auto_cpu_offload_inference_consistent(self):
pass
@pytest.mark.skip(reason=_SAMPLING_SKIP)
def test_unload_components_auto_cpu_offload(self):
pass

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

@@ -27,7 +27,14 @@ from diffusers.modular_pipelines import (
)
from ...testing_utils import torch_device
from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularGuiderTesterMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS = {
@@ -127,11 +134,10 @@ QWEN_IMAGE_TEXT2IMAGE_WORKFLOWS = {
}
class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
class QwenImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = QwenImageModularPipeline
pipeline_blocks_class = QwenImageAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-modular"
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
@@ -150,10 +156,28 @@ class TestQwenImageModularPipelineFast(ModularPipelineTesterMixin, ModularGuider
}
return inputs
class TestQwenImageModularPipelineFast(QwenImageModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-4)
class TestQwenImageModularPipelineLoading(QwenImageModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestQwenImageModularPipelineWorkflow(QwenImageModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestQwenImageModularPipelineMemory(QwenImageModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestQwenImageModularPipelineGuider(QwenImageModularPipelineTesterConfig, ModularGuiderTesterMixin):
pass
QWEN_IMAGE_EDIT_WORKFLOWS = {
"image_conditioned": [
("text_encoder.resize", "QwenImageEditResizeStep"),
@@ -192,11 +216,10 @@ QWEN_IMAGE_EDIT_WORKFLOWS = {
}
class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
class QwenImageEditModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = QwenImageEditModularPipeline
pipeline_blocks_class = QwenImageEditAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-edit-modular"
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
@@ -215,15 +238,36 @@ class TestQwenImageEditModularPipelineFast(ModularPipelineTesterMixin, ModularGu
inputs["image"] = PIL.Image.new("RGB", (32, 32), 0)
return inputs
class TestQwenImageEditModularPipelineFast(QwenImageEditModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_num_images_per_prompt(self):
# `QwenImageEditResizeStep` blows the conditioning image up to a ~1024x1024 area before the VL encoder
# regardless of how tiny the test image is, so every forward here is expensive — the default four
# combinations run past the 60s CI timeout. One combination with both multipliers > 1 pins the same thing.
super().test_num_images_per_prompt(batch_sizes=[2], num_images_per_prompts=[2])
class TestQwenImageEditModularPipelineLoading(QwenImageEditModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestQwenImageEditModularPipelineWorkflow(QwenImageEditModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestQwenImageEditModularPipelineMemory(QwenImageEditModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestQwenImageEditModularPipelineGuider(QwenImageEditModularPipelineTesterConfig, ModularGuiderTesterMixin):
def test_guider_cfg(self):
super().test_guider_cfg(7e-5)
class TestQwenImageEditPlusModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
class QwenImageEditPlusModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = QwenImageEditPlusModularPipeline
pipeline_blocks_class = QwenImageEditPlusAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-qwenimage-edit-plus-modular"
# No `mask_image` yet.
params = frozenset(["prompt", "height", "width", "negative_prompt", "attention_kwargs", "image"])
batch_params = frozenset(["prompt", "negative_prompt", "image"])
@@ -242,6 +286,10 @@ class TestQwenImageEditPlusModularPipelineFast(ModularPipelineTesterMixin, Modul
inputs["image"] = PIL.Image.new("RGB", (32, 32), 0)
return inputs
class TestQwenImageEditPlusModularPipelineFast(
QwenImageEditPlusModularPipelineTesterConfig, ModularPipelineTesterMixin
):
def test_multi_images_as_input(self):
inputs = self.get_dummy_inputs()
image = inputs.pop("image")
@@ -264,5 +312,27 @@ class TestQwenImageEditPlusModularPipelineFast(ModularPipelineTesterMixin, Modul
def test_inference_batch_single_identical():
super().test_inference_batch_single_identical()
class TestQwenImageEditPlusModularPipelineLoading(
QwenImageEditPlusModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestQwenImageEditPlusModularPipelineWorkflow(
QwenImageEditPlusModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestQwenImageEditPlusModularPipelineMemory(
QwenImageEditPlusModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass
class TestQwenImageEditPlusModularPipelineGuider(
QwenImageEditPlusModularPipelineTesterConfig, ModularGuiderTesterMixin
):
def test_guider_cfg(self):
super().test_guider_cfg(1e-6)

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

@@ -26,7 +26,13 @@ from diffusers.modular_pipelines.stable_diffusion_3 import (
)
from ...testing_utils import floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
SD3_TEXT2IMAGE_WORKFLOWS = {
@@ -41,20 +47,14 @@ SD3_TEXT2IMAGE_WORKFLOWS = {
}
class TestStableDiffusion3ModularPipelineFast(ModularPipelineTesterMixin):
class StableDiffusion3ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = StableDiffusion3ModularPipeline
pipeline_blocks_class = StableDiffusion3AutoBlocks
pretrained_model_name_or_path = "AlanPonnachan/tiny-sd3-modular"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = SD3_TEXT2IMAGE_WORKFLOWS
def test_pipeline_call_signature(self):
# Override to prevent signature check failure for guider configurations
# (guidance_scale) which are intentionally omitted from pipeline inputs.
pass
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
return {
@@ -67,9 +67,20 @@ class TestStableDiffusion3ModularPipelineFast(ModularPipelineTesterMixin):
"output_type": "pt",
}
def get_pipeline(self, components_manager=None, dtype=torch.float32):
return super().get_pipeline(components_manager, dtype)
class TestStableDiffusion3ModularPipelineFast(StableDiffusion3ModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_pipeline_call_signature(self):
# Override to prevent signature check failure for guider configurations
# (guidance_scale) which are intentionally omitted from pipeline inputs.
pass
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestStableDiffusion3ModularPipelineLoading(
StableDiffusion3ModularPipelineTesterConfig, ModularLoadingTesterMixin
):
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
@@ -98,8 +109,15 @@ class TestStableDiffusion3ModularPipelineFast(ModularPipelineTesterMixin):
assert set(base_pipe.components.keys()) == set(pipe.components.keys())
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestStableDiffusion3ModularPipelineWorkflow(
StableDiffusion3ModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestStableDiffusion3ModularPipelineMemory(StableDiffusion3ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
SD3_IMAGE2IMAGE_WORKFLOWS = {
@@ -121,20 +139,14 @@ SD3_IMAGE2IMAGE_WORKFLOWS = {
}
class TestStableDiffusion3Img2ImgModularPipelineFast(ModularPipelineTesterMixin):
class StableDiffusion3Img2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = StableDiffusion3ModularPipeline
pipeline_blocks_class = StableDiffusion3AutoBlocks
pretrained_model_name_or_path = "AlanPonnachan/tiny-sd3-modular"
params = frozenset(["prompt", "height", "width", "image"])
batch_params = frozenset(["prompt", "image"])
expected_workflow_blocks = SD3_IMAGE2IMAGE_WORKFLOWS
def test_pipeline_call_signature(self):
# Override to prevent signature check failure for guider configurations
# (guidance_scale) which are intentionally omitted from pipeline inputs.
pass
def get_pipeline(self, components_manager=None, dtype=torch.float32):
pipeline = super().get_pipeline(components_manager, dtype)
pipeline.image_processor = VaeImageProcessor(vae_scale_factor=8)
@@ -158,6 +170,22 @@ class TestStableDiffusion3Img2ImgModularPipelineFast(ModularPipelineTesterMixin)
inputs["strength"] = 0.5
return inputs
class TestStableDiffusion3Img2ImgModularPipelineFast(
StableDiffusion3Img2ImgModularPipelineTesterConfig, ModularPipelineTesterMixin
):
def test_pipeline_call_signature(self):
# Override to prevent signature check failure for guider configurations
# (guidance_scale) which are intentionally omitted from pipeline inputs.
pass
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestStableDiffusion3Img2ImgModularPipelineLoading(
StableDiffusion3Img2ImgModularPipelineTesterConfig, ModularLoadingTesterMixin
):
def test_save_from_pretrained(self, tmp_path):
pipes = []
base_pipe = self.get_pipeline().to(torch_device)
@@ -187,5 +215,14 @@ class TestStableDiffusion3Img2ImgModularPipelineFast(ModularPipelineTesterMixin)
assert set(base_pipe.components.keys()) == set(pipe.components.keys())
def test_float16_inference(self):
super().test_float16_inference(9e-2)
class TestStableDiffusion3Img2ImgModularPipelineWorkflow(
StableDiffusion3Img2ImgModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestStableDiffusion3Img2ImgModularPipelineMemory(
StableDiffusion3Img2ImgModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass

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

@@ -24,14 +24,22 @@ from diffusers import ClassifierFreeGuidance, StableDiffusionXLAutoBlocks, Stabl
from diffusers.loaders import ModularIPAdapterMixin
from ...models.unets.test_models_unet_2d_condition import create_ip_adapter_state_dict
from ...testing_utils import enable_full_determinism, floats_tensor, torch_device
from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin
from ...testing_utils import enable_full_determinism, floats_tensor, is_ip_adapter, torch_device
from ..testing_utils import (
BaseModularPipelineOutputMixin,
BaseModularPipelineTesterConfig,
ModularGuiderTesterMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
enable_full_determinism()
class SDXLModularTesterMixin:
class SDXLModularTesterMixin(BaseModularPipelineOutputMixin):
"""
This mixin defines method to create pipeline, base input and base test across all SDXL modular tests.
"""
@@ -48,6 +56,7 @@ class SDXLModularTesterMixin:
assert max_diff < expected_max_diff, f"Image slice does not match expected slice. Max Difference: {max_diff}"
@is_ip_adapter
class SDXLModularIPAdapterTesterMixin:
"""
This mixin is designed to test IP Adapter.
@@ -183,7 +192,7 @@ class SDXLModularIPAdapterTesterMixin:
)
class SDXLModularControlNetTesterMixin:
class SDXLModularControlNetTesterMixin(BaseModularPipelineOutputMixin):
"""
This mixin is designed to test ControlNet.
"""
@@ -321,14 +330,8 @@ TEXT2IMAGE_WORKFLOWS = {
}
class TestSDXLModularPipelineFast(
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
SDXLModularControlNetTesterMixin,
ModularGuiderTesterMixin,
ModularPipelineTesterMixin,
):
"""Test cases for Stable Diffusion XL modular pipeline fast tests."""
class SDXLModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""Shared configuration for the Stable Diffusion XL text-to-image modular pipeline tests."""
pipeline_class = StableDiffusionXLModularPipeline
pipeline_blocks_class = StableDiffusionXLAutoBlocks
@@ -357,6 +360,8 @@ class TestSDXLModularPipelineFast(
}
return inputs
class TestSDXLModularPipelineFast(SDXLModularPipelineTesterConfig, SDXLModularTesterMixin, ModularPipelineTesterMixin):
def test_stable_diffusion_xl_euler(self):
self._test_stable_diffusion_xl_euler(
expected_image_shape=self.expected_image_output_shape,
@@ -370,6 +375,30 @@ class TestSDXLModularPipelineFast(
super().test_inference_batch_single_identical(expected_max_diff=3e-3)
class TestSDXLModularPipelineIPAdapter(SDXLModularPipelineTesterConfig, SDXLModularIPAdapterTesterMixin):
pass
class TestSDXLModularPipelineControlNet(SDXLModularPipelineTesterConfig, SDXLModularControlNetTesterMixin):
pass
class TestSDXLModularPipelineGuider(SDXLModularPipelineTesterConfig, ModularGuiderTesterMixin):
pass
class TestSDXLModularPipelineLoading(SDXLModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestSDXLModularPipelineWorkflow(SDXLModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestSDXLModularPipelineMemory(SDXLModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
IMAGE2IMAGE_WORKFLOWS = {
"image2image": [
("text_encoder", "StableDiffusionXLTextEncoderStep"),
@@ -429,14 +458,8 @@ IMAGE2IMAGE_WORKFLOWS = {
}
class TestSDXLImg2ImgModularPipelineFast(
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
SDXLModularControlNetTesterMixin,
ModularGuiderTesterMixin,
ModularPipelineTesterMixin,
):
"""Test cases for Stable Diffusion XL image-to-image modular pipeline fast tests."""
class SDXLImg2ImgModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
"""Shared configuration for the Stable Diffusion XL image-to-image modular pipeline tests."""
pipeline_class = StableDiffusionXLModularPipeline
pipeline_blocks_class = StableDiffusionXLAutoBlocks
@@ -472,6 +495,10 @@ class TestSDXLImg2ImgModularPipelineFast(
return inputs
class TestSDXLImg2ImgModularPipelineFast(
SDXLImg2ImgModularPipelineTesterConfig, SDXLModularTesterMixin, ModularPipelineTesterMixin
):
def test_stable_diffusion_xl_euler(self):
self._test_stable_diffusion_xl_euler(
expected_image_shape=self.expected_image_output_shape,
@@ -483,6 +510,32 @@ class TestSDXLImg2ImgModularPipelineFast(
super().test_inference_batch_single_identical(expected_max_diff=3e-3)
class TestSDXLImg2ImgModularPipelineIPAdapter(SDXLImg2ImgModularPipelineTesterConfig, SDXLModularIPAdapterTesterMixin):
pass
class TestSDXLImg2ImgModularPipelineControlNet(
SDXLImg2ImgModularPipelineTesterConfig, SDXLModularControlNetTesterMixin
):
pass
class TestSDXLImg2ImgModularPipelineGuider(SDXLImg2ImgModularPipelineTesterConfig, ModularGuiderTesterMixin):
pass
class TestSDXLImg2ImgModularPipelineLoading(SDXLImg2ImgModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestSDXLImg2ImgModularPipelineWorkflow(SDXLImg2ImgModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestSDXLImg2ImgModularPipelineMemory(SDXLImg2ImgModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
INPAINTING_WORKFLOWS = {
"inpainting": [
("text_encoder", "StableDiffusionXLTextEncoderStep"),
@@ -543,6 +596,7 @@ INPAINTING_WORKFLOWS = {
class SDXLInpaintingModularPipelineFastTests(
BaseModularPipelineTesterConfig,
SDXLModularTesterMixin,
SDXLModularIPAdapterTesterMixin,
SDXLModularControlNetTesterMixin,

165
tests/modular_pipelines/test_components_manager.py
View File

@@ -15,7 +15,6 @@
import gc
from unittest import mock
import pytest
import torch
from diffusers import ComponentsManager
@@ -23,6 +22,7 @@ from diffusers.models import ModelMixin
from diffusers.utils import is_accelerate_available
from ..testing_utils import backend_empty_cache, require_accelerate, require_accelerator, torch_device
from .testing_utils import patch_free_memory
if is_accelerate_available():
@@ -68,20 +68,10 @@ def _patch_cuda_mem_get_info(free_bytes: int, total_bytes: int = 80 * UNIT):
return mock.patch.object(torch.cuda, "mem_get_info", return_value=(free_bytes, total_bytes))
def _patch_free_memory(free_bytes: int, total_bytes: int = 80 * UNIT):
# Integration tests run on the real `torch_device`; patch `mem_get_info` on
# whichever backend module (cuda/xpu/...) actually backs it. `mem_get_info` returns
# `(free, total)` and is the single point where the strategy learns how much memory
# is available, so patching it simulates arbitrary memory pressure.
device_type = torch.device(torch_device).type
device_module = getattr(torch, device_type, torch.cuda)
return mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes))
@require_accelerate
class ComponentsManagerTesterMixin:
class TestComponentsManager:
"""
Common tests for `ComponentsManager` and its auto-offload strategy.
Tests for `ComponentsManager` and its auto-offload strategy.
"""
# A `cuda:0` device descriptor is enough to drive the strategy's device-type and
@@ -89,8 +79,8 @@ class ComponentsManagerTesterMixin:
strategy_execution_device = torch.device("cuda:0")
def setup_method(self):
# Mirror `ModularPipelineTesterMixin` cleanup so this mixin stays interchangeable
# in the MRO when stacked into a pipeline test class.
# Free VRAM before/after each test; the auto-offload integration tests below reason
# about device residency, so they must not inherit another test's allocations.
torch.compiler.reset()
gc.collect()
backend_empty_cache(torch_device)
@@ -273,13 +263,13 @@ class ComponentsManagerTesterMixin:
# Ample free memory: running m1 just moves it onto the device, evicting
# nothing (m2 is not resident, so it is not even a candidate).
with _patch_free_memory(70 * UNIT):
with patch_free_memory(70 * UNIT):
m1(x)
assert next(m1.parameters()).device.type == device_type
# Memory pressure: usable = 4 - 1 = 3 but m2 needs 4, so the only resident
# model (m1) must be evicted back to the CPU to make room for m2.
with _patch_free_memory(4 * UNIT):
with patch_free_memory(4 * UNIT):
m2(x)
assert next(m2.parameters()).device.type == device_type
assert next(m1.parameters()).device.type == "cpu"
@@ -297,7 +287,7 @@ class ComponentsManagerTesterMixin:
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=UNIT)
try:
x = torch.randn(2, 4, device=torch_device)
with _patch_free_memory(70 * UNIT):
with patch_free_memory(70 * UNIT):
m1(x)
m2(x)
# Both fit comfortably, so neither gets evicted.
@@ -305,142 +295,3 @@ class ComponentsManagerTesterMixin:
assert next(m2.parameters()).device.type == device_type
finally:
cm.disable_auto_cpu_offload()
class TestComponentsManager(ComponentsManagerTesterMixin):
pass
# More free memory than any tiny test checkpoint could ever need, so the strategy never
# decides to offload. Used to assert the *negative*: no eviction without memory pressure.
_AMPLE_FREE_BYTES = 1024**4
class ModularPipelineOffloadTesterMixin:
"""
Auto-CPU-offload tests for a modular pipeline's components.
"""
@staticmethod
def _managed_models(cm):
"""The registered components that the offloader actually manages (parameterized
`nn.Module`s)."""
models = []
for component in cm.components.values():
if isinstance(component, torch.nn.Module) and next(component.parameters(), None) is not None:
models.append(component)
return models
@staticmethod
def _is_resident(model):
return next(model.parameters()).device.type == torch.device(torch_device).type
def _run_offloaded(self, free_bytes):
"""
Run the pipeline with auto offload on and `free_bytes` of *simulated* device
memory, recording every offload decision the strategy makes.
Each record is `{"incoming", "resident_before", "offloaded"}` (lists of model
ids), captured by spying on `AutoOffloadStrategy.__call__`, which the hooks call
each time a model is about to be moved onto the device.
"""
cm = ComponentsManager()
pipe = self.get_pipeline(components_manager=cm)
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=0)
records = []
original_call = AutoOffloadStrategy.__call__
def spy_call(strategy, hooks, model_id, model, execution_device):
selected = original_call(
strategy, hooks=hooks, model_id=model_id, model=model, execution_device=execution_device
)
records.append(
{
"incoming": model_id,
"resident_before": [hook.model_id for hook in hooks],
"offloaded": [hook.model_id for hook in selected],
}
)
return selected
with _patch_free_memory(free_bytes), mock.patch.object(AutoOffloadStrategy, "__call__", spy_call):
output = pipe(**self.get_dummy_inputs(), output=self.output_name)
return cm, records, output
@staticmethod
def _peak_co_residency(records):
"""
Largest number of models simultaneously on the device, reconstructed from the
strategy's view of residency just before each load.
"""
peak = 0
for record in records:
resident = (set(record["resident_before"]) - set(record["offloaded"])) | {record["incoming"]}
peak = max(peak, len(resident))
return peak
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_serializes_models_under_memory_pressure(self):
# Zero simulated free memory: every model that runs must first evict whatever is
# currently resident (comfy-style serialized execution).
cm, records, _ = self._run_offloaded(free_bytes=0)
try:
distinct_models = {record["incoming"] for record in records}
if len(distinct_models) < 2:
pytest.skip("pipeline has fewer than two offloadable model components")
# Offloading actually fired (at least one eviction happened).
assert any(record["offloaded"] for record in records), "expected at least one eviction"
# Sequencing: models run one at a time, never two co-resident on the device.
peak = self._peak_co_residency(records)
assert peak == 1, f"expected serialized execution under pressure, saw {peak} models co-resident"
# Device placement after the run: at most the last-run model stays on the
# accelerator, and at least one managed model was pushed back to the CPU.
models = self._managed_models(cm)
resident = [m for m in models if self._is_resident(m)]
assert len(resident) <= 1
assert any(not self._is_resident(m) for m in models), "expected some model offloaded to CPU"
finally:
cm.disable_auto_cpu_offload()
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self):
# Negative case: with ample simulated memory the strategy is still consulted on
# every load, but it must never decide to evict anything.
cm, records, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES)
try:
distinct_models = {record["incoming"] for record in records}
if len(distinct_models) < 2:
pytest.skip("pipeline has fewer than two offloadable model components")
# Nothing was ever offloaded...
assert all(record["offloaded"] == [] for record in records), "no model should be evicted"
# ...and models accumulate on the device instead of being serialized.
peak = self._peak_co_residency(records)
assert peak >= 2, f"expected models to co-reside without pressure, saw peak {peak}"
models = self._managed_models(cm)
assert sum(self._is_resident(m) for m in models) >= 2, "expected multiple models resident on device"
finally:
cm.disable_auto_cpu_offload()
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_inference_consistent_under_memory_pressure(self, expected_max_diff=1e-3):
# Sensible results: forcing offload (zero simulated free memory) must not change
# the output relative to an ordinary, non-offloaded run.
base_pipe = self.get_pipeline().to(torch_device)
baseline = base_pipe(**self.get_dummy_inputs(), output=self.output_name)
cm, _, offloaded = self._run_offloaded(free_bytes=0)
try:
max_diff = torch.abs(baseline - offloaded).max()
assert max_diff < expected_max_diff, f"offloaded output diverged from baseline (max diff {max_diff})"
finally:
cm.disable_auto_cpu_offload()

232
tests/modular_pipelines/test_modular_model_card.py Normal file
View File

@@ -0,0 +1,232 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from diffusers.modular_pipelines.modular_pipeline_utils import (
ComponentSpec,
ConfigSpec,
InputParam,
OutputParam,
generate_modular_model_card_content,
)
class TestModularModelCardContent:
def create_mock_block(self, name="TestBlock", description="Test block description"):
class MockBlock:
def __init__(self, name, description):
self.__class__.__name__ = name
self.description = description
self.sub_blocks = {}
return MockBlock(name, description)
def create_mock_blocks(
self,
class_name="TestBlocks",
description="Test pipeline description",
num_blocks=2,
components=None,
configs=None,
inputs=None,
outputs=None,
trigger_inputs=None,
model_name=None,
):
class MockBlocks:
def __init__(self):
self.__class__.__name__ = class_name
self.description = description
self.sub_blocks = {}
self.expected_components = components or []
self.expected_configs = configs or []
self.inputs = inputs or []
self.outputs = outputs or []
self.trigger_inputs = trigger_inputs
self.model_name = model_name
blocks = MockBlocks()
# Add mock sub-blocks
for i in range(num_blocks):
block_name = f"block_{i}"
blocks.sub_blocks[block_name] = self.create_mock_block(f"Block{i}", f"Description for block {i}")
return blocks
def test_basic_model_card_content_structure(self):
"""Test that all expected keys are present in the output."""
blocks = self.create_mock_blocks()
content = generate_modular_model_card_content(blocks)
expected_keys = [
"pipeline_name",
"model_description",
"blocks_description",
"components_description",
"configs_section",
"io_specification_section",
"trigger_inputs_section",
"tags",
]
for key in expected_keys:
assert key in content, f"Expected key '{key}' not found in model card content"
assert isinstance(content["tags"], list), "Tags should be a list"
def test_pipeline_name_generation(self):
"""Test that pipeline name is correctly generated from blocks class name."""
blocks = self.create_mock_blocks(class_name="StableDiffusionBlocks")
content = generate_modular_model_card_content(blocks)
assert content["pipeline_name"] == "StableDiffusion Pipeline"
def test_tags_generation_text_to_image(self):
"""Test that text-to-image tags are correctly generated."""
blocks = self.create_mock_blocks(trigger_inputs=None)
content = generate_modular_model_card_content(blocks)
assert "modular-diffusers" in content["tags"]
assert "diffusers" in content["tags"]
assert "text-to-image" in content["tags"]
def test_tags_generation_with_trigger_inputs(self):
"""Test that tags are correctly generated based on trigger inputs."""
# Test inpainting
blocks = self.create_mock_blocks(trigger_inputs=["mask", "prompt"])
content = generate_modular_model_card_content(blocks)
assert "inpainting" in content["tags"]
# Test image-to-image
blocks = self.create_mock_blocks(trigger_inputs=["image", "prompt"])
content = generate_modular_model_card_content(blocks)
assert "image-to-image" in content["tags"]
# Test controlnet
blocks = self.create_mock_blocks(trigger_inputs=["control_image", "prompt"])
content = generate_modular_model_card_content(blocks)
assert "controlnet" in content["tags"]
def test_tags_with_model_name(self):
"""Test that model name is included in tags when present."""
blocks = self.create_mock_blocks(model_name="stable-diffusion-xl")
content = generate_modular_model_card_content(blocks)
assert "stable-diffusion-xl" in content["tags"]
def test_components_description_formatting(self):
"""Test that components are correctly formatted."""
components = [
ComponentSpec(name="vae", description="VAE component"),
ComponentSpec(name="text_encoder", description="Text encoder component"),
]
blocks = self.create_mock_blocks(components=components)
content = generate_modular_model_card_content(blocks)
assert "vae" in content["components_description"]
assert "text_encoder" in content["components_description"]
# Should be enumerated
assert "1." in content["components_description"]
def test_components_description_empty(self):
"""Test handling of pipelines without components."""
blocks = self.create_mock_blocks(components=None)
content = generate_modular_model_card_content(blocks)
assert "No specific components required" in content["components_description"]
def test_configs_section_with_configs(self):
"""Test that configs section is generated when configs are present."""
configs = [
ConfigSpec(name="num_train_timesteps", default=1000, description="Number of training timesteps"),
]
blocks = self.create_mock_blocks(configs=configs)
content = generate_modular_model_card_content(blocks)
assert "## Configuration Parameters" in content["configs_section"]
def test_configs_section_empty(self):
"""Test that configs section is empty when no configs are present."""
blocks = self.create_mock_blocks(configs=None)
content = generate_modular_model_card_content(blocks)
assert content["configs_section"] == ""
def test_inputs_description_required_and_optional(self):
"""Test that required and optional inputs are correctly formatted."""
inputs = [
InputParam(name="prompt", type_hint=str, required=True, description="The input prompt"),
InputParam(name="num_steps", type_hint=int, required=False, default=50, description="Number of steps"),
]
blocks = self.create_mock_blocks(inputs=inputs)
content = generate_modular_model_card_content(blocks)
io_section = content["io_specification_section"]
assert "**Inputs:**" in io_section
assert "prompt" in io_section
assert "num_steps" in io_section
assert "*optional*" in io_section
assert "defaults to `50`" in io_section
def test_inputs_description_empty(self):
"""Test handling of pipelines without specific inputs."""
blocks = self.create_mock_blocks(inputs=[])
content = generate_modular_model_card_content(blocks)
assert "No specific inputs defined" in content["io_specification_section"]
def test_outputs_description_formatting(self):
"""Test that outputs are correctly formatted."""
outputs = [
OutputParam(name="images", type_hint=torch.Tensor, description="Generated images"),
]
blocks = self.create_mock_blocks(outputs=outputs)
content = generate_modular_model_card_content(blocks)
io_section = content["io_specification_section"]
assert "images" in io_section
assert "Generated images" in io_section
def test_outputs_description_empty(self):
"""Test handling of pipelines without specific outputs."""
blocks = self.create_mock_blocks(outputs=[])
content = generate_modular_model_card_content(blocks)
assert "Standard pipeline outputs" in content["io_specification_section"]
def test_trigger_inputs_section_with_triggers(self):
"""Test that trigger inputs section is generated when present."""
blocks = self.create_mock_blocks(trigger_inputs=["mask", "image"])
content = generate_modular_model_card_content(blocks)
assert "### Conditional Execution" in content["trigger_inputs_section"]
assert "`mask`" in content["trigger_inputs_section"]
assert "`image`" in content["trigger_inputs_section"]
def test_trigger_inputs_section_empty(self):
"""Test that trigger inputs section is empty when not present."""
blocks = self.create_mock_blocks(trigger_inputs=None)
content = generate_modular_model_card_content(blocks)
assert content["trigger_inputs_section"] == ""
def test_model_description_includes_block_count(self):
"""Test that model description includes the number of blocks."""
blocks = self.create_mock_blocks(num_blocks=5)
content = generate_modular_model_card_content(blocks)
assert "5-block architecture" in content["model_description"]

241
tests/modular_pipelines/test_modular_pipeline_loading.py Normal file
View File

@@ -0,0 +1,241 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
import torch
from diffusers import AutoModel, ControlNetModel, ModularPipeline, UNet2DConditionModel
from diffusers.modular_pipelines.modular_pipeline_utils import ComponentSpec
class TestAutoModelLoadIdTagging:
def test_automodel_tags_load_id(self):
model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
assert hasattr(model, "_diffusers_load_id"), "Model should have _diffusers_load_id attribute"
assert model._diffusers_load_id != "null", "_diffusers_load_id should not be 'null'"
# Verify load_id contains the expected fields
load_id = model._diffusers_load_id
assert "hf-internal-testing/tiny-stable-diffusion-xl-pipe" in load_id
assert "unet" in load_id
def test_automodel_update_components(self):
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
auto_model = AutoModel.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet")
pipe.update_components(unet=auto_model)
assert pipe.unet is auto_model
assert "unet" in pipe._component_specs
spec = pipe._component_specs["unet"]
assert spec.pretrained_model_name_or_path == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
assert spec.subfolder == "unet"
def test_load_components_loads_local_single_file_path(self, tmp_path):
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
model = ControlNetModel.from_pretrained("hf-internal-testing/tiny-controlnet")
model.save_pretrained(tmp_path)
local_ckpt_path = str(tmp_path / "diffusion_pytorch_model.safetensors")
pipe._component_specs["controlnet"] = ComponentSpec(
name="controlnet",
type_hint=ControlNetModel,
pretrained_model_name_or_path=local_ckpt_path,
)
pipe.load_components(names="controlnet", config=str(tmp_path))
assert pipe.controlnet is not None
assert isinstance(pipe.controlnet, ControlNetModel)
assert pipe._component_specs["controlnet"].pretrained_model_name_or_path == local_ckpt_path
assert getattr(pipe.controlnet, "_diffusers_load_id", None) not in (None, "null")
class TestLoadComponentsSkipBehavior:
def test_load_components_skips_already_loaded(self):
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
original_unet = pipe.unet
pipe.load_components()
# Verify that the unet is the same object (not reloaded)
assert pipe.unet is original_unet, "load_components should skip already loaded components"
def test_load_components_selective_loading(self):
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(names="unet", dtype=torch.float32)
# Verify only requested component was loaded.
assert hasattr(pipe, "unet")
assert pipe.unet is not None
assert getattr(pipe, "vae", None) is None
def test_load_components_selective_loading_incremental(self):
"""Loading a subset of components should not affect already-loaded components."""
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(names="unet", dtype=torch.float32)
pipe.load_components(names="text_encoder", dtype=torch.float32)
assert hasattr(pipe, "unet")
assert pipe.unet is not None
assert hasattr(pipe, "text_encoder")
assert pipe.text_encoder is not None
def test_load_components_skips_invalid_pretrained_path(self):
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe._component_specs["test_component"] = ComponentSpec(
name="test_component",
type_hint=torch.nn.Module,
pretrained_model_name_or_path=None,
default_creation_method="from_pretrained",
)
pipe.load_components(dtype=torch.float32)
# Verify test_component was not loaded
assert not hasattr(pipe, "test_component") or pipe.test_component is None
class TestCustomModelSavePretrained:
def test_save_pretrained_updates_index_for_local_model(self, tmp_path):
"""When a component without _diffusers_load_id (custom/local model) is saved,
modular_model_index.json should point to the save directory."""
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
pipe.unet._diffusers_load_id = "null"
save_dir = str(tmp_path / "my-pipeline")
pipe.save_pretrained(save_dir)
with open(os.path.join(save_dir, "modular_model_index.json")) as f:
index = json.load(f)
_library, _cls, unet_spec = index["unet"]
assert unet_spec["pretrained_model_name_or_path"] == save_dir
assert unet_spec["subfolder"] == "unet"
_library, _cls, vae_spec = index["vae"]
assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
def test_save_pretrained_roundtrip_with_local_model(self, tmp_path):
"""A pipeline with a custom/local model should be saveable and re-loadable with identical outputs."""
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
pipe.unet._diffusers_load_id = "null"
original_state_dict = pipe.unet.state_dict()
save_dir = str(tmp_path / "my-pipeline")
pipe.save_pretrained(save_dir)
loaded_pipe = ModularPipeline.from_pretrained(save_dir)
loaded_pipe.load_components(dtype=torch.float32)
assert loaded_pipe.unet is not None
assert loaded_pipe.unet.__class__.__name__ == pipe.unet.__class__.__name__
loaded_state_dict = loaded_pipe.unet.state_dict()
assert set(original_state_dict.keys()) == set(loaded_state_dict.keys())
for key in original_state_dict:
assert torch.equal(original_state_dict[key], loaded_state_dict[key]), f"Mismatch in {key}"
def test_save_pretrained_updates_index_for_model_with_no_load_id(self, tmp_path):
"""testing the workflow of update the pipeline with a custom model and save the pipeline,
the modular_model_index.json should point to the save directory."""
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
unet = UNet2DConditionModel.from_pretrained(
"hf-internal-testing/tiny-stable-diffusion-xl-pipe", subfolder="unet"
)
assert not hasattr(unet, "_diffusers_load_id")
pipe.update_components(unet=unet)
save_dir = str(tmp_path / "my-pipeline")
pipe.save_pretrained(save_dir)
with open(os.path.join(save_dir, "modular_model_index.json")) as f:
index = json.load(f)
_library, _cls, unet_spec = index["unet"]
assert unet_spec["pretrained_model_name_or_path"] == save_dir
assert unet_spec["subfolder"] == "unet"
_library, _cls, vae_spec = index["vae"]
assert vae_spec["pretrained_model_name_or_path"] == "hf-internal-testing/tiny-stable-diffusion-xl-pipe"
def test_save_pretrained_overwrite_modular_index(self, tmp_path):
"""With overwrite_modular_index=True, all component references should point to the save directory."""
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
pipe.load_components(dtype=torch.float32)
save_dir = str(tmp_path / "my-pipeline")
pipe.save_pretrained(save_dir, overwrite_modular_index=True)
with open(os.path.join(save_dir, "modular_model_index.json")) as f:
index = json.load(f)
for component_name in ["unet", "vae", "text_encoder", "text_encoder_2"]:
if component_name not in index:
continue
_library, _cls, spec = index[component_name]
assert spec["pretrained_model_name_or_path"] == save_dir, (
f"{component_name} should point to save dir but got {spec['pretrained_model_name_or_path']}"
)
assert spec["subfolder"] == component_name
loaded_pipe = ModularPipeline.from_pretrained(save_dir)
loaded_pipe.load_components(dtype=torch.float32)
assert loaded_pipe.unet is not None
assert loaded_pipe.vae is not None
class TestModularPipelineInitFallback:
"""Test that ModularPipeline.__init__ falls back to default_blocks_name when
_blocks_class_name is a base class (e.g. SequentialPipelineBlocks saved by from_blocks_dict)."""
def test_init_fallback_when_blocks_class_name_is_base_class(self, tmp_path):
# 1. Load pipeline and get a workflow (returns a base SequentialPipelineBlocks)
pipe = ModularPipeline.from_pretrained("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
t2i_blocks = pipe.blocks.get_workflow("text2image")
assert t2i_blocks.__class__.__name__ == "SequentialPipelineBlocks"
# 2. Use init_pipeline to create a new pipeline from the workflow blocks
t2i_pipe = t2i_blocks.init_pipeline("hf-internal-testing/tiny-stable-diffusion-xl-pipe")
# 3. Save and reload — the saved config will have _blocks_class_name="SequentialPipelineBlocks"
save_dir = str(tmp_path / "pipeline")
t2i_pipe.save_pretrained(save_dir)
loaded_pipe = ModularPipeline.from_pretrained(save_dir)
# 4. Verify it fell back to default_blocks_name and has correct blocks
assert loaded_pipe.__class__.__name__ == pipe.__class__.__name__
assert loaded_pipe._blocks.__class__.__name__ == pipe._blocks.__class__.__name__
assert len(loaded_pipe._blocks.sub_blocks) == len(pipe._blocks.sub_blocks)

1215
tests/modular_pipelines/test_modular_pipelines_common.py
View File

File diff suppressed because it is too large Load Diff

32
tests/modular_pipelines/testing_utils/__init__.py Normal file
View File

@@ -0,0 +1,32 @@
from .common import (
BaseModularPipelineOutputMixin,
BaseModularPipelineTesterConfig,
ModularPipelineTesterMixin,
)
from .guider import ModularGuiderTesterMixin
from .loading import ModularLoadingTesterMixin
from .memory import (
ModularAutoOffloadTesterMixin,
ModularGroupOffloadTesterMixin,
ModularMemoryTesterMixin,
ModularOffloadTesterMixin,
)
from .utils import backend_memory_allocated, get_specified_components, patch_free_memory
from .workflow import ModularWorkflowTesterMixin
__all__ = [
"BaseModularPipelineOutputMixin",
"BaseModularPipelineTesterConfig",
"ModularAutoOffloadTesterMixin",
"ModularGroupOffloadTesterMixin",
"ModularGuiderTesterMixin",
"ModularLoadingTesterMixin",
"ModularMemoryTesterMixin",
"ModularOffloadTesterMixin",
"ModularPipelineTesterMixin",
"ModularWorkflowTesterMixin",
"backend_memory_allocated",
"get_specified_components",
"patch_free_memory",
]

363
tests/modular_pipelines/testing_utils/common.py Normal file
View File

@@ -0,0 +1,363 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import gc
from typing import Callable
import pytest
import torch
import diffusers
from diffusers import ModularPipeline, ModularPipelineBlocks
from diffusers.utils import logging
from ...testing_utils import (
backend_empty_cache,
numpy_cosine_similarity_distance,
require_accelerator,
torch_device,
)
class BaseModularPipelineTesterConfig:
"""
Base class defining the configuration interface for modular pipeline testing.
A concrete config must set `pipeline_blocks_class` and `pretrained_model_name_or_path` and implement
`get_dummy_inputs()`; `params` and `batch_params` declare which inputs the blocks are expected to accept and which
of them are batched. This class only declares the testing contract; the pipeline builder and the cached reference
output live on `BaseModularPipelineOutputMixin` (mirroring the non-modular `BasePipelineTesterConfig`).
"""
# Canonical parameters that are passed to `__call__` regardless of the type of pipeline. They are always
# optional and have common sense default values.
optional_params = frozenset(["num_inference_steps", "num_images_per_prompt", "latents", "output_type"])
# Parameters the pipeline deliberately does NOT accept — e.g. `negative_prompt` on a
# guidance-distilled pipeline. `test_pipeline_call_signature` asserts they are absent,
# so accidentally (re)introducing one fails the test.
not_params = frozenset()
# this is modular specific: generator needs to be a intermediate input because it's mutable
intermediate_params = frozenset(["generator"])
# 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"
# ==================== Required interface ====================
@property
def pipeline_class(self) -> Callable | ModularPipeline:
raise NotImplementedError(
"You need to set the attribute `pipeline_class = ClassNameOfPipeline` in the child test class. "
"See existing pipeline tests for reference."
)
@property
def pretrained_model_name_or_path(self) -> str:
raise NotImplementedError(
"You need to set the attribute `pretrained_model_name_or_path` in the child test class. See existing pipeline tests for reference."
)
@property
def pipeline_blocks_class(self) -> Callable | ModularPipelineBlocks:
raise NotImplementedError(
"You need to set the attribute `pipeline_blocks_class = ClassNameOfPipelineBlocks` in the child test class. "
"See existing pipeline tests for reference."
)
def get_dummy_inputs(self, seed=0):
raise NotImplementedError(
"You need to implement `get_dummy_inputs(self, device, seed)` in the child test class. "
"See existing pipeline tests for reference."
)
@property
def params(self) -> frozenset:
raise NotImplementedError(
"You need to set the attribute `params` in the child test class. "
"`params` are checked for if all values are present in `__call__`'s signature."
" You can set `params` using one of the common set of parameters defined in `pipeline_params.py`"
" e.g., `TEXT_TO_IMAGE_PARAMS` defines the common parameters used in text to "
"image pipelines, including prompts and prompt embedding overrides."
"If your pipeline's set of arguments has minor changes from one of the common sets of arguments, "
"do not make modifications to the existing common sets of arguments. I.e. a text to image pipeline "
"with non-configurable height and width arguments should set the attribute as "
"`params = TEXT_TO_IMAGE_PARAMS - {'height', 'width'}`. "
"See existing pipeline tests for reference."
)
@property
def batch_params(self) -> frozenset:
raise NotImplementedError(
"You need to set the attribute `batch_params` in the child test class. "
"`batch_params` are the parameters required to be batched when passed to the pipeline's "
"`__call__` method. `pipeline_params.py` provides some common sets of parameters such as "
"`TEXT_TO_IMAGE_BATCH_PARAMS`, `IMAGE_VARIATION_BATCH_PARAMS`, etc... If your pipeline's "
"set of batch arguments has minor changes from one of the common sets of batch arguments, "
"do not make modifications to the existing common sets of batch arguments. I.e. a text to "
"image pipeline `negative_prompt` is not batched should set the attribute as "
"`batch_params = TEXT_TO_IMAGE_BATCH_PARAMS - {'negative_prompt'}`. "
"See existing pipeline tests for reference."
)
@property
def expected_workflow_blocks(self) -> dict:
raise NotImplementedError(
"You need to set the attribute `expected_workflow_blocks` in the child test class. "
"`expected_workflow_blocks` is a dictionary that maps workflow names to list of block names. "
"See existing pipeline tests for reference."
)
# ==================== Shared helpers ====================
def get_generator(self, seed=0):
# Always build the generator on CPU: a CPU generator works with a pipeline placed on any device (the tensor
# is created on CPU and moved), whereas an accelerator generator cannot seed a CPU tensor, which the tests
# that run the pipeline on CPU rely on.
return torch.Generator("cpu").manual_seed(seed)
# ==================== Fixtures ====================
@pytest.fixture(autouse=True)
def cleanup(self):
"""Free VRAM before/after each test (replaces unittest setUp/tearDown)."""
torch.compiler.reset()
gc.collect()
backend_empty_cache(torch_device)
yield
torch.compiler.reset()
gc.collect()
backend_empty_cache(torch_device)
class BaseModularPipelineOutputMixin:
"""Provides the `get_pipeline` builder and the class-scoped `base_pipe_output` fixture shared across tester
mixins.
Kept separate from `BaseModularPipelineTesterConfig` — which only declares the testing contract and performs no
computation — so any mixin that needs to build a pipeline or read the cached reference output
(`ModularPipelineTesterMixin`, the loading and memory mixins, ...) can inherit it without duplicating the
build-and-forward.
"""
def get_pipeline(self, components_manager=None, dtype=torch.float32):
"""Build the pipeline under test from `pipeline_blocks_class` and load its components.
The pipeline is left wherever `load_components` put it — callers that need it elsewhere should chain
`.to(torch_device)`.
"""
pipeline = self.pipeline_blocks_class().init_pipeline(
self.pretrained_model_name_or_path, components_manager=components_manager
)
pipeline.load_components(dtype=dtype)
pipeline.set_progress_bar_config(disable=None)
return pipeline
@pytest.fixture(scope="class")
def base_pipe_output(self):
"""Output of a freshly built pipeline on the standard dummy inputs, computed once per test class."""
pipe = self.get_pipeline().to(torch_device)
return pipe(**self.get_dummy_inputs(), output=self.output_name)
class ModularPipelineTesterMixin(BaseModularPipelineOutputMixin):
"""
Common inference tests for each modular pipeline: call signature, batching, dtype/device handling and NaN-free
outputs.
Designed to be composed with `BaseModularPipelineTesterConfig` (which provides `pipeline_blocks_class`,
`pretrained_model_name_or_path`, `get_dummy_inputs()` and the shared fixtures).
"""
def test_pipeline_call_signature(self):
pipe = self.get_pipeline()
input_parameters = pipe.blocks.input_names
optional_parameters = pipe.default_call_parameters
def _check_for_parameters(parameters, expected_parameters, param_type):
remaining_parameters = {param for param in parameters if param not in expected_parameters}
assert len(remaining_parameters) == 0, (
f"Required {param_type} parameters not present: {remaining_parameters}"
)
_check_for_parameters(self.params, input_parameters, "input")
_check_for_parameters(self.optional_params, optional_parameters, "optional")
unsupported_parameters = {param for param in self.not_params if param in input_parameters}
assert len(unsupported_parameters) == 0, (
f"Parameters declared in `not_params` unexpectedly present in the pipeline inputs: {unsupported_parameters}"
)
def test_inference_batch_consistent(self, batch_sizes=[2], batch_generator=True):
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs()
inputs["generator"] = self.get_generator(0)
logger = logging.get_logger(pipe.__module__)
logger.setLevel(level=diffusers.logging.FATAL)
# prepare batched inputs
batched_inputs = []
for batch_size in batch_sizes:
batched_input = {}
batched_input.update(inputs)
for name in self.batch_params:
if name not in inputs:
continue
value = inputs[name]
batched_input[name] = batch_size * [value]
if batch_generator and "generator" in inputs:
batched_input["generator"] = [self.get_generator(i) for i in range(batch_size)]
if "batch_size" in inputs:
batched_input["batch_size"] = batch_size
batched_inputs.append(batched_input)
logger.setLevel(level=diffusers.logging.WARNING)
for batch_size, batched_input in zip(batch_sizes, batched_inputs):
output = pipe(**batched_input, output=self.output_name)
assert len(output) == batch_size, "Output is different from expected batch size"
def test_inference_batch_single_identical(
self,
batch_size=2,
expected_max_diff=1e-4,
):
pipe = self.get_pipeline().to(torch_device)
inputs = self.get_dummy_inputs()
# Reset generator in case it is has been used in self.get_dummy_inputs
inputs["generator"] = self.get_generator(0)
logger = logging.get_logger(pipe.__module__)
logger.setLevel(level=diffusers.logging.FATAL)
# batchify inputs
batched_inputs = {}
batched_inputs.update(inputs)
for name in self.batch_params:
if name not in inputs:
continue
value = inputs[name]
batched_inputs[name] = batch_size * [value]
if "generator" in inputs:
batched_inputs["generator"] = [self.get_generator(i) for i in range(batch_size)]
if "batch_size" in inputs:
batched_inputs["batch_size"] = batch_size
output = pipe(**inputs, output=self.output_name)
output_batch = pipe(**batched_inputs, output=self.output_name)
assert output_batch.shape[0] == batch_size
# For batch comparison, we only need to compare the first item
if output_batch.shape[0] == batch_size and output.shape[0] == 1:
output_batch = output_batch[0:1]
max_diff = torch.abs(output_batch - output).max()
assert max_diff < expected_max_diff, "Batch inference results different from single inference results"
@require_accelerator
def test_float16_inference(self, expected_max_diff=5e-2):
pipe = self.get_pipeline()
pipe.to(torch_device, torch.float32)
pipe_fp16 = self.get_pipeline()
pipe_fp16.to(torch_device, torch.float16)
inputs = self.get_dummy_inputs()
# Reset generator in case it is used inside dummy inputs
if "generator" in inputs:
inputs["generator"] = self.get_generator(0)
output = pipe(**inputs, output=self.output_name)
fp16_inputs = self.get_dummy_inputs()
# Reset generator in case it is used inside dummy inputs
if "generator" in fp16_inputs:
fp16_inputs["generator"] = self.get_generator(0)
output_fp16 = pipe_fp16(**fp16_inputs, output=self.output_name)
output_tensor = output.float().cpu()
output_fp16_tensor = output_fp16.float().cpu()
# Check for NaNs in outputs (can happen with tiny models in FP16)
if torch.isnan(output_tensor).any() or torch.isnan(output_fp16_tensor).any():
pytest.skip("FP16 inference produces NaN values - this is a known issue with tiny models")
max_diff = numpy_cosine_similarity_distance(
output_tensor.flatten().numpy(), output_fp16_tensor.flatten().numpy()
)
# Check if cosine similarity is NaN (which can happen if vectors are zero or very small)
if torch.isnan(torch.tensor(max_diff)):
pytest.skip("Cosine similarity is NaN - outputs may be too small for reliable comparison")
assert max_diff < expected_max_diff, f"FP16 inference is different from FP32 inference (max_diff: {max_diff})"
@require_accelerator
def test_to_device(self):
pipe = self.get_pipeline().to("cpu")
model_devices = [
component.device.type for component in pipe.components.values() if hasattr(component, "device")
]
assert all(device == "cpu" for device in model_devices), "All pipeline components are not on CPU"
pipe.to(torch_device)
model_devices = [
component.device.type for component in pipe.components.values() if hasattr(component, "device")
]
assert all(device == torch_device for device in model_devices), (
"All pipeline components are not on accelerator device"
)
def test_inference_is_not_nan_cpu(self):
pipe = self.get_pipeline().to("cpu")
inputs = self.get_dummy_inputs()
output = pipe(**inputs, output=self.output_name)
assert torch.isnan(output).sum() == 0, "CPU Inference returns NaN"
@require_accelerator
def test_inference_is_not_nan(self, base_pipe_output):
assert torch.isnan(base_pipe_output).sum() == 0, "Accelerator Inference returns NaN"
def test_num_images_per_prompt(self, batch_sizes=[1, 2], num_images_per_prompts=[1, 2]):
pipe = self.get_pipeline().to(torch_device)
if "num_images_per_prompt" not in pipe.blocks.input_names:
pytest.skip("Skipping test as `num_images_per_prompt` is not present in input names.")
for batch_size in batch_sizes:
for num_images_per_prompt in num_images_per_prompts:
inputs = self.get_dummy_inputs()
for key in inputs.keys():
if key in self.batch_params:
inputs[key] = batch_size * [inputs[key]]
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

45
tests/modular_pipelines/testing_utils/guider.py Normal file
View File

@@ -0,0 +1,45 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from diffusers.guiders import ClassifierFreeGuidance
from ...testing_utils import torch_device
from .common import BaseModularPipelineOutputMixin
class ModularGuiderTesterMixin(BaseModularPipelineOutputMixin):
"""Guidance tests for pipelines that expose a `guider` component."""
def test_guider_cfg(self, expected_max_diff=1e-2):
pipe = self.get_pipeline().to(torch_device)
# forward pass with CFG not applied
guider = ClassifierFreeGuidance(guidance_scale=1.0)
pipe.update_components(guider=guider)
inputs = self.get_dummy_inputs()
out_no_cfg = pipe(**inputs, output=self.output_name)
# forward pass with CFG applied
guider = ClassifierFreeGuidance(guidance_scale=7.5)
pipe.update_components(guider=guider)
inputs = self.get_dummy_inputs()
out_cfg = pipe(**inputs, output=self.output_name)
assert out_cfg.shape == out_no_cfg.shape
max_diff = torch.abs(out_cfg - out_no_cfg).max()
assert max_diff > expected_max_diff, "Output with CFG must be different from normal inference"

182
tests/modular_pipelines/testing_utils/loading.py Normal file
View File

@@ -0,0 +1,182 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import weakref
import pytest
import torch
import diffusers
from diffusers import ComponentsManager, ModularPipeline
from diffusers.utils import logging
from ...testing_utils import CaptureLogger, torch_device
from .common import BaseModularPipelineOutputMixin
from .utils import get_specified_components
class ModularLoadingTesterMixin(BaseModularPipelineOutputMixin):
"""
Serialization and component lifecycle tests for a modular pipeline: `save_pretrained`/`from_pretrained`
round-trips, what `modular_model_index.json` records, and `load_components`/`unload_components`.
Tests that assert on *device memory* after unloading live in `memory.py` instead.
"""
def test_save_from_pretrained(self, tmp_path, base_pipe_output):
base_pipe = self.get_pipeline().to(torch_device)
base_pipe.save_pretrained(str(tmp_path))
pipe = ModularPipeline.from_pretrained(tmp_path)
pipe.load_components(dtype=torch.float32)
pipe.to(torch_device)
image = pipe(**self.get_dummy_inputs(), output=self.output_name)
expected_slice = base_pipe_output[0, -3:, -3:, -1].flatten()
image_slice = image[0, -3:, -3:, -1].flatten()
assert torch.abs(expected_slice - image_slice).max() < 1e-3
def test_load_expected_components_from_pretrained(self, tmp_path):
pipe = self.get_pipeline()
expected = get_specified_components(self.pretrained_model_name_or_path, cache_dir=tmp_path)
if not expected:
pytest.skip("Skipping test as we couldn't fetch the expected components.")
actual = {
name
for name in pipe.components
if getattr(pipe, name, None) is not None
and getattr(getattr(pipe, name), "_diffusers_load_id", None) not in (None, "null")
}
assert expected == actual, f"Component mismatch: missing={expected - actual}, unexpected={actual - expected}"
def test_load_expected_components_from_save_pretrained(self, tmp_path):
pipe = self.get_pipeline()
save_dir = str(tmp_path / "saved-pipeline")
pipe.save_pretrained(save_dir)
expected = get_specified_components(save_dir)
loaded_pipe = ModularPipeline.from_pretrained(save_dir)
loaded_pipe.load_components(dtype=torch.float32)
actual = {
name
for name in loaded_pipe.components
if getattr(loaded_pipe, name, None) is not None
and getattr(getattr(loaded_pipe, name), "_diffusers_load_id", None) not in (None, "null")
}
assert expected == actual, (
f"Component mismatch after save/load: missing={expected - actual}, unexpected={actual - expected}"
)
def test_modular_index_consistency(self, tmp_path):
pipe = self.get_pipeline()
components_spec = pipe._component_specs
components = sorted(components_spec.keys())
pipe.save_pretrained(str(tmp_path))
index_file = tmp_path / "modular_model_index.json"
assert index_file.exists()
with open(index_file) as f:
index_contents = json.load(f)
compulsory_keys = {"_blocks_class_name", "_class_name", "_diffusers_version"}
for k in compulsory_keys:
assert k in index_contents
to_check_attrs = {"pretrained_model_name_or_path", "revision", "subfolder"}
for component in components:
spec = components_spec[component]
for attr in to_check_attrs:
if getattr(spec, "pretrained_model_name_or_path", None) is not None:
for attr in to_check_attrs:
assert component in index_contents, f"{component} should be present in index but isn't."
attr_value_from_index = index_contents[component][2][attr]
assert getattr(spec, attr) == attr_value_from_index
def test_unload_components(self):
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components()
name = next(name for name in pipe.pretrained_component_names if pipe.components.get(name) is not None)
spec_before = pipe._component_specs[name]
pipe.unload_components(name)
assert getattr(pipe, name) is None
# `components` is the mapping most callers iterate over: the entry stays, its value becomes None
assert pipe.components[name] is None
# unloading an already unloaded component is a no-op, not an error
pipe.unload_components(name)
assert pipe.components[name] is None
# the spec survives, so the component can be loaded again
assert pipe._component_specs[name] is spec_before
pipe.load_components(names=name)
assert getattr(pipe, name) is not None
# with a ComponentsManager attached, unloading also removes the component from the manager
manager = ComponentsManager()
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path, components_manager=manager)
pipe.load_components(names=name)
assert len(manager._lookup_ids(name=name)) == 1
pipe.unload_components(name)
assert getattr(pipe, name) is None
assert len(manager._lookup_ids(name=name)) == 0
def test_unload_components_multiple_names(self):
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components()
names = [name for name in pipe.pretrained_component_names if pipe.components.get(name) is not None]
if len(names) < 2:
pytest.skip("Skipping test as the pipeline has fewer than two loaded pretrained components.")
pipe.unload_components(names)
assert all(pipe.components[name] is None for name in names)
pipe.load_components(names=names)
assert all(pipe.components[name] is not None for name in names)
def test_unload_components_invalid_names(self):
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components()
name = next(name for name in pipe.pretrained_component_names if pipe.components.get(name) is not None)
with pytest.raises(ValueError, match="Invalid type for names"):
pipe.unload_components((name,))
assert pipe.components[name] is not None
# an unknown name is warned about and skipped; the known names are still unloaded
logger = logging.get_logger("diffusers.modular_pipelines.modular_pipeline")
logger.setLevel(diffusers.logging.WARNING)
with CaptureLogger(logger) as cap_logger:
pipe.unload_components([name, "not_a_component"])
assert "not_a_component" in cap_logger.out
assert pipe.components[name] is None
def test_unload_components_releases_component(self):
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components()
name = next(
name for name in pipe.pretrained_component_names if isinstance(pipe.components.get(name), torch.nn.Module)
)
# a weakref keeps no strong reference, so it goes dead only if nothing in the pipeline holds the
# component anymore — which is what makes the memory actually reclaimable
component_ref = weakref.ref(pipe.components[name])
pipe.unload_components(name)
assert component_ref() is None

279
tests/modular_pipelines/testing_utils/memory.py Normal file
View File

@@ -0,0 +1,279 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import gc
from unittest import mock
import pytest
import torch
from diffusers import ComponentsManager, ModularPipeline
from diffusers.hooks import apply_group_offloading
from diffusers.utils import is_accelerate_available
from ...testing_utils import (
backend_empty_cache,
is_cpu_offload,
is_group_offload,
is_memory,
require_accelerate,
require_accelerator,
torch_device,
)
from .common import BaseModularPipelineOutputMixin
from .utils import backend_memory_allocated, patch_free_memory
if is_accelerate_available():
from diffusers.modular_pipelines.components_manager import AutoOffloadStrategy
# More free memory than any tiny test checkpoint could ever need, so the strategy never
# decides to offload. Used to assert the *negative*: no eviction without memory pressure.
_AMPLE_FREE_BYTES = 1024**4
@is_cpu_offload
class ModularOffloadTesterMixin(BaseModularPipelineOutputMixin):
"""Auto CPU offload driven by a `ComponentsManager`: inference stays correct, and unloading a component
re-applies the offload hooks to the ones that are left."""
@require_accelerator
def test_components_auto_cpu_offload_inference_consistent(self, base_pipe_output):
cm = ComponentsManager()
cm.enable_auto_cpu_offload(device=torch_device)
offload_pipe = self.get_pipeline(components_manager=cm)
image = offload_pipe(**self.get_dummy_inputs(), output=self.output_name)
expected_slice = base_pipe_output[0, -3:, -3:, -1].flatten()
image_slice = image[0, -3:, -3:, -1].flatten()
assert torch.abs(expected_slice - image_slice).max() < 1e-3
@require_accelerator
def test_unload_components_auto_cpu_offload(self, base_pipe_output):
cm = ComponentsManager()
cm.enable_auto_cpu_offload(device=torch_device)
pipe = self.get_pipeline(components_manager=cm)
name = next(
name for name in pipe.pretrained_component_names if isinstance(pipe.components.get(name), torch.nn.Module)
)
component_id = f"{name}_{id(pipe.components[name])}"
pipe.unload_components(name)
# removing a component re-applies auto offload to the ones that are left
assert component_id not in cm.components
assert component_id not in {hook.model_id for hook in cm.model_hooks}
remaining = [component for component in cm.components.values() if isinstance(component, torch.nn.Module)]
assert all(hasattr(component, "_hf_hook") for component in remaining)
# the reloaded component is hooked up again, so the pipeline still runs
pipe.load_components(names=name, dtype=torch.float32)
image = pipe(**self.get_dummy_inputs(), output=self.output_name)
assert torch.abs(base_pipe_output - image).max() < 1e-3
@is_group_offload
class ModularGroupOffloadTesterMixin(BaseModularPipelineOutputMixin):
"""Group offloading applied to the pipeline's components."""
@require_accelerator
def test_group_offloading_execution_device(self):
pipe = self.get_pipeline().to("cpu")
assert pipe._execution_device.type == "cpu"
offloaded = None
for name, component in pipe.components.items():
if not isinstance(component, torch.nn.Module):
continue
if not getattr(component, "_supports_group_offloading", True):
continue
apply_group_offloading(
component,
onload_device=torch.device(torch_device),
offload_device=torch.device("cpu"),
offload_type="leaf_level",
)
offloaded = name
break
if offloaded is None:
pytest.skip("No component supports group offloading.")
assert pipe._execution_device.type == torch.device(torch_device).type
@is_memory
class ModularMemoryTesterMixin(ModularOffloadTesterMixin, ModularGroupOffloadTesterMixin):
"""Combined mixin for the memory optimizations every modular pipeline is expected to support: auto CPU offload,
group offload, and reclaiming device memory on `unload_components`."""
@require_accelerator
def test_unload_components_frees_device_memory(self):
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
name = next(
name for name in pipe.pretrained_component_names if isinstance(pipe.components.get(name), torch.nn.Module)
)
pipe.to(torch_device)
component = pipe.components[name]
footprint = sum(t.numel() * t.element_size() for t in [*component.parameters(), *component.buffers()])
del component
gc.collect()
backend_empty_cache(torch_device)
allocated_before = backend_memory_allocated(torch_device)
pipe.unload_components(name)
freed = allocated_before - backend_memory_allocated(torch_device)
assert freed >= 0.9 * footprint, (
f"Unloading '{name}' freed {freed} bytes on {torch_device}, expected around {footprint}"
)
@is_cpu_offload
class ModularAutoOffloadTesterMixin(BaseModularPipelineOutputMixin):
"""
Auto-CPU-offload *decisions* for a modular pipeline's components, driven by simulated device memory.
Opt-in on top of `ModularMemoryTesterMixin`: these tests spy on `AutoOffloadStrategy` to assert how models are
sequenced onto the device, which is only meaningful for pipelines with several offloadable model components.
"""
@staticmethod
def _managed_models(cm):
"""The registered components that the offloader actually manages (parameterized
`nn.Module`s)."""
models = []
for component in cm.components.values():
if isinstance(component, torch.nn.Module) and next(component.parameters(), None) is not None:
models.append(component)
return models
@staticmethod
def _is_resident(model):
return next(model.parameters()).device.type == torch.device(torch_device).type
def _run_offloaded(self, free_bytes):
"""
Run the pipeline with auto offload on and `free_bytes` of *simulated* device
memory, recording every offload decision the strategy makes.
Each record is `{"incoming", "resident_before", "offloaded"}` (lists of model
ids), captured by spying on `AutoOffloadStrategy.__call__`, which the hooks call
each time a model is about to be moved onto the device.
"""
cm = ComponentsManager()
pipe = self.get_pipeline(components_manager=cm)
cm.enable_auto_cpu_offload(device=torch_device, memory_reserve_margin=0)
records = []
original_call = AutoOffloadStrategy.__call__
def spy_call(strategy, hooks, model_id, model, execution_device):
selected = original_call(
strategy, hooks=hooks, model_id=model_id, model=model, execution_device=execution_device
)
records.append(
{
"incoming": model_id,
"resident_before": [hook.model_id for hook in hooks],
"offloaded": [hook.model_id for hook in selected],
}
)
return selected
with patch_free_memory(free_bytes), mock.patch.object(AutoOffloadStrategy, "__call__", spy_call):
output = pipe(**self.get_dummy_inputs(), output=self.output_name)
return cm, records, output
@staticmethod
def _peak_co_residency(records):
"""
Largest number of models simultaneously on the device, reconstructed from the
strategy's view of residency just before each load.
"""
peak = 0
for record in records:
resident = (set(record["resident_before"]) - set(record["offloaded"])) | {record["incoming"]}
peak = max(peak, len(resident))
return peak
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_serializes_models_under_memory_pressure(self):
# Zero simulated free memory: every model that runs must first evict whatever is
# currently resident (comfy-style serialized execution).
cm, records, _ = self._run_offloaded(free_bytes=0)
try:
distinct_models = {record["incoming"] for record in records}
if len(distinct_models) < 2:
pytest.skip("pipeline has fewer than two offloadable model components")
# Offloading actually fired (at least one eviction happened).
assert any(record["offloaded"] for record in records), "expected at least one eviction"
# Sequencing: models run one at a time, never two co-resident on the device.
peak = self._peak_co_residency(records)
assert peak == 1, f"expected serialized execution under pressure, saw {peak} models co-resident"
# Device placement after the run: at most the last-run model stays on the
# accelerator, and at least one managed model was pushed back to the CPU.
models = self._managed_models(cm)
resident = [m for m in models if self._is_resident(m)]
assert len(resident) <= 1
assert any(not self._is_resident(m) for m in models), "expected some model offloaded to CPU"
finally:
cm.disable_auto_cpu_offload()
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_keeps_models_resident_without_memory_pressure(self):
# Negative case: with ample simulated memory the strategy is still consulted on
# every load, but it must never decide to evict anything.
cm, records, _ = self._run_offloaded(free_bytes=_AMPLE_FREE_BYTES)
try:
distinct_models = {record["incoming"] for record in records}
if len(distinct_models) < 2:
pytest.skip("pipeline has fewer than two offloadable model components")
# Nothing was ever offloaded...
assert all(record["offloaded"] == [] for record in records), "no model should be evicted"
# ...and models accumulate on the device instead of being serialized.
peak = self._peak_co_residency(records)
assert peak >= 2, f"expected models to co-reside without pressure, saw peak {peak}"
models = self._managed_models(cm)
assert sum(self._is_resident(m) for m in models) >= 2, "expected multiple models resident on device"
finally:
cm.disable_auto_cpu_offload()
@require_accelerate
@require_accelerator
def test_auto_cpu_offload_inference_consistent_under_memory_pressure(
self, base_pipe_output, expected_max_diff=1e-3
):
# Sensible results: forcing offload (zero simulated free memory) must not change
# the output relative to an ordinary, non-offloaded run.
cm, _, offloaded = self._run_offloaded(free_bytes=0)
try:
max_diff = torch.abs(base_pipe_output - offloaded).max()
assert max_diff < expected_max_diff, f"offloaded output diverged from baseline (max diff {max_diff})"
finally:
cm.disable_auto_cpu_offload()

79
tests/modular_pipelines/testing_utils/utils.py Normal file
View File

@@ -0,0 +1,79 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import json
import os
from unittest import mock
import pytest
import torch
from huggingface_hub import hf_hub_download
from ...testing_utils import torch_device
def backend_memory_allocated(device: str) -> int:
"""
Bytes currently allocated on `device`. `tests/testing_utils.py` only exposes the *peak* allocation, which cannot
show memory being released. Skips on backends that do not implement `memory_allocated()` (e.g. mps).
"""
device_module = getattr(torch, torch.device(device).type)
if not hasattr(device_module, "memory_allocated"):
pytest.skip(f"`memory_allocated()` is not implemented for {device}.")
return device_module.memory_allocated()
def patch_free_memory(free_bytes: int, total_bytes: int = 80 * 1024):
"""
Simulate `free_bytes` of free device memory on whichever backend module (cuda/xpu/...) backs `torch_device`.
`mem_get_info` returns `(free, total)` and is the single point where `AutoOffloadStrategy` learns how much memory
is available, so patching it makes offloading decisions deterministic instead of dependent on the real free memory
of the test hardware (an 80GB GPU never runs low on a handful of KB-sized models).
"""
device_type = torch.device(torch_device).type
device_module = getattr(torch, device_type, torch.cuda)
return mock.patch.object(device_module, "mem_get_info", return_value=(free_bytes, total_bytes))
def get_specified_components(path_or_repo_id, cache_dir=None):
"""
The component names a `modular_model_index.json` actually points at a checkpoint for. Returns `None` when the
index cannot be fetched, which callers treat as "skip the comparison".
"""
if os.path.isdir(path_or_repo_id):
config_path = os.path.join(path_or_repo_id, "modular_model_index.json")
else:
try:
config_path = hf_hub_download(
repo_id=path_or_repo_id,
filename="modular_model_index.json",
local_dir=cache_dir,
)
except Exception:
return None
with open(config_path) as f:
config = json.load(f)
components = set()
for k, v in config.items():
if isinstance(v, (str, int, float, bool)):
continue
for entry in v:
if isinstance(entry, dict) and (entry.get("repo") or entry.get("pretrained_model_name_or_path")):
components.add(k)
break
return components

183
tests/modular_pipelines/testing_utils/workflow.py Normal file
View File

@@ -0,0 +1,183 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import pytest
from diffusers import ModularPipeline
class ModularWorkflowTesterMixin:
"""
Tests for the workflows a blocks class declares through `_workflow_map`: which blocks each workflow resolves to,
the components/configs/inputs it ends up expecting, and the two entry points that prune by workflow name
(`ModularPipeline.from_pretrained(..., workflow=...)` and `load_components(workflow=...)`).
Every test skips when the blocks class declares no `_workflow_map`, so this mixin is safe to compose onto any
modular pipeline test.
"""
def test_workflow_map(self):
blocks = self.pipeline_blocks_class()
if blocks._workflow_map is None:
pytest.skip("Skipping test as _workflow_map is not set")
assert hasattr(self, "expected_workflow_blocks") and self.expected_workflow_blocks, (
"expected_workflow_blocks must be defined in the test class"
)
for workflow_name, expected_blocks in self.expected_workflow_blocks.items():
workflow_blocks = blocks.get_workflow(workflow_name)
actual_blocks = list(workflow_blocks.sub_blocks.items())
# Check that the number of blocks matches
assert len(actual_blocks) == len(expected_blocks), (
f"Workflow '{workflow_name}' has {len(actual_blocks)} blocks, expected {len(expected_blocks)}"
)
# Check that each block name and type matches
for i, ((actual_name, actual_block), (expected_name, expected_class_name)) in enumerate(
zip(actual_blocks, expected_blocks)
):
assert actual_name == expected_name
assert actual_block.__class__.__name__ == expected_class_name, (
f"Workflow '{workflow_name}': block '{actual_name}' has type "
f"{actual_block.__class__.__name__}, expected {expected_class_name}"
)
# a `_workflow_map` value is either one trigger dict or a tuple of trigger dicts — alternative spellings
# of the same workflow, which must all resolve the same blocks since `get_workflow` prunes with the first
for workflow_name, trigger_inputs in blocks._workflow_map.items():
if isinstance(trigger_inputs, dict):
continue
assert isinstance(trigger_inputs, tuple) and all(
isinstance(alternative, dict) for alternative in trigger_inputs
), (
f"Workflow '{workflow_name}': a `_workflow_map` value must be a trigger dict or a tuple of "
f"trigger dicts, got {trigger_inputs!r}"
)
resolved = [list(blocks.get_execution_blocks(**alternative).sub_blocks) for alternative in trigger_inputs]
for alternative_blocks in resolved[1:]:
assert alternative_blocks == resolved[0], (
f"Workflow '{workflow_name}': its trigger spellings resolve different blocks: "
f"{resolved[0]} vs {alternative_blocks}"
)
def test_workflow_defaults(self):
if not getattr(self, "expected_workflow_defaults", None):
pytest.skip("Skipping test as expected_workflow_defaults is not set")
blocks = self.pipeline_blocks_class()
for workflow_name, expected_defaults in self.expected_workflow_defaults.items():
# a pipeline without workflows is tested as the single unnamed workflow `None` over the full blockset
workflow_blocks = blocks if workflow_name is None else blocks.get_workflow(workflow_name)
# components: every one of the workflow is named with its class, so one appearing, disappearing or
# changing type fails loudly
component_specs = {spec.name: spec for spec in workflow_blocks.expected_components}
expected_components = expected_defaults["components"]
assert set(component_specs) == set(expected_components), (
f"Workflow '{workflow_name}' expects components {sorted(component_specs)}, "
f"the test expects {sorted(expected_components)}"
)
for component_name, expected_class_name in expected_components.items():
actual_class_name = component_specs[component_name].type_hint.__name__
assert actual_class_name == expected_class_name, (
f"Workflow '{workflow_name}': component '{component_name}' is a {actual_class_name}, "
f"expected {expected_class_name}"
)
# configs: the pipeline-level configs the workflow declares, with their defaults
config_specs = {spec.name: spec.default for spec in workflow_blocks.expected_configs}
expected_configs = expected_defaults.get("configs", {})
assert config_specs == expected_configs, (
f"Workflow '{workflow_name}' declares configs {config_specs}, the test expects {expected_configs}"
)
# inputs: `inputs` names every optional input with its default and `required_inputs` the required ones,
# and together they are exactly the workflow's inputs. kwargs-style inputs (e.g.
# **denoiser_input_fields) have no name and no default to pin
input_params = {param.name: param for param in workflow_blocks.inputs if param.name is not None}
expected_inputs = expected_defaults["inputs"]
expected_required = expected_defaults.get("required_inputs", [])
assert set(input_params) == set(expected_inputs) | set(expected_required), (
f"Workflow '{workflow_name}' takes inputs {sorted(input_params)}, "
f"the test expects {sorted(set(expected_inputs) | set(expected_required))}"
)
for input_name in expected_required:
assert input_params[input_name].required, (
f"Workflow '{workflow_name}': input '{input_name}' should be required"
)
for input_name, expected_default in expected_inputs.items():
param = input_params[input_name]
assert not param.required and param.default == expected_default, (
f"Workflow '{workflow_name}': input '{input_name}' default is "
f"{param.default!r} (required={param.required}), expected {expected_default!r}"
)
# component configs: the values a workflow pins on a component it creates itself — the guidance scale
# of its guider, say — which is what tells otherwise identical presets apart
for component_name, expected_config in expected_defaults.get("component_configs", {}).items():
actual_config = dict(component_specs[component_name].config or {})
for config_name, expected_value in expected_config.items():
assert actual_config.get(config_name) == expected_value, (
f"Workflow '{workflow_name}': component '{component_name}' config '{config_name}' is "
f"{actual_config.get(config_name)!r}, expected {expected_value!r}"
)
def test_from_pretrained_workflow(self):
blocks = self.pipeline_blocks_class()
if blocks._workflow_map is None:
pytest.skip("Skipping test as _workflow_map is not set")
for workflow_name in blocks.available_workflows:
# the workflow argument should be equivalent to pruning the blocks by hand
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path, workflow=workflow_name)
ref_pipe = blocks.get_workflow(workflow_name).init_pipeline(self.pretrained_model_name_or_path)
assert set(pipe.component_names) == set(ref_pipe.component_names), (
f"Workflow '{workflow_name}': pipeline expects components {sorted(pipe.component_names)}, "
f"the workflow blocks expect {sorted(ref_pipe.component_names)}"
)
for name in pipe.pretrained_component_names:
assert pipe.get_component_spec(name) == ref_pipe.get_component_spec(name), (
f"Workflow '{workflow_name}': component '{name}' has a different spec than the one "
f"created from the workflow blocks"
)
with pytest.raises(ValueError, match="Available workflows"):
ModularPipeline.from_pretrained(self.pretrained_model_name_or_path, workflow="not_a_workflow")
def test_load_components_workflow(self):
blocks = self.pipeline_blocks_class()
if blocks._workflow_map is None:
pytest.skip("Skipping test as _workflow_map is not set")
workflow_name = blocks.available_workflows[0]
# a full pipeline restricted at load time should load the same components as a pipeline
# created from the pruned workflow blocks
pipe = ModularPipeline.from_pretrained(self.pretrained_model_name_or_path)
pipe.load_components(workflow=workflow_name)
ref_pipe = blocks.get_workflow(workflow_name).init_pipeline(self.pretrained_model_name_or_path)
ref_pipe.load_components()
loaded = {name for name in pipe.pretrained_component_names if pipe.components[name] is not None}
ref_loaded = {name for name in ref_pipe.pretrained_component_names if ref_pipe.components[name] is not None}
assert loaded == ref_loaded, (
f"Workflow '{workflow_name}': load_components(workflow=...) loaded {sorted(loaded)}, "
f"the pipeline created from the workflow blocks loaded {sorted(ref_loaded)}"
)
with pytest.raises(ValueError, match="not both"):
pipe.load_components(names="unet", workflow=workflow_name)

13
tests/modular_pipelines/utils.py
View File

@@ -1,13 +0,0 @@
import pytest
import torch
def backend_memory_allocated(device: str) -> int:
"""
Bytes currently allocated on `device`. `tests/testing_utils.py` only exposes the *peak* allocation, which cannot
show memory being released. Skips on backends that do not implement `memory_allocated()` (e.g. mps).
"""
device_module = getattr(torch, torch.device(device).type)
if not hasattr(device_module, "memory_allocated"):
pytest.skip(f"`memory_allocated()` is not implemented for {device}.")
return device_module.memory_allocated()

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

@@ -17,14 +17,19 @@ import pytest
from diffusers.modular_pipelines import WanBlocks, WanModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
class TestWanModularPipelineFast(ModularPipelineTesterMixin):
class WanModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = WanModularPipeline
pipeline_blocks_class = WanBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-modular-pipe"
params = frozenset(["prompt", "height", "width", "num_frames"])
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"])
@@ -44,6 +49,20 @@ class TestWanModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestWanModularPipelineFast(WanModularPipelineTesterConfig, ModularPipelineTesterMixin):
@pytest.mark.skip(reason="num_videos_per_prompt")
def test_num_images_per_prompt(self):
pass
class TestWanModularPipelineLoading(WanModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestWanModularPipelineWorkflow(WanModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestWanModularPipelineMemory(WanModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

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

@@ -24,7 +24,14 @@ from diffusers.modular_pipelines import (
WanAnimate2ModularPipeline,
)
from ..test_modular_pipelines_common import ModularGuiderTesterMixin, ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularGuiderTesterMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
# Every component with its class, every input — optional ones with their defaults — and the config values worth
@@ -73,7 +80,7 @@ WAN_ANIMATE_2_DISTILLED_DEFAULTS = {
}
class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTesterMixin):
class WanAnimate2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = WanAnimate2ModularPipeline
pipeline_blocks_class = WanAnimate2Blocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-modular"
@@ -103,6 +110,10 @@ class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuid
}
return inputs
class WanAnimate2ModularPipelineFastTesterMixin(ModularPipelineTesterMixin):
"""`ModularPipelineTesterMixin` minus the tests that assume a batchable pipeline; shared by both presets."""
@pytest.mark.skip(reason="Wan-Animate-2 is unbatched: one character image and driving video per call")
def test_inference_batch_consistent(self):
pass
@@ -115,6 +126,26 @@ class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuid
def test_num_images_per_prompt(self):
pass
class TestWanAnimate2ModularPipelineFast(
WanAnimate2ModularPipelineTesterConfig, WanAnimate2ModularPipelineFastTesterMixin
):
pass
class TestWanAnimate2ModularPipelineLoading(WanAnimate2ModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestWanAnimate2ModularPipelineWorkflow(WanAnimate2ModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestWanAnimate2ModularPipelineMemory(WanAnimate2ModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass
class TestWanAnimate2ModularPipelineGuider(WanAnimate2ModularPipelineTesterConfig, ModularGuiderTesterMixin):
def test_guider_cfg(self):
# The tiny random transformer responds only weakly to the text embeddings (cond vs uncond
# predictions differ by ~1e-4), so the CFG effect on the decoded pixels is real but small.
@@ -122,12 +153,33 @@ class TestWanAnimate2ModularPipelineFast(ModularPipelineTesterMixin, ModularGuid
super().test_guider_cfg(expected_max_diff=1e-6)
class TestWanAnimate2DistilledModularPipelineFast(TestWanAnimate2ModularPipelineFast):
class WanAnimate2DistilledModularPipelineTesterConfig(WanAnimate2ModularPipelineTesterConfig):
pipeline_class = WanAnimate2DistilledModularPipeline
pipeline_blocks_class = WanAnimate2DistilledBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-wan-animate-2-distilled-modular"
expected_workflow_defaults = WAN_ANIMATE_2_DISTILLED_DEFAULTS
@pytest.mark.skip(reason="The distilled preset pins its guider to guidance_scale=1.0")
def test_guider_cfg(self):
pass
# No guider test class for the distilled preset: it pins its guider to guidance_scale=1.0.
class TestWanAnimate2DistilledModularPipelineFast(
WanAnimate2DistilledModularPipelineTesterConfig, WanAnimate2ModularPipelineFastTesterMixin
):
pass
class TestWanAnimate2DistilledModularPipelineLoading(
WanAnimate2DistilledModularPipelineTesterConfig, ModularLoadingTesterMixin
):
pass
class TestWanAnimate2DistilledModularPipelineWorkflow(
WanAnimate2DistilledModularPipelineTesterConfig, ModularWorkflowTesterMixin
):
pass
class TestWanAnimate2DistilledModularPipelineMemory(
WanAnimate2DistilledModularPipelineTesterConfig, ModularMemoryTesterMixin
):
pass

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

@@ -16,7 +16,13 @@
from diffusers.modular_pipelines import ZImageAutoBlocks, ZImageModularPipeline
from ..test_modular_pipelines_common import ModularPipelineTesterMixin
from ..testing_utils import (
BaseModularPipelineTesterConfig,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
ZIMAGE_WORKFLOWS = {
@@ -43,11 +49,10 @@ ZIMAGE_WORKFLOWS = {
}
class TestZImageModularPipelineFast(ModularPipelineTesterMixin):
class ZImageModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = ZImageModularPipeline
pipeline_blocks_class = ZImageAutoBlocks
pretrained_model_name_or_path = "hf-internal-testing/tiny-zimage-modular-pipe"
params = frozenset(["prompt", "height", "width"])
batch_params = frozenset(["prompt"])
expected_workflow_blocks = ZIMAGE_WORKFLOWS
@@ -65,5 +70,19 @@ class TestZImageModularPipelineFast(ModularPipelineTesterMixin):
}
return inputs
class TestZImageModularPipelineFast(ZImageModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=5e-3)
class TestZImageModularPipelineLoading(ZImageModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestZImageModularPipelineWorkflow(ZImageModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestZImageModularPipelineMemory(ZImageModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass