FilesExpand file tree

master

/

cuda_malloc.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

115 lines (99 loc) · 4.08 KB

master

/

cuda_malloc.py

Copy path

Top

File metadata and controls

• Code
• Blame

115 lines (99 loc) · 4.08 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

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

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

import os

import importlib.util

from comfy.cli_args import args, PerformanceFeature

import subprocess

import re

#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():

if os.name == 'nt':

import ctypes

# Define necessary C structures and types

class DISPLAY_DEVICEA(ctypes.Structure):

_fields_ = [

('cb', ctypes.c_ulong),

('DeviceName', ctypes.c_char * 32),

('DeviceString', ctypes.c_char * 128),

('StateFlags', ctypes.c_ulong),

('DeviceID', ctypes.c_char * 128),

('DeviceKey', ctypes.c_char * 128)

]

# Load user32.dll

user32 = ctypes.windll.user32

# Call EnumDisplayDevicesA

def enum_display_devices():

device_info = DISPLAY_DEVICEA()

device_info.cb = ctypes.sizeof(device_info)

device_index = 0

gpu_names = []

while user32.EnumDisplayDevicesA(None, device_index, ctypes.byref(device_info), 0):

device_index += 1

gpu_names.append(device_info.DeviceString.decode('utf-8'))

return gpu_names

return enum_display_devices()

else:

gpu_names = []

out = subprocess.check_output(['nvidia-smi', '-L'])

for l in out.split(b'\n'):

if len(l) > 0:

gpu_names.append(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",

"Quadro K1200", "Quadro K2200", "Quadro M500", "Quadro M520", "Quadro M600", "Quadro M620", "Quadro M1000",

"Quadro M1200", "Quadro M2000", "Quadro M2200", "Quadro M3000", "Quadro M4000", "Quadro M5000", "Quadro M5500", "Quadro M6000",

"GeForce MX110", "GeForce MX130", "GeForce 830M", "GeForce 840M", "GeForce GTX 850M", "GeForce GTX 860M",

"GeForce GTX 1650", "GeForce GTX 1630", "Tesla M4", "Tesla M6", "Tesla M10", "Tesla M40", "Tesla M60"

}

def cuda_malloc_supported():

try:

names = get_gpu_names()

except:

names = set()

for x in names:

if "NVIDIA" in x:

for b in blacklist:

if b in x:

return False

return True

version = ""

try:

torch_spec = importlib.util.find_spec("torch")

for folder in torch_spec.submodule_search_locations:

ver_file = os.path.join(folder, "version.py")

if os.path.isfile(ver_file):

spec = importlib.util.spec_from_file_location("torch_version_import", ver_file)

module = importlib.util.module_from_spec(spec)

spec.loader.exec_module(module)

version = module.__version__

except:

pass

def get_raw_cuda_version(version_str):

match = re.search(r'\+cu(\d+)', version_str)

if match:

try:

return int(match.group(1))

except:

pass

return None

if not args.cuda_malloc:

try:

if int(version[0]) >= 2 and "+cu" in version: # enable by default for torch version 2.0 and up only on cuda torch

if PerformanceFeature.AutoTune not in args.fast: # Autotune has issues with cuda malloc

cuda_version = get_raw_cuda_version(version)

if cuda_version is not None and cuda_version >= 130:

args.cuda_malloc = True

else:

args.cuda_malloc = cuda_malloc_supported()

except:

pass

if args.disable_cuda_malloc:

args.cuda_malloc = False

if args.cuda_malloc:

env_var = os.environ.get('PYTORCH_CUDA_ALLOC_CONF', None)

if env_var is None:

env_var = "backend:cudaMallocAsync"

else:

env_var += ",backend:cudaMallocAsync"

os.environ['PYTORCH_CUDA_ALLOC_CONF'] = env_var

def get_torch_version_noimport():

return str(version)