up

Author: sayakpaul

src/diffusers/pipelines/pipeline_utils.py

@@ -67,7 +67,6 @@
     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,8 +108,7 @@
 for library in LOADABLE_CLASSES:
     LIBRARIES.append(library)
 
-# TODO: support single-device namings
-SUPPORTED_DEVICE_MAP = ["balanced"] + [torch_device]
+SUPPORTED_DEVICE_MAP = ["balanced"] + [get_device()]
 
 logger = logging.get_logger(__name__)
 

src/diffusers/utils/torch_utils.py

@@ -15,6 +15,7 @@
 PyTorch utilities: Utilities related to PyTorch
 """
 
+import functools
 from typing import List, Optional, Tuple, Union
 
 from . import logging
@@ -168,6 +169,7 @@ def get_torch_cuda_device_capability():
         return None
 
 
+@functools.lru_cache
 def get_device():
     if torch.cuda.is_available():
         return "cuda"