From 5352999e14c4b9945553836ef0c58c18837603ce Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 11 Mar 2026 10:39:09 +0530 Subject: [PATCH 01/13] update --- .../test_models_transformer_z_image.py | 194 ++++++++++-------- 1 file changed, 105 insertions(+), 89 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 79054019f2d2..5748f5a52a98 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -1,4 +1,3 @@ -# coding=utf-8 # Copyright 2025 HuggingFace Inc. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,16 +12,21 @@ # See the License for the specific language governing permissions and # limitations under the License. -import gc import os -import unittest +import pytest import torch from diffusers import ZImageTransformer2DModel from ...testing_utils import IS_GITHUB_ACTIONS, torch_device -from ..test_modeling_common import ModelTesterMixin, TorchCompileTesterMixin +from ..testing_utils import ( + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, +) # Z-Image requires torch.use_deterministic_algorithms(False) due to complex64 RoPE operations @@ -36,44 +40,39 @@ torch.backends.cuda.matmul.allow_tf32 = False -@unittest.skipIf( +pytestmark = pytest.mark.skipif( IS_GITHUB_ACTIONS, reason="Skipping test-suite inside the CI because the model has `torch.empty()` inside of it during init and we don't have a clear way to override it in the modeling tests.", ) -class ZImageTransformerTests(ModelTesterMixin, unittest.TestCase): - model_class = ZImageTransformer2DModel - main_input_name = "x" - # We override the items here because the transformer under consideration is small. - model_split_percents = [0.9, 0.9, 0.9] - - def prepare_dummy_input(self, height=16, width=16): - batch_size = 1 - num_channels = 16 - embedding_dim = 16 - sequence_length = 16 - - hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)] - encoder_hidden_states = [ - torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size) - ] - timestep = torch.tensor([0.0]).to(torch_device) - return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep} +class ZImageTransformerTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - return self.prepare_dummy_input() + def model_class(self): + return ZImageTransformer2DModel @property - def input_shape(self): + def output_shape(self) -> tuple[int, ...]: return (4, 32, 32) @property - def output_shape(self): + def input_shape(self) -> tuple[int, ...]: return (4, 32, 32) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + @property + def model_split_percents(self) -> list: + return [0.9, 0.9, 0.9] + + @property + def main_input_name(self) -> str: + return "x" + + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | float]: + return { "all_patch_size": (2,), "all_f_patch_size": (1,), "in_channels": 16, @@ -89,83 +88,100 @@ def prepare_init_args_and_inputs_for_common(self): "axes_dims": [8, 4, 4], "axes_lens": [256, 32, 32], } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - - def setUp(self): - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) - - def tearDown(self): - super().tearDown() - gc.collect() - if torch.cuda.is_available(): - torch.cuda.empty_cache() - torch.cuda.synchronize() - torch.manual_seed(0) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(0) - def test_gradient_checkpointing_is_applied(self): - expected_set = {"ZImageTransformer2DModel"} - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + def get_dummy_inputs(self) -> dict[str, torch.Tensor | list]: + batch_size = 1 + num_channels = 16 + embedding_dim = 16 + sequence_length = 16 + height = 16 + width = 16 - @unittest.skip("Test is not supported for handling main inputs that are lists.") - def test_training(self): - super().test_training() + hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)] + encoder_hidden_states = [ + torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size) + ] + timestep = torch.tensor([0.0]).to(torch_device) + + return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep} + + +class TestZImageTransformer(ZImageTransformerTesterConfig, ModelTesterMixin): + """Core model tests for Z-Image Transformer.""" + + @pytest.mark.skip("Test is not supported for handling main inputs that are lists.") + def test_outputs_equivalence(self, atol=1e-5, rtol=0): + pass - @unittest.skip("Test is not supported for handling main inputs that are lists.") - def test_ema_training(self): - super().test_ema_training() - @unittest.skip("Test is not supported for handling main inputs that are lists.") - def test_effective_gradient_checkpointing(self): - super().test_effective_gradient_checkpointing() +class TestZImageTransformerMemory(ZImageTransformerTesterConfig, MemoryTesterMixin): + """Memory optimization tests for Z-Image Transformer.""" - @unittest.skip( - "Test needs to be revisited. But we need to ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices." + @pytest.mark.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") + def test_group_offloading(self, record_stream, atol=1e-5, rtol=0): + pass + + @pytest.mark.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") + def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type, atol=1e-5, rtol=0): + pass + + @pytest.mark.skip( + "Test needs to be revisited. Ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices." ) def test_layerwise_casting_training(self): - super().test_layerwise_casting_training() + pass - @unittest.skip("Test is not supported for handling main inputs that are lists.") - def test_outputs_equivalence(self): - super().test_outputs_equivalence() - @unittest.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") - def test_group_offloading(self): - super().test_group_offloading() +class TestZImageTransformerTraining(ZImageTransformerTesterConfig, TrainingTesterMixin): + """Training tests for Z-Image Transformer.""" - @unittest.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") - def test_group_offloading_with_disk(self): - super().test_group_offloading_with_disk() + def test_gradient_checkpointing_is_applied(self): + super().test_gradient_checkpointing_is_applied(expected_set={"ZImageTransformer2DModel"}) + @pytest.mark.skip("Test is not supported for handling main inputs that are lists.") + def test_training(self): + pass -class ZImageTransformerCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = ZImageTransformer2DModel - different_shapes_for_compilation = [(4, 4), (4, 8), (8, 8)] + @pytest.mark.skip("Test is not supported for handling main inputs that are lists.") + def test_training_with_ema(self): + pass - def prepare_init_args_and_inputs_for_common(self): - return ZImageTransformerTests().prepare_init_args_and_inputs_for_common() + @pytest.mark.skip("Test is not supported for handling main inputs that are lists.") + def test_gradient_checkpointing_equivalence(self, loss_tolerance=1e-5, param_grad_tol=5e-5, skip=None): + pass - def prepare_dummy_input(self, height, width): - return ZImageTransformerTests().prepare_dummy_input(height=height, width=width) - @unittest.skip( - "The repeated block in this model is ZImageTransformerBlock, which is used for noise_refiner, context_refiner, and layers. As a consequence of this, the inputs recorded for the block would vary during compilation and full compilation with fullgraph=True would trigger recompilation at least thrice." +class TestZImageTransformerCompile(ZImageTransformerTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for Z-Image Transformer.""" + + @property + def different_shapes_for_compilation(self): + return [(4, 4), (4, 8), (8, 8)] + + def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict[str, torch.Tensor | list]: + batch_size = 1 + num_channels = 16 + embedding_dim = 16 + sequence_length = 16 + + hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)] + encoder_hidden_states = [ + torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size) + ] + timestep = torch.tensor([0.0]).to(torch_device) + + return {"x": hidden_states, "cap_feats": encoder_hidden_states, "t": timestep} + + @pytest.mark.skip( + "The repeated block in this model is ZImageTransformerBlock, which is used for noise_refiner, context_refiner, and layers. The inputs recorded for the block would vary during compilation and full compilation with fullgraph=True would trigger recompilation at least thrice." ) def test_torch_compile_recompilation_and_graph_break(self): - super().test_torch_compile_recompilation_and_graph_break() + pass - @unittest.skip("Fullgraph AoT is broken") - def test_compile_works_with_aot(self): - super().test_compile_works_with_aot() + @pytest.mark.skip("Fullgraph AoT is broken") + def test_compile_works_with_aot(self, tmp_path): + pass - @unittest.skip("Fullgraph is broken") + @pytest.mark.skip("Fullgraph is broken") def test_compile_on_different_shapes(self): - super().test_compile_on_different_shapes() + pass From d15761686a8651cf2b309da06a59eaac6310da53 Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 11 Mar 2026 11:23:01 +0530 Subject: [PATCH 02/13] update --- .../test_models_transformer_z_image.py | 116 +++++++++++++++++- 1 file changed, 110 insertions(+), 6 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 5748f5a52a98..a2139bdad569 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -19,7 +19,7 @@ from diffusers import ZImageTransformer2DModel -from ...testing_utils import IS_GITHUB_ACTIONS, torch_device +from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( BaseModelTesterConfig, MemoryTesterMixin, @@ -40,10 +40,9 @@ torch.backends.cuda.matmul.allow_tf32 = False -pytestmark = pytest.mark.skipif( - IS_GITHUB_ACTIONS, - reason="Skipping test-suite inside the CI because the model has `torch.empty()` inside of it during init and we don't have a clear way to override it in the modeling tests.", -) +def _concat_list_output(output): + """Model output `sample` is a list of tensors. Concatenate them for comparison.""" + return torch.cat([t.flatten() for t in output]) class ZImageTransformerTesterConfig(BaseModelTesterConfig): @@ -109,10 +108,115 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor | list]: class TestZImageTransformer(ZImageTransformerTesterConfig, ModelTesterMixin): """Core model tests for Z-Image Transformer.""" - @pytest.mark.skip("Test is not supported for handling main inputs that are lists.") + @torch.no_grad() + def test_determinism(self, atol=1e-5, rtol=0): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + inputs_dict = self.get_dummy_inputs() + first = _concat_list_output(model(**inputs_dict, return_dict=False)[0]) + second = _concat_list_output(model(**inputs_dict, return_dict=False)[0]) + + mask = ~(torch.isnan(first) | torch.isnan(second)) + assert_tensors_close( + first[mask], second[mask], atol=atol, rtol=rtol, msg="Model outputs are not deterministic" + ) + + def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + model.save_pretrained(tmp_path) + new_model = self.model_class.from_pretrained(tmp_path) + new_model.to(torch_device) + + for param_name in model.state_dict().keys(): + param_1 = model.state_dict()[param_name] + param_2 = new_model.state_dict()[param_name] + assert param_1.shape == param_2.shape + + inputs_dict = self.get_dummy_inputs() + image = _concat_list_output(model(**inputs_dict, return_dict=False)[0]) + new_image = _concat_list_output(new_model(**inputs_dict, return_dict=False)[0]) + + assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") + + @torch.no_grad() + def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0): + model = self.model_class(**self.get_init_dict()) + model.to(torch_device) + model.eval() + + model.save_pretrained(tmp_path, variant="fp16") + new_model = self.model_class.from_pretrained(tmp_path, variant="fp16") + + with pytest.raises(OSError) as exc_info: + self.model_class.from_pretrained(tmp_path) + + assert "Error no file named diffusion_pytorch_model.bin found in directory" in str(exc_info.value) + + new_model.to(torch_device) + + inputs_dict = self.get_dummy_inputs() + image = _concat_list_output(model(**inputs_dict, return_dict=False)[0]) + new_image = _concat_list_output(new_model(**inputs_dict, return_dict=False)[0]) + + assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") + + @pytest.mark.skip("Model output `sample` is a list of tensors, not a single tensor.") def test_outputs_equivalence(self, atol=1e-5, rtol=0): pass + def test_sharded_checkpoints_with_parallel_loading(self, tmp_path, atol=1e-5, rtol=0): + from diffusers.utils import SAFE_WEIGHTS_INDEX_NAME, constants + + from ..testing_utils.common import calculate_expected_num_shards, compute_module_persistent_sizes + + torch.manual_seed(0) + config = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**config).eval() + model = model.to(torch_device) + + base_output = _concat_list_output(model(**inputs_dict, return_dict=False)[0]) + + model_size = compute_module_persistent_sizes(model)[""] + max_shard_size = int((model_size * 0.75) / (2**10)) + + original_parallel_loading = constants.HF_ENABLE_PARALLEL_LOADING + original_parallel_workers = getattr(constants, "HF_PARALLEL_WORKERS", None) + + try: + model.cpu().save_pretrained(tmp_path, max_shard_size=f"{max_shard_size}KB") + assert os.path.exists(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)) + + expected_num_shards = calculate_expected_num_shards(os.path.join(tmp_path, SAFE_WEIGHTS_INDEX_NAME)) + actual_num_shards = len([file for file in os.listdir(tmp_path) if file.endswith(".safetensors")]) + assert actual_num_shards == expected_num_shards + + constants.HF_ENABLE_PARALLEL_LOADING = False + self.model_class.from_pretrained(tmp_path).eval().to(torch_device) + + constants.HF_ENABLE_PARALLEL_LOADING = True + constants.DEFAULT_HF_PARALLEL_LOADING_WORKERS = 2 + + torch.manual_seed(0) + model_parallel = self.model_class.from_pretrained(tmp_path).eval() + model_parallel = model_parallel.to(torch_device) + + output_parallel = _concat_list_output(model_parallel(**inputs_dict, return_dict=False)[0]) + + assert_tensors_close( + base_output, output_parallel, atol=atol, rtol=rtol, msg="Output should match with parallel loading" + ) + finally: + constants.HF_ENABLE_PARALLEL_LOADING = original_parallel_loading + if original_parallel_workers is not None: + constants.HF_PARALLEL_WORKERS = original_parallel_workers + class TestZImageTransformerMemory(ZImageTransformerTesterConfig, MemoryTesterMixin): """Memory optimization tests for Z-Image Transformer.""" From c15472d2c4a5668962dab20a076e446840c97a9a Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 11 Mar 2026 11:32:28 +0530 Subject: [PATCH 03/13] update --- .../test_models_transformer_z_image.py | 29 ++++++++++--------- 1 file changed, 15 insertions(+), 14 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index a2139bdad569..8df539a037ba 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -18,6 +18,7 @@ import torch from diffusers import ZImageTransformer2DModel +from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( @@ -70,7 +71,7 @@ def main_input_name(self) -> str: def generator(self): return torch.Generator("cpu").manual_seed(0) - def get_init_dict(self) -> dict[str, int | list[int] | tuple | str | bool | float]: + def get_init_dict(self): return { "all_patch_size": (2,), "all_f_patch_size": (1,), @@ -96,9 +97,13 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor | list]: height = 16 width = 16 - hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)] + hidden_states = [ + randn_tensor((num_channels, 1, height, width), generator=self.generator, device=torch_device) + for _ in range(batch_size) + ] encoder_hidden_states = [ - torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size) + randn_tensor((sequence_length, embedding_dim), generator=self.generator, device=torch_device) + for _ in range(batch_size) ] timestep = torch.tensor([0.0]).to(torch_device) @@ -221,16 +226,8 @@ def test_sharded_checkpoints_with_parallel_loading(self, tmp_path, atol=1e-5, rt class TestZImageTransformerMemory(ZImageTransformerTesterConfig, MemoryTesterMixin): """Memory optimization tests for Z-Image Transformer.""" - @pytest.mark.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") - def test_group_offloading(self, record_stream, atol=1e-5, rtol=0): - pass - - @pytest.mark.skip("Test will pass if we change to deterministic values instead of empty in the DiT.") - def test_group_offloading_with_disk(self, tmp_path, record_stream, offload_type, atol=1e-5, rtol=0): - pass - @pytest.mark.skip( - "Test needs to be revisited. Ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices." + "Ensure `x_pad_token` and `cap_pad_token` are cast to the same dtype as the destination tensor before they are assigned to the padding indices." ) def test_layerwise_casting_training(self): pass @@ -268,9 +265,13 @@ def get_dummy_inputs(self, height: int = 16, width: int = 16) -> dict[str, torch embedding_dim = 16 sequence_length = 16 - hidden_states = [torch.randn((num_channels, 1, height, width)).to(torch_device) for _ in range(batch_size)] + hidden_states = [ + randn_tensor((num_channels, 1, height, width), generator=self.generator, device=torch_device) + for _ in range(batch_size) + ] encoder_hidden_states = [ - torch.randn((sequence_length, embedding_dim)).to(torch_device) for _ in range(batch_size) + randn_tensor((sequence_length, embedding_dim), generator=self.generator, device=torch_device) + for _ in range(batch_size) ] timestep = torch.tensor([0.0]).to(torch_device) From 73b23dc92e8182b8448dd695ea118b42db3845c5 Mon Sep 17 00:00:00 2001 From: DN6 Date: Wed, 11 Mar 2026 14:02:31 +0530 Subject: [PATCH 04/13] update --- .../test_models_transformer_z_image.py | 19 +++++++++++++++++++ 1 file changed, 19 insertions(+) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index 8df539a037ba..c2e7990e0424 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -23,6 +23,7 @@ from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( BaseModelTesterConfig, + LoraTesterMixin, MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, @@ -252,6 +253,24 @@ def test_gradient_checkpointing_equivalence(self, loss_tolerance=1e-5, param_gra pass +class TestZImageTransformerLoRA(ZImageTransformerTesterConfig, LoraTesterMixin): + """LoRA adapter tests for Z-Image Transformer.""" + + @pytest.mark.skip("Model output `sample` is a list of tensors, not a single tensor.") + def test_save_load_lora_adapter(self, tmp_path, rank=4, lora_alpha=4, use_dora=False, atol=1e-4, rtol=1e-4): + pass + + +# TODO: Add pretrained_model_name_or_path once a tiny Z-Image model is available on the Hub +# class TestZImageTransformerBitsAndBytes(ZImageTransformerTesterConfig, BitsAndBytesTesterMixin): +# """BitsAndBytes quantization tests for Z-Image Transformer.""" + + +# TODO: Add pretrained_model_name_or_path once a tiny Z-Image model is available on the Hub +# class TestZImageTransformerTorchAo(ZImageTransformerTesterConfig, TorchAoTesterMixin): +# """TorchAo quantization tests for Z-Image Transformer.""" + + class TestZImageTransformerCompile(ZImageTransformerTesterConfig, TorchCompileTesterMixin): """Torch compile tests for Z-Image Transformer.""" From 3ff29b20df2f9a125ed81933a3cbea93cd72f8f0 Mon Sep 17 00:00:00 2001 From: Dhruv Nair Date: Tue, 9 Jun 2026 18:00:23 +0530 Subject: [PATCH 05/13] [CI] Refactor SD3 Transformer Test (#13340) * update * update --------- Co-authored-by: Sayak Paul --- .../test_models_transformer_sd3.py | 230 ++++++++++-------- 1 file changed, 128 insertions(+), 102 deletions(-) diff --git a/tests/models/transformers/test_models_transformer_sd3.py b/tests/models/transformers/test_models_transformer_sd3.py index c4ee7017a380..50590fab8be1 100644 --- a/tests/models/transformers/test_models_transformer_sd3.py +++ b/tests/models/transformers/test_models_transformer_sd3.py @@ -13,58 +13,63 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - import torch from diffusers import SD3Transformer2DModel -from diffusers.utils.import_utils import is_xformers_available - -from ...testing_utils import ( - enable_full_determinism, - torch_device, +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + BaseModelTesterConfig, + BitsAndBytesTesterMixin, + ModelTesterMixin, + TorchAoTesterMixin, + TorchCompileTesterMixin, + TrainingTesterMixin, ) -from ..test_modeling_common import ModelTesterMixin enable_full_determinism() -class SD3TransformerTests(ModelTesterMixin, unittest.TestCase): - model_class = SD3Transformer2DModel - main_input_name = "hidden_states" - model_split_percents = [0.8, 0.8, 0.9] +# ======================== SD3 Transformer ======================== + +class SD3TransformerTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 2 - num_channels = 4 - height = width = embedding_dim = 32 - pooled_embedding_dim = embedding_dim * 2 - sequence_length = 154 + def model_class(self): + return SD3Transformer2DModel + + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-sd3-pipe" - hidden_states = torch.randn((batch_size, num_channels, height, width)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) - pooled_prompt_embeds = torch.randn((batch_size, pooled_embedding_dim)).to(torch_device) - timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device) + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "timestep": timestep, - } + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def model_split_percents(self) -> list: + return [0.8, 0.8, 0.9] @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 32, 32) @property - def output_shape(self): + def input_shape(self) -> tuple: return (4, 32, 32) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { "sample_size": 32, "patch_size": 1, "in_channels": 4, @@ -79,67 +84,79 @@ def prepare_init_args_and_inputs_for_common(self): "dual_attention_layers": (), "qk_norm": None, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_enable_works(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) + def get_dummy_inputs(self, batch_size: int = 2) -> dict[str, torch.Tensor]: + num_channels = 4 + height = width = embedding_dim = 32 + pooled_embedding_dim = embedding_dim * 2 + sequence_length = 154 + + return { + "hidden_states": randn_tensor( + (batch_size, num_channels, height, width), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device + ), + "pooled_projections": randn_tensor( + (batch_size, pooled_embedding_dim), generator=self.generator, device=torch_device + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + } - model.enable_xformers_memory_efficient_attention() - assert model.transformer_blocks[0].attn.processor.__class__.__name__ == "XFormersJointAttnProcessor", ( - "xformers is not enabled" - ) +class TestSD3Transformer(SD3TransformerTesterConfig, ModelTesterMixin): + pass - @unittest.skip("SD3Transformer2DModel uses a dedicated attention processor. This test doesn't apply") - def test_set_attn_processor_for_determinism(self): - pass +class TestSD3TransformerTraining(SD3TransformerTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): expected_set = {"SD3Transformer2DModel"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) -class SD35TransformerTests(ModelTesterMixin, unittest.TestCase): - model_class = SD3Transformer2DModel - main_input_name = "hidden_states" - model_split_percents = [0.8, 0.8, 0.9] +class TestSD3TransformerCompile(SD3TransformerTesterConfig, TorchCompileTesterMixin): + pass + + +# ======================== SD3.5 Transformer ======================== + +class SD35TransformerTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 2 - num_channels = 4 - height = width = embedding_dim = 32 - pooled_embedding_dim = embedding_dim * 2 - sequence_length = 154 + def model_class(self): + return SD3Transformer2DModel - hidden_states = torch.randn((batch_size, num_channels, height, width)).to(torch_device) - encoder_hidden_states = torch.randn((batch_size, sequence_length, embedding_dim)).to(torch_device) - pooled_prompt_embeds = torch.randn((batch_size, pooled_embedding_dim)).to(torch_device) - timestep = torch.randint(0, 1000, size=(batch_size,)).to(torch_device) + @property + def pretrained_model_name_or_path(self): + return "hf-internal-testing/tiny-sd35-pipe" - return { - "hidden_states": hidden_states, - "encoder_hidden_states": encoder_hidden_states, - "pooled_projections": pooled_prompt_embeds, - "timestep": timestep, - } + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + @property + def main_input_name(self) -> str: + return "hidden_states" + + @property + def model_split_percents(self) -> list: + return [0.8, 0.8, 0.9] @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 32, 32) @property - def output_shape(self): + def input_shape(self) -> tuple: return (4, 32, 32) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) + + def get_init_dict(self) -> dict: + return { "sample_size": 32, "patch_size": 1, "in_channels": 4, @@ -154,47 +171,56 @@ def prepare_init_args_and_inputs_for_common(self): "dual_attention_layers": (0,), "qk_norm": "rms_norm", } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_enable_works(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.enable_xformers_memory_efficient_attention() - - assert model.transformer_blocks[0].attn.processor.__class__.__name__ == "XFormersJointAttnProcessor", ( - "xformers is not enabled" - ) + def get_dummy_inputs(self, batch_size: int = 2) -> dict[str, torch.Tensor]: + num_channels = 4 + height = width = embedding_dim = 32 + pooled_embedding_dim = embedding_dim * 2 + sequence_length = 154 - @unittest.skip("SD3Transformer2DModel uses a dedicated attention processor. This test doesn't apply") - def test_set_attn_processor_for_determinism(self): - pass + return { + "hidden_states": randn_tensor( + (batch_size, num_channels, height, width), generator=self.generator, device=torch_device + ), + "encoder_hidden_states": randn_tensor( + (batch_size, sequence_length, embedding_dim), generator=self.generator, device=torch_device + ), + "pooled_projections": randn_tensor( + (batch_size, pooled_embedding_dim), generator=self.generator, device=torch_device + ), + "timestep": torch.randint(0, 1000, size=(batch_size,), generator=self.generator).to(torch_device), + } - def test_gradient_checkpointing_is_applied(self): - expected_set = {"SD3Transformer2DModel"} - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) +class TestSD35Transformer(SD35TransformerTesterConfig, ModelTesterMixin): def test_skip_layers(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() model = self.model_class(**init_dict).to(torch_device) - # Forward pass without skipping layers output_full = model(**inputs_dict).sample - # Forward pass with skipping layers 0 (since there's only one layer in this test setup) inputs_dict_with_skip = inputs_dict.copy() inputs_dict_with_skip["skip_layers"] = [0] output_skip = model(**inputs_dict_with_skip).sample - # Check that the outputs are different - self.assertFalse( - torch.allclose(output_full, output_skip, atol=1e-5), "Outputs should differ when layers are skipped" - ) + assert not torch.allclose(output_full, output_skip, atol=1e-5), "Outputs should differ when layers are skipped" + assert output_full.shape == output_skip.shape, "Outputs should have the same shape" + + +class TestSD35TransformerTraining(SD35TransformerTesterConfig, TrainingTesterMixin): + def test_gradient_checkpointing_is_applied(self): + expected_set = {"SD3Transformer2DModel"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestSD35TransformerCompile(SD35TransformerTesterConfig, TorchCompileTesterMixin): + pass + + +class TestSD35TransformerBitsAndBytes(SD35TransformerTesterConfig, BitsAndBytesTesterMixin): + """BitsAndBytes quantization tests for SD3.5 Transformer.""" + - # Check that the outputs have the same shape - self.assertEqual(output_full.shape, output_skip.shape, "Outputs should have the same shape") +class TestSD35TransformerTorchAo(SD35TransformerTesterConfig, TorchAoTesterMixin): + """TorchAO quantization tests for SD3.5 Transformer.""" From b4381fa455ccfada1d548e93de69b3e5763ce2f4 Mon Sep 17 00:00:00 2001 From: Akshan Krithick <97239696+akshan-main@users.noreply.github.com> Date: Tue, 9 Jun 2026 20:13:46 -0700 Subject: [PATCH 06/13] refactor unet tests (3d_condition, motion, controlnetxs) (#13897) * refactor unet_3d_condition tests * refactor unet_motion tests * refactor unet_controlnetxs tests --- .../unets/test_models_unet_3d_condition.py | 171 ++++------- .../unets/test_models_unet_controlnetxs.py | 184 +++++------ tests/models/unets/test_models_unet_motion.py | 285 ++++++------------ 3 files changed, 228 insertions(+), 412 deletions(-) diff --git a/tests/models/unets/test_models_unet_3d_condition.py b/tests/models/unets/test_models_unet_3d_condition.py index f73e3461c38e..9ebdd659ee4b 100644 --- a/tests/models/unets/test_models_unet_3d_condition.py +++ b/tests/models/unets/test_models_unet_3d_condition.py @@ -13,52 +13,43 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import numpy as np import torch -from diffusers.models import ModelMixin, UNet3DConditionModel -from diffusers.utils import logging -from diffusers.utils.import_utils import is_xformers_available +from diffusers import UNet3DConditionModel +from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, floats_tensor, skip_mps, torch_device -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, +) enable_full_determinism() -logger = logging.get_logger(__name__) - - -@skip_mps -class UNet3DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet3DConditionModel - main_input_name = "sample" +class UNet3DConditionModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_channels = 4 - num_frames = 4 - sizes = (16, 16) + def model_class(self): + return UNet3DConditionModel - noise = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) - encoder_hidden_states = floats_tensor((batch_size, 4, 8)).to(torch_device) - - return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states} + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 4, 16, 16) @property - def output_shape(self): - return (4, 4, 16, 16) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "block_out_channels": (4, 8), "norm_num_groups": 4, "down_block_types": ( @@ -73,111 +64,57 @@ def prepare_init_args_and_inputs_for_common(self): "layers_per_block": 1, "sample_size": 16, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_enable_works(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.enable_xformers_memory_efficient_attention() + def get_dummy_inputs(self) -> dict: + batch_size = 4 + num_channels = 4 + num_frames = 4 + sizes = (16, 16) + noise = randn_tensor( + (batch_size, num_channels, num_frames, *sizes), generator=self.generator, device=torch_device + ) + timestep = torch.tensor([10], device=torch_device) + encoder_hidden_states = randn_tensor((batch_size, 4, 8), generator=self.generator, device=torch_device) + return {"sample": noise, "timestep": timestep, "encoder_hidden_states": encoder_hidden_states} - assert ( - model.mid_block.attentions[0].transformer_blocks[0].attn1.processor.__class__.__name__ - == "XFormersAttnProcessor" - ), "xformers is not enabled" - # Overriding to set `norm_num_groups` needs to be different for this model. +class TestUNet3DConditionModel(UNet3DConditionModelTesterConfig, ModelTesterMixin): + # Overridden because UNet3DConditionModel needs a different `norm_num_groups`. def test_forward_with_norm_groups(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() init_dict["block_out_channels"] = (32, 64) init_dict["norm_num_groups"] = 32 - - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() + model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): - output = model(**inputs_dict) + output = model(**self.get_dummy_inputs()).sample - if isinstance(output, dict): - output = output.sample + assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" - self.assertIsNotNone(output) - expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") - - # Overriding since the UNet3D outputs a different structure. - def test_determinism(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() + def test_feed_forward_chunking(self): + init_dict = self.get_init_dict() + init_dict["block_out_channels"] = (32, 64) + init_dict["norm_num_groups"] = 32 + model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): - # Warmup pass when using mps (see #372) - if torch_device == "mps" and isinstance(model, ModelMixin): - model(**self.dummy_input) - - first = model(**inputs_dict) - if isinstance(first, dict): - first = first.sample - - second = model(**inputs_dict) - if isinstance(second, dict): - second = second.sample - - out_1 = first.cpu().numpy() - out_2 = second.cpu().numpy() - out_1 = out_1[~np.isnan(out_1)] - out_2 = out_2[~np.isnan(out_2)] - max_diff = np.amax(np.abs(out_1 - out_2)) - self.assertLessEqual(max_diff, 1e-5) + output = model(**self.get_dummy_inputs())[0] - def test_model_attention_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - init_dict["block_out_channels"] = (16, 32) - init_dict["attention_head_dim"] = 8 - - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - model.set_attention_slice("auto") + model.enable_forward_chunking() with torch.no_grad(): - output = model(**inputs_dict) - assert output is not None + output_2 = model(**self.get_dummy_inputs())[0] - model.set_attention_slice("max") - with torch.no_grad(): - output = model(**inputs_dict) - assert output is not None + assert output.shape == output_2.shape, "Shape doesn't match" + assert (output - output_2).abs().max() < 1e-2 - model.set_attention_slice(2) - with torch.no_grad(): - output = model(**inputs_dict) - assert output is not None - def test_feed_forward_chunking(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - init_dict["block_out_channels"] = (32, 64) - init_dict["norm_num_groups"] = 32 +class TestUNet3DConditionModelTraining(UNet3DConditionModelTesterConfig, TrainingTesterMixin): + """Training tests for UNet3DConditionModel.""" - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - with torch.no_grad(): - output = model(**inputs_dict)[0] +class TestUNet3DConditionModelMemory(UNet3DConditionModelTesterConfig, MemoryTesterMixin): + """Memory optimization tests for UNet3DConditionModel.""" - model.enable_forward_chunking() - with torch.no_grad(): - output_2 = model(**inputs_dict)[0] - self.assertEqual(output.shape, output_2.shape, "Shape doesn't match") - assert np.abs(output.cpu() - output_2.cpu()).max() < 1e-2 +class TestUNet3DConditionModelAttention(UNet3DConditionModelTesterConfig, AttentionTesterMixin): + """Attention processor tests for UNet3DConditionModel.""" diff --git a/tests/models/unets/test_models_unet_controlnetxs.py b/tests/models/unets/test_models_unet_controlnetxs.py index 40773536df70..a12eb57228da 100644 --- a/tests/models/unets/test_models_unet_controlnetxs.py +++ b/tests/models/unets/test_models_unet_controlnetxs.py @@ -13,59 +13,44 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import numpy as np import torch from torch import nn from diffusers import ControlNetXSAdapter, UNet2DConditionModel, UNetControlNetXSModel -from diffusers.utils import logging - -from ...testing_utils import enable_full_determinism, floats_tensor, is_flaky, torch_device -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from diffusers.utils.torch_utils import randn_tensor +from ...testing_utils import enable_full_determinism, is_flaky, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, +) -logger = logging.get_logger(__name__) enable_full_determinism() -class UNetControlNetXSModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNetControlNetXSModel - main_input_name = "sample" - +class UNetControlNetXSModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_channels = 4 - sizes = (16, 16) - conditioning_image_size = (3, 32, 32) # size of additional, unprocessed image for control-conditioning - - noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) - encoder_hidden_states = floats_tensor((batch_size, 4, 8)).to(torch_device) - controlnet_cond = floats_tensor((batch_size, *conditioning_image_size)).to(torch_device) - conditioning_scale = 1 + def model_class(self): + return UNetControlNetXSModel - return { - "sample": noise, - "timestep": time_step, - "encoder_hidden_states": encoder_hidden_states, - "controlnet_cond": controlnet_cond, - "conditioning_scale": conditioning_scale, - } + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 16, 16) @property - def output_shape(self): - return (4, 16, 16) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "sample_size": 16, "down_block_types": ("DownBlock2D", "CrossAttnDownBlock2D"), "up_block_types": ("CrossAttnUpBlock2D", "UpBlock2D"), @@ -80,11 +65,27 @@ def prepare_init_args_and_inputs_for_common(self): "ctrl_max_norm_num_groups": 2, "ctrl_conditioning_embedding_out_channels": (2, 2), } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict: + batch_size = 4 + num_channels = 4 + sizes = (16, 16) + noise = randn_tensor((batch_size, num_channels, *sizes), generator=self.generator, device=torch_device) + timestep = torch.tensor([10], device=torch_device) + encoder_hidden_states = randn_tensor((batch_size, 4, 8), generator=self.generator, device=torch_device) + controlnet_cond = randn_tensor((batch_size, 3, 32, 32), generator=self.generator, device=torch_device) + return { + "sample": noise, + "timestep": timestep, + "encoder_hidden_states": encoder_hidden_states, + "controlnet_cond": controlnet_cond, + "conditioning_scale": 1, + } + + +class TestUNetControlNetXSModel(UNetControlNetXSModelTesterConfig, ModelTesterMixin): def get_dummy_unet(self): - """For some tests we also need the underlying UNet. For these, we'll build the UNetControlNetXSModel from the UNet and ControlNetXS-Adapter""" + """The underlying UNet, used to build the UNetControlNetXSModel from a UNet and a ControlNetXS-Adapter.""" return UNet2DConditionModel( block_out_channels=(4, 8), layers_per_block=2, @@ -99,8 +100,7 @@ def get_dummy_unet(self): ) def get_dummy_controlnet_from_unet(self, unet, **kwargs): - """For some tests we also need the underlying ControlNetXS-Adapter. For these, we'll build the UNetControlNetXSModel from the UNet and ControlNetXS-Adapter""" - # size_ratio and conditioning_embedding_out_channels chosen to keep model small + """The underlying ControlNetXS-Adapter. size_ratio and conditioning_embedding_out_channels keep the model small.""" return ControlNetXSAdapter.from_unet(unet, size_ratio=1, conditioning_embedding_out_channels=(2, 2), **kwargs) def test_from_unet(self): @@ -114,21 +114,11 @@ def assert_equal_weights(module, weight_dict_prefix): for param_name, param_value in module.named_parameters(): assert torch.equal(model_state_dict[weight_dict_prefix + "." + param_name], param_value) - # # check unet - # everything expect down,mid,up blocks - modules_from_unet = [ - "time_embedding", - "conv_in", - "conv_norm_out", - "conv_out", - ] + # check unet: everything except down, mid, up blocks + modules_from_unet = ["time_embedding", "conv_in", "conv_norm_out", "conv_out"] for p in modules_from_unet: assert_equal_weights(getattr(unet, p), "base_" + p) - optional_modules_from_unet = [ - "class_embedding", - "add_time_proj", - "add_embedding", - ] + optional_modules_from_unet = ["class_embedding", "add_time_proj", "add_embedding"] for p in optional_modules_from_unet: if hasattr(unet, p) and getattr(unet, p) is not None: assert_equal_weights(getattr(unet, p), "base_" + p) @@ -151,8 +141,7 @@ def assert_equal_weights(module, weight_dict_prefix): if hasattr(u, "upsamplers") and getattr(u, "upsamplers") is not None: assert_equal_weights(u.upsamplers[0], f"up_blocks.{i}.upsamplers") - # # check controlnet - # everything expect down,mid,up blocks + # check controlnet: everything except down, mid, up blocks modules_from_controlnet = { "controlnet_cond_embedding": "controlnet_cond_embedding", "conv_in": "ctrl_conv_in", @@ -161,7 +150,6 @@ def assert_equal_weights(module, weight_dict_prefix): optional_modules_from_controlnet = {"time_embedding": "ctrl_time_embedding"} for name_in_controlnet, name_in_unetcnxs in modules_from_controlnet.items(): assert_equal_weights(getattr(controlnet, name_in_controlnet), name_in_unetcnxs) - for name_in_controlnet, name_in_unetcnxs in optional_modules_from_controlnet.items(): if hasattr(controlnet, name_in_controlnet) and getattr(controlnet, name_in_controlnet) is not None: assert_equal_weights(getattr(controlnet, name_in_controlnet), name_in_unetcnxs) @@ -193,12 +181,10 @@ def assert_unfrozen(module): for p in module.parameters(): assert p.requires_grad - init_dict, _ = self.prepare_init_args_and_inputs_for_common() - model = UNetControlNetXSModel(**init_dict) + model = UNetControlNetXSModel(**self.get_init_dict()) model.freeze_unet_params() - # # check unet - # everything expect down,mid,up blocks + # check unet: everything except down, mid, up blocks modules_from_unet = [ model.base_time_embedding, model.base_conv_in, @@ -207,49 +193,39 @@ def assert_unfrozen(module): ] for m in modules_from_unet: assert_frozen(m) - - optional_modules_from_unet = [ - model.base_add_time_proj, - model.base_add_embedding, - ] + optional_modules_from_unet = [model.base_add_time_proj, model.base_add_embedding] for m in optional_modules_from_unet: if m is not None: assert_frozen(m) - # down blocks - for i, d in enumerate(model.down_blocks): + for d in model.down_blocks: assert_frozen(d.base_resnets) if isinstance(d.base_attentions, nn.ModuleList): # attentions can be list of Nones assert_frozen(d.base_attentions) if d.base_downsamplers is not None: assert_frozen(d.base_downsamplers) - # mid block assert_frozen(model.mid_block.base_midblock) - # up blocks - for i, u in enumerate(model.up_blocks): + for u in model.up_blocks: assert_frozen(u.resnets) if isinstance(u.attentions, nn.ModuleList): # attentions can be list of Nones assert_frozen(u.attentions) if u.upsamplers is not None: assert_frozen(u.upsamplers) - # # check controlnet - # everything expect down,mid,up blocks + # check controlnet: everything except down, mid, up blocks modules_from_controlnet = [ model.controlnet_cond_embedding, model.ctrl_conv_in, model.control_to_base_for_conv_in, ] optional_modules_from_controlnet = [model.ctrl_time_embedding] - for m in modules_from_controlnet: assert_unfrozen(m) for m in optional_modules_from_controlnet: if m is not None: assert_unfrozen(m) - # down blocks for d in model.down_blocks: assert_unfrozen(d.ctrl_resnets) @@ -267,36 +243,24 @@ def assert_unfrozen(module): for u in model.up_blocks: assert_unfrozen(u.ctrl_to_base) - def test_gradient_checkpointing_is_applied(self): - expected_set = { - "Transformer2DModel", - "UNetMidBlock2DCrossAttn", - "ControlNetXSCrossAttnDownBlock2D", - "ControlNetXSCrossAttnMidBlock2D", - "ControlNetXSCrossAttnUpBlock2D", - } - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - @is_flaky def test_forward_no_control(self): unet = self.get_dummy_unet() controlnet = self.get_dummy_controlnet_from_unet(unet) model = UNetControlNetXSModel.from_unet(unet, controlnet) - unet = unet.to(torch_device) model = model.to(torch_device) - input_ = self.dummy_input - + inputs = self.get_dummy_inputs() control_specific_input = ["controlnet_cond", "conditioning_scale"] - input_for_unet = {k: v for k, v in input_.items() if k not in control_specific_input} + input_for_unet = {k: v for k, v in inputs.items() if k not in control_specific_input} with torch.no_grad(): unet_output = unet(**input_for_unet).sample.cpu() - unet_controlnet_output = model(**input_, apply_control=False).sample.cpu() + unet_controlnet_output = model(**inputs, apply_control=False).sample.cpu() - assert np.abs(unet_output.flatten() - unet_controlnet_output.flatten()).max() < 3e-4 + assert (unet_output.flatten() - unet_controlnet_output.flatten()).abs().max() < 3e-4 def test_time_embedding_mixing(self): unet = self.get_dummy_unet() @@ -305,22 +269,34 @@ def test_time_embedding_mixing(self): unet, time_embedding_mix=0.5, learn_time_embedding=True ) - model = UNetControlNetXSModel.from_unet(unet, controlnet) - model_mix_time = UNetControlNetXSModel.from_unet(unet, controlnet_mix_time) - - unet = unet.to(torch_device) - model = model.to(torch_device) - model_mix_time = model_mix_time.to(torch_device) - - input_ = self.dummy_input + model = UNetControlNetXSModel.from_unet(unet, controlnet).to(torch_device) + model_mix_time = UNetControlNetXSModel.from_unet(unet, controlnet_mix_time).to(torch_device) + inputs = self.get_dummy_inputs() with torch.no_grad(): - output = model(**input_).sample - output_mix_time = model_mix_time(**input_).sample + output = model(**inputs).sample + output_mix_time = model_mix_time(**inputs).sample assert output.shape == output_mix_time.shape - @unittest.skip("Test not supported.") - def test_forward_with_norm_groups(self): - # UNetControlNetXSModel currently only supports StableDiffusion and StableDiffusion-XL, both of which have norm_num_groups fixed at 32. So we don't need to test different values for norm_num_groups. - pass + +class TestUNetControlNetXSModelTraining(UNetControlNetXSModelTesterConfig, TrainingTesterMixin): + """Training tests for UNetControlNetXSModel.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = { + "Transformer2DModel", + "UNetMidBlock2DCrossAttn", + "ControlNetXSCrossAttnDownBlock2D", + "ControlNetXSCrossAttnMidBlock2D", + "ControlNetXSCrossAttnUpBlock2D", + } + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestUNetControlNetXSModelMemory(UNetControlNetXSModelTesterConfig, MemoryTesterMixin): + """Memory optimization tests for UNetControlNetXSModel.""" + + +class TestUNetControlNetXSModelAttention(UNetControlNetXSModelTesterConfig, AttentionTesterMixin): + """Attention processor tests for UNetControlNetXSModel.""" diff --git a/tests/models/unets/test_models_unet_motion.py b/tests/models/unets/test_models_unet_motion.py index d931b345fd09..21bc9bd62de9 100644 --- a/tests/models/unets/test_models_unet_motion.py +++ b/tests/models/unets/test_models_unet_motion.py @@ -15,56 +15,44 @@ import copy import os -import tempfile -import unittest -import numpy as np import torch from diffusers import MotionAdapter, UNet2DConditionModel, UNetMotionModel -from diffusers.utils import logging -from diffusers.utils.import_utils import is_xformers_available - -from ...testing_utils import ( - enable_full_determinism, - floats_tensor, - torch_device, +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + MemoryTesterMixin, + ModelTesterMixin, + TrainingTesterMixin, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin - -logger = logging.get_logger(__name__) enable_full_determinism() -class UNetMotionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNetMotionModel - main_input_name = "sample" - +class UNetMotionModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_channels = 4 - num_frames = 4 - sizes = (16, 16) - - noise = floats_tensor((batch_size, num_channels, num_frames) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) - encoder_hidden_states = floats_tensor((batch_size * num_frames, 4, 16)).to(torch_device) + def model_class(self): + return UNetMotionModel - return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states} + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 4, 16, 16) @property - def output_shape(self): - return (4, 4, 16, 16) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "block_out_channels": (16, 32), "norm_num_groups": 16, "down_block_types": ("CrossAttnDownBlockMotion", "DownBlockMotion"), @@ -76,9 +64,23 @@ def prepare_init_args_and_inputs_for_common(self): "layers_per_block": 1, "sample_size": 16, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict: + batch_size = 4 + num_channels = 4 + num_frames = 4 + sizes = (16, 16) + noise = randn_tensor( + (batch_size, num_channels, num_frames, *sizes), generator=self.generator, device=torch_device + ) + timestep = torch.tensor([10], device=torch_device) + encoder_hidden_states = randn_tensor( + (batch_size * num_frames, 4, 16), generator=self.generator, device=torch_device + ) + return {"sample": noise, "timestep": timestep, "encoder_hidden_states": encoder_hidden_states} + + +class TestUNetMotionModel(UNetMotionModelTesterConfig, ModelTesterMixin): def test_from_unet2d(self): torch.manual_seed(0) unet2d = UNet2DConditionModel() @@ -88,19 +90,17 @@ def test_from_unet2d(self): model_state_dict = model.state_dict() for param_name, param_value in unet2d.named_parameters(): - self.assertTrue(torch.equal(model_state_dict[param_name], param_value)) + assert torch.equal(model_state_dict[param_name], param_value) def test_freeze_unet2d(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) + model = self.model_class(**self.get_init_dict()) model.freeze_unet2d_params() for param_name, param_value in model.named_parameters(): if "motion_modules" not in param_name: - self.assertFalse(param_value.requires_grad) - + assert not param_value.requires_grad else: - self.assertTrue(param_value.requires_grad) + assert param_value.requires_grad def test_loading_motion_adapter(self): model = self.model_class() @@ -110,210 +110,113 @@ def test_loading_motion_adapter(self): for idx, down_block in enumerate(model.down_blocks): adapter_state_dict = adapter.down_blocks[idx].motion_modules.state_dict() for param_name, param_value in down_block.motion_modules.named_parameters(): - self.assertTrue(torch.equal(adapter_state_dict[param_name], param_value)) + assert torch.equal(adapter_state_dict[param_name], param_value) for idx, up_block in enumerate(model.up_blocks): adapter_state_dict = adapter.up_blocks[idx].motion_modules.state_dict() for param_name, param_value in up_block.motion_modules.named_parameters(): - self.assertTrue(torch.equal(adapter_state_dict[param_name], param_value)) + assert torch.equal(adapter_state_dict[param_name], param_value) mid_block_adapter_state_dict = adapter.mid_block.motion_modules.state_dict() for param_name, param_value in model.mid_block.motion_modules.named_parameters(): - self.assertTrue(torch.equal(mid_block_adapter_state_dict[param_name], param_value)) + assert torch.equal(mid_block_adapter_state_dict[param_name], param_value) - def test_saving_motion_modules(self): + def test_saving_motion_modules(self, tmp_path): torch.manual_seed(0) - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_motion_modules(tmpdirname) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "diffusion_pytorch_model.safetensors"))) + init_dict = self.get_init_dict() + model = self.model_class(**init_dict).to(torch_device) - adapter_loaded = MotionAdapter.from_pretrained(tmpdirname) + model.save_motion_modules(tmp_path) + assert os.path.isfile(os.path.join(tmp_path, "diffusion_pytorch_model.safetensors")) - torch.manual_seed(0) - model_loaded = self.model_class(**init_dict) - model_loaded.load_motion_modules(adapter_loaded) - model_loaded.to(torch_device) + adapter_loaded = MotionAdapter.from_pretrained(tmp_path) + torch.manual_seed(0) + model_loaded = self.model_class(**init_dict) + model_loaded.load_motion_modules(adapter_loaded) + model_loaded.to(torch_device) with torch.no_grad(): - output = model(**inputs_dict)[0] - output_loaded = model_loaded(**inputs_dict)[0] - - max_diff = (output - output_loaded).abs().max().item() - self.assertLessEqual(max_diff, 1e-4, "Models give different forward passes") + output = model(**self.get_dummy_inputs())[0] + output_loaded = model_loaded(**self.get_dummy_inputs())[0] - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_enable_works(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - - model.enable_xformers_memory_efficient_attention() - - assert ( - model.mid_block.attentions[0].transformer_blocks[0].attn1.processor.__class__.__name__ - == "XFormersAttnProcessor" - ), "xformers is not enabled" - - def test_gradient_checkpointing_is_applied(self): - expected_set = { - "CrossAttnUpBlockMotion", - "CrossAttnDownBlockMotion", - "UNetMidBlockCrossAttnMotion", - "UpBlockMotion", - "Transformer2DModel", - "DownBlockMotion", - } - super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + assert (output - output_loaded).abs().max().item() <= 1e-4, "Models give different forward passes" def test_feed_forward_chunking(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() init_dict["block_out_channels"] = (32, 64) init_dict["norm_num_groups"] = 32 - - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() + model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): - output = model(**inputs_dict)[0] + output = model(**self.get_dummy_inputs())[0] model.enable_forward_chunking() with torch.no_grad(): - output_2 = model(**inputs_dict)[0] + output_2 = model(**self.get_dummy_inputs())[0] - self.assertEqual(output.shape, output_2.shape, "Shape doesn't match") - assert np.abs(output.cpu() - output_2.cpu()).max() < 1e-2 + assert output.shape == output_2.shape, "Shape doesn't match" + assert (output - output_2).abs().max() < 1e-2 def test_pickle(self): - # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) + model = self.model_class(**self.get_init_dict()).to(torch_device) with torch.no_grad(): - sample = model(**inputs_dict).sample + sample = model(**self.get_dummy_inputs()).sample sample_copy = copy.copy(sample) - assert (sample - sample_copy).abs().max() < 1e-4 - def test_from_save_pretrained(self, expected_max_diff=5e-5): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname, safe_serialization=False) - torch.manual_seed(0) - new_model = self.model_class.from_pretrained(tmpdirname) - new_model.to(torch_device) - - with torch.no_grad(): - image = model(**inputs_dict) - if isinstance(image, dict): - image = image.to_tuple()[0] - - new_image = new_model(**inputs_dict) - - if isinstance(new_image, dict): - new_image = new_image.to_tuple()[0] - - max_diff = (image - new_image).abs().max().item() - self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes") - - def test_from_save_pretrained_variant(self, expected_max_diff=5e-5): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - torch.manual_seed(0) - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_pretrained(tmpdirname, variant="fp16", safe_serialization=False) - - torch.manual_seed(0) - new_model = self.model_class.from_pretrained(tmpdirname, variant="fp16") - # non-variant cannot be loaded - with self.assertRaises(OSError) as error_context: - self.model_class.from_pretrained(tmpdirname) - - # make sure that error message states what keys are missing - assert "Error no file named diffusion_pytorch_model.bin found in directory" in str(error_context.exception) - - new_model.to(torch_device) - - with torch.no_grad(): - image = model(**inputs_dict) - if isinstance(image, dict): - image = image.to_tuple()[0] - - new_image = new_model(**inputs_dict) - - if isinstance(new_image, dict): - new_image = new_image.to_tuple()[0] - - max_diff = (image - new_image).abs().max().item() - self.assertLessEqual(max_diff, expected_max_diff, "Models give different forward passes") - def test_forward_with_norm_groups(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - + init_dict = self.get_init_dict() init_dict["norm_num_groups"] = 16 init_dict["block_out_channels"] = (16, 32) - - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() + model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): - output = model(**inputs_dict) + output = model(**self.get_dummy_inputs()).sample - if isinstance(output, dict): - output = output.to_tuple()[0] - - self.assertIsNotNone(output) - expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" def test_asymmetric_motion_model(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - + init_dict = self.get_init_dict() init_dict["layers_per_block"] = (2, 3) init_dict["transformer_layers_per_block"] = ((1, 2), (3, 4, 5)) init_dict["reverse_transformer_layers_per_block"] = ((7, 6, 7, 4), (4, 2, 2)) - init_dict["temporal_transformer_layers_per_block"] = ((2, 5), (2, 3, 5)) init_dict["reverse_temporal_transformer_layers_per_block"] = ((5, 4, 3, 4), (3, 2, 2)) - init_dict["num_attention_heads"] = (2, 4) init_dict["motion_num_attention_heads"] = (4, 4) init_dict["reverse_motion_num_attention_heads"] = (2, 2) - init_dict["use_motion_mid_block"] = True init_dict["mid_block_layers"] = 2 init_dict["transformer_layers_per_mid_block"] = (1, 5) init_dict["temporal_transformer_layers_per_mid_block"] = (2, 4) - - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() + model = self.model_class(**init_dict).to(torch_device).eval() with torch.no_grad(): - output = model(**inputs_dict) + output = model(**self.get_dummy_inputs()).sample + + assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" + + +class TestUNetMotionModelTraining(UNetMotionModelTesterConfig, TrainingTesterMixin): + """Training tests for UNetMotionModel.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = { + "CrossAttnUpBlockMotion", + "CrossAttnDownBlockMotion", + "UNetMidBlockCrossAttnMotion", + "UpBlockMotion", + "Transformer2DModel", + "DownBlockMotion", + } + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestUNetMotionModelMemory(UNetMotionModelTesterConfig, MemoryTesterMixin): + """Memory optimization tests for UNetMotionModel.""" - if isinstance(output, dict): - output = output.to_tuple()[0] - self.assertIsNotNone(output) - expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") +class TestUNetMotionModelAttention(UNetMotionModelTesterConfig, AttentionTesterMixin): + """Attention processor tests for UNetMotionModel.""" From 1ead334f9a20397f62d6fb8e6f20a6af08969cb5 Mon Sep 17 00:00:00 2001 From: Akshan Krithick <97239696+akshan-main@users.noreply.github.com> Date: Tue, 9 Jun 2026 20:39:57 -0700 Subject: [PATCH 07/13] refactor unet_1d tests (#13898) * refactor unet_1d tests * use per-sample output_shape for unet_1d tests --------- Co-authored-by: Sayak Paul --- tests/models/unets/test_models_unet_1d.py | 262 ++++++---------------- 1 file changed, 69 insertions(+), 193 deletions(-) diff --git a/tests/models/unets/test_models_unet_1d.py b/tests/models/unets/test_models_unet_1d.py index bac017e7e7d3..ee11cba3dd2c 100644 --- a/tests/models/unets/test_models_unet_1d.py +++ b/tests/models/unets/test_models_unet_1d.py @@ -13,77 +13,37 @@ # See the License for the specific language governing permissions and # limitations under the License. -import unittest - -import pytest import torch from diffusers import UNet1DModel +from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import ( - backend_manual_seed, - floats_tensor, - slow, - torch_device, -) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin - +from ...testing_utils import backend_manual_seed, enable_full_determinism, slow, torch_device +from ..testing_utils import BaseModelTesterConfig, ModelTesterMixin -class UNet1DModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet1DModel - main_input_name = "sample" - - @property - def dummy_input(self): - batch_size = 4 - num_features = 14 - seq_len = 16 - noise = floats_tensor((batch_size, num_features, seq_len)).to(torch_device) - time_step = torch.tensor([10] * batch_size).to(torch_device) +enable_full_determinism() - return {"sample": noise, "timestep": time_step} +class UNet1DModelTesterConfig(BaseModelTesterConfig): @property - def input_shape(self): - return (4, 14, 16) + def model_class(self): + return UNet1DModel @property - def output_shape(self): - return (4, 14, 16) - - @unittest.skip("Test not supported.") - def test_ema_training(self): - pass - - @unittest.skip("Test not supported.") - def test_training(self): - pass - - @unittest.skip("Test not supported.") - def test_layerwise_casting_training(self): - pass - - def test_determinism(self): - super().test_determinism() - - def test_outputs_equivalence(self): - super().test_outputs_equivalence() - - def test_from_save_pretrained(self): - super().test_from_save_pretrained() - - def test_from_save_pretrained_variant(self): - super().test_from_save_pretrained_variant() + def main_input_name(self) -> str: + return "sample" - def test_model_from_pretrained(self): - super().test_model_from_pretrained() + @property + def output_shape(self) -> tuple: + return (14, 16) - def test_output(self): - super().test_output() + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "block_out_channels": (8, 8, 16, 16), "in_channels": 14, "out_channels": 14, @@ -97,19 +57,26 @@ def prepare_init_args_and_inputs_for_common(self): "up_block_types": ("UpResnetBlock1D", "UpResnetBlock1D", "UpResnetBlock1D"), "act_fn": "swish", } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict: + batch_size = 4 + num_features = 14 + seq_len = 16 + noise = randn_tensor((batch_size, num_features, seq_len), generator=self.generator, device=torch_device) + timestep = torch.tensor([10] * batch_size, device=torch_device) + return {"sample": noise, "timestep": timestep} + + +class TestUNet1DModel(UNet1DModelTesterConfig, ModelTesterMixin): def test_from_pretrained_hub(self): model, loading_info = UNet1DModel.from_pretrained( "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="unet" ) - self.assertIsNotNone(model) - self.assertEqual(len(loading_info["missing_keys"]), 0) + assert model is not None + assert len(loading_info["missing_keys"]) == 0 model.to(torch_device) - image = model(**self.dummy_input) - + image = model(**self.get_dummy_inputs()) assert image is not None, "Make sure output is not None" def test_output_pretrained(self): @@ -119,9 +86,7 @@ def test_output_pretrained(self): num_features = model.config.in_channels seq_len = 16 - noise = torch.randn((1, seq_len, num_features)).permute( - 0, 2, 1 - ) # match original, we can update values and remove + noise = torch.randn((1, seq_len, num_features)).permute(0, 2, 1) time_step = torch.full((num_features,), 0) with torch.no_grad(): @@ -131,12 +96,7 @@ def test_output_pretrained(self): # fmt: off expected_output_slice = torch.tensor([-2.137172, 1.1426016, 0.3688687, -0.766922, 0.7303146, 0.11038864, -0.4760633, 0.13270172, 0.02591348]) # fmt: on - self.assertTrue(torch.allclose(output_slice, expected_output_slice, rtol=1e-3)) - - @unittest.skip("Test not supported.") - def test_forward_with_norm_groups(self): - # Not implemented yet for this UNet - pass + assert torch.allclose(output_slice, expected_output_slice, rtol=1e-3) @slow def test_unet_1d_maestro(self): @@ -157,98 +117,26 @@ def test_unet_1d_maestro(self): assert (output_sum - 224.0896).abs() < 0.5 assert (output_max - 0.0607).abs() < 4e-4 - @pytest.mark.xfail( - reason=( - "RuntimeError: 'fill_out' not implemented for 'Float8_e4m3fn'. The error is caused due to certain torch.float8_e4m3fn and torch.float8_e5m2 operations " - "not being supported when using deterministic algorithms (which is what the tests run with). To fix:\n" - "1. Wait for next PyTorch release: https://github.com/pytorch/pytorch/issues/137160.\n" - "2. Unskip this test." - ), - ) - def test_layerwise_casting_inference(self): - super().test_layerwise_casting_inference() - - @pytest.mark.xfail( - reason=( - "RuntimeError: 'fill_out' not implemented for 'Float8_e4m3fn'. The error is caused due to certain torch.float8_e4m3fn and torch.float8_e5m2 operations " - "not being supported when using deterministic algorithms (which is what the tests run with). To fix:\n" - "1. Wait for next PyTorch release: https://github.com/pytorch/pytorch/issues/137160.\n" - "2. Unskip this test." - ), - ) - def test_layerwise_casting_memory(self): - pass - - -class UNetRLModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet1DModel - main_input_name = "sample" +class UNetRLModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_features = 14 - seq_len = 16 - - noise = floats_tensor((batch_size, num_features, seq_len)).to(torch_device) - time_step = torch.tensor([10] * batch_size).to(torch_device) - - return {"sample": noise, "timestep": time_step} + def model_class(self): + return UNet1DModel @property - def input_shape(self): - return (4, 14, 16) + def main_input_name(self) -> str: + return "sample" @property - def output_shape(self): - return (4, 14, 1) - - def test_determinism(self): - super().test_determinism() - - def test_outputs_equivalence(self): - super().test_outputs_equivalence() - - def test_from_save_pretrained(self): - super().test_from_save_pretrained() - - def test_from_save_pretrained_variant(self): - super().test_from_save_pretrained_variant() - - def test_model_from_pretrained(self): - super().test_model_from_pretrained() - - def test_output(self): - # UNetRL is a value-function is different output shape - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - with torch.no_grad(): - output = model(**inputs_dict) - - if isinstance(output, dict): - output = output.sample - - self.assertIsNotNone(output) - expected_shape = torch.Size((inputs_dict["sample"].shape[0], 1)) - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") - - @unittest.skip("Test not supported.") - def test_ema_training(self): - pass + def output_shape(self) -> tuple: + return (1,) - @unittest.skip("Test not supported.") - def test_training(self): - pass - - @unittest.skip("Test not supported.") - def test_layerwise_casting_training(self): - pass + @property + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "in_channels": 14, "out_channels": 14, "down_block_types": ["DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D", "DownResnetBlock1D"], @@ -264,19 +152,36 @@ def prepare_init_args_and_inputs_for_common(self): "time_embedding_type": "positional", "act_fn": "mish", } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + + def get_dummy_inputs(self) -> dict: + batch_size = 4 + num_features = 14 + seq_len = 16 + noise = randn_tensor((batch_size, num_features, seq_len), generator=self.generator, device=torch_device) + timestep = torch.tensor([10] * batch_size, device=torch_device) + return {"sample": noise, "timestep": timestep} + + +class TestUNetRLModel(UNetRLModelTesterConfig, ModelTesterMixin): + # UNetRL is a value function, so it has a different output shape. + def test_output(self): + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + + inputs = self.get_dummy_inputs() + with torch.no_grad(): + output = model(**inputs).sample + + assert output.shape == (inputs["sample"].shape[0], 1), "Input and output shapes do not match" def test_from_pretrained_hub(self): value_function, vf_loading_info = UNet1DModel.from_pretrained( "bglick13/hopper-medium-v2-value-function-hor32", output_loading_info=True, subfolder="value_function" ) - self.assertIsNotNone(value_function) - self.assertEqual(len(vf_loading_info["missing_keys"]), 0) + assert value_function is not None + assert len(vf_loading_info["missing_keys"]) == 0 value_function.to(torch_device) - image = value_function(**self.dummy_input) - + image = value_function(**self.get_dummy_inputs()) assert image is not None, "Make sure output is not None" def test_output_pretrained(self): @@ -288,9 +193,7 @@ def test_output_pretrained(self): num_features = value_function.config.in_channels seq_len = 14 - noise = torch.randn((1, seq_len, num_features)).permute( - 0, 2, 1 - ) # match original, we can update values and remove + noise = torch.randn((1, seq_len, num_features)).permute(0, 2, 1) time_step = torch.full((num_features,), 0) with torch.no_grad(): @@ -299,31 +202,4 @@ def test_output_pretrained(self): # fmt: off expected_output_slice = torch.tensor([165.25] * seq_len) # fmt: on - self.assertTrue(torch.allclose(output, expected_output_slice, rtol=1e-3)) - - @unittest.skip("Test not supported.") - def test_forward_with_norm_groups(self): - # Not implemented yet for this UNet - pass - - @pytest.mark.xfail( - reason=( - "RuntimeError: 'fill_out' not implemented for 'Float8_e4m3fn'. The error is caused due to certain torch.float8_e4m3fn and torch.float8_e5m2 operations " - "not being supported when using deterministic algorithms (which is what the tests run with). To fix:\n" - "1. Wait for next PyTorch release: https://github.com/pytorch/pytorch/issues/137160.\n" - "2. Unskip this test." - ), - ) - def test_layerwise_casting_inference(self): - pass - - @pytest.mark.xfail( - reason=( - "RuntimeError: 'fill_out' not implemented for 'Float8_e4m3fn'. The error is caused due to certain torch.float8_e4m3fn and torch.float8_e5m2 operations " - "not being supported when using deterministic algorithms (which is what the tests run with). To fix:\n" - "1. Wait for next PyTorch release: https://github.com/pytorch/pytorch/issues/137160.\n" - "2. Unskip this test." - ), - ) - def test_layerwise_casting_memory(self): - pass + assert torch.allclose(output, expected_output_slice, rtol=1e-3) From d4fecddb1e59670a5904891e0c0444d1e0fd28fd Mon Sep 17 00:00:00 2001 From: Akshan Krithick <97239696+akshan-main@users.noreply.github.com> Date: Tue, 9 Jun 2026 21:30:25 -0700 Subject: [PATCH 08/13] refactor unet_2d tests (#13901) Co-authored-by: Sayak Paul --- tests/models/unets/test_models_unet_2d.py | 296 +++++++++------------- 1 file changed, 113 insertions(+), 183 deletions(-) diff --git a/tests/models/unets/test_models_unet_2d.py b/tests/models/unets/test_models_unet_2d.py index e289f44303f2..a5cd8abd873a 100644 --- a/tests/models/unets/test_models_unet_2d.py +++ b/tests/models/unets/test_models_unet_2d.py @@ -15,12 +15,12 @@ import gc import math -import unittest +import pytest import torch from diffusers import UNet2DModel -from diffusers.utils import logging +from diffusers.utils.torch_utils import randn_tensor from ...testing_utils import ( backend_empty_cache, @@ -31,39 +31,31 @@ torch_all_close, torch_device, ) -from ..test_modeling_common import ModelTesterMixin, UNetTesterMixin +from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin -logger = logging.get_logger(__name__) - enable_full_determinism() -class Unet2DModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet2DModel - main_input_name = "sample" - +class Unet2DModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_channels = 3 - sizes = (32, 32) - - noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) + def model_class(self): + return UNet2DModel - return {"sample": noise, "timestep": time_step} + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (3, 32, 32) @property - def output_shape(self): - return (3, 32, 32) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "block_out_channels": (4, 8), "norm_num_groups": 2, "down_block_types": ("DownBlock2D", "AttnDownBlock2D"), @@ -74,110 +66,77 @@ def prepare_init_args_and_inputs_for_common(self): "layers_per_block": 2, "sample_size": 32, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - def test_mid_block_attn_groups(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + def get_dummy_inputs(self) -> dict: + noise = randn_tensor((4, 3, 32, 32), generator=self.generator, device=torch_device) + timestep = torch.tensor([10], device=torch_device) + return {"sample": noise, "timestep": timestep} + +class TestUnet2DModel(Unet2DModelTesterConfig, ModelTesterMixin): + def test_mid_block_attn_groups(self): + init_dict = self.get_init_dict() init_dict["add_attention"] = True init_dict["attn_norm_num_groups"] = 4 + model = self.model_class(**init_dict).to(torch_device).eval() - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - self.assertIsNotNone( - model.mid_block.attentions[0].group_norm, "Mid block Attention group norm should exist but does not." + assert model.mid_block.attentions[0].group_norm is not None, ( + "Mid block Attention group norm should exist but does not." ) - self.assertEqual( - model.mid_block.attentions[0].group_norm.num_groups, - init_dict["attn_norm_num_groups"], - "Mid block Attention group norm does not have the expected number of groups.", + assert model.mid_block.attentions[0].group_norm.num_groups == init_dict["attn_norm_num_groups"], ( + "Mid block Attention group norm does not have the expected number of groups." ) with torch.no_grad(): - output = model(**inputs_dict) - - if isinstance(output, dict): - output = output.to_tuple()[0] + output = model(**self.get_dummy_inputs()).sample - self.assertIsNotNone(output) - expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == self.get_dummy_inputs()["sample"].shape, "Input and output shapes do not match" def test_mid_block_none(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - mid_none_init_dict, mid_none_inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + mid_none_init_dict = self.get_init_dict() mid_none_init_dict["mid_block_type"] = None - model = self.model_class(**init_dict) - model.to(torch_device) - model.eval() - - mid_none_model = self.model_class(**mid_none_init_dict) - mid_none_model.to(torch_device) - mid_none_model.eval() - - self.assertIsNone(mid_none_model.mid_block, "Mid block should not exist.") - - with torch.no_grad(): - output = model(**inputs_dict) - - if isinstance(output, dict): - output = output.to_tuple()[0] + model = self.model_class(**init_dict).to(torch_device).eval() + mid_none_model = self.model_class(**mid_none_init_dict).to(torch_device).eval() + assert mid_none_model.mid_block is None, "Mid block should not exist." with torch.no_grad(): - mid_none_output = mid_none_model(**mid_none_inputs_dict) + output = model(**self.get_dummy_inputs()).sample + mid_none_output = mid_none_model(**self.get_dummy_inputs()).sample - if isinstance(mid_none_output, dict): - mid_none_output = mid_none_output.to_tuple()[0] + assert not torch.allclose(output, mid_none_output, rtol=1e-3), "outputs should be different." - self.assertFalse(torch.allclose(output, mid_none_output, rtol=1e-3), "outputs should be different.") +class TestUnet2DModelTraining(Unet2DModelTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): - expected_set = { - "AttnUpBlock2D", - "AttnDownBlock2D", - "UNetMidBlock2D", - "UpBlock2D", - "DownBlock2D", - } - - # NOTE: unlike UNet2DConditionModel, UNet2DModel does not currently support tuples for `attention_head_dim` - attention_head_dim = 8 - block_out_channels = (16, 32) + expected_set = {"AttnUpBlock2D", "AttnDownBlock2D", "UNetMidBlock2D", "UpBlock2D", "DownBlock2D"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - super().test_gradient_checkpointing_is_applied( - expected_set=expected_set, attention_head_dim=attention_head_dim, block_out_channels=block_out_channels - ) +class TestUnet2DModelMemory(Unet2DModelTesterConfig, MemoryTesterMixin): + """Memory optimization tests for UNet2DModel.""" -class UNetLDMModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet2DModel - main_input_name = "sample" +class UNetLDMModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self): - batch_size = 4 - num_channels = 4 - sizes = (32, 32) - - noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) + def model_class(self): + return UNet2DModel - return {"sample": noise, "timestep": time_step} + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (4, 32, 32) @property - def output_shape(self): - return (4, 32, 32) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "sample_size": 32, "in_channels": 4, "out_channels": 4, @@ -187,26 +146,28 @@ def prepare_init_args_and_inputs_for_common(self): "down_block_types": ("DownBlock2D", "DownBlock2D"), "up_block_types": ("UpBlock2D", "UpBlock2D"), } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict: + noise = randn_tensor((4, 4, 32, 32), generator=self.generator, device=torch_device) + timestep = torch.tensor([10], device=torch_device) + return {"sample": noise, "timestep": timestep} + + +class TestUNetLDMModel(UNetLDMModelTesterConfig, ModelTesterMixin): def test_from_pretrained_hub(self): model, loading_info = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True) - - self.assertIsNotNone(model) - self.assertEqual(len(loading_info["missing_keys"]), 0) + assert model is not None + assert len(loading_info["missing_keys"]) == 0 model.to(torch_device) - image = model(**self.dummy_input).sample - + image = model(**self.get_dummy_inputs()).sample assert image is not None, "Make sure output is not None" @require_torch_accelerator def test_from_pretrained_accelerate(self): model, _ = UNet2DModel.from_pretrained("fusing/unet-ldm-dummy-update", output_loading_info=True) model.to(torch_device) - image = model(**self.dummy_input).sample - + image = model(**self.get_dummy_inputs()).sample assert image is not None, "Make sure output is not None" @require_torch_accelerator @@ -264,45 +225,38 @@ def test_output_pretrained(self): # fmt: off expected_output_slice = torch.tensor([-13.3258, -20.1100, -15.9873, -17.6617, -23.0596, -17.9419, -13.3675, -16.1889, -12.3800]) # fmt: on + assert torch_all_close(output_slice, expected_output_slice, rtol=1e-3) - self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-3)) +class TestUNetLDMModelTraining(UNetLDMModelTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): expected_set = {"DownBlock2D", "UNetMidBlock2D", "UpBlock2D"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - # NOTE: unlike UNet2DConditionModel, UNet2DModel does not currently support tuples for `attention_head_dim` - attention_head_dim = 32 - block_out_channels = (32, 64) - - super().test_gradient_checkpointing_is_applied( - expected_set=expected_set, attention_head_dim=attention_head_dim, block_out_channels=block_out_channels - ) +class TestUNetLDMModelMemory(UNetLDMModelTesterConfig, MemoryTesterMixin): + """Memory optimization tests for the LDM UNet2DModel config.""" -class NCSNppModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet2DModel - main_input_name = "sample" +class NCSNppModelTesterConfig(BaseModelTesterConfig): @property - def dummy_input(self, sizes=(32, 32)): - batch_size = 4 - num_channels = 3 - - noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor(batch_size * [10]).to(dtype=torch.int32, device=torch_device) + def model_class(self): + return UNet2DModel - return {"sample": noise, "timestep": time_step} + @property + def main_input_name(self) -> str: + return "sample" @property - def input_shape(self): + def output_shape(self) -> tuple: return (3, 32, 32) @property - def output_shape(self): - return (3, 32, 32) + def generator(self): + return torch.Generator("cpu").manual_seed(0) - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + return { "block_out_channels": [32, 64, 64, 64], "in_channels": 3, "layers_per_block": 1, @@ -311,34 +265,27 @@ def prepare_init_args_and_inputs_for_common(self): "norm_eps": 1e-6, "mid_block_scale_factor": math.sqrt(2.0), "norm_num_groups": None, - "down_block_types": [ - "SkipDownBlock2D", - "AttnSkipDownBlock2D", - "SkipDownBlock2D", - "SkipDownBlock2D", - ], - "up_block_types": [ - "SkipUpBlock2D", - "SkipUpBlock2D", - "AttnSkipUpBlock2D", - "SkipUpBlock2D", - ], + "down_block_types": ["SkipDownBlock2D", "AttnSkipDownBlock2D", "SkipDownBlock2D", "SkipDownBlock2D"], + "up_block_types": ["SkipUpBlock2D", "SkipUpBlock2D", "AttnSkipUpBlock2D", "SkipUpBlock2D"], } - inputs_dict = self.dummy_input - return init_dict, inputs_dict + def get_dummy_inputs(self) -> dict: + noise = randn_tensor((4, 3, 32, 32), generator=self.generator, device=torch_device) + timestep = torch.tensor(4 * [10], dtype=torch.int32, device=torch_device) + return {"sample": noise, "timestep": timestep} + + +class TestNCSNppModel(NCSNppModelTesterConfig, ModelTesterMixin): @slow def test_from_pretrained_hub(self): model, loading_info = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256", output_loading_info=True) - self.assertIsNotNone(model) - self.assertEqual(len(loading_info["missing_keys"]), 0) + assert model is not None + assert len(loading_info["missing_keys"]) == 0 model.to(torch_device) - inputs = self.dummy_input - noise = floats_tensor((4, 3) + (256, 256)).to(torch_device) - inputs["sample"] = noise + inputs = self.get_dummy_inputs() + inputs["sample"] = floats_tensor((4, 3) + (256, 256)).to(torch_device) image = model(**inputs) - assert image is not None, "Make sure output is not None" @slow @@ -346,12 +293,8 @@ def test_output_pretrained_ve_mid(self): model = UNet2DModel.from_pretrained("google/ncsnpp-celebahq-256") model.to(torch_device) - batch_size = 4 - num_channels = 3 - sizes = (256, 256) - - noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor(batch_size * [1e-4]).to(torch_device) + noise = torch.ones((4, 3) + (256, 256)).to(torch_device) + time_step = torch.tensor(4 * [1e-4]).to(torch_device) with torch.no_grad(): output = model(noise, time_step).sample @@ -360,19 +303,14 @@ def test_output_pretrained_ve_mid(self): # fmt: off expected_output_slice = torch.tensor([-4836.2178, -6487.1470, -3816.8196, -7964.9302, -10966.3037, -20043.5957, 8137.0513, 2340.3328, 544.6056]) # fmt: on - - self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2)) + assert torch_all_close(output_slice, expected_output_slice, rtol=1e-2) def test_output_pretrained_ve_large(self): model = UNet2DModel.from_pretrained("fusing/ncsnpp-ffhq-ve-dummy-update") model.to(torch_device) - batch_size = 4 - num_channels = 3 - sizes = (32, 32) - - noise = torch.ones((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor(batch_size * [1e-4]).to(torch_device) + noise = torch.ones((4, 3) + (32, 32)).to(torch_device) + time_step = torch.tensor(4 * [1e-4]).to(torch_device) with torch.no_grad(): output = model(noise, time_step).sample @@ -381,36 +319,28 @@ def test_output_pretrained_ve_large(self): # fmt: off expected_output_slice = torch.tensor([-0.0325, -0.0900, -0.0869, -0.0332, -0.0725, -0.0270, -0.0101, 0.0227, 0.0256]) # fmt: on + assert torch_all_close(output_slice, expected_output_slice, rtol=1e-2) - self.assertTrue(torch_all_close(output_slice, expected_output_slice, rtol=1e-2)) - - @unittest.skip("Test not supported.") - def test_forward_with_norm_groups(self): - # not required for this model - pass +class TestNCSNppModelTraining(NCSNppModelTesterConfig, TrainingTesterMixin): def test_gradient_checkpointing_is_applied(self): - expected_set = { - "UNetMidBlock2D", - } + expected_set = {"UNetMidBlock2D"} + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - block_out_channels = (32, 64, 64, 64) + def test_gradient_checkpointing_equivalence(self): + super().test_gradient_checkpointing_equivalence(skip={"time_proj.weight"}) - super().test_gradient_checkpointing_is_applied( - expected_set=expected_set, block_out_channels=block_out_channels - ) - def test_effective_gradient_checkpointing(self): - super().test_effective_gradient_checkpointing(skip={"time_proj.weight"}) +class TestNCSNppModelMemory(NCSNppModelTesterConfig, MemoryTesterMixin): + # Layerwise casting is not supported for this model. + @pytest.mark.skip("Layerwise casting is not supported for this model.") + def test_layerwise_casting_memory(self): + pass - @unittest.skip( - "To make layerwise casting work with this model, we will have to update the implementation. Due to potentially low usage, we don't support it here." - ) - def test_layerwise_casting_inference(self): + @pytest.mark.skip("Layerwise casting is not supported for this model.") + def test_layerwise_casting_training(self): pass - @unittest.skip( - "To make layerwise casting work with this model, we will have to update the implementation. Due to potentially low usage, we don't support it here." - ) - def test_layerwise_casting_memory(self): + @pytest.mark.skip("Layerwise casting is not supported for this model.") + def test_group_offloading_with_layerwise_casting(self, *args, **kwargs): pass From f1d8d7f8e877cfd0e323bd72b7165672f89129cf Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 10 Jun 2026 11:04:55 +0530 Subject: [PATCH 09/13] [chore] log quant config to the user_agent (#13850) log quant config to the user_agent --- src/diffusers/models/modeling_utils.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/diffusers/models/modeling_utils.py b/src/diffusers/models/modeling_utils.py index 0423b7287193..f0b8e4a58c69 100644 --- a/src/diffusers/models/modeling_utils.py +++ b/src/diffusers/models/modeling_utils.py @@ -1100,6 +1100,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None "diffusers": __version__, "file_type": "model", "framework": "pytorch", + "model_class": str(cls.__name__), } unused_kwargs = {} @@ -1146,8 +1147,9 @@ def from_pretrained(cls, pretrained_model_name_or_path: str | os.PathLike | None torch_dtype = hf_quantizer.update_torch_dtype(torch_dtype) device_map = hf_quantizer.update_device_map(device_map) - # In order to ensure popular quantization methods are supported. Can be disable with `disable_telemetry` + # In order to ensure popular quantization methods are supported. Can be disabled with `disable_telemetry` user_agent["quant"] = hf_quantizer.quantization_config.quant_method.value + user_agent["quant_config"] = json.dumps(quantization_config.to_dict(), sort_keys=True) # Force-set to `True` for more mem efficiency if low_cpu_mem_usage is None: From 43f48b60ef1204f883adf1f4e73763152eb86f88 Mon Sep 17 00:00:00 2001 From: Xin He Date: Wed, 10 Jun 2026 14:11:18 +0800 Subject: [PATCH 10/13] Integrate AutoRound into Diffusers (#13552) * support auto_round Signed-off-by: Xin He * add document and unit tests Signed-off-by: Xin He * fix CI Signed-off-by: Xin He * Apply suggestions from code review Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> * update document and overwrite the default quantization_config with specified backend. Signed-off-by: Xin He * add UT and fix bug Signed-off-by: Xin He * update per comments Signed-off-by: Xin He * update per comments Signed-off-by: Xin He * fix compile error in doc Signed-off-by: Xin He * Apply style fixes * small nits * Add auto_round dependency to the versions table Signed-off-by: Xin He * fix make deps_table_check_updated Signed-off-by: Xin He * fix CI Signed-off-by: Xin He --------- Signed-off-by: Xin He Co-authored-by: Steven Liu <59462357+stevhliu@users.noreply.github.com> Co-authored-by: github-actions[bot] Co-authored-by: Sayak Paul --- docs/source/en/_toctree.yml | 2 + docs/source/en/quantization/autoround.md | 206 ++++++++++++++++++ setup.py | 1 + src/diffusers/__init__.py | 21 ++ src/diffusers/dependency_versions_table.py | 1 + src/diffusers/quantizers/auto.py | 17 ++ .../quantizers/autoround/__init__.py | 1 + .../autoround/autoround_quantizer.py | 122 +++++++++++ .../quantizers/quantization_config.py | 79 +++++++ src/diffusers/utils/__init__.py | 1 + .../utils/dummy_auto_round_objects.py | 17 ++ src/diffusers/utils/import_utils.py | 5 + tests/models/testing_utils/__init__.py | 5 + tests/models/testing_utils/quantization.py | 127 ++++++++++- .../test_models_transformer_z_image.py | 121 ++++++++++ tests/others/test_dependencies.py | 2 + tests/testing_utils.py | 23 ++ 17 files changed, 748 insertions(+), 3 deletions(-) create mode 100644 docs/source/en/quantization/autoround.md create mode 100644 src/diffusers/quantizers/autoround/__init__.py create mode 100644 src/diffusers/quantizers/autoround/autoround_quantizer.py create mode 100644 src/diffusers/utils/dummy_auto_round_objects.py diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 1d6dbb4a301c..25b9f0ec2fbe 100644 --- a/docs/source/en/_toctree.yml +++ b/docs/source/en/_toctree.yml @@ -180,6 +180,8 @@ title: quanto - local: quantization/modelopt title: NVIDIA ModelOpt + - local: quantization/autoround + title: AutoRound title: Quantization - isExpanded: false sections: diff --git a/docs/source/en/quantization/autoround.md b/docs/source/en/quantization/autoround.md new file mode 100644 index 000000000000..f4fcf1a780c3 --- /dev/null +++ b/docs/source/en/quantization/autoround.md @@ -0,0 +1,206 @@ + + +# AutoRound + +[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit. It achieves high accuracy at ultra-low bit widths (2-4 bits) with minimal tuning by leveraging sign-gradient descent and providing broad hardware compatibility. See our papers [SignRoundV1](https://arxiv.org/pdf/2309.05516) and [SignRoundV2](https://arxiv.org/abs/2512.04746) for more details. + + +Install `auto-round`(version ≥ 0.13.0): + +```bash +pip install "auto-round>=0.13.0" +``` + +To use the Marlin kernel for faster CUDA inference, install `gptqmodel`: + +```bash +pip install "gptqmodel>=5.8.0" +``` + +## Load a quantized model + +Load a pre-quantized AutoRound model by passing [`AutoRoundConfig`] to [`~ModelMixin.from_pretrained`]. The method works with any model that loads via [Accelerate](https://hf.co/docs/accelerate/index) and has `torch.nn.Linear` layers. + +You can use [`PipelineQuantizationConfig`] to quantize specific components of a pipeline: + +```python +import torch +from diffusers import DiffusionPipeline, PipelineQuantizationConfig, AutoRoundConfig + +pipeline_quant_config = PipelineQuantizationConfig( + quant_mapping={"transformer": AutoRoundConfig(backend="auto")} +) +pipe = DiffusionPipeline.from_pretrained( + "INCModel/Z-Image-W4A16-AutoRound", + quantization_config=pipeline_quant_config, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").images[0] +image.save("output.png") +``` + +Or load a quantized model component directly: + +```python +import torch +from diffusers import ZImageTransformer2DModel, ZImagePipeline, AutoRoundConfig + +model_id = "INCModel/Z-Image-W4A16-AutoRound" + +quantization_config = AutoRoundConfig(backend="auto") +transformer = ZImageTransformer2DModel.from_pretrained( + model_id, + subfolder="transformer", + quantization_config=quantization_config, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +pipe = ZImagePipeline.from_pretrained( + model_id, + transformer=transformer, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").images[0] +image.save("output.png") +``` + +> [!NOTE] +> AutoRound in Diffusers only supports loading *pre-quantized* models. To quantize a model from scratch, use the [AutoRound CLI or Python API](https://github.com/intel/auto-round) directly, then load the result with Diffusers. + +## torch.compile + +AutoRound is compatible with [`torch.compile`](../optimization/fp16#torchcompile) for faster inference. You can compile the quantized transformer (DiT) for better performance: + +```python +import torch +from diffusers import DiffusionPipeline, PipelineQuantizationConfig, AutoRoundConfig + +pipeline_quant_config = PipelineQuantizationConfig( + quant_mapping={"transformer": AutoRoundConfig(backend="auto")} +) +pipe = DiffusionPipeline.from_pretrained( + "INCModel/Z-Image-W4A16-AutoRound", + quantization_config=pipeline_quant_config, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +pipe.transformer = torch.compile(pipe.transformer, mode="default", fullgraph=False) +``` + +## Backends + +AutoRound supports multiple inference backends for Weight-only quantized model. The backend controls which kernel handles dequantization during the forward pass. Set the `backend` parameter in [`AutoRoundConfig`] to choose one: + +| Backend | Value | Device | Requirements | Notes | +|---------|-------|--------|--------------|-------| +| **Auto** | `"auto"` | Any | — | Default. Automatically selects the best available backend. | +| **PyTorch** | `"torch"` | CPU / CUDA | — | Pure PyTorch implementation. Broadest compatibility. | +| **Triton** | `"tritonv2"` | CUDA | `triton` | Triton-based kernel for GPU inference. | +| **ExllamaV2** | `"exllamav2"` | CUDA | `gptqmodel>=5.8.0` | Good CUDA performance via the ExllamaV2 kernel. | +| **Marlin** | `"marlin"` | CUDA | `gptqmodel>=5.8.0` | Best CUDA performance via the Marlin kernel. | + + +```python +from diffusers import AutoRoundConfig + +# Auto-select (default) +config = AutoRoundConfig() + +# Explicit Triton backend for CUDA +config = AutoRoundConfig(backend="tritonv2") + +# Marlin backend for best CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(backend="marlin") + +# ExllamaV2 backend for good CUDA performance (requires gptqmodel>=5.8.0) +config = AutoRoundConfig(backend="exllamav2") + +# PyTorch backend for CPU/CUDA inference +config = AutoRoundConfig(backend="torch") +``` + + +## Save and load + + + + +AutoRound requires data calibration to quantize a model. This is done outside of Diffusers using the [AutoRound library](https://github.com/intel/auto-round) directly: + +```python +from auto_round import AutoRound + +autoround = AutoRound( + "Tongyi-MAI/Z-Image", + scheme="W4A16", # W4G128 symmetric + enable_torch_compile=True, + num_inference_steps=3, + guidance_scale=7.5, + dataset="coco2014", +) +autoround.quantize_and_save("Z-Image-W4A16-AutoRound") +``` + +For more details on calibration options, see the [AutoRound documentation](https://github.com/intel/auto-round). + + + + + +```python +import torch +from diffusers import ZImageTransformer2DModel, ZImagePipeline + +model_id = "INCModel/Z-Image-W4A16-AutoRound" + +# The inference backend will be automatically selected. +pipe = ZImagePipeline.from_pretrained( + model_id, + torch_dtype=torch.bfloat16, + device_map="cuda", +) + +image = pipe("a cat holding a sign that says hello").images[0] +image.save("output.png") +``` + + + + +### Supported Quantization Schemes + +AutoRound supports several Schemes: + +- **W4A16**(bits:4,group_size:128,sym:True,act_bits:16) +- **W8A16**(bits:8,group_size:128,sym:True,act_bits:16) +- **W3A16**(bits:3,group_size:128,sym:True,act_bits:16) +- **W2A16**(bits:2,group_size:128,sym:True,act_bits:16) +- **GGUF:Q4_K_M**(all Q*_K,Q*_0,Q*_1 provided by llamacpp are supported) +- **NVFP4**(Experimental feature, recommend exporting to `llm_compressor` format.data_type nvfp4,act_data_type nvfp4,static_global_scale,group_size 16) +- **MXFP4**(**Research feature, no real kernel**, Standard MXFP4, data_type mxfp,act_data_type mxfp,bits 4, act_bits 4, group_size 32) +- **MXINT4**(**Research feature, no real kernel**, Standard MXINT4, data_type mxint,act_data_type mxint,bits 4, act_bits 4, group_size 32) +- **MXFP4_RCEIL**(**Research feature,no real kernel**, NVIDIA's variant, data_type mxfp,act_data_type mxfp_rceil,bits 4, act_bits 4, group_size 32) +- **MXFP8**(**Research feature, no real kernel**, data_type mxfp,act_data_type mxfp_rceil,group_size 32) +- **FPW8A16**(**Research feature, no real kernel**, data_type fp8,group_size 0->per tensor ) +- **FP8_STATIC**(**Research feature, no real kernel**, data_type:fp8,act_data_type:fp8,group_size -1 ->per channel, act_group_size=0->per tensor) + +Besides, you could modify the `group_size`, `bits`, `sym` and many other configs you want, though there are maybe no real kernels. + +## Resources + +- [Pre-quantized AutoRound models on the Hub](https://huggingface.co/models?search=autoround) diff --git a/setup.py b/setup.py index bc8110bbc594..a9bafaff5399 100644 --- a/setup.py +++ b/setup.py @@ -130,6 +130,7 @@ "onnx", "optimum_quanto>=0.2.6", "gguf>=0.10.0", + "auto-round>=0.13.0", "torchao>=0.7.0", "bitsandbytes>=0.43.3", "nvidia_modelopt[hf]>=0.33.1", diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 4a2c3bca5bcc..0f4eb50a709a 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -7,6 +7,7 @@ OptionalDependencyNotAvailable, _LazyModule, is_accelerate_available, + is_auto_round_available, is_bitsandbytes_available, is_flax_available, is_gguf_available, @@ -123,6 +124,18 @@ else: _import_structure["quantizers.quantization_config"].append("NVIDIAModelOptConfig") +try: + if not is_auto_round_available(): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from .utils import dummy_auto_round_objects + + _import_structure["utils.dummy_auto_round_objects"] = [ + name for name in dir(dummy_auto_round_objects) if not name.startswith("_") + ] +else: + _import_structure["quantizers.quantization_config"].append("AutoRoundConfig") + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() @@ -982,6 +995,14 @@ else: from .quantizers.quantization_config import NVIDIAModelOptConfig + try: + if not is_auto_round_available(): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from .utils.dummy_auto_round_objects import * + else: + from .quantizers.quantization_config import AutoRoundConfig + try: if not is_onnx_available(): raise OptionalDependencyNotAvailable() diff --git a/src/diffusers/dependency_versions_table.py b/src/diffusers/dependency_versions_table.py index 747d1011aa40..3aac2f280af6 100644 --- a/src/diffusers/dependency_versions_table.py +++ b/src/diffusers/dependency_versions_table.py @@ -37,6 +37,7 @@ "onnx": "onnx", "optimum_quanto": "optimum_quanto>=0.2.6", "gguf": "gguf>=0.10.0", + "auto-round": "auto-round>=0.13.0", "torchao": "torchao>=0.7.0", "bitsandbytes": "bitsandbytes>=0.43.3", "nvidia_modelopt[hf]": "nvidia_modelopt[hf]>=0.33.1", diff --git a/src/diffusers/quantizers/auto.py b/src/diffusers/quantizers/auto.py index 6cd24c459c9d..a10bf0cdcb3f 100644 --- a/src/diffusers/quantizers/auto.py +++ b/src/diffusers/quantizers/auto.py @@ -18,10 +18,12 @@ import warnings +from .autoround import AutoRoundQuantizer from .bitsandbytes import BnB4BitDiffusersQuantizer, BnB8BitDiffusersQuantizer from .gguf import GGUFQuantizer from .modelopt import NVIDIAModelOptQuantizer from .quantization_config import ( + AutoRoundConfig, BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, @@ -41,6 +43,7 @@ "quanto": QuantoQuantizer, "torchao": TorchAoHfQuantizer, "modelopt": NVIDIAModelOptQuantizer, + "auto-round": AutoRoundQuantizer, } AUTO_QUANTIZATION_CONFIG_MAPPING = { @@ -50,6 +53,7 @@ "quanto": QuantoConfig, "torchao": TorchAoConfig, "modelopt": NVIDIAModelOptConfig, + "auto-round": AutoRoundConfig, } @@ -143,6 +147,19 @@ def merge_quantization_configs( if isinstance(quantization_config, NVIDIAModelOptConfig): quantization_config.check_model_patching() + if quantization_config_from_args is not None and isinstance(quantization_config, AutoRoundConfig): + # For AutoRound, allow overriding fields like `backend` from user args, + # since the model config may store a default value (e.g. backend="auto"). + for key, value in quantization_config_from_args.__dict__.items(): + if key in ("quant_method",): + continue + if hasattr(quantization_config, key) and getattr(quantization_config, key) != value: + warnings.warn( + f"Overriding `{key}` in the model's quantization_config with value {value!r} " + f"from the user-provided `quantization_config`." + ) + setattr(quantization_config, key, value) + if warning_msg != "": warnings.warn(warning_msg) diff --git a/src/diffusers/quantizers/autoround/__init__.py b/src/diffusers/quantizers/autoround/__init__.py new file mode 100644 index 000000000000..2fe2083d4a5f --- /dev/null +++ b/src/diffusers/quantizers/autoround/__init__.py @@ -0,0 +1 @@ +from .autoround_quantizer import AutoRoundQuantizer diff --git a/src/diffusers/quantizers/autoround/autoround_quantizer.py b/src/diffusers/quantizers/autoround/autoround_quantizer.py new file mode 100644 index 000000000000..f80563fed406 --- /dev/null +++ b/src/diffusers/quantizers/autoround/autoround_quantizer.py @@ -0,0 +1,122 @@ +# Copyright 2025 The Intel and The HuggingFace Inc. teams. All rights reserved. +# +# 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. + +from typing import TYPE_CHECKING + +from ...utils import is_auto_round_available, logging +from ..base import DiffusersQuantizer + + +if TYPE_CHECKING: + from ...models.modeling_utils import ModelMixin + + +logger = logging.get_logger(__name__) + + +class AutoRoundQuantizer(DiffusersQuantizer): + r""" + Diffusers Quantizer for AutoRound (https://github.com/intel/auto-round). + + AutoRound is a weight-only quantization method that uses sign gradient descent to jointly optimize rounding values + and min-max ranges for weights. It supports W4A16 (4-bit weight, 16-bit activation) quantization for efficient + inference. + + This quantizer only supports loading pre-quantized AutoRound models. On-the-fly quantization (calibration) is not + supported through this interface. + """ + + # AutoRound requires data calibration — we only support loading pre-quantized checkpoints. + requires_calibration = True + required_packages = ["auto_round"] + + def __init__(self, quantization_config, **kwargs): + super().__init__(quantization_config, **kwargs) + + def validate_environment(self, *args, **kwargs): + """ + Validates that the auto-round library (>= 0.5) is installed and captures the device_map for later use during + model conversion. + """ + self.device_map = kwargs.get("device_map", None) + if not is_auto_round_available(): + raise ImportError( + "Loading an AutoRound quantized model requires the auto-round library " + "(`pip install 'auto-round>=0.13.0'`)" + ) + if not self.pre_quantized: + raise ValueError( + "AutoRound quantizer in diffusers only supports loading pre-quantized models. " + "To quantize a model from scratch, use the AutoRound CLI or Python API " + "(https://github.com/intel/auto-round) directly, then load the result with Diffusers." + ) + + def _process_model_before_weight_loading( + self, + model: "ModelMixin", + device_map, + keep_in_fp32_modules: list[str] = [], + **kwargs, + ): + """ + Replaces target nn.Linear layers with AutoRound's quantized QuantLinear layers before weights are loaded from + the checkpoint. + + Uses `auto_round.inference.convert_model.convert_hf_model` which: + - Inspects the model architecture and the quantization config (bits, group_size, sym, backend). + - Replaces eligible nn.Linear modules with the appropriate QuantLinear variant (the packed-weight layer that + stores qweight, scales, qzeros). + - Returns the converted model and a set of used backend names. + + `infer_target_device` resolves the device_map into a single target device string that AutoRound uses to select + the correct kernel backend (e.g. "cuda", "cpu"). + """ + from auto_round.inference.convert_model import convert_hf_model, infer_target_device + + target_device = infer_target_device(self.device_map) + model, used_backends = convert_hf_model(model, target_device) + self.used_backends = used_backends + + def _process_model_after_weight_loading(self, model, **kwargs): + """ + Finalizes the model after all quantized weights (qweight, scales, qzeros, etc.) have been loaded into the + QuantLinear layers. + + Uses `auto_round.inference.convert_model.post_init` which: + - Performs backend-specific finalization (e.g. repacking weights into the kernel's expected memory layout, + moving buffers to the correct device). + - Freezes quantized parameters (requires_grad=False). + - Prepares the model for inference. + + """ + from auto_round.inference.convert_model import post_init + + post_init(model, self.used_backends) + + return model + + @property + def is_trainable(self) -> bool: + """AutoRound W4A16 pre-quantized models do not support training.""" + return False + + @property + def is_serializable(self): + """AutoRound quantized models can be serialized (the quantization config may be + updated by the backend, e.g. for GPTQ/AWQ-compatible formats).""" + return True + + @property + def is_compileable(self) -> bool: + return True diff --git a/src/diffusers/quantizers/quantization_config.py b/src/diffusers/quantizers/quantization_config.py index c3d829fde8cf..0c98e40ba962 100644 --- a/src/diffusers/quantizers/quantization_config.py +++ b/src/diffusers/quantizers/quantization_config.py @@ -48,6 +48,7 @@ class QuantizationMethod(str, Enum): TORCHAO = "torchao" QUANTO = "quanto" MODELOPT = "modelopt" + AUTOROUND = "auto-round" @dataclass @@ -749,3 +750,81 @@ def get_config_from_quant_type(self) -> dict[str, Any]: ) return BASE_CONFIG + + +@dataclass +class AutoRoundConfig(QuantizationConfigMixin): + """Configuration class for AutoRound quantization. + + AutoRound is a weight-only quantization algorithm that uses sign gradient descent to jointly optimize weight + rounding and min-max values. This config targets the W4A16 (4-bit weights, 16-bit activations) setting. + + Reference: https://github.com/intel/auto-round + + Args: + bits (`int`, *optional*, defaults to `4`): + The number of bits to quantize weights to. For W4A16 this should be 4. + group_size (`int`, *optional*, defaults to `128`): + The group size for weight quantization. Weights in each group share the same scale and zero-point. Common + choices: 32, 64, 128, -1 (per-channel). + sym (`bool`, *optional*, defaults to `True`): + Whether to use symmetric quantization (zero-point fixed at 0) or asymmetric quantization (zero-point is + learned). + backend (`str`, *optional*, defaults to `"auto"`): + The backend kernel to use for quantized inference. Available backends: + - `"auto"`: Automatically select the best available backend for the current device. + - `"torch"`: Pure PyTorch kernel — works on CPU and CUDA. + - `"tritonv2"`: Triton-based kernel — requires CUDA. + - `"exllamav2"`: Exllamav2 kernel via GPTQModel — requires CUDA and `gptqmodel>=5.8.0`. Offers good CUDA + inference performance. + - `"marlin"`: Marlin kernel via GPTQModel — requires CUDA and `gptqmodel>=5.8.0`. Offers the best CUDA + inference performance. + kwargs (`dict[str, Any]`, *optional*): + Additional keyword arguments forwarded to AutoRound (e.g. `iters`, `seqlen`, `batch_size`, `lr`, + `minmax_lr` for calibration when quantizing from scratch). + """ + + VALID_BACKENDS = ["auto", "torch", "tritonv2", "exllamav2", "marlin"] + + def __init__( + self, + bits: int = 4, + group_size: int = 128, + sym: bool = True, + backend: str = "auto", + **kwargs, + ) -> None: + self.quant_method = QuantizationMethod.AUTOROUND + self._validate_backend(backend) + self.bits = bits + self.group_size = group_size + self.sym = sym + self.backend = backend + for k, v in kwargs.items(): + setattr(self, k, v) + + def _validate_backend(self, backend): + if backend not in self.VALID_BACKENDS: + raise ValueError(f"Invalid backend '{backend}'. Valid options are: {self.VALID_BACKENDS}") + + def to_dict(self) -> dict: + """Serialize the config to a JSON-compatible dict. + + Output: A dict containing all config fields. The `quant_method` is stored as its string value so it can be + round-tripped through JSON. + """ + output = super().to_dict() + output["quant_method"] = output["quant_method"].value + return output + + @classmethod + def from_dict(cls, config_dict: dict, return_unused_kwargs: bool = False, **kwargs): + """Instantiate an AutoRoundConfig from a dictionary. + + Input: config_dict with keys like bits, group_size, sym, etc. Output: An AutoRoundConfig instance (and + optionally unused kwargs). + """ + # Filter out keys that are not constructor parameters + # (e.g. quant_method is set automatically) + config_dict = {k: v for k, v in config_dict.items() if k != "quant_method"} + return super().from_dict(config_dict, return_unused_kwargs=return_unused_kwargs, **kwargs) diff --git a/src/diffusers/utils/__init__.py b/src/diffusers/utils/__init__.py index 10ad75d92f17..5cd6885e0364 100644 --- a/src/diffusers/utils/__init__.py +++ b/src/diffusers/utils/__init__.py @@ -69,6 +69,7 @@ is_accelerate_version, is_aiter_available, is_aiter_version, + is_auto_round_available, is_av_available, is_better_profanity_available, is_bitsandbytes_available, diff --git a/src/diffusers/utils/dummy_auto_round_objects.py b/src/diffusers/utils/dummy_auto_round_objects.py new file mode 100644 index 000000000000..be7a6b8403cb --- /dev/null +++ b/src/diffusers/utils/dummy_auto_round_objects.py @@ -0,0 +1,17 @@ +# This file is autogenerated by the command `make fix-copies`, do not edit. +from ..utils import DummyObject, requires_backends + + +class AutoRoundConfig(metaclass=DummyObject): + _backends = ["auto_round"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["auto_round"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["auto_round"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["auto_round"]) diff --git a/src/diffusers/utils/import_utils.py b/src/diffusers/utils/import_utils.py index ce439bfecbf2..a0fa882d2705 100644 --- a/src/diffusers/utils/import_utils.py +++ b/src/diffusers/utils/import_utils.py @@ -232,6 +232,7 @@ def _is_package_available(pkg_name: str, get_dist_name: bool = False) -> tuple[b _aiter_available, _aiter_version = _is_package_available("aiter", get_dist_name=True) _kornia_available, _kornia_version = _is_package_available("kornia") _nvidia_modelopt_available, _nvidia_modelopt_version = _is_package_available("modelopt", get_dist_name=True) +_auto_round_available, _auto_round_version = _is_package_available("auto_round") _flashpack_available, _flashpack_version = _is_package_available("flashpack") _av_available, _av_version = _is_package_available("av") @@ -404,6 +405,10 @@ def is_nvidia_modelopt_available(): return _nvidia_modelopt_available +def is_auto_round_available(): + return _auto_round_available + + def is_timm_available(): return _timm_available diff --git a/tests/models/testing_utils/__init__.py b/tests/models/testing_utils/__init__.py index 0b31342ffd4a..728a7ac80248 100644 --- a/tests/models/testing_utils/__init__.py +++ b/tests/models/testing_utils/__init__.py @@ -19,6 +19,9 @@ from .memory import CPUOffloadTesterMixin, GroupOffloadTesterMixin, LayerwiseCastingTesterMixin, MemoryTesterMixin from .parallelism import ContextParallelAttentionBackendsTesterMixin, ContextParallelTesterMixin from .quantization import ( + AutoRoundCompileTesterMixin, + AutoRoundConfigMixin, + AutoRoundTesterMixin, BitsAndBytesCompileTesterMixin, BitsAndBytesConfigMixin, BitsAndBytesTesterMixin, @@ -44,6 +47,8 @@ __all__ = [ "AttentionBackendTesterMixin", "AttentionTesterMixin", + "AutoRoundConfigMixin", + "AutoRoundTesterMixin", "BaseModelTesterConfig", "BitsAndBytesCompileTesterMixin", "BitsAndBytesConfigMixin", diff --git a/tests/models/testing_utils/quantization.py b/tests/models/testing_utils/quantization.py index ded5cab52268..cae2397dfd9b 100644 --- a/tests/models/testing_utils/quantization.py +++ b/tests/models/testing_utils/quantization.py @@ -18,7 +18,14 @@ import pytest import torch -from diffusers import BitsAndBytesConfig, GGUFQuantizationConfig, NVIDIAModelOptConfig, QuantoConfig, TorchAoConfig +from diffusers import ( + AutoRoundConfig, + BitsAndBytesConfig, + GGUFQuantizationConfig, + NVIDIAModelOptConfig, + QuantoConfig, + TorchAoConfig, +) from diffusers.utils.import_utils import ( is_bitsandbytes_available, is_gguf_available, @@ -31,6 +38,7 @@ backend_empty_cache, backend_max_memory_allocated, backend_reset_peak_memory_stats, + is_autoround, is_bitsandbytes, is_gguf, is_modelopt, @@ -40,6 +48,7 @@ is_torchao, require_accelerate, require_accelerator, + require_auto_round_version_greater_or_equal, require_bitsandbytes_version_greater, require_gguf_version_greater_or_equal, require_modelopt_version_greater_or_equal, @@ -1183,7 +1192,7 @@ def teardown_method(self): torch.compiler.reset() @torch.no_grad() - def _test_torch_compile(self, config_kwargs): + def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile=True): """ Test that torch.compile works correctly with a quantized model. @@ -1196,7 +1205,7 @@ def _test_torch_compile(self, config_kwargs): model.compile(fullgraph=True) - with torch._dynamo.config.patch(error_on_recompile=True): + with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): inputs = self.get_dummy_inputs() output = model(**inputs, return_dict=False)[0] assert output is not None, "Model output is None" @@ -1375,3 +1384,115 @@ def test_modelopt_torch_compile(self, config_name): @pytest.mark.parametrize("config_name", ["fp8"], ids=["fp8"]) def test_modelopt_torch_compile_with_group_offload(self, config_name): self._test_torch_compile_with_group_offload(ModelOptConfigMixin.MODELOPT_CONFIGS[config_name]) + + +@is_quantization +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundConfigMixin: + """ + Base mixin providing AutoRound quantization config and model creation. + + AutoRound is a weight-only quantization method (W4A16). It supports multiple inference + + When `backend="auto"`, AutoRound selects the best available backend automatically. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - quantized_model_name_or_path: Hub repository ID for the quantized model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained + """ + + config_dict = {"backend": "auto"} + + def _load_unquantized_model(self): + kwargs = getattr(self, "pretrained_model_kwargs", {}) + return self.model_class.from_pretrained(self.pretrained_model_name_or_path, **kwargs) + + def _create_quantized_model(self, config_kwargs, **extra_kwargs): + config = AutoRoundConfig(**config_kwargs) + kwargs = getattr(self, "pretrained_model_kwargs", {}).copy() + kwargs["quantization_config"] = config + kwargs["torch_dtype"] = torch.bfloat16 + if "device_map" not in kwargs: + kwargs["device_map"] = torch_device + kwargs.update(extra_kwargs) + return self.model_class.from_pretrained(self.quantized_model_name_or_path, **kwargs) + + def _verify_if_layer_quantized(self, name, module, config_kwargs): + # AutoRound replaces linear layers with quantized linear layers + assert isinstance(module, torch.nn.Linear), f"Layer {name} is not Linear, got {type(module)}" + + +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundTesterMixin(AutoRoundConfigMixin, QuantizationTesterMixin): + """ + Mixin class for testing AutoRound quantization on models. + + Expected class attributes: + - model_class: The model class to test + - pretrained_model_name_or_path: Hub repository ID for the pretrained model + - quantized_model_name_or_path: Hub repository ID for the quantized model + - pretrained_model_kwargs: (Optional) Dict of kwargs to pass to from_pretrained (e.g., {"subfolder": "transformer"}) + + Expected methods to be implemented by subclasses: + - get_dummy_inputs(): Returns dict of inputs to pass to the model forward pass + + Optional class attributes: + - AUTOROUND_CONFIGS: Dict of config name -> AutoRoundConfig kwargs to test + + Pytest mark: autoround + Use `pytest -m "not autoround"` to skip these tests + """ + + config_dict = {"backend": "auto"} + + def test_autoround_quantization_memory_footprint(self): + expected = 1.5 # AutoRound is a W4A16 method, so we expect around 1.5x memory reduction + self._test_quantization_memory_footprint(self.config_dict, expected_memory_reduction=expected) + + def test_autoround_quantization_inference(self): + self._test_quantization_inference(self.config_dict) + + def test_autoround_device_map(self): + """Test that device_map='auto' works correctly with quantization.""" + self._test_quantization_device_map(self.config_dict) + + +@is_autoround +@require_accelerator +@require_accelerate +@require_auto_round_version_greater_or_equal("0.13.0") +class AutoRoundCompileTesterMixin(AutoRoundConfigMixin, QuantizationCompileTesterMixin): + """ + Mixin class for testing `torch.compile` with AutoRound-quantized models. + + This mixin provides tests that verify `torch.compile` works correctly with models + quantized using AutoRound. Subclasses are expected to inherit from + `AutoRoundConfigMixin` (which defines `config_dict`) and to provide the + following class attributes: `model_class`, `pretrained_model_name_or_path`, and + `quantized_model_name_or_path`. + + The mixin uses `config_dict` (defaults to {"backend": "auto"}) as the + quantization configuration passed into `_create_quantized_model` when + invoking the compile-related tests. + + Provided tests: + - `test_autoround_torch_compile`: Ensures `torch.compile` runs and produces + valid, non-NaN outputs for an AutoRound-quantized model. + - `test_autoround_torch_compile_with_group_offload`: Ensures `torch.compile` + works together with group offloading when supported by the quantized + model implementation. + """ + + def test_autoround_torch_compile(self): + self._test_torch_compile(self.config_dict, fullgraph=False, error_on_recompile=False) + + def test_autoround_torch_compile_with_group_offload(self): + self._test_torch_compile_with_group_offload(self.config_dict) diff --git a/tests/models/transformers/test_models_transformer_z_image.py b/tests/models/transformers/test_models_transformer_z_image.py index c2e7990e0424..3a0fe18bc692 100644 --- a/tests/models/transformers/test_models_transformer_z_image.py +++ b/tests/models/transformers/test_models_transformer_z_image.py @@ -22,6 +22,8 @@ from ...testing_utils import assert_tensors_close, torch_device from ..testing_utils import ( + AutoRoundCompileTesterMixin, + AutoRoundTesterMixin, BaseModelTesterConfig, LoraTesterMixin, MemoryTesterMixin, @@ -309,3 +311,122 @@ def test_compile_works_with_aot(self, tmp_path): @pytest.mark.skip("Fullgraph is broken") def test_compile_on_different_shapes(self): pass + + +class ZImageTransformerAutoRoundTesterConfig: + """Configuration class for Z-Image Transformer AutoRound quantization tests.""" + + @property + def model_class(self): + return ZImageTransformer2DModel + + @property + def pretrained_model_name_or_path(self): + return "INCModel/Z-Image-tiny-for-testing" + + @property + def quantized_model_name_or_path(self): + return "INCModel/Z-Image-tiny-for-testing-W4A16-AutoRound" + + @property + def pretrained_model_kwargs(self): + return {"subfolder": "transformer"} + + def get_dummy_inputs(self): + batch_size = 1 + in_channels = 16 + cap_feat_dim = 512 + height = width = 8 + frames = 1 + seq_len = 16 + + torch.manual_seed(0) + x = [ + torch.randn((in_channels, frames, height, width)).to(torch_device, dtype=torch.bfloat16) + for _ in range(batch_size) + ] + cap_feats = [ + torch.randn((seq_len, cap_feat_dim)).to(torch_device, dtype=torch.bfloat16) for _ in range(batch_size) + ] + t = torch.tensor([0.5]).to(torch_device, dtype=torch.bfloat16) + + return {"x": x, "cap_feats": cap_feats, "t": t} + + +class TestZImageTransformerAutoRound(ZImageTransformerAutoRoundTesterConfig, AutoRoundTesterMixin): + """AutoRound quantization tests for Z-Image Transformer.""" + + @torch.no_grad() + def _test_quantization_inference(self, config_kwargs): + model_quantized = self._create_quantized_model(config_kwargs) + model_quantized.to(torch_device) + + inputs = self.get_dummy_inputs() + output = model_quantized(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_quantization_device_map(self, config_kwargs): + model = self._create_quantized_model(config_kwargs, device_map="auto") + + assert hasattr(model, "hf_device_map"), "Model should have hf_device_map attribute" + assert model.hf_device_map is not None, "hf_device_map should not be None" + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + +class TestZImageTransformerAutoRoundCompile(ZImageTransformerAutoRoundTesterConfig, AutoRoundCompileTesterMixin): + """AutoRound quantization + torch.compile tests for Z-Image Transformer.""" + + @torch.no_grad() + def _test_torch_compile(self, config_kwargs, fullgraph=True, error_on_recompile=True): + model = self._create_quantized_model(config_kwargs) + model.to(torch_device) + model.eval() + + model = torch.compile(model, fullgraph=fullgraph) + + with torch._dynamo.config.patch(error_on_recompile=error_on_recompile): + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" + + @torch.no_grad() + def _test_torch_compile_with_group_offload(self, config_kwargs, use_stream=False): + import pytest + + torch._dynamo.config.cache_size_limit = 1000 + + model = self._create_quantized_model(config_kwargs) + model.eval() + + if not hasattr(model, "enable_group_offload"): + pytest.skip("Model does not support group offloading") + + group_offload_kwargs = { + "onload_device": torch.device(torch_device), + "offload_device": torch.device("cpu"), + "offload_type": "leaf_level", + "use_stream": use_stream, + } + model.enable_group_offload(**group_offload_kwargs) + model = torch.compile(model) + + inputs = self.get_dummy_inputs() + output = model(**inputs, return_dict=False)[0] + # Z-Image returns a list of tensors from unpatchify + output = output[0] if isinstance(output, (list, tuple)) else output + assert output is not None, "Model output is None" + assert not torch.isnan(output).any(), "Model output contains NaN" diff --git a/tests/others/test_dependencies.py b/tests/others/test_dependencies.py index b2e28077b131..4a33bd529c15 100644 --- a/tests/others/test_dependencies.py +++ b/tests/others/test_dependencies.py @@ -39,6 +39,8 @@ def test_backend_registration(self): backend = "opencv-python" elif backend == "nvidia_modelopt": backend = "nvidia_modelopt[hf]" + elif backend == "auto_round": + backend = "auto-round" assert backend in deps, f"{backend} is not in the deps table!" def test_pipeline_imports(self): diff --git a/tests/testing_utils.py b/tests/testing_utils.py index 6d6df8b24d1e..a8306b3d65f8 100644 --- a/tests/testing_utils.py +++ b/tests/testing_utils.py @@ -32,6 +32,7 @@ from diffusers.utils.import_utils import ( BACKENDS_MAPPING, is_accelerate_available, + is_auto_round_available, is_bitsandbytes_available, is_compel_available, is_flashpack_available, @@ -449,6 +450,15 @@ def is_gguf(test_case): return pytest.mark.gguf(test_case) +def is_autoround(test_case): + """ + Decorator marking a test as an AutoRound quantization test. These tests can be filtered using: + pytest -m "not autoround" to skip + pytest -m autoround to run only these tests + """ + return pytest.mark.autoround(test_case) + + def is_modelopt(test_case): """ Decorator marking a test as a NVIDIA ModelOpt quantization test. These tests can be filtered using: @@ -836,6 +846,19 @@ def decorator(test_case): return decorator +def require_auto_round_version_greater_or_equal(auto_round_version): + def decorator(test_case): + correct_auto_round_version = is_auto_round_available() and version.parse( + version.parse(importlib.metadata.version("auto_round")).base_version + ) >= version.parse(auto_round_version) + return pytest.mark.skipif( + not correct_auto_round_version, + reason=f"Test requires auto-round with version greater than {auto_round_version}.", + )(test_case) + + return decorator + + def require_kernels_version_greater_or_equal(kernels_version): def decorator(test_case): correct_kernels_version = is_kernels_available() and version.parse( From 33457566e3db25121f9ec4a1c6337a4b2c9e8d25 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 10 Jun 2026 12:33:34 +0530 Subject: [PATCH 11/13] [tests] refactor UNet model tests to align with the new pattern (#13153) * refactor unet2d condition model tests. * fix tests * up * fix * Revert "fix" This reverts commit 46d44b73d8d703070912896ee47ff1b60f385305. * up * recompile limit * [tests] refactor test_models_unet_1d.py to use modular testing mixins Refactor UNet1D model tests to follow the modern testing pattern using BaseModelTesterConfig and focused mixin classes (ModelTesterMixin, MemoryTesterMixin, TrainingTesterMixin, LoraTesterMixin). Both UNet1D standard and RL variants now have separate config classes and dedicated test classes organized by concern (core, memory, training, LoRA, hub loading). Co-Authored-By: Claude Opus 4.6 * [tests] refactor test_models_unet_2d.py to use modular testing mixins Refactor UNet2D model tests (standard, LDM, NCSN++) to follow the modern testing pattern. Each variant gets its own config class and dedicated test classes organized by concern (core, memory, training, LoRA, hub loading). Co-Authored-By: Claude Opus 4.6 * [tests] refactor test_models_unet_3d_condition.py to use modular testing mixins Refactor UNet3DConditionModel tests to follow the modern testing pattern with separate classes for core, attention, memory, training, and LoRA. Co-Authored-By: Claude Opus 4.6 * [tests] refactor test_models_unet_controlnetxs.py to use modular testing mixins Refactor UNetControlNetXSModel tests to follow the modern testing pattern with separate classes for core, memory, training, and LoRA. Specialized tests (from_unet, freeze_unet, forward_no_control, time_embedding_mixing) remain in the core test class. Co-Authored-By: Claude Opus 4.6 * [tests] refactor test_models_unet_spatiotemporal.py to use modular testing mixins Refactored the spatiotemporal UNet test file to follow the modern modular testing pattern with BaseModelTesterConfig and focused test classes: - UNetSpatioTemporalTesterConfig: Base configuration with model setup - TestUNetSpatioTemporal: Core model tests (ModelTesterMixin, UNetTesterMixin) - TestUNetSpatioTemporalAttention: Attention-related tests (AttentionTesterMixin) - TestUNetSpatioTemporalMemory: Memory/offloading tests (MemoryTesterMixin) - TestUNetSpatioTemporalTraining: Training tests (TrainingTesterMixin) - TestUNetSpatioTemporalLoRA: LoRA adapter tests (LoraTesterMixin) Co-Authored-By: Claude Opus 4.6 * remove test suites that are passed. * fix consistencydecodervae tests * Revert "fix consistencydecodervae tests" This reverts commit 41b036b9891ab8209b73be4c15e3967d5885f6e9. --------- Co-authored-by: Claude Opus 4.6 Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> --- tests/models/test_modeling_common.py | 7 +- tests/models/testing_utils/common.py | 38 +- tests/models/testing_utils/compile.py | 3 - .../unets/test_models_unet_2d_condition.py | 691 ++++++++++-------- 4 files changed, 398 insertions(+), 341 deletions(-) diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index dc961c70c0fe..8575439649d7 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -465,7 +465,8 @@ def _accepts_norm_num_groups(model_class): def test_forward_with_norm_groups(self): if not self._accepts_norm_num_groups(self.model_class): pytest.skip(f"Test not supported for {self.model_class.__name__}") - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["norm_num_groups"] = 16 init_dict["block_out_channels"] = (16, 32) @@ -480,9 +481,9 @@ def test_forward_with_norm_groups(self): if isinstance(output, dict): output = output.to_tuple()[0] - self.assertIsNotNone(output) + assert output is not None expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == expected_shape, "Input and output shapes do not match" class ModelTesterMixin: diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py index ba060b3b120d..ba119b9a212f 100644 --- a/tests/models/testing_utils/common.py +++ b/tests/models/testing_utils/common.py @@ -292,8 +292,9 @@ def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5): f"Parameter shape mismatch for {param_name}. Original: {param_1.shape}, loaded: {param_2.shape}" ) - image = model(**self.get_dummy_inputs(), return_dict=False)[0] - new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] + inputs_dict = self.get_dummy_inputs() + image = model(**inputs_dict, return_dict=False)[0] + new_image = new_model(**inputs_dict, return_dict=False)[0] assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") @@ -313,8 +314,9 @@ def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0): new_model.to(torch_device) - image = model(**self.get_dummy_inputs(), return_dict=False)[0] - new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] + inputs_dict = self.get_dummy_inputs() + image = model(**inputs_dict, return_dict=False)[0] + new_image = new_model(**inputs_dict, return_dict=False)[0] assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") @@ -342,8 +344,9 @@ def test_determinism(self, atol=1e-5, rtol=0): model.to(torch_device) model.eval() - first = model(**self.get_dummy_inputs(), return_dict=False)[0] - second = model(**self.get_dummy_inputs(), return_dict=False)[0] + inputs_dict = self.get_dummy_inputs() + first = model(**inputs_dict, return_dict=False)[0] + second = model(**inputs_dict, return_dict=False)[0] first_flat = first.flatten() second_flat = second.flatten() @@ -400,8 +403,9 @@ def recursive_check(tuple_object, dict_object): model.to(torch_device) model.eval() - outputs_dict = model(**self.get_dummy_inputs()) - outputs_tuple = model(**self.get_dummy_inputs(), return_dict=False) + inputs_dict = self.get_dummy_inputs() + outputs_dict = model(**inputs_dict) + outputs_tuple = model(**inputs_dict, return_dict=False) recursive_check(outputs_tuple, outputs_dict) @@ -528,8 +532,10 @@ def test_sharded_checkpoints(self, tmp_path, atol=1e-5, rtol=0): new_model = new_model.to(torch_device) torch.manual_seed(0) - inputs_dict_new = self.get_dummy_inputs() - new_output = new_model(**inputs_dict_new, return_dict=False)[0] + # Re-create inputs only if they contain a generator (which needs to be reset) + if "generator" in inputs_dict: + inputs_dict = self.get_dummy_inputs() + new_output = new_model(**inputs_dict, return_dict=False)[0] assert_tensors_close( base_output, new_output, atol=atol, rtol=rtol, msg="Output should match after sharded save/load" @@ -568,8 +574,10 @@ def test_sharded_checkpoints_with_variant(self, tmp_path, atol=1e-5, rtol=0): new_model = new_model.to(torch_device) torch.manual_seed(0) - inputs_dict_new = self.get_dummy_inputs() - new_output = new_model(**inputs_dict_new, return_dict=False)[0] + # Re-create inputs only if they contain a generator (which needs to be reset) + if "generator" in inputs_dict: + inputs_dict = self.get_dummy_inputs() + new_output = new_model(**inputs_dict, return_dict=False)[0] assert_tensors_close( base_output, new_output, atol=atol, rtol=rtol, msg="Output should match after variant sharded save/load" @@ -619,8 +627,10 @@ def test_sharded_checkpoints_with_parallel_loading(self, tmp_path, atol=1e-5, rt model_parallel = model_parallel.to(torch_device) torch.manual_seed(0) - inputs_dict_parallel = self.get_dummy_inputs() - output_parallel = model_parallel(**inputs_dict_parallel, return_dict=False)[0] + # Re-create inputs only if they contain a generator (which needs to be reset) + if "generator" in inputs_dict: + inputs_dict = self.get_dummy_inputs() + output_parallel = model_parallel(**inputs_dict, return_dict=False)[0] assert_tensors_close( base_output, output_parallel, atol=atol, rtol=rtol, msg="Output should match with parallel loading" diff --git a/tests/models/testing_utils/compile.py b/tests/models/testing_utils/compile.py index 998b88fb469e..4787d0742b18 100644 --- a/tests/models/testing_utils/compile.py +++ b/tests/models/testing_utils/compile.py @@ -92,9 +92,6 @@ def test_torch_compile_repeated_blocks(self, recompile_limit=1): model.eval() model.compile_repeated_blocks(fullgraph=True) - if self.model_class.__name__ == "UNet2DConditionModel": - recompile_limit = 2 - with ( torch._inductor.utils.fresh_inductor_cache(), torch._dynamo.config.patch(recompile_limit=recompile_limit), diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index 4dbb8ca7c075..a7293208d370 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/tests/models/unets/test_models_unet_2d_condition.py @@ -20,6 +20,7 @@ import unittest from collections import OrderedDict +import pytest import torch from huggingface_hub import snapshot_download from parameterized import parameterized @@ -52,17 +53,24 @@ torch_all_close, torch_device, ) -from ..test_modeling_common import ( +from ..test_modeling_common import UNetTesterMixin +from ..testing_utils import ( + AttentionTesterMixin, + BaseModelTesterConfig, + IPAdapterTesterMixin, LoraHotSwappingForModelTesterMixin, + LoraTesterMixin, + MemoryTesterMixin, ModelTesterMixin, TorchCompileTesterMixin, - UNetTesterMixin, + TrainingTesterMixin, ) if is_peft_available(): from peft import LoraConfig - from peft.tuners.tuners_utils import BaseTunerLayer + + from ..testing_utils.lora import check_if_lora_correctly_set logger = logging.get_logger(__name__) @@ -82,16 +90,6 @@ def get_unet_lora_config(): return unet_lora_config -def check_if_lora_correctly_set(model) -> bool: - """ - Checks if the LoRA layers are correctly set with peft - """ - for module in model.modules(): - if isinstance(module, BaseTunerLayer): - return True - return False - - def create_ip_adapter_state_dict(model): # "ip_adapter" (cross-attention weights) ip_cross_attn_state_dict = {} @@ -354,34 +352,28 @@ def create_custom_diffusion_layers(model, mock_weights: bool = True): return custom_diffusion_attn_procs -class UNet2DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.TestCase): - model_class = UNet2DConditionModel - main_input_name = "sample" - # We override the items here because the unet under consideration is small. - model_split_percents = [0.5, 0.34, 0.4] +class UNet2DConditionTesterConfig(BaseModelTesterConfig): + """Base configuration for UNet2DConditionModel testing.""" @property - def dummy_input(self): - batch_size = 4 - num_channels = 4 - sizes = (16, 16) - - noise = floats_tensor((batch_size, num_channels) + sizes).to(torch_device) - time_step = torch.tensor([10]).to(torch_device) - encoder_hidden_states = floats_tensor((batch_size, 4, 8)).to(torch_device) - - return {"sample": noise, "timestep": time_step, "encoder_hidden_states": encoder_hidden_states} + def model_class(self): + return UNet2DConditionModel @property - def input_shape(self): + def output_shape(self) -> tuple[int, int, int]: return (4, 16, 16) @property - def output_shape(self): - return (4, 16, 16) + def model_split_percents(self) -> list[float]: + return [0.5, 0.34, 0.4] + + @property + def main_input_name(self) -> str: + return "sample" - def prepare_init_args_and_inputs_for_common(self): - init_dict = { + def get_init_dict(self) -> dict: + """Return UNet2D model initialization arguments.""" + return { "block_out_channels": (4, 8), "norm_num_groups": 4, "down_block_types": ("CrossAttnDownBlock2D", "DownBlock2D"), @@ -393,26 +385,24 @@ def prepare_init_args_and_inputs_for_common(self): "layers_per_block": 1, "sample_size": 16, } - inputs_dict = self.dummy_input - return init_dict, inputs_dict - @unittest.skipIf( - torch_device != "cuda" or not is_xformers_available(), - reason="XFormers attention is only available with CUDA and `xformers` installed", - ) - def test_xformers_enable_works(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) + def get_dummy_inputs(self) -> dict[str, torch.Tensor]: + """Return dummy inputs for UNet2D model.""" + batch_size = 4 + num_channels = 4 + sizes = (16, 16) - model.enable_xformers_memory_efficient_attention() + return { + "sample": floats_tensor((batch_size, num_channels) + sizes).to(torch_device), + "timestep": torch.tensor([10]).to(torch_device), + "encoder_hidden_states": floats_tensor((batch_size, 4, 8)).to(torch_device), + } - assert ( - model.mid_block.attentions[0].transformer_blocks[0].attn1.processor.__class__.__name__ - == "XFormersAttnProcessor" - ), "xformers is not enabled" +class TestUNet2DCondition(UNet2DConditionTesterConfig, ModelTesterMixin, UNetTesterMixin): def test_model_with_attention_head_dim_tuple(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -427,12 +417,13 @@ def test_model_with_attention_head_dim_tuple(self): if isinstance(output, dict): output = output.sample - self.assertIsNotNone(output) + assert output is not None expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == expected_shape, "Input and output shapes do not match" def test_model_with_use_linear_projection(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["use_linear_projection"] = True @@ -446,12 +437,13 @@ def test_model_with_use_linear_projection(self): if isinstance(output, dict): output = output.sample - self.assertIsNotNone(output) + assert output is not None expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == expected_shape, "Input and output shapes do not match" def test_model_with_cross_attention_dim_tuple(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["cross_attention_dim"] = (8, 8) @@ -465,12 +457,13 @@ def test_model_with_cross_attention_dim_tuple(self): if isinstance(output, dict): output = output.sample - self.assertIsNotNone(output) + assert output is not None expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == expected_shape, "Input and output shapes do not match" def test_model_with_simple_projection(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() batch_size, _, _, sample_size = inputs_dict["sample"].shape @@ -489,12 +482,13 @@ def test_model_with_simple_projection(self): if isinstance(output, dict): output = output.sample - self.assertIsNotNone(output) + assert output is not None expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + assert output.shape == expected_shape, "Input and output shapes do not match" def test_model_with_class_embeddings_concat(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() batch_size, _, _, sample_size = inputs_dict["sample"].shape @@ -514,12 +508,287 @@ def test_model_with_class_embeddings_concat(self): if isinstance(output, dict): output = output.sample - self.assertIsNotNone(output) + assert output is not None + expected_shape = inputs_dict["sample"].shape + assert output.shape == expected_shape, "Input and output shapes do not match" + + # see diffusers.models.attention_processor::Attention#prepare_attention_mask + # note: we may not need to fix mask padding to work for stable-diffusion cross-attn masks. + # since the use-case (somebody passes in a too-short cross-attn mask) is pretty small, + # maybe it's fine that this only works for the unclip use-case. + @mark.skip( + reason="we currently pad mask by target_length tokens (what unclip needs), whereas stable-diffusion's cross-attn needs to instead pad by remaining_length." + ) + def test_model_xattn_padding(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + model = self.model_class(**{**init_dict, "attention_head_dim": (8, 16)}) + model.to(torch_device) + model.eval() + + cond = inputs_dict["encoder_hidden_states"] + with torch.no_grad(): + full_cond_out = model(**inputs_dict).sample + assert full_cond_out is not None + + batch, tokens, _ = cond.shape + keeplast_mask = (torch.arange(tokens) == tokens - 1).expand(batch, -1).to(cond.device, torch.bool) + keeplast_out = model(**{**inputs_dict, "encoder_attention_mask": keeplast_mask}).sample + assert not keeplast_out.allclose(full_cond_out), "a 'keep last token' mask should change the result" + + trunc_mask = torch.zeros(batch, tokens - 1, device=cond.device, dtype=torch.bool) + trunc_mask_out = model(**{**inputs_dict, "encoder_attention_mask": trunc_mask}).sample + assert trunc_mask_out.allclose(keeplast_out), ( + "a mask with fewer tokens than condition, will be padded with 'keep' tokens. a 'discard-all' mask missing the final token is thus equivalent to a 'keep last' mask." + ) + + def test_pickle(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + + init_dict["block_out_channels"] = (16, 32) + init_dict["attention_head_dim"] = (8, 16) + + model = self.model_class(**init_dict) + model.to(torch_device) + + with torch.no_grad(): + sample = model(**inputs_dict).sample + + sample_copy = copy.copy(sample) + + assert (sample - sample_copy).abs().max() < 1e-4 + + def test_asymmetrical_unet(self): + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + # Add asymmetry to configs + init_dict["transformer_layers_per_block"] = [[3, 2], 1] + init_dict["reverse_transformer_layers_per_block"] = [[3, 4], 1] + + torch.manual_seed(0) + model = self.model_class(**init_dict) + model.to(torch_device) + + output = model(**inputs_dict).sample expected_shape = inputs_dict["sample"].shape - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + + # Check if input and output shapes are the same + assert output.shape == expected_shape, "Input and output shapes do not match" + + +class TestUNet2DConditionHubLoading(UNet2DConditionTesterConfig): + """Hub checkpoint loading tests for UNet2DConditionModel.""" + + @parameterized.expand( + [ + ("hf-internal-testing/unet2d-sharded-dummy", None), + ("hf-internal-testing/tiny-sd-unet-sharded-latest-format", "fp16"), + ] + ) + @require_torch_accelerator + def test_load_sharded_checkpoint_from_hub(self, repo_id, variant): + inputs_dict = self.get_dummy_inputs() + loaded_model = self.model_class.from_pretrained(repo_id, variant=variant) + loaded_model = loaded_model.to(torch_device) + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @parameterized.expand( + [ + ("hf-internal-testing/unet2d-sharded-dummy-subfolder", None), + ("hf-internal-testing/tiny-sd-unet-sharded-latest-format-subfolder", "fp16"), + ] + ) + @require_torch_accelerator + def test_load_sharded_checkpoint_from_hub_subfolder(self, repo_id, variant): + inputs_dict = self.get_dummy_inputs() + loaded_model = self.model_class.from_pretrained(repo_id, subfolder="unet", variant=variant) + loaded_model = loaded_model.to(torch_device) + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + def test_load_sharded_checkpoint_from_hub_local(self): + inputs_dict = self.get_dummy_inputs() + ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy") + loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True) + loaded_model = loaded_model.to(torch_device) + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + def test_load_sharded_checkpoint_from_hub_local_subfolder(self): + inputs_dict = self.get_dummy_inputs() + ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder") + loaded_model = self.model_class.from_pretrained(ckpt_path, subfolder="unet", local_files_only=True) + loaded_model = loaded_model.to(torch_device) + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + @parameterized.expand( + [ + ("hf-internal-testing/unet2d-sharded-dummy", None), + ("hf-internal-testing/tiny-sd-unet-sharded-latest-format", "fp16"), + ] + ) + def test_load_sharded_checkpoint_device_map_from_hub(self, repo_id, variant): + inputs_dict = self.get_dummy_inputs() + loaded_model = self.model_class.from_pretrained(repo_id, variant=variant, device_map="auto") + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + @parameterized.expand( + [ + ("hf-internal-testing/unet2d-sharded-dummy-subfolder", None), + ("hf-internal-testing/tiny-sd-unet-sharded-latest-format-subfolder", "fp16"), + ] + ) + def test_load_sharded_checkpoint_device_map_from_hub_subfolder(self, repo_id, variant): + inputs_dict = self.get_dummy_inputs() + loaded_model = self.model_class.from_pretrained(repo_id, variant=variant, subfolder="unet", device_map="auto") + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + def test_load_sharded_checkpoint_device_map_from_hub_local(self): + inputs_dict = self.get_dummy_inputs() + ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy") + loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True, device_map="auto") + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + @require_torch_accelerator + def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self): + inputs_dict = self.get_dummy_inputs() + ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder") + loaded_model = self.model_class.from_pretrained( + ckpt_path, local_files_only=True, subfolder="unet", device_map="auto" + ) + new_output = loaded_model(**inputs_dict) + + assert loaded_model + assert new_output.sample.shape == (4, 4, 16, 16) + + +class TestUNet2DConditionLoRA(UNet2DConditionTesterConfig, LoraTesterMixin): + """LoRA adapter tests for UNet2DConditionModel.""" + + @require_peft_backend + def test_load_attn_procs_raise_warning(self): + """Test that deprecated load_attn_procs method raises FutureWarning.""" + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() + model = self.model_class(**init_dict) + model.to(torch_device) + + # forward pass without LoRA + with torch.no_grad(): + non_lora_sample = model(**inputs_dict).sample + + unet_lora_config = get_unet_lora_config() + model.add_adapter(unet_lora_config) + + assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." + + # forward pass with LoRA + with torch.no_grad(): + lora_sample_1 = model(**inputs_dict).sample + + with tempfile.TemporaryDirectory() as tmpdirname: + model.save_attn_procs(tmpdirname) + model.unload_lora() + + with pytest.warns(FutureWarning, match="Using the `load_attn_procs\\(\\)` method has been deprecated"): + model.load_attn_procs(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) + + # import to still check for the rest of the stuff. + assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." + + with torch.no_grad(): + lora_sample_2 = model(**inputs_dict).sample + + assert not torch.allclose(non_lora_sample, lora_sample_1, atol=1e-4, rtol=1e-4), ( + "LoRA injected UNet should produce different results." + ) + assert torch.allclose(lora_sample_1, lora_sample_2, atol=1e-4, rtol=1e-4), ( + "Loading from a saved checkpoint should produce identical results." + ) + + @require_peft_backend + def test_save_attn_procs_raise_warning(self): + """Test that deprecated save_attn_procs method raises FutureWarning.""" + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + model.to(torch_device) + + unet_lora_config = get_unet_lora_config() + model.add_adapter(unet_lora_config) + + assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." + + with tempfile.TemporaryDirectory() as tmpdirname: + with pytest.warns(FutureWarning, match="Using the `save_attn_procs\\(\\)` method has been deprecated"): + model.save_attn_procs(os.path.join(tmpdirname)) + + +class TestUNet2DConditionMemory(UNet2DConditionTesterConfig, MemoryTesterMixin): + """Memory optimization tests for UNet2DConditionModel.""" + + +class TestUNet2DConditionTraining(UNet2DConditionTesterConfig, TrainingTesterMixin): + """Training tests for UNet2DConditionModel.""" + + def test_gradient_checkpointing_is_applied(self): + expected_set = { + "CrossAttnUpBlock2D", + "CrossAttnDownBlock2D", + "UNetMidBlock2DCrossAttn", + "UpBlock2D", + "Transformer2DModel", + "DownBlock2D", + } + super().test_gradient_checkpointing_is_applied(expected_set=expected_set) + + +class TestUNet2DConditionAttention(UNet2DConditionTesterConfig, AttentionTesterMixin): + """Attention processor tests for UNet2DConditionModel.""" + + @unittest.skipIf( + torch_device != "cuda" or not is_xformers_available(), + reason="XFormers attention is only available with CUDA and `xformers` installed", + ) + def test_xformers_enable_works(self): + init_dict = self.get_init_dict() + model = self.model_class(**init_dict) + + model.enable_xformers_memory_efficient_attention() + + assert ( + model.mid_block.attentions[0].transformer_blocks[0].attn1.processor.__class__.__name__ + == "XFormersAttnProcessor" + ), "xformers is not enabled" def test_model_attention_slicing(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -544,7 +813,7 @@ def test_model_attention_slicing(self): assert output is not None def test_model_sliceable_head_dim(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -562,21 +831,6 @@ def check_sliceable_dim_attr(module: torch.nn.Module): for module in model.children(): check_sliceable_dim_attr(module) - def test_gradient_checkpointing_is_applied(self): - expected_set = { - "CrossAttnUpBlock2D", - "CrossAttnDownBlock2D", - "UNetMidBlock2DCrossAttn", - "UpBlock2D", - "Transformer2DModel", - "DownBlock2D", - } - attention_head_dim = (8, 16) - block_out_channels = (16, 32) - super().test_gradient_checkpointing_is_applied( - expected_set=expected_set, attention_head_dim=attention_head_dim, block_out_channels=block_out_channels - ) - def test_special_attn_proc(self): class AttnEasyProc(torch.nn.Module): def __init__(self, num): @@ -618,7 +872,8 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_ma return hidden_states # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -645,7 +900,8 @@ def __call__(self, attn, hidden_states, encoder_hidden_states=None, attention_ma ] ) def test_model_xattn_mask(self, mask_dtype): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() model = self.model_class(**{**init_dict, "attention_head_dim": (8, 16), "block_out_channels": (16, 32)}) model.to(torch_device) @@ -675,39 +931,13 @@ def test_model_xattn_mask(self, mask_dtype): "masking the last token from our cond should be equivalent to truncating that token out of the condition" ) - # see diffusers.models.attention_processor::Attention#prepare_attention_mask - # note: we may not need to fix mask padding to work for stable-diffusion cross-attn masks. - # since the use-case (somebody passes in a too-short cross-attn mask) is pretty esoteric. - # maybe it's fine that this only works for the unclip use-case. - @mark.skip( - reason="we currently pad mask by target_length tokens (what unclip needs), whereas stable-diffusion's cross-attn needs to instead pad by remaining_length." - ) - def test_model_xattn_padding(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - model = self.model_class(**{**init_dict, "attention_head_dim": (8, 16)}) - model.to(torch_device) - model.eval() - - cond = inputs_dict["encoder_hidden_states"] - with torch.no_grad(): - full_cond_out = model(**inputs_dict).sample - assert full_cond_out is not None - batch, tokens, _ = cond.shape - keeplast_mask = (torch.arange(tokens) == tokens - 1).expand(batch, -1).to(cond.device, torch.bool) - keeplast_out = model(**{**inputs_dict, "encoder_attention_mask": keeplast_mask}).sample - assert not keeplast_out.allclose(full_cond_out), "a 'keep last token' mask should change the result" - - trunc_mask = torch.zeros(batch, tokens - 1, device=cond.device, dtype=torch.bool) - trunc_mask_out = model(**{**inputs_dict, "encoder_attention_mask": trunc_mask}).sample - assert trunc_mask_out.allclose(keeplast_out), ( - "a mask with fewer tokens than condition, will be padded with 'keep' tokens. a 'discard-all' mask missing the final token is thus equivalent to a 'keep last' mask." - ) +class TestUNet2DConditionCustomDiffusion(UNet2DConditionTesterConfig): + """Custom Diffusion processor tests for UNet2DConditionModel.""" def test_custom_diffusion_processors(self): - # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -733,8 +963,8 @@ def test_custom_diffusion_processors(self): assert (sample1 - sample2).abs().max() < 3e-3 def test_custom_diffusion_save_load(self): - # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -754,7 +984,7 @@ def test_custom_diffusion_save_load(self): with tempfile.TemporaryDirectory() as tmpdirname: model.save_attn_procs(tmpdirname, safe_serialization=False) - self.assertTrue(os.path.isfile(os.path.join(tmpdirname, "pytorch_custom_diffusion_weights.bin"))) + assert os.path.isfile(os.path.join(tmpdirname, "pytorch_custom_diffusion_weights.bin")) torch.manual_seed(0) new_model = self.model_class(**init_dict) new_model.load_attn_procs(tmpdirname, weight_name="pytorch_custom_diffusion_weights.bin") @@ -773,8 +1003,8 @@ def test_custom_diffusion_save_load(self): reason="XFormers attention is only available with CUDA and `xformers` installed", ) def test_custom_diffusion_xformers_on_off(self): - # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -798,41 +1028,28 @@ def test_custom_diffusion_xformers_on_off(self): assert (sample - on_sample).abs().max() < 1e-4 assert (sample - off_sample).abs().max() < 1e-4 - def test_pickle(self): - # enable deterministic behavior for gradient checkpointing - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - - init_dict["block_out_channels"] = (16, 32) - init_dict["attention_head_dim"] = (8, 16) - - model = self.model_class(**init_dict) - model.to(torch_device) - - with torch.no_grad(): - sample = model(**inputs_dict).sample - sample_copy = copy.copy(sample) +class TestUNet2DConditionIPAdapter(UNet2DConditionTesterConfig, IPAdapterTesterMixin): + """IP Adapter tests for UNet2DConditionModel.""" - assert (sample - sample_copy).abs().max() < 1e-4 - - def test_asymmetrical_unet(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - # Add asymmetry to configs - init_dict["transformer_layers_per_block"] = [[3, 2], 1] - init_dict["reverse_transformer_layers_per_block"] = [[3, 4], 1] - - torch.manual_seed(0) - model = self.model_class(**init_dict) - model.to(torch_device) + @property + def ip_adapter_processor_cls(self): + return (IPAdapterAttnProcessor, IPAdapterAttnProcessor2_0) - output = model(**inputs_dict).sample - expected_shape = inputs_dict["sample"].shape + def create_ip_adapter_state_dict(self, model): + return create_ip_adapter_state_dict(model) - # Check if input and output shapes are the same - self.assertEqual(output.shape, expected_shape, "Input and output shapes do not match") + def modify_inputs_for_ip_adapter(self, model, inputs_dict): + batch_size = inputs_dict["encoder_hidden_states"].shape[0] + # for ip-adapter image_embeds has shape [batch_size, num_image, embed_dim] + cross_attention_dim = getattr(model.config, "cross_attention_dim", 8) + image_embeds = floats_tensor((batch_size, 1, cross_attention_dim)).to(torch_device) + inputs_dict["added_cond_kwargs"] = {"image_embeds": [image_embeds]} + return inputs_dict def test_ip_adapter(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -905,7 +1122,8 @@ def test_ip_adapter(self): assert sample2.allclose(sample6, atol=1e-4, rtol=1e-4) def test_ip_adapter_plus(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() + init_dict = self.get_init_dict() + inputs_dict = self.get_dummy_inputs() init_dict["block_out_channels"] = (16, 32) init_dict["attention_head_dim"] = (8, 16) @@ -977,185 +1195,16 @@ def test_ip_adapter_plus(self): assert sample2.allclose(sample5, atol=1e-4, rtol=1e-4) assert sample2.allclose(sample6, atol=1e-4, rtol=1e-4) - @parameterized.expand( - [ - ("hf-internal-testing/unet2d-sharded-dummy", None), - ("hf-internal-testing/tiny-sd-unet-sharded-latest-format", "fp16"), - ] - ) - @require_torch_accelerator - def test_load_sharded_checkpoint_from_hub(self, repo_id, variant): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained(repo_id, variant=variant) - loaded_model = loaded_model.to(torch_device) - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @parameterized.expand( - [ - ("hf-internal-testing/unet2d-sharded-dummy-subfolder", None), - ("hf-internal-testing/tiny-sd-unet-sharded-latest-format-subfolder", "fp16"), - ] - ) - @require_torch_accelerator - def test_load_sharded_checkpoint_from_hub_subfolder(self, repo_id, variant): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained(repo_id, subfolder="unet", variant=variant) - loaded_model = loaded_model.to(torch_device) - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - def test_load_sharded_checkpoint_from_hub_local(self): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy") - loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True) - loaded_model = loaded_model.to(torch_device) - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - def test_load_sharded_checkpoint_from_hub_local_subfolder(self): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder") - loaded_model = self.model_class.from_pretrained(ckpt_path, subfolder="unet", local_files_only=True) - loaded_model = loaded_model.to(torch_device) - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - @parameterized.expand( - [ - ("hf-internal-testing/unet2d-sharded-dummy", None), - ("hf-internal-testing/tiny-sd-unet-sharded-latest-format", "fp16"), - ] - ) - def test_load_sharded_checkpoint_device_map_from_hub(self, repo_id, variant): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained(repo_id, variant=variant, device_map="auto") - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - @parameterized.expand( - [ - ("hf-internal-testing/unet2d-sharded-dummy-subfolder", None), - ("hf-internal-testing/tiny-sd-unet-sharded-latest-format-subfolder", "fp16"), - ] - ) - def test_load_sharded_checkpoint_device_map_from_hub_subfolder(self, repo_id, variant): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - loaded_model = self.model_class.from_pretrained(repo_id, variant=variant, subfolder="unet", device_map="auto") - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - def test_load_sharded_checkpoint_device_map_from_hub_local(self): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy") - loaded_model = self.model_class.from_pretrained(ckpt_path, local_files_only=True, device_map="auto") - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_torch_accelerator - def test_load_sharded_checkpoint_device_map_from_hub_local_subfolder(self): - _, inputs_dict = self.prepare_init_args_and_inputs_for_common() - ckpt_path = snapshot_download("hf-internal-testing/unet2d-sharded-dummy-subfolder") - loaded_model = self.model_class.from_pretrained( - ckpt_path, local_files_only=True, subfolder="unet", device_map="auto" - ) - new_output = loaded_model(**inputs_dict) - - assert loaded_model - assert new_output.sample.shape == (4, 4, 16, 16) - - @require_peft_backend - def test_load_attn_procs_raise_warning(self): - init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) - - # forward pass without LoRA - with torch.no_grad(): - non_lora_sample = model(**inputs_dict).sample - - unet_lora_config = get_unet_lora_config() - model.add_adapter(unet_lora_config) - - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - # forward pass with LoRA - with torch.no_grad(): - lora_sample_1 = model(**inputs_dict).sample - - with tempfile.TemporaryDirectory() as tmpdirname: - model.save_attn_procs(tmpdirname) - model.unload_lora() - - with self.assertWarns(FutureWarning) as warning: - model.load_attn_procs(os.path.join(tmpdirname, "pytorch_lora_weights.safetensors")) - - warning_message = str(warning.warnings[0].message) - assert "Using the `load_attn_procs()` method has been deprecated" in warning_message - - # import to still check for the rest of the stuff. - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - with torch.no_grad(): - lora_sample_2 = model(**inputs_dict).sample - - assert not torch.allclose(non_lora_sample, lora_sample_1, atol=1e-4, rtol=1e-4), ( - "LoRA injected UNet should produce different results." - ) - assert torch.allclose(lora_sample_1, lora_sample_2, atol=1e-4, rtol=1e-4), ( - "Loading from a saved checkpoint should produce identical results." - ) - - @require_peft_backend - def test_save_attn_procs_raise_warning(self): - init_dict, _ = self.prepare_init_args_and_inputs_for_common() - model = self.model_class(**init_dict) - model.to(torch_device) - - unet_lora_config = get_unet_lora_config() - model.add_adapter(unet_lora_config) - - assert check_if_lora_correctly_set(model), "Lora not correctly set in UNet." - - with tempfile.TemporaryDirectory() as tmpdirname: - with self.assertWarns(FutureWarning) as warning: - model.save_attn_procs(tmpdirname) - - warning_message = str(warning.warnings[0].message) - assert "Using the `save_attn_procs()` method has been deprecated" in warning_message - - -class UNet2DConditionModelCompileTests(TorchCompileTesterMixin, unittest.TestCase): - model_class = UNet2DConditionModel - def prepare_init_args_and_inputs_for_common(self): - return UNet2DConditionModelTests().prepare_init_args_and_inputs_for_common() +class TestUNet2DConditionModelCompile(UNet2DConditionTesterConfig, TorchCompileTesterMixin): + """Torch compile tests for UNet2DConditionModel.""" + def test_torch_compile_repeated_blocks(self): + return super().test_torch_compile_repeated_blocks(recompile_limit=2) -class UNet2DConditionModelLoRAHotSwapTests(LoraHotSwappingForModelTesterMixin, unittest.TestCase): - model_class = UNet2DConditionModel - def prepare_init_args_and_inputs_for_common(self): - return UNet2DConditionModelTests().prepare_init_args_and_inputs_for_common() +class TestUNet2DConditionModelLoRAHotSwap(UNet2DConditionTesterConfig, LoraHotSwappingForModelTesterMixin): + """LoRA hot-swapping tests for UNet2DConditionModel.""" @slow From 0fe108d81e10f09617b7f7d053dfb1eca091f423 Mon Sep 17 00:00:00 2001 From: Sayak Paul Date: Wed, 10 Jun 2026 13:35:29 +0530 Subject: [PATCH 12/13] [tests] fix vidtok tests (#13894) * fix vidtok tests * style * Update tests/models/autoencoders/test_models_autoencoder_vidtok.py Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> * Apply style fixes --------- Co-authored-by: dg845 <58458699+dg845@users.noreply.github.com> Co-authored-by: github-actions[bot] --- .../test_models_autoencoder_vidtok.py | 29 ++++++++++++++++--- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/models/autoencoders/test_models_autoencoder_vidtok.py b/tests/models/autoencoders/test_models_autoencoder_vidtok.py index eb2863121a21..087dca5debfa 100644 --- a/tests/models/autoencoders/test_models_autoencoder_vidtok.py +++ b/tests/models/autoencoders/test_models_autoencoder_vidtok.py @@ -19,7 +19,7 @@ from diffusers import AutoencoderVidTok from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import IS_GITHUB_ACTIONS, enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, torch_device from ..testing_utils import BaseModelTesterConfig, MemoryTesterMixin, ModelTesterMixin, TrainingTesterMixin from .testing_utils import NewAutoencoderTesterMixin @@ -27,6 +27,16 @@ enable_full_determinism() +def _run_nondeterministic(fn): + # avg_pool3d_backward_cuda has no deterministic CUDA implementation; + # temporarily relax the requirement for tests that do backward passes. + torch.use_deterministic_algorithms(False) + try: + fn() + finally: + torch.use_deterministic_algorithms(True) + + class AutoencoderVidTokTesterConfig(BaseModelTesterConfig): @property def model_class(self): @@ -82,14 +92,25 @@ def test_gradient_checkpointing_is_applied(self): expected_set = {"VidTokEncoder3D", "VidTokDecoder3D"} super().test_gradient_checkpointing_is_applied(expected_set=expected_set) - @pytest.mark.skipif(IS_GITHUB_ACTIONS, reason="Skipping test inside GitHub Actions environment") - def test_layerwise_casting_training(self): - super().test_layerwise_casting_training() + def test_training(self): + _run_nondeterministic(super().test_training) + + def test_training_with_ema(self): + _run_nondeterministic(super().test_training_with_ema) + + def test_mixed_precision_training(self): + _run_nondeterministic(super().test_mixed_precision_training) + + def test_gradient_checkpointing_equivalence(self): + _run_nondeterministic(super().test_gradient_checkpointing_equivalence) class TestAutoencoderVidTokMemory(AutoencoderVidTokTesterConfig, MemoryTesterMixin): """Memory optimization tests for AutoencoderVidTok.""" + def test_layerwise_casting_training(self): + _run_nondeterministic(super().test_layerwise_casting_training) + class TestAutoencoderVidTokSlicingTiling(AutoencoderVidTokTesterConfig, NewAutoencoderTesterMixin): """Slicing and tiling tests for AutoencoderVidTok.""" From bb26b9078bd2839469e67f17c99e583d97fee6fb Mon Sep 17 00:00:00 2001 From: DN6 Date: Tue, 16 Jun 2026 10:39:13 +0530 Subject: [PATCH 13/13] clean up --- tests/models/testing_utils/common.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/tests/models/testing_utils/common.py b/tests/models/testing_utils/common.py index 8f1ee222d2c1..5726dba9c600 100644 --- a/tests/models/testing_utils/common.py +++ b/tests/models/testing_utils/common.py @@ -295,9 +295,8 @@ def test_from_save_pretrained(self, tmp_path, atol=5e-5, rtol=5e-5): f"Parameter shape mismatch for {param_name}. Original: {param_1.shape}, loaded: {param_2.shape}" ) - inputs_dict = self.get_dummy_inputs() - image = model(**inputs_dict, return_dict=False)[0] - new_image = new_model(**inputs_dict, return_dict=False)[0] + image = model(**self.get_dummy_inputs(), return_dict=False)[0] + new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") @@ -317,9 +316,8 @@ def test_from_save_pretrained_variant(self, tmp_path, atol=5e-5, rtol=0): new_model.to(torch_device) - inputs_dict = self.get_dummy_inputs() - image = model(**inputs_dict, return_dict=False)[0] - new_image = new_model(**inputs_dict, return_dict=False)[0] + image = model(**self.get_dummy_inputs(), return_dict=False)[0] + new_image = new_model(**self.get_dummy_inputs(), return_dict=False)[0] assert_tensors_close(image, new_image, atol=atol, rtol=rtol, msg="Models give different forward passes.") @@ -347,9 +345,8 @@ def test_determinism(self, atol=1e-5, rtol=0): model.to(torch_device) model.eval() - inputs_dict = self.get_dummy_inputs() - first = model(**inputs_dict, return_dict=False)[0] - second = model(**inputs_dict, return_dict=False)[0] + first = model(**self.get_dummy_inputs(), return_dict=False)[0] + second = model(**self.get_dummy_inputs(), return_dict=False)[0] first_flat = first.flatten() second_flat = second.flatten() @@ -406,9 +403,8 @@ def recursive_check(tuple_object, dict_object): model.to(torch_device) model.eval() - inputs_dict = self.get_dummy_inputs() - outputs_dict = model(**inputs_dict) - outputs_tuple = model(**inputs_dict, return_dict=False) + outputs_dict = model(**self.get_dummy_inputs()) + outputs_tuple = model(**self.get_dummy_inputs(), return_dict=False) recursive_check(outputs_tuple, outputs_dict)