up

Author: sayakpaul

src/diffusers/pipelines/pipeline_utils.py

@@ -67,6 +67,7 @@
     numpy_to_pil,
 )
 from ..utils.hub_utils import _check_legacy_sharding_variant_format, load_or_create_model_card, populate_model_card
+from ..utils.testing_utils import torch_device
 from ..utils.torch_utils import empty_device_cache, get_device, is_compiled_module
 
 
@@ -109,7 +110,7 @@
     LIBRARIES.append(library)
 
 # TODO: support single-device namings
-SUPPORTED_DEVICE_MAP = ["balanced", "cuda"]
+SUPPORTED_DEVICE_MAP = ["balanced"] + [torch_device]
 
 logger = logging.get_logger(__name__)
 

tests/pipelines/test_pipelines_common.py

@@ -2339,6 +2339,26 @@ def test_torch_dtype_dict(self):
                     f"Component '{name}' has dtype {component.dtype} but expected {expected_dtype}",
                 )
 
+    @require_torch_accelerator
+    def test_pipeline_with_accelerator_device_map(self, expected_max_difference=1e-4):
+        components = self.get_dummy_components()
+        pipe = self.pipeline_class(**components)
+        pipe = pipe.to(torch_device)
+        pipe.set_progress_bar_config(disable=None)
+
+        torch.manual_seed(0)
+        inputs = self.get_dummy_inputs(torch_device)
+        inputs["generator"] = torch.manual_seed(0)
+        out = pipe(**inputs)[0]
+
+        with tempfile.TemporaryDirectory() as tmpdir:
+            pipe.save_pretrained(tmpdir)
+            loaded_pipe = self.pipeline_class.from_pretrained(tmpdir, device_map=torch_device)
+        inputs = self.get_dummy_inputs(torch_device)
+        loaded_out = loaded_pipe(**inputs)[0]
+        max_diff = np.abs(to_np(out) - to_np(loaded_out)).max()
+        self.assertLess(max_diff, expected_max_difference)
+
 
 @is_staging_test
 class PipelinePushToHubTester(unittest.TestCase):