From e94da484b9aaa70a25a178cc03c36d5390c560ec Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Tue, 11 Aug 2026 11:47:45 +0000 Subject: [PATCH 1/4] notes from .ai --- .ai/testing.md | 12 +- .../anima/test_modular_pipeline_anima.py | 49 +- .../cosmos/test_modular_pipeline_cosmos3.py | 51 +- ...test_modular_pipeline_cosmos3_distilled.py | 53 +- .../test_modular_pipeline_ernie_image.py | 25 +- .../flux/test_modular_pipeline_flux.py | 101 +- .../flux2/test_modular_pipeline_flux2.py | 50 +- .../test_modular_pipeline_flux2_klein.py | 50 +- .../test_modular_pipeline_flux2_klein_base.py | 50 +- .../helios/test_modular_pipeline_helios.py | 46 +- .../test_modular_pipeline_hunyuan_video1_5.py | 25 +- .../krea2/test_modular_pipeline_krea2.py | 25 +- .../test_modular_pipeline_krea2_turbo.py | 25 +- .../ltx/test_modular_pipeline_ltx.py | 25 +- .../test_modular_pipeline_minimax_h3.py | 102 +- .../qwen/test_modular_pipeline_qwenimage.py | 80 +- ...est_modular_pipeline_stable_diffusion_3.py | 79 +- ...st_modular_pipeline_stable_diffusion_xl.py | 94 +- .../test_components_manager.py | 165 +-- .../test_modular_model_card.py | 232 ++++ .../test_modular_pipeline_loading.py | 241 ++++ .../test_modular_pipelines_common.py | 1203 ----------------- .../testing_utils/__init__.py | 32 + .../modular_pipelines/testing_utils/common.py | 366 +++++ .../modular_pipelines/testing_utils/guider.py | 45 + .../testing_utils/loading.py | 182 +++ .../modular_pipelines/testing_utils/memory.py | 279 ++++ .../modular_pipelines/testing_utils/utils.py | 79 ++ .../testing_utils/workflow.py | 172 +++ tests/modular_pipelines/utils.py | 13 - .../wan/test_modular_pipeline_wan.py | 25 +- .../z_image/test_modular_pipeline_z_image.py | 25 +- 32 files changed, 2428 insertions(+), 1573 deletions(-) create mode 100644 tests/modular_pipelines/test_modular_model_card.py create mode 100644 tests/modular_pipelines/test_modular_pipeline_loading.py delete mode 100644 tests/modular_pipelines/test_modular_pipelines_common.py create mode 100644 tests/modular_pipelines/testing_utils/__init__.py create mode 100644 tests/modular_pipelines/testing_utils/common.py create mode 100644 tests/modular_pipelines/testing_utils/guider.py create mode 100644 tests/modular_pipelines/testing_utils/loading.py create mode 100644 tests/modular_pipelines/testing_utils/memory.py create mode 100644 tests/modular_pipelines/testing_utils/utils.py create mode 100644 tests/modular_pipelines/testing_utils/workflow.py delete mode 100644 tests/modular_pipelines/utils.py diff --git a/.ai/testing.md b/.ai/testing.md index 24d39da3bd68..ac0fda2d440b 100644 --- a/.ai/testing.md +++ b/.ai/testing.md @@ -34,9 +34,15 @@ Follow the style introduced in [#14113](https://github.com/huggingface/diffusers ### Modular pipelines -- Location: `tests/modular_pipelines//test_modular_pipeline_.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. +- Location: `tests/modular_pipelines//test_modular_pipeline_.py` (one config class + set of test classes per blockset / pipeline variant). +- **Define one config class**, `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. diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index c35def821d10..c3ae0d9f2e0d 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py @@ -26,8 +26,16 @@ 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 @@ def test_conditioner_output_shape_and_padding(self): 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 @@ def get_dummy_inputs(self, seed=0): "output_type": "pt", } + +class TestAnimaModularPipelineFast(AnimaModularPipelineTesterConfig, ModularPipelineTesterMixin): def test_inference_empty_negative_prompt(self): pipe = self.get_pipeline() @@ -130,6 +140,8 @@ def test_inference_empty_negative_prompt(self): 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 @@ def test_save_load_components(self): 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 @@ def test_load_lora_weights(self): 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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_inference_empty_negative_prompt(self): 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 diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py index 846b7d1b1c3f..9c4addb69566 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3.py @@ -34,7 +34,13 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_num_images_per_prompt(self): 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 @@ def test_set_timesteps_native_flow_schedule(self): 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 diff --git a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py index 00a7bd5390bb..97bb2b478989 100644 --- a/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py +++ b/tests/modular_pipelines/cosmos/test_modular_pipeline_cosmos3_distilled.py @@ -21,7 +21,13 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_rejects_guidance_scale_override(self): 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 diff --git a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py index 0c1b82700044..10df99406326 100644 --- a/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py +++ b/tests/modular_pipelines/ernie_image/test_modular_pipeline_ernie_image.py @@ -17,7 +17,13 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): "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 diff --git a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py index 1b7423f45004..502df3f447bf 100644 --- a/tests/modular_pipelines/flux/test_modular_pipeline_flux.py +++ b/tests/modular_pipelines/flux/test_modular_pipeline_flux.py @@ -29,8 +29,14 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_float16_inference(self): } -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,30 +153,36 @@ def get_dummy_inputs(self, seed=0): return inputs - def test_save_from_pretrained(self, tmp_path): - pipes = [] - base_pipe = self.get_pipeline().to(torch_device) - pipes.append(base_pipe) +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).to(torch_device) + + 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) - 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 TestFluxImg2ImgModularPipelineWorkflow(FluxImg2ImgModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass - def test_float16_inference(self): - super().test_float16_inference(8e-2) + +class TestFluxImg2ImgModularPipelineMemory(FluxImg2ImgModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass FLUX_KONTEXT_WORKFLOWS = { @@ -181,7 +211,7 @@ def test_float16_inference(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py index 2acb2d4f6dae..eb2ed35c46b7 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2.py @@ -25,7 +25,13 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_float16_inference(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py index 299809b02ca7..4932efa7dd74 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein.py @@ -25,7 +25,13 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_float16_inference(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 diff --git a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py index 20328a8e4d99..30460105ad34 100644 --- a/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py +++ b/tests/modular_pipelines/flux2/test_modular_pipeline_flux2_klein_base.py @@ -25,7 +25,13 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_float16_inference(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 diff --git a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py index cdd7cfa19fc9..5c270c20d9de 100644 --- a/tests/modular_pipelines/helios/test_modular_pipeline_helios.py +++ b/tests/modular_pipelines/helios/test_modular_pipeline_helios.py @@ -22,7 +22,13 @@ HeliosPyramidModularPipeline, ) -from ..test_modular_pipelines_common import ModularPipelineTesterMixin +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) HELIOS_WORKFLOWS = { @@ -61,11 +67,10 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_num_images_per_prompt(self): } -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 @@ def get_dummy_inputs(self, seed=0): } 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 diff --git a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py index a6db1289bb35..f6368d91b011 100644 --- a/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py +++ b/tests/modular_pipelines/hunyuan_video1_5/test_modular_pipeline_hunyuan_video1_5.py @@ -17,7 +17,13 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_inference_batch_single_identical(self): 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 diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py index 15caa5abe45f..133b3eb1bf66 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2.py @@ -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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 diff --git a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py index 9143e67d003a..40d278156ce9 100644 --- a/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py +++ b/tests/modular_pipelines/krea2/test_modular_pipeline_krea2_turbo.py @@ -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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 diff --git a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py index 733b79e98de6..1a58465c23c6 100644 --- a/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py +++ b/tests/modular_pipelines/ltx/test_modular_pipeline_ltx.py @@ -17,7 +17,13 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 diff --git a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py index 366cb366b220..f4b6904c5727 100644 --- a/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py +++ b/tests/modular_pipelines/minimax_h3/test_modular_pipeline_minimax_h3.py @@ -35,7 +35,13 @@ ) 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 @@ -146,13 +152,12 @@ } -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() @@ -163,6 +168,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 @@ -188,20 +209,6 @@ def test_duration_ceiling_holds_for_the_aligned_count(self): 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() @@ -395,13 +402,24 @@ def test_check_inputs(self, overrides, message): 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() @@ -430,6 +448,23 @@ def get_pipeline(self, components_manager=None, dtype=torch.float32): 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 @@ -458,21 +493,6 @@ def test_duration_ceiling_holds_for_the_aligned_count(self): 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() @@ -741,6 +761,20 @@ def test_check_inputs_references(self, references, message): 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. diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py index 71345ace1d6b..979c06fec2ed 100644 --- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py +++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py @@ -27,7 +27,14 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self): } 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 @@ def test_inference_batch_single_identical(self): } -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,32 @@ def get_dummy_inputs(self): inputs["image"] = PIL.Image.new("RGB", (32, 32), 0) return inputs + +class TestQwenImageEditModularPipelineFast(QwenImageEditModularPipelineTesterConfig, ModularPipelineTesterMixin): + pass + + +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 +282,10 @@ def get_dummy_inputs(self): 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 +308,27 @@ def test_inference_batch_consistent(): 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) diff --git a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py index 52965a653935..03d0c5687794 100644 --- a/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py +++ b/tests/modular_pipelines/stable_diffusion_3/test_modular_pipeline_stable_diffusion_3.py @@ -26,7 +26,13 @@ ) 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_load_expected_components_from_save_pretrained(self, tmp_path): 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 @@ def test_float16_inference(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 @@ def test_load_expected_components_from_save_pretrained(self, tmp_path): 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 diff --git a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py index d6a0523c2d11..1e1777b90790 100644 --- a/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py +++ b/tests/modular_pipelines/stable_diffusion_xl/test_modular_pipeline_stable_diffusion_xl.py @@ -24,14 +24,22 @@ 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 @@ def _test_stable_diffusion_xl_euler(self, expected_image_shape, expected_slice, 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 @@ def test_ip_adapter(self, expected_max_diff: float = 1e-4, expected_pipe_slice=N ) -class SDXLModularControlNetTesterMixin: +class SDXLModularControlNetTesterMixin(BaseModularPipelineOutputMixin): """ This mixin is designed to test ControlNet. """ @@ -321,14 +330,8 @@ def test_controlnet_cfg(self): } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_inference_batch_single_identical(self): 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 @@ def test_inference_batch_single_identical(self): } -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 @@ def get_dummy_inputs(self, seed=0): 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 @@ def test_inference_batch_single_identical(self): 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 @@ def test_inference_batch_single_identical(self): class SDXLInpaintingModularPipelineFastTests( + BaseModularPipelineTesterConfig, SDXLModularTesterMixin, SDXLModularIPAdapterTesterMixin, SDXLModularControlNetTesterMixin, diff --git a/tests/modular_pipelines/test_components_manager.py b/tests/modular_pipelines/test_components_manager.py index 42929b04e3d8..be2a247995f1 100644 --- a/tests/modular_pipelines/test_components_manager.py +++ b/tests/modular_pipelines/test_components_manager.py @@ -15,7 +15,6 @@ import gc from unittest import mock -import pytest import torch from diffusers import ComponentsManager @@ -23,6 +22,7 @@ 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 @@ def test_auto_offload_evicts_resident_model_under_memory_pressure(self): # 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 @@ def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): 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 @@ def test_auto_offload_keeps_models_resident_when_memory_is_ample(self): 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() diff --git a/tests/modular_pipelines/test_modular_model_card.py b/tests/modular_pipelines/test_modular_model_card.py new file mode 100644 index 000000000000..bebcadeb4b35 --- /dev/null +++ b/tests/modular_pipelines/test_modular_model_card.py @@ -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"] diff --git a/tests/modular_pipelines/test_modular_pipeline_loading.py b/tests/modular_pipelines/test_modular_pipeline_loading.py new file mode 100644 index 000000000000..3b9ebcc1cdf3 --- /dev/null +++ b/tests/modular_pipelines/test_modular_pipeline_loading.py @@ -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) diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py deleted file mode 100644 index fadaef9d3e58..000000000000 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ /dev/null @@ -1,1203 +0,0 @@ -import gc -import json -import os -import weakref -from typing import Callable - -import pytest -import torch -from huggingface_hub import hf_hub_download - -import diffusers -from diffusers import AutoModel, ComponentsManager, ControlNetModel, ModularPipeline, ModularPipelineBlocks -from diffusers.guiders import ClassifierFreeGuidance -from diffusers.modular_pipelines.modular_pipeline_utils import ( - ComponentSpec, - ConfigSpec, - InputParam, - OutputParam, - generate_modular_model_card_content, -) -from diffusers.utils import logging - -from ..testing_utils import ( - CaptureLogger, - backend_empty_cache, - numpy_cosine_similarity_distance, - require_accelerator, - torch_device, -) -from .utils import backend_memory_allocated - - -def _get_specified_components(path_or_repo_id, cache_dir=None): - 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 - - -class ModularPipelineTesterMixin: - """ - It provides a set of common tests for each modular pipeline, - including: - - test_pipeline_call_signature: check if the pipeline's __call__ method has all required parameters - - test_inference_batch_consistent: check if the pipeline's __call__ method can handle batch inputs - - test_inference_batch_single_identical: check if the pipeline's __call__ method can handle single input - - test_float16_inference: check if the pipeline's __call__ method can handle float16 inputs - - test_to_device: check if the pipeline's __call__ method can handle different devices - """ - - # 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" - - def get_generator(self, seed=0): - generator = torch.Generator("cpu").manual_seed(seed) - return generator - - @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." - ) - - def setup_method(self): - # clean up the VRAM before each test - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - def teardown_method(self): - # clean up the VRAM after each test in case of CUDA runtime errors - torch.compiler.reset() - gc.collect() - backend_empty_cache(torch_device) - - def get_pipeline(self, components_manager=None, dtype=torch.float32): - 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 - - 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): - pipe = self.get_pipeline().to(torch_device) - - inputs = self.get_dummy_inputs() - output = pipe(**inputs, output=self.output_name) - assert torch.isnan(output).sum() == 0, "Accelerator Inference returns NaN" - - def test_num_images_per_prompt(self): - pipe = self.get_pipeline().to(torch_device) - - if "num_images_per_prompt" not in pipe.blocks.input_names: - pytest.mark.skip("Skipping test as `num_images_per_prompt` is not present in input names.") - - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] - - 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 - - @require_accelerator - def test_components_auto_cpu_offload_inference_consistent(self): - base_pipe = self.get_pipeline().to(torch_device) - - cm = ComponentsManager() - cm.enable_auto_cpu_offload(device=torch_device) - offload_pipe = self.get_pipeline(components_manager=cm) - - image_slices = [] - for pipe in [base_pipe, offload_pipe]: - inputs = self.get_dummy_inputs() - image = pipe(**inputs, output=self.output_name) - image_slices.append(image[0, -3:, -3:, -1].flatten()) - - assert torch.abs(image_slices[0] - image_slices[1]).max() < 1e-3 - - @require_accelerator - def test_group_offloading_execution_device(self): - from diffusers.hooks import apply_group_offloading - - 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 - - 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) - - pipes.append(pipe) - - image_slices = [] - for pipe in pipes: - inputs = self.get_dummy_inputs() - image = pipe(**inputs, output=self.output_name) - image_slices.append(image[0, -3:, -3:, -1].flatten()) - - assert torch.abs(image_slices[0] - image_slices[1]).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_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(): - workflow_blocks = 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}" - ) - - 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) - - 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 - - @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}" - ) - - @require_accelerator - def test_unload_components_auto_cpu_offload(self): - base_pipe = self.get_pipeline().to(torch_device) - expected_image = base_pipe(**self.get_dummy_inputs(), output=self.output_name) - - 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(expected_image - image).max() < 1e-3 - - -class ModularGuiderTesterMixin: - 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" - - -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"] - - -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.""" - import json - - 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.""" - import json - - from diffusers import UNet2DConditionModel - - 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.""" - import json - - 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) diff --git a/tests/modular_pipelines/testing_utils/__init__.py b/tests/modular_pipelines/testing_utils/__init__.py new file mode 100644 index 000000000000..ce6ab1b8f216 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/__init__.py @@ -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", +] diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py new file mode 100644 index 000000000000..887efa4de0ab --- /dev/null +++ b/tests/modular_pipelines/testing_utils/common.py @@ -0,0 +1,366 @@ +# 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): + pipe = self.get_pipeline().to(torch_device) + + if "num_images_per_prompt" not in pipe.blocks.input_names: + pytest.mark.skip("Skipping test as `num_images_per_prompt` is not present in input names.") + + batch_sizes = [1, 2] + num_images_per_prompts = [1, 2] + + 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 diff --git a/tests/modular_pipelines/testing_utils/guider.py b/tests/modular_pipelines/testing_utils/guider.py new file mode 100644 index 000000000000..8c6a82d435d8 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/guider.py @@ -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" diff --git a/tests/modular_pipelines/testing_utils/loading.py b/tests/modular_pipelines/testing_utils/loading.py new file mode 100644 index 000000000000..64a9dd345008 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/loading.py @@ -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 diff --git a/tests/modular_pipelines/testing_utils/memory.py b/tests/modular_pipelines/testing_utils/memory.py new file mode 100644 index 000000000000..1185567a5ed6 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/memory.py @@ -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() diff --git a/tests/modular_pipelines/testing_utils/utils.py b/tests/modular_pipelines/testing_utils/utils.py new file mode 100644 index 000000000000..32a69f3dd0d7 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/utils.py @@ -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 diff --git a/tests/modular_pipelines/testing_utils/workflow.py b/tests/modular_pipelines/testing_utils/workflow.py new file mode 100644 index 000000000000..9c4cb4196683 --- /dev/null +++ b/tests/modular_pipelines/testing_utils/workflow.py @@ -0,0 +1,172 @@ +# 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(): + workflow_blocks = 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}" + ) + + 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) diff --git a/tests/modular_pipelines/utils.py b/tests/modular_pipelines/utils.py deleted file mode 100644 index 21d1a74f7f00..000000000000 --- a/tests/modular_pipelines/utils.py +++ /dev/null @@ -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() diff --git a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py index 1141633f2b7d..d35c21455ba9 100644 --- a/tests/modular_pipelines/wan/test_modular_pipeline_wan.py +++ b/tests/modular_pipelines/wan/test_modular_pipeline_wan.py @@ -17,14 +17,19 @@ 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 @@ def get_dummy_inputs(self, seed=0): } 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 diff --git a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py index 49de376024d6..b67be5bc88e9 100644 --- a/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py +++ b/tests/modular_pipelines/z_image/test_modular_pipeline_z_image.py @@ -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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 From aaecb78784857659ffdc49d1c4cf6197f8c80651 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 12 Aug 2026 07:29:07 +0000 Subject: [PATCH 2/4] ltx2 --- .../ltx2/test_modular_pipeline_ltx2.py | 212 ++++++++++++------ 1 file changed, 139 insertions(+), 73 deletions(-) diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py index dafb59768b1e..e73944f4534a 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py @@ -22,19 +22,88 @@ 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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_inference_batch_single_identical(self): 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 @@ def test_auto_duration_predicts_a_grid_valid_frame_count(self): 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 @@ def get_dummy_inputs(self, seed=0): 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 @@ def get_dummy_inputs(self, seed=0): 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 From 2a1bffb33230a7f1ed3726d0b540d123e519e79b Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 14 Aug 2026 10:33:04 +0200 Subject: [PATCH 3/4] cover the remaining tests --- .../ltx2/test_modular_pipeline_ltx25.py | 214 ++++++++++++------ .../test_modular_pipeline_minimax_music3.py | 101 +++++---- .../test_modular_pipeline_wan_animate_2.py | 64 +++++- 3 files changed, 261 insertions(+), 118 deletions(-) diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index eca4765e9dcd..69080095c627 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py @@ -22,19 +22,88 @@ 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 @@ def get_dummy_inputs(self, seed=0): "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 @@ def test_inference_batch_single_identical(self): 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 @@ def test_auto_duration_predicts_a_grid_valid_frame_count(self): 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 @@ def get_dummy_inputs(self, seed=0): 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 @@ def get_dummy_inputs(self, seed=0): 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 diff --git a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py index 0197a6d34809..923c4704dd7b 100644 --- a/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py +++ b/tests/modular_pipelines/minimax_music3/test_modular_pipeline_minimax_music3.py @@ -21,15 +21,30 @@ 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 @@ def test_condition_encoder_output_shape(self): 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 @@ def get_dummy_inputs(self, seed=0): "output_type": "pt", } + +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): - pytest.skip("MiniMax Music 3 generates a single waveform per call; `prompt` and `lyrics` are single strings.") + 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 diff --git a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py index 44e636de0e94..057bf11af5d5 100644 --- a/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py +++ b/tests/modular_pipelines/wan_animate_2/test_modular_pipeline_wan_animate_2.py @@ -24,7 +24,14 @@ 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 @@ } -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 @@ def get_dummy_inputs(self, seed=0): } 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 @@ def test_inference_batch_single_identical(self): 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 @@ def test_guider_cfg(self): 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 From 72e02fcc72e92befc238552754e3ac2993ee6a9f Mon Sep 17 00:00:00 2001 From: sayakpaul Date: Fri, 14 Aug 2026 11:18:51 +0200 Subject: [PATCH 4/4] fixes --- .../qwen/test_modular_pipeline_qwenimage.py | 6 +++++- tests/modular_pipelines/testing_utils/common.py | 7 ++----- tests/modular_pipelines/testing_utils/workflow.py | 13 ++++++++++++- 3 files changed, 19 insertions(+), 7 deletions(-) diff --git a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py index 979c06fec2ed..26a1cba8b30a 100644 --- a/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py +++ b/tests/modular_pipelines/qwen/test_modular_pipeline_qwenimage.py @@ -240,7 +240,11 @@ def get_dummy_inputs(self): class TestQwenImageEditModularPipelineFast(QwenImageEditModularPipelineTesterConfig, ModularPipelineTesterMixin): - pass + 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): diff --git a/tests/modular_pipelines/testing_utils/common.py b/tests/modular_pipelines/testing_utils/common.py index 887efa4de0ab..47614fd51005 100644 --- a/tests/modular_pipelines/testing_utils/common.py +++ b/tests/modular_pipelines/testing_utils/common.py @@ -344,14 +344,11 @@ def test_inference_is_not_nan_cpu(self): 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): + 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.mark.skip("Skipping test as `num_images_per_prompt` is not present in input names.") - - batch_sizes = [1, 2] - num_images_per_prompts = [1, 2] + 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: diff --git a/tests/modular_pipelines/testing_utils/workflow.py b/tests/modular_pipelines/testing_utils/workflow.py index 9c4cb4196683..70c4f6b84b2b 100644 --- a/tests/modular_pipelines/testing_utils/workflow.py +++ b/tests/modular_pipelines/testing_utils/workflow.py @@ -80,7 +80,8 @@ def test_workflow_defaults(self): blocks = self.pipeline_blocks_class() for workflow_name, expected_defaults in self.expected_workflow_defaults.items(): - workflow_blocks = blocks.get_workflow(workflow_name) + # 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 @@ -125,6 +126,16 @@ def test_workflow_defaults(self): 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: