update

Author: sayakpaul

src/diffusers/models/modeling_utils.py

@@ -987,7 +987,7 @@ def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.P
             else:
                 device_map = {"": device_map}
         # {"": device} case.
-        elif isinstance(device_map, dict) and len(dict) == 1:
+        elif isinstance(device_map, dict) and len(device_map) == 1:
             device_value = list(device_map.values())[0]
             if isinstance(device_value, str):
                 try:

tests/models/unets/test_models_unet_2d_condition.py

@@ -1108,14 +1108,14 @@ def test_passing_non_dict_device_map_works(self, device_map):
         output = loaded_model(**inputs_dict)
         assert output.sample.shape == (4, 4, 16, 16)
 
-    @parameterized.expand([{"": "cuda"}, {"": torch.device("cuda")}, {"": "cpu"}, {"": torch.device("cpu")}])
+    @parameterized.expand([("", "cuda"), ("", torch.device("cuda"))])
     @require_torch_gpu
-    def test_passing_dict_device_map_works(self, device_map):
+    def test_passing_dict_device_map_works(self, name, device_map):
         # 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.
         _, inputs_dict = self.prepare_init_args_and_inputs_for_common()
         loaded_model = self.model_class.from_pretrained(
-            "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map=device_map
+            "hf-internal-testing/unet2d-sharded-dummy-subfolder", subfolder="unet", device_map={name: device_map}
         )
         output = loaded_model(**inputs_dict)
         assert output.sample.shape == (4, 4, 16, 16)