model_management.py 7.93 KB
Newer Older
1

2
3
4
5
CPU = 0
NO_VRAM = 1
LOW_VRAM = 2
NORMAL_VRAM = 3
6
HIGH_VRAM = 4
Yurii Mazurevich's avatar
Yurii Mazurevich committed
7
MPS = 5
8
9
10
11

accelerate_enabled = False
vram_state = NORMAL_VRAM

12
total_vram = 0
13
14
total_vram_available_mb = -1

15
import sys
16
import psutil
17
18
19

set_vram_to = NORMAL_VRAM

20
21
22
try:
    import torch
    total_vram = torch.cuda.mem_get_info(torch.cuda.current_device())[1] / (1024 * 1024)
23
24
    total_ram = psutil.virtual_memory().total / (1024 * 1024)
    forced_normal_vram = "--normalvram" in sys.argv
25
26
    forced_cpu = "--cpu" in sys.argv
    if not forced_normal_vram and not forced_cpu:
27
28
29
        if total_vram <= 4096:
            print("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
            set_vram_to = LOW_VRAM
comfyanonymous's avatar
comfyanonymous committed
30
        elif total_vram > total_ram * 1.1 and total_vram > 14336:
31
32
            print("Enabling highvram mode because your GPU has more vram than your computer has ram. If you don't want this use: --normalvram")
            vram_state = HIGH_VRAM
33
34
35
except:
    pass

36
37
38
39
40
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

41
42
if "--disable-xformers" in sys.argv:
    XFORMERS_IS_AVAILBLE = False
43
44
45
46
47
48
49
50
else:
    try:
        import xformers
        import xformers.ops
        XFORMERS_IS_AVAILBLE = True
    except:
        XFORMERS_IS_AVAILBLE = False

51
52
53
54
55
56
57
58
ENABLE_PYTORCH_ATTENTION = False
if "--use-pytorch-cross-attention" in sys.argv:
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
    ENABLE_PYTORCH_ATTENTION = True
    XFORMERS_IS_AVAILBLE = False

59

60
61
62
63
if "--lowvram" in sys.argv:
    set_vram_to = LOW_VRAM
if "--novram" in sys.argv:
    set_vram_to = NO_VRAM
64
65
if "--highvram" in sys.argv:
    vram_state = HIGH_VRAM
66

67

68
if set_vram_to == LOW_VRAM or set_vram_to == NO_VRAM:
69
70
71
72
73
74
75
76
    try:
        import accelerate
        accelerate_enabled = True
        vram_state = set_vram_to
    except Exception as e:
        import traceback
        print(traceback.format_exc())
        print("ERROR: COULD NOT ENABLE LOW VRAM MODE.")
77
78

    total_vram_available_mb = (total_vram - 1024) // 2
79
    total_vram_available_mb = int(max(256, total_vram_available_mb))
80

81
82
83
84
85
86
try:
    if torch.backends.mps.is_available():
        vram_state = MPS
except:
    pass

87
88
if "--cpu" in sys.argv:
    vram_state = CPU
89

Yurii Mazurevich's avatar
Yurii Mazurevich committed
90
print("Set vram state to:", ["CPU", "NO VRAM", "LOW VRAM", "NORMAL VRAM", "HIGH VRAM", "MPS"][vram_state])
91

92
93

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
94
current_gpu_controlnets = []
95

96
97
98
model_accelerated = False


99
100
def unload_model():
    global current_loaded_model
101
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
102
    global current_gpu_controlnets
103
104
    global vram_state

105
    if current_loaded_model is not None:
106
107
108
109
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

110
111
112
        #never unload models from GPU on high vram
        if vram_state != HIGH_VRAM:
            current_loaded_model.model.cpu()
113
114
        current_loaded_model.unpatch_model()
        current_loaded_model = None
115
116
117
118
119
120

    if vram_state != HIGH_VRAM:
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
121
122
123
124


def load_model_gpu(model):
    global current_loaded_model
125
126
127
    global vram_state
    global model_accelerated

128
129
130
131
132
133
134
135
136
    if model is current_loaded_model:
        return
    unload_model()
    try:
        real_model = model.patch_model()
    except Exception as e:
        model.unpatch_model()
        raise e
    current_loaded_model = model
137
138
    if vram_state == CPU:
        pass
Yurii Mazurevich's avatar
Yurii Mazurevich committed
139
140
141
142
    elif vram_state == MPS:
        mps_device = torch.device("mps")
        real_model.to(mps_device)
        pass
143
    elif vram_state == NORMAL_VRAM or vram_state == HIGH_VRAM:
144
145
146
147
148
149
        model_accelerated = False
        real_model.cuda()
    else:
        if vram_state == NO_VRAM:
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
        elif vram_state == LOW_VRAM:
150
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "{}MiB".format(total_vram_available_mb), "cpu": "16GiB"})
comfyanonymous's avatar
comfyanonymous committed
151

152
153
        accelerate.dispatch_model(real_model, device_map=device_map, main_device="cuda")
        model_accelerated = True
154
    return current_loaded_model
155

comfyanonymous's avatar
comfyanonymous committed
156
157
def load_controlnet_gpu(models):
    global current_gpu_controlnets
158
    global vram_state
159
160
    if vram_state == CPU:
        return
161
162
163
164
165

    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
        return

comfyanonymous's avatar
comfyanonymous committed
166
167
168
169
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

170
    device = get_torch_device()
comfyanonymous's avatar
comfyanonymous committed
171
172
    current_gpu_controlnets = []
    for m in models:
173
        current_gpu_controlnets.append(m.to(device))
comfyanonymous's avatar
comfyanonymous committed
174

175

176
177
178
179
180
181
182
183
184
185
186
187
def load_if_low_vram(model):
    global vram_state
    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        return model.cuda()
    return model

def unload_if_low_vram(model):
    global vram_state
    if vram_state == LOW_VRAM or vram_state == NO_VRAM:
        return model.cpu()
    return model

188
def get_torch_device():
Yurii Mazurevich's avatar
Yurii Mazurevich committed
189
190
    if vram_state == MPS:
        return torch.device("mps")
191
192
193
194
195
196
197
198
199
    if vram_state == CPU:
        return torch.device("cpu")
    else:
        return torch.cuda.current_device()

def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
200

201
202
203
204
205
def xformers_enabled():
    if vram_state == CPU:
        return False
    return XFORMERS_IS_AVAILBLE

206
207
208
def pytorch_attention_enabled():
    return ENABLE_PYTORCH_ATTENTION

209
210
def get_free_memory(dev=None, torch_free_too=False):
    if dev is None:
211
        dev = get_torch_device()
212

Yurii Mazurevich's avatar
Yurii Mazurevich committed
213
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
214
215
216
217
218
219
220
221
222
223
224
225
226
227
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
        stats = torch.cuda.memory_stats(dev)
        mem_active = stats['active_bytes.all.current']
        mem_reserved = stats['reserved_bytes.all.current']
        mem_free_cuda, _ = torch.cuda.mem_get_info(dev)
        mem_free_torch = mem_reserved - mem_active
        mem_free_total = mem_free_cuda + mem_free_torch

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
228
229
230
231
232
233
234
235
236

def maximum_batch_area():
    global vram_state
    if vram_state == NO_VRAM:
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
    area = ((memory_free - 1024) * 0.9) / (0.6)
    return int(max(area, 0))
237
238
239
240
241

def cpu_mode():
    global vram_state
    return vram_state == CPU

Yurii Mazurevich's avatar
Yurii Mazurevich committed
242
243
244
245
def mps_mode():
    global vram_state
    return vram_state == MPS

246
def should_use_fp16():
Yurii Mazurevich's avatar
Yurii Mazurevich committed
247
    if cpu_mode() or mps_mode():
248
249
250
251
252
        return False #TODO ?

    if torch.cuda.is_bf16_supported():
        return True

comfyanonymous's avatar
comfyanonymous committed
253
    props = torch.cuda.get_device_properties("cuda")
254
255
256
257
    if props.major < 7:
        return False

    #FP32 is faster on those cards?
258
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600"]
259
260
261
262
263
264
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
#TODO: might be cleaner to put this somewhere else
import threading

class InterruptProcessingException(Exception):
    pass

interrupt_processing_mutex = threading.RLock()

interrupt_processing = False
def interrupt_current_processing(value=True):
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        interrupt_processing = value

def processing_interrupted():
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        return interrupt_processing

def throw_exception_if_processing_interrupted():
    global interrupt_processing
    global interrupt_processing_mutex
    with interrupt_processing_mutex:
        if interrupt_processing:
            interrupt_processing = False
            raise InterruptProcessingException()