more

Author: sayakpaul

src/diffusers/models/modeling_utils.py

@@ -986,6 +986,14 @@ 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:
+            device_value = list(device_map.values())[0]
+            if isinstance(device_value, str):
+                try:
+                    device_map = {"": torch.device(device_value)}
+                except RuntimeError:
+                    raise ValueError(f"Invalid value ({device_value}) specified in the {device_map=}.")
 
         if device_map is not None:
             if low_cpu_mem_usage is None:

tests/models/unets/test_models_unet_2d_condition.py

@@ -1108,6 +1108,18 @@ 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")}])
+    @require_torch_gpu
+    def test_passing_dict_device_map_works(self, 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
+        )
+        output = loaded_model(**inputs_dict)
+        assert 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()