Spaces:
Running
Running
File size: 1,341 Bytes
e4316f1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 |
import shutil
import subprocess
import sys
def pick_device() -> str:
"""Return the best available device string for MinerU."""
# 1. NVIDIA GPU
if shutil.which("nvidia-smi") is not None:
try:
# check if there's any free GPU memory
out = subprocess.check_output(
[
"nvidia-smi",
"--query-gpu=memory.free",
"--format=csv,noheader,nounits",
],
text=True,
)
if any(int(line) > 0 for line in out.strip().splitlines()):
return "cuda:0"
except Exception: # pylint: disable=broad-except
pass
# 2. Apple Silicon
if sys.platform == "darwin" and shutil.which("sysctl"):
try:
brand = subprocess.check_output(
["sysctl", "-n", "machdep.cpu.brand_string"], text=True
)
if "Apple" in brand:
return "mps"
except Exception: # pylint: disable=broad-except
pass
# 3. Ascend NPU
if shutil.which("npu-smi") is not None:
try:
subprocess.check_call(["npu-smi", "info"], stdout=subprocess.DEVNULL)
return "npu"
except Exception: # pylint: disable=broad-except
pass
return "cpu"
|