diff --git a/src/diffusers/models/unets/unet_2d_condition.py b/src/diffusers/models/unets/unet_2d_condition.py index 0f789d3961fc..736deb28c376 100644 --- a/src/diffusers/models/unets/unet_2d_condition.py +++ b/src/diffusers/models/unets/unet_2d_condition.py @@ -165,7 +165,7 @@ class conditioning with `class_embed_type` equal to `None`. """ _supports_gradient_checkpointing = True - _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D"] + _no_split_modules = ["BasicTransformerBlock", "ResnetBlock2D", "CrossAttnUpBlock2D", "UpBlock2D"] _skip_layerwise_casting_patterns = ["norm"] _repeated_blocks = ["BasicTransformerBlock"] diff --git a/tests/models/test_modeling_common.py b/tests/models/test_modeling_common.py index 01dea057def3..8309700ce106 100644 --- a/tests/models/test_modeling_common.py +++ b/tests/models/test_modeling_common.py @@ -75,7 +75,6 @@ require_torch_2, require_torch_accelerator, require_torch_accelerator_with_training, - require_torch_gpu, require_torch_multi_accelerator, require_torch_version_greater, run_test_in_subprocess, @@ -1829,8 +1828,8 @@ def test_wrong_device_map_raises_error(self, device_map, msg_substring): assert msg_substring in str(err_ctx.exception) - @parameterized.expand([0, "cuda", torch.device("cuda")]) - @require_torch_gpu + @parameterized.expand([0, torch_device, torch.device(torch_device)]) + @require_torch_accelerator def test_passing_non_dict_device_map_works(self, device_map): init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common() model = self.model_class(**init_dict).eval() @@ -1839,8 +1838,8 @@ def test_passing_non_dict_device_map_works(self, device_map): loaded_model = self.model_class.from_pretrained(tmpdir, device_map=device_map) _ = loaded_model(**inputs_dict) - @parameterized.expand([("", "cuda"), ("", torch.device("cuda"))]) - @require_torch_gpu + @parameterized.expand([("", torch_device), ("", torch.device(torch_device))]) + @require_torch_accelerator def test_passing_dict_device_map_works(self, name, device): # There are other valid dict-based `device_map` values too. It's best to refer to # the docs for those: https://huggingface.co/docs/accelerate/en/concept_guides/big_model_inference#the-devicemap. @@ -1945,7 +1944,7 @@ def test_push_to_hub_library_name(self): delete_repo(self.repo_id, token=TOKEN) -@require_torch_gpu +@require_torch_accelerator @require_torch_2 @is_torch_compile @slow @@ -2013,7 +2012,7 @@ def test_compile_with_group_offloading(self): model.eval() # TODO: Can test for other group offloading kwargs later if needed. group_offload_kwargs = { - "onload_device": "cuda", + "onload_device": torch_device, "offload_device": "cpu", "offload_type": "block_level", "num_blocks_per_group": 1, diff --git a/tests/models/unets/test_models_unet_2d_condition.py b/tests/models/unets/test_models_unet_2d_condition.py index abf44aa7447b..123dff16f8b0 100644 --- a/tests/models/unets/test_models_unet_2d_condition.py +++ b/tests/models/unets/test_models_unet_2d_condition.py @@ -358,7 +358,7 @@ class UNet2DConditionModelTests(ModelTesterMixin, UNetTesterMixin, unittest.Test 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.3, 0.4] + model_split_percents = [0.5, 0.34, 0.4] @property def dummy_input(self):