fluxdev/stable-diffusion-webui-forge

1import os2import importlib.util3 4 5# https://github.com/comfyanonymous/ComfyUI/blob/master/cuda_malloc.py6def get_gpu_names():7    if os.name == 'nt':8        import ctypes9 10        # Define necessary C structures and types11        class DISPLAY_DEVICEA(ctypes.Structure):12            _fields_ = [13                ('cb', ctypes.c_ulong),14                ('DeviceName', ctypes.c_char * 32),15                ('DeviceString', ctypes.c_char * 128),16                ('StateFlags', ctypes.c_ulong),17                ('DeviceID', ctypes.c_char * 128),18                ('DeviceKey', ctypes.c_char * 128)19            ]20 21        # Load user32.dll22        user32 = ctypes.windll.user3223 24        # Call EnumDisplayDevicesA25        def enum_display_devices():26            device_info = DISPLAY_DEVICEA()27            device_info.cb = ctypes.sizeof(device_info)28            device_index = 029            gpu_names = set()30 31            while user32.EnumDisplayDevicesA(None, device_index, ctypes.byref(device_info), 0):32                device_index += 133                gpu_names.add(device_info.DeviceString.decode('utf-8'))34            return gpu_names35        return enum_display_devices()36    else:37        return set()38 39 40blacklist = {"GeForce GTX TITAN X", "GeForce GTX 980", "GeForce GTX 970", "GeForce GTX 960", "GeForce GTX 950", "GeForce 945M",41                "GeForce 940M", "GeForce 930M", "GeForce 920M", "GeForce 910M", "GeForce GTX 750", "GeForce GTX 745", "Quadro K620",42                "Quadro K1200", "Quadro K2200", "Quadro M500", "Quadro M520", "Quadro M600", "Quadro M620", "Quadro M1000",43                "Quadro M1200", "Quadro M2000", "Quadro M2200", "Quadro M3000", "Quadro M4000", "Quadro M5000", "Quadro M5500", "Quadro M6000",44                "GeForce MX110", "GeForce MX130", "GeForce 830M", "GeForce 840M", "GeForce GTX 850M", "GeForce GTX 860M",45                "GeForce GTX 1650", "GeForce GTX 1630"46                }47 48 49def cuda_malloc_supported():50    try:51        names = get_gpu_names()52    except:53        names = set()54    for x in names:55        if "NVIDIA" in x:56            for b in blacklist:57                if b in x:58                    return False59    return True60 61 62def try_cuda_malloc():63    do_cuda_malloc = False64 65    try:66        version = ""67        torch_spec = importlib.util.find_spec("torch")68        for folder in torch_spec.submodule_search_locations:69            ver_file = os.path.join(folder, "version.py")70            if os.path.isfile(ver_file):71                spec = importlib.util.spec_from_file_location("torch_version_import", ver_file)72                module = importlib.util.module_from_spec(spec)73                spec.loader.exec_module(module)74                version = module.__version__75        if int(version[0]) >= 2:76            do_cuda_malloc = cuda_malloc_supported()77    except:78        pass79 80    if do_cuda_malloc:81        env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)82        if env_var is None:83            env_var = "backend:cudaMallocAsync"84        else:85            env_var += ",backend:cudaMallocAsync"86 87        os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var88 89        print('Using cudaMallocAsync backend.')90    else:91        print('Failed to use cudaMallocAsync backend.')92    return93