diff --git a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py b/tests/modular_pipelines/anima/test_modular_pipeline_anima.py index d1c4387331d8..c35def821d10 100644 --- a/tests/modular_pipelines/anima/test_modular_pipeline_anima.py +++ b/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 { diff --git a/tests/modular_pipelines/test_modular_pipelines_common.py b/tests/modular_pipelines/test_modular_pipelines_common.py index 5c51295a7e79..fadaef9d3e58 100644 --- a/tests/modular_pipelines/test_modular_pipelines_common.py +++ b/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): diff --git a/tests/modular_pipelines/utils.py b/tests/modular_pipelines/utils.py new file mode 100644 index 000000000000..21d1a74f7f00 --- /dev/null +++ b/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()