Consolidate torch device backend dispatch

docs/source/en/api/utilities.md

@@ -50,6 +50,10 @@ Utility and helper functions for working with 🤗 Diffusers.
 
 [[autodoc]] utils.torch_utils.randn_tensor
 
+## TorchDeviceBackend
+
+[[autodoc]] utils.torch_utils.TorchDeviceBackend
+
 ## apply_layerwise_casting
 
 [[autodoc]] hooks.layerwise_casting.apply_layerwise_casting

src/diffusers/hooks/group_offloading.py

@@ -23,6 +23,7 @@
 import torch
 
 from ..utils import get_logger, is_accelerate_available, is_torchao_available
+from ..utils.torch_utils import TorchDeviceBackend
 from ._common import _GO_LC_SUPPORTED_PYTORCH_LAYERS
 from .hooks import HookRegistry, ModelHook
 
@@ -166,11 +167,7 @@ def __init__(
         else:
             self.cpu_param_dict = self._init_cpu_param_dict()
 
-        self._torch_accelerator_module = (
-            getattr(torch, torch.accelerator.current_accelerator().type)
-            if hasattr(torch, "accelerator")
-            else torch.cuda
-        )
+        self._torch_accelerator_module = TorchDeviceBackend(self.onload_device)
 
     @staticmethod
     def _to_cpu(tensor, low_cpu_mem_usage):
@@ -671,12 +668,13 @@ def apply_group_offloading(
 
     stream = None
     if use_stream:
-        if torch.cuda.is_available():
-            stream = torch.cuda.Stream()
-        elif hasattr(torch, "xpu") and torch.xpu.is_available():
-            stream = torch.Stream()
-        else:
-            raise ValueError("Using streams for data transfer requires a CUDA device, or an Intel XPU device.")
+        backend = TorchDeviceBackend(onload_device)
+        if onload_device.type == "cpu" or not hasattr(backend, "Stream"):
+            raise ValueError(
+                "Using streams for data transfer requires an onload device whose backend implements streams, "
+                f"got `{onload_device.type}`. Pass `use_stream=False`."
+            )
+        stream = backend.Stream()
 
     if not use_stream and record_stream:
         raise ValueError("`record_stream` cannot be True when `use_stream=False`.")

src/diffusers/loaders/lora_pipeline.py

@@ -30,6 +30,7 @@
     logging,
     require_peft_backend,
 )
+from ..utils.torch_utils import get_device
 from .lora_base import (  # noqa
     LORA_WEIGHT_NAME,
     LORA_WEIGHT_NAME_SAFE,
@@ -109,7 +110,7 @@ def _maybe_dequantize_weight_for_expanded_lora(model, module):
     if module.weight.device.type == "cpu":
         weight_on_cpu = True
 
-    device = torch.accelerator.current_accelerator().type if hasattr(torch, "accelerator") else "cuda"
+    device = get_device()
     if is_bnb_4bit_quantized or is_bnb_8bit_quantized:
         module_weight = dequantize_bnb_weight(
             module.weight.to(device) if weight_on_cpu else module.weight,

src/diffusers/modular_pipelines/components_manager.py

@@ -28,7 +28,7 @@
     is_accelerate_available,
     logging,
 )
-from ..utils.torch_utils import get_device
+from ..utils.torch_utils import TorchDeviceBackend, empty_device_cache, get_device
 
 
 if is_accelerate_available():
@@ -179,13 +179,7 @@ def __call__(self, hooks, model_id, model, execution_device):
         except AttributeError:
             raise AttributeError(f"Do not know how to compute memory footprint of `{model.__class__.__name__}.")
 
-        device_type = execution_device.type
-        device_module = getattr(torch, device_type, torch.cuda)
-        try:
-            mem_on_device = device_module.mem_get_info(execution_device.index)[0]
-        except AttributeError:
-            raise AttributeError(f"Do not know how to obtain obtain memory info for {str(device_module)}.")
-
+        mem_on_device = TorchDeviceBackend(execution_device).mem_get_info()[0]
         mem_on_device = mem_on_device - self.memory_reserve_margin
         if current_module_size < mem_on_device:
             return []
@@ -513,10 +507,7 @@ def remove(self, component_id: str = None):
             import gc
 
             gc.collect()
-            if torch.cuda.is_available():
-                torch.cuda.empty_cache()
-            if torch.xpu.is_available():
-                torch.xpu.empty_cache()
+            empty_device_cache()
 
     # YiYi TODO: rename to search_components for now, may remove this method
     def search_components(
@@ -743,12 +734,8 @@ def enable_auto_cpu_offload(
         if not isinstance(device, torch.device):
             device = torch.device(device)
 
-        device_type = device.type
-        device_module = getattr(torch, device_type, torch.cuda)
-        if not hasattr(device_module, "mem_get_info"):
-            raise NotImplementedError(
-                f"`enable_auto_cpu_offload() relies on the `mem_get_info()` method. It's not implemented for {str(device.type)}."
-            )
+        # Fail here rather than on the first forward: the strategy cannot run without a free-memory query.
+        TorchDeviceBackend(device).mem_get_info()
 
         if device.index is None:
             device = torch.device(f"{device.type}:{0}")

src/diffusers/pipelines/pag/pipeline_pag_sana.py

@@ -35,7 +35,7 @@
     logging,
     replace_example_docstring,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline, ImagePipelineOutput
 from ..pixart_alpha.pipeline_pixart_alpha import (
     ASPECT_RATIO_512_BIN,
@@ -892,15 +892,9 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             try:
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/pipeline_utils.py

@@ -1287,7 +1287,7 @@ def maybe_free_model_hooks(self):
             return
 
         # make sure the model is in the same state as before calling it
-        self.enable_model_cpu_offload(device=getattr(self, "_offload_device", "cuda"))
+        self.enable_model_cpu_offload(device=getattr(self, "_offload_device", get_device()))
 
     def enable_sequential_cpu_offload(self, gpu_id: int | None = None, device: torch.device | str = None):
         r"""

src/diffusers/pipelines/sana/pipeline_sana.py

@@ -38,7 +38,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from ..pixart_alpha.pipeline_pixart_alpha import (
     ASPECT_RATIO_512_BIN,
@@ -957,15 +957,9 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             try:
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/sana/pipeline_sana_controlnet.py

@@ -38,7 +38,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from ..pixart_alpha.pipeline_pixart_alpha import (
     ASPECT_RATIO_512_BIN,
@@ -1053,15 +1053,9 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             try:
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/sana/pipeline_sana_sprint.py

@@ -38,7 +38,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN
 from .pipeline_output import SanaPipelineOutput
@@ -839,15 +839,9 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             try:
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/sana/pipeline_sana_sprint_img2img.py

@@ -39,7 +39,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ..pipeline_utils import DiffusionPipeline
 from ..pixart_alpha.pipeline_pixart_alpha import ASPECT_RATIO_1024_BIN
 from .pipeline_output import SanaPipelineOutput
@@ -929,15 +929,9 @@ def __call__(
             image = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             try:
                 image = self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/sana_video/pipeline_sana_video.py

@@ -37,7 +37,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
 from ..pipeline_utils import DiffusionPipeline
 from .pipeline_output import SanaVideoPipelineOutput
@@ -990,12 +990,6 @@ def __call__(
             video = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             if isinstance(self.vae, AutoencoderKLLTX2Video):
                 latents_mean = self.vae.latents_mean
                 latents_std = self.vae.latents_std
@@ -1014,7 +1008,7 @@ def __call__(
             latents = latents / latents_std + latents_mean
             try:
                 video = self.vae.decode(latents, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/pipelines/sana_video/pipeline_sana_video_i2v.py

@@ -39,7 +39,7 @@
     scale_lora_layers,
     unscale_lora_layers,
 )
-from ...utils.torch_utils import get_device, is_torch_version, randn_tensor
+from ...utils.torch_utils import randn_tensor
 from ...video_processor import VideoProcessor
 from ..pipeline_utils import DiffusionPipeline
 from .pipeline_output import SanaVideoPipelineOutput
@@ -1043,12 +1043,6 @@ def __call__(
             video = latents
         else:
             latents = latents.to(self.vae.dtype)
-            torch_accelerator_module = getattr(torch, get_device(), torch.cuda)
-            oom_error = (
-                torch.OutOfMemoryError
-                if is_torch_version(">=", "2.5.0")
-                else torch_accelerator_module.OutOfMemoryError
-            )
             if isinstance(self.vae, AutoencoderKLLTX2Video):
                 latents_mean = self.vae.latents_mean
                 latents_std = self.vae.latents_std
@@ -1067,7 +1061,7 @@ def __call__(
             latents = latents / latents_std + latents_mean
             try:
                 video = self.vae.decode(latents, return_dict=False)[0]
-            except oom_error as e:
+            except torch.OutOfMemoryError as e:
                 warnings.warn(
                     f"{e}. \n"
                     f"Try to use VAE tiling for large images. For example: \n"

src/diffusers/quantizers/gguf/gguf_quantizer.py

@@ -18,6 +18,7 @@
     is_torch_available,
     logging,
 )
+from ...utils.torch_utils import get_device
 
 
 if is_torch_available() and is_gguf_available():
@@ -177,12 +178,7 @@ def _dequantize(self, model):
             logger.info(
                 "Model was found to be on CPU (could happen as a result of `enable_model_cpu_offload()`). So, moving it to accelerator. After dequantization, will move the model back to CPU again to preserve the previous device."
             )
-            device = (
-                torch.accelerator.current_accelerator()
-                if hasattr(torch, "accelerator")
-                else torch.cuda.current_device()
-            )
-            model.to(device)
+            model.to(get_device())
 
         model = _dequantize_gguf_and_restore_linear(model, self.modules_to_not_convert)
         if is_model_on_cpu:

src/diffusers/training_utils.py

@@ -37,6 +37,7 @@
     is_torchvision_available,
     is_transformers_available,
 )
+from .utils.torch_utils import empty_device_cache
 
 
 if is_transformers_available():
@@ -412,15 +413,7 @@ def free_memory():
     Runs garbage collection. Then clears the cache of the available accelerator.
     """
     gc.collect()
-
-    if torch.cuda.is_available():
-        torch.cuda.empty_cache()
-    elif torch.backends.mps.is_available():
-        torch.mps.empty_cache()
-    elif is_torch_npu_available():
-        torch_npu.npu.empty_cache()
-    elif hasattr(torch, "xpu") and torch.xpu.is_available():
-        torch.xpu.empty_cache()
+    empty_device_cache()
 
 
 @contextmanager