[tests] tighten tests for unload_components()

tests/modular_pipelines/anima/test_modular_pipeline_anima.py

@@ -17,17 +17,13 @@
 
 import numpy as np
 import PIL.Image
-import pytest
 import torch
-from transformers import Qwen2Tokenizer, Qwen3Config, Qwen3Model, T5TokenizerFast
 
 from diffusers import (
     AnimaAutoBlocks,
     AnimaModularPipeline,
     AnimaTextConditioner,
-    AutoencoderKLQwenImage,
     CosmosTransformer3DModel,
-    FlowMatchEulerDiscreteScheduler,
 )
 
 from ...testing_utils import enable_full_determinism, require_peft_backend
@@ -37,12 +33,6 @@
 enable_full_determinism()
 
 
-# TODO: `hf-internal-testing/tiny-anima-modular-pipe` carries no `modular_model_index.json`, so the three tests that
-# build a pipeline from it cannot run. Anima assembles its dummy components in `get_pipeline` instead, which is why
-# nothing noticed until those tests arrived. Publish the tiny repository and drop the overrides below.
-_NO_MODULAR_REPO = "TODO: no tiny Anima modular repository to load from yet."
-
-
 ANIMA_TEXT2IMAGE_WORKFLOWS = {
     "text2image": [
         ("text_encoder", "AnimaTextEncoderStep"),
@@ -72,76 +62,6 @@
 }
 
 
-def get_dummy_components():
-    torch.manual_seed(0)
-    transformer = CosmosTransformer3DModel(
-        in_channels=4,
-        out_channels=4,
-        num_attention_heads=2,
-        attention_head_dim=16,
-        num_layers=2,
-        mlp_ratio=2,
-        text_embed_dim=16,
-        adaln_lora_dim=4,
-        max_size=(4, 32, 32),
-        patch_size=(1, 2, 2),
-        rope_scale=(1.0, 4.0, 4.0),
-        concat_padding_mask=True,
-        extra_pos_embed_type=None,
-    )
-
-    torch.manual_seed(0)
-    vae = AutoencoderKLQwenImage(
-        base_dim=24,
-        z_dim=4,
-        dim_mult=[1, 2, 4],
-        num_res_blocks=1,
-        temperal_downsample=[False, True],
-        latents_mean=[0.0] * 4,
-        latents_std=[1.0] * 4,
-    )
-
-    torch.manual_seed(0)
-    text_conditioner = AnimaTextConditioner(
-        source_dim=16,
-        target_dim=16,
-        model_dim=16,
-        num_layers=2,
-        num_attention_heads=4,
-        target_vocab_size=32128,
-        min_sequence_length=16,
-    )
-
-    torch.manual_seed(0)
-    text_encoder_config = Qwen3Config(
-        vocab_size=152064,
-        hidden_size=16,
-        intermediate_size=32,
-        num_hidden_layers=2,
-        num_attention_heads=4,
-        num_key_value_heads=2,
-        max_position_embeddings=128,
-        rms_norm_eps=1e-6,
-        rope_theta=1000000.0,
-        head_dim=4,
-        attention_bias=False,
-    )
-    text_encoder = Qwen3Model(text_encoder_config).eval()
-    tokenizer = Qwen2Tokenizer.from_pretrained("hf-internal-testing/tiny-random-Qwen2VLForConditionalGeneration")
-    t5_tokenizer = T5TokenizerFast.from_pretrained("hf-internal-testing/tiny-random-t5")
-    scheduler = FlowMatchEulerDiscreteScheduler(shift=3.0)
-
-    return {
-        "transformer": transformer,
-        "vae": vae,
-        "scheduler": scheduler,
-        "text_encoder": text_encoder,
-        "tokenizer": tokenizer,
-        "t5_tokenizer": t5_tokenizer,
-        "text_conditioner": text_conditioner,
-    }
-
-
 def get_dummy_image(height=32, width=32):
     image_array = np.random.randint(0, 256, (height, width, 3), dtype=np.uint8)
     return PIL.Image.fromarray(image_array)
@@ -184,25 +104,6 @@ class TestAnimaModularPipelineFast(ModularPipelineTesterMixin, ModularGuiderTest
     batch_params = frozenset(["prompt", "negative_prompt"])
     expected_workflow_blocks = ANIMA_TEXT2IMAGE_WORKFLOWS
 
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_from_pretrained_workflow(self):
-        pass
-
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_load_components_workflow(self):
-        pass
-
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_unload_components(self):
-        pass
-
-    def get_pipeline(self, components_manager=None, dtype=torch.float32):
-        pipe = self.pipeline_blocks_class().init_pipeline(components_manager=components_manager)
-        pipe.update_components(**get_dummy_components())
-        pipe.to(dtype=dtype)
-        pipe.set_progress_bar_config(disable=None)
-        return pipe
-
     def get_dummy_inputs(self, seed=0):
         generator = torch.Generator(device="cpu").manual_seed(seed)
         return {
@@ -280,25 +181,6 @@ class TestAnimaImg2ImgModularPipelineFast(ModularPipelineTesterMixin):
     batch_params = frozenset(["prompt", "negative_prompt"])
     expected_workflow_blocks = ANIMA_IMG2IMG_WORKFLOWS
 
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_from_pretrained_workflow(self):
-        pass
-
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_load_components_workflow(self):
-        pass
-
-    @pytest.mark.skip(reason=_NO_MODULAR_REPO)
-    def test_unload_components(self):
-        pass
-
-    def get_pipeline(self, components_manager=None, torch_dtype=torch.float32):
-        pipe = self.pipeline_blocks_class().init_pipeline(components_manager=components_manager)
-        pipe.update_components(**get_dummy_components())
-        pipe.to(dtype=torch_dtype)
-        pipe.set_progress_bar_config(disable=None)
-        return pipe
-
     def get_dummy_inputs(self, seed=0):
         generator = torch.Generator(device="cpu").manual_seed(seed)
         return {

tests/modular_pipelines/test_modular_pipelines_common.py

@@ -1,6 +1,7 @@
 import gc
 import json
 import os
+import weakref
 from typing import Callable
 
 import pytest
@@ -20,11 +21,13 @@
 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):
@@ -628,6 +631,11 @@ def test_unload_components(self):
 
         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)
@@ -642,6 +650,101 @@ def test_unload_components(self):
         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):

tests/modular_pipelines/utils.py

@@ -0,0 +1,13 @@
+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()