From caddef8d8863adaa3d85faa03e3e2f1c6a04fb58 Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Mon, 4 Mar 2024 09:03:59 -0500 Subject: [PATCH] Auto disable cuda malloc on unsupported GPUs on Linux. --- cuda_malloc.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/cuda_malloc.py b/cuda_malloc.py index 144cdacd..70e7ecf9 100644 --- a/cuda_malloc.py +++ b/cuda_malloc.py @@ -1,6 +1,7 @@ import os import importlib.util 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. def get_gpu_names(): @@ -34,7 +35,12 @@ def get_gpu_names(): return gpu_names return enum_display_devices() 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", "GeForce 940M", "GeForce 930M", "GeForce 920M", "GeForce 910M", "GeForce GTX 750", "GeForce GTX 745", "Quadro K620",