oobabooga/text-generation-webui
1
0
Fork 0
mirror of https://github.com/oobabooga/text-generation-webui.git synced 2026-04-04 14:17:28 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add Ascend NPU support (basic) (#5541)

Browse source
This commit is contained in:
wangshuai09 2024-04-12 05:42:20 +08:00 • committed by GitHub
parent a90509d82e
commit fd4e46bce2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 35 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

5
modules/ui_model_menu.py
View file

@ -8,7 +8,7 @@ from pathlib import Path
import gradio as gr
import psutil
import torch
from transformers import is_torch_xpu_available
from transformers import is_torch_npu_available, is_torch_xpu_available
from modules import loaders, shared, ui, utils
from modules.logging_colors import logger
@ -32,6 +32,9 @@ def create_ui():
if is_torch_xpu_available():
for i in range(torch.xpu.device_count()):
total_mem.append(math.floor(torch.xpu.get_device_properties(i).total_memory / (1024 * 1024)))
elif is_torch_npu_available():
for i in range(torch.npu.device_count()):
total_mem.append(math.floor(torch.npu.get_device_properties(i).total_memory / (1024 * 1024)))
else:
for i in range(torch.cuda.device_count()):
total_mem.append(math.floor(torch.cuda.get_device_properties(i).total_memory / (1024 * 1024)))