oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2025-12-06 07:12:10 +01:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add fallbacks after 93aa7b3ed3

Browse source
This commit is contained in:
oobabooga 2025-10-09 10:59:34 -07:00
parent 1aa2b924d2
commit 218dc01b51
1 changed files with 16 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
modules/torch_utils.py
View file

@ -2,12 +2,27 @@ import gc
import torch import torch
from accelerate.utils import is_npu_available, is_xpu_available from accelerate.utils import is_npu_available, is_xpu_available
from transformers import is_torch_npu_available, is_torch_xpu_available
from modules import shared from modules import shared
def get_device(): def get_device():
return getattr(shared.model, 'device', None) if hasattr(shared.model, 'device'):
return shared.model.device
elif torch.cuda.is_available():
return torch.device('cuda')
elif shared.args.deepspeed:
import deepspeed
return deepspeed.get_accelerator().current_device_name()
elif torch.backends.mps.is_available():
return torch.device('mps')
elif is_torch_xpu_available():
return torch.device('xpu:0')
elif is_torch_npu_available():
return torch.device('npu:0')
else:
return None
def clear_torch_cache(): def clear_torch_cache():