AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-01-27 17:22:53 +03:00
Code Activity

Adding NPU for get device function (#11617)

Browse Source 
* Adding device choice for npu

* Adding device choice for npu

---------

Co-authored-by: J石页 <jiangshuo9@h-partners.com>
Co-authored-by: Sayak Paul <spsayakpaul@gmail.com>
This commit is contained in:
Leo Jiang
2025-05-27 07:59:15 -06:00
committed by GitHub
parent cc59505e26
commit 5939ace91b
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
src/diffusers/utils/torch_utils.py
View File

@@ -18,7 +18,7 @@ PyTorch utilities: Utilities related to PyTorch
from typing import List, Optional, Tuple, Union
from . import logging
from .import_utils import is_torch_available, is_torch_version
from .import_utils import is_torch_available, is_torch_npu_available, is_torch_version
if is_torch_available():
@@ -166,6 +166,8 @@ def get_torch_cuda_device_capability():
def get_device():
if torch.cuda.is_available():
return "cuda"
elif is_torch_npu_available():
return "npu"
elif hasattr(torch, "xpu") and torch.xpu.is_available():
return "xpu"
else: