Auto disable cuda malloc on unsupported GPUs on Linux.

This commit is contained in:
comfyanonymous 2024-03-04 09:03:59 -05:00
parent 478f71a249
commit caddef8d88
1 changed files with 7 additions and 1 deletions

View File

@ -1,6 +1,7 @@
import os import os
import importlib.util import importlib.util
from comfy.cli_args import args from comfy.cli_args import args
import subprocess
#Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import. #Can't use pytorch to get the GPU names because the cuda malloc has to be set before the first import.
def get_gpu_names(): def get_gpu_names():
@ -34,7 +35,12 @@ def get_gpu_names():
return gpu_names return gpu_names
return enum_display_devices() return enum_display_devices()
else: else:
return set() gpu_names = set()
out = subprocess.check_output(['nvidia-smi', '-L'])
for l in out.split(b'\n'):
if len(l) > 0:
gpu_names.add(l.decode('utf-8').split(' (UUID')[0])
return gpu_names
blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeForce GTX 960", "GeForce GTX 950", "GeForce 945M", blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeForce GTX 960", "GeForce GTX 950", "GeForce 945M",
"GeForce 940M", "GeForce 930M", "GeForce 920M", "GeForce 910M", "GeForce GTX 750", "GeForce GTX 745", "Quadro K620", "GeForce 940M", "GeForce 930M", "GeForce 920M", "GeForce 910M", "GeForce GTX 750", "GeForce GTX 745", "Quadro K620",