Merge branch 'main' into t5-fixes-ci

Author: sayakpaul

.ai/testing.md

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

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

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

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

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