model_management.py 16.9 KB
Newer Older
1
2
import psutil
from enum import Enum
comfyanonymous's avatar
comfyanonymous committed
3
from comfy.cli_args import args
4
import torch
5

6
class VRAMState(Enum):
7
8
    DISABLED = 0    #No vram present: no need to move models to vram
    NO_VRAM = 1     #Very low vram: enable all the options to save vram
9
10
11
    LOW_VRAM = 2
    NORMAL_VRAM = 3
    HIGH_VRAM = 4
12
    SHARED = 5      #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.
13
14
15
16
17

class CPUState(Enum):
    GPU = 0
    CPU = 1
    MPS = 2
18

19
20
21
# Determine VRAM State
vram_state = VRAMState.NORMAL_VRAM
set_vram_to = VRAMState.NORMAL_VRAM
22
cpu_state = CPUState.GPU
23

24
total_vram = 0
25

26
lowvram_available = True
藍+85CD's avatar
藍+85CD committed
27
xpu_available = False
28

29
directml_enabled = False
30
if args.directml is not None:
31
32
    import torch_directml
    directml_enabled = True
33
34
35
36
37
38
    device_index = args.directml
    if device_index < 0:
        directml_device = torch_directml.device()
    else:
        directml_device = torch_directml.device(device_index)
    print("Using directml with device:", torch_directml.device_name(device_index))
39
    # torch_directml.disable_tiled_resources(True)
40
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
41

42
try:
43
44
45
    import intel_extension_for_pytorch as ipex
    if torch.xpu.is_available():
        xpu_available = True
46
47
48
except:
    pass

49
50
51
52
53
54
55
56
57
try:
    if torch.backends.mps.is_available():
        cpu_state = CPUState.MPS
except:
    pass

if args.cpu:
    cpu_state = CPUState.CPU

58
59
60
def get_torch_device():
    global xpu_available
    global directml_enabled
61
    global cpu_state
62
63
64
    if directml_enabled:
        global directml_device
        return directml_device
65
    if cpu_state == CPUState.MPS:
66
        return torch.device("mps")
67
    if cpu_state == CPUState.CPU:
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
        return torch.device("cpu")
    else:
        if xpu_available:
            return torch.device("xpu")
        else:
            return torch.device(torch.cuda.current_device())

def get_total_memory(dev=None, torch_total_too=False):
    global xpu_available
    global directml_enabled
    if dev is None:
        dev = get_torch_device()

    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
        mem_total = psutil.virtual_memory().total
        mem_total_torch = mem_total
    else:
        if directml_enabled:
            mem_total = 1024 * 1024 * 1024 #TODO
            mem_total_torch = mem_total
        elif xpu_available:
            mem_total = torch.xpu.get_device_properties(dev).total_memory
            mem_total_torch = mem_total
        else:
            stats = torch.cuda.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
            _, mem_total_cuda = torch.cuda.mem_get_info(dev)
            mem_total_torch = mem_reserved
            mem_total = mem_total_cuda

    if torch_total_too:
        return (mem_total, mem_total_torch)
    else:
        return mem_total

total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024)
print("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
if not args.normalvram and not args.cpu:
    if lowvram_available and 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 = VRAMState.LOW_VRAM
    elif total_vram > total_ram * 1.1 and total_vram > 14336:
        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 = VRAMState.HIGH_VRAM

114
115
116
117
118
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

119
120
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
121
122
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
123
124
125
126
else:
    try:
        import xformers
        import xformers.ops
127
        XFORMERS_IS_AVAILABLE = True
128
129
130
131
132
133
134
135
136
137
138
        try:
            XFORMERS_VERSION = xformers.version.__version__
            print("xformers version:", XFORMERS_VERSION)
            if XFORMERS_VERSION.startswith("0.0.18"):
                print()
                print("WARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
                print("Please downgrade or upgrade xformers to a different version.")
                print()
                XFORMERS_ENABLED_VAE = False
        except:
            pass
139
    except:
140
        XFORMERS_IS_AVAILABLE = False
141

142
143
144
145
146
147
def is_nvidia():
    global cpu_state
    if cpu_state == CPUState.GPU:
        if torch.version.cuda:
            return True

148
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
149
150
151
152
153
154
155
156
157
158

if ENABLE_PYTORCH_ATTENTION == False and XFORMERS_IS_AVAILABLE == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
    try:
        if is_nvidia():
            torch_version = torch.version.__version__
            if int(torch_version[0]) >= 2:
                ENABLE_PYTORCH_ATTENTION = True
    except:
        pass

159
if ENABLE_PYTORCH_ATTENTION:
160
161
162
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
163
    XFORMERS_IS_AVAILABLE = False
164

165
166
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
167
    lowvram_available = True
168
169
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
170
elif args.highvram or args.gpu_only:
171
    vram_state = VRAMState.HIGH_VRAM
172

173
FORCE_FP32 = False
174
FORCE_FP16 = False
175
176
177
178
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

179
180
181
182
if args.force_fp16:
    print("Forcing FP16.")
    FORCE_FP16 = True

183
if lowvram_available:
184
185
    try:
        import accelerate
186
187
        if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
            vram_state = set_vram_to
188
189
190
    except Exception as e:
        import traceback
        print(traceback.format_exc())
191
192
        print("ERROR: LOW VRAM MODE NEEDS accelerate.")
        lowvram_available = False
193

194

195
196
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
197

198
199
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
200

201
print(f"Set vram state to: {vram_state.name}")
202

203

204
205
def get_torch_device_name(device):
    if hasattr(device, 'type'):
206
207
208
209
210
211
        if device.type == "cuda":
            return "{} {}".format(device, torch.cuda.get_device_name(device))
        else:
            return "{}".format(device.type)
    else:
        return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
212
213

try:
214
    print("Device:", get_torch_device_name(get_torch_device()))
215
216
217
except:
    print("Could not pick default device.")

218
219

current_loaded_model = None
comfyanonymous's avatar
comfyanonymous committed
220
current_gpu_controlnets = []
221

222
223
224
model_accelerated = False


225
226
def unload_model():
    global current_loaded_model
227
    global model_accelerated
comfyanonymous's avatar
comfyanonymous committed
228
    global current_gpu_controlnets
229
230
    global vram_state

231
    if current_loaded_model is not None:
232
233
234
235
        if model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(current_loaded_model.model)
            model_accelerated = False

236

237
238
        current_loaded_model.model.to(current_loaded_model.offload_device)
        current_loaded_model.model_patches_to(current_loaded_model.offload_device)
239
240
        current_loaded_model.unpatch_model()
        current_loaded_model = None
241

242
    if vram_state != VRAMState.HIGH_VRAM:
243
244
245
246
        if len(current_gpu_controlnets) > 0:
            for n in current_gpu_controlnets:
                n.cpu()
            current_gpu_controlnets = []
247

248
249
def minimum_inference_memory():
    return (768 * 1024 * 1024)
250
251
252

def load_model_gpu(model):
    global current_loaded_model
253
254
255
    global vram_state
    global model_accelerated

256
257
258
259
260
261
262
263
    if model is current_loaded_model:
        return
    unload_model()
    try:
        real_model = model.patch_model()
    except Exception as e:
        model.unpatch_model()
        raise e
264

265
    torch_dev = model.load_device
266
    model.model_patches_to(torch_dev)
267
    model.model_patches_to(model.model_dtype())
268

269
270
271
272
273
    if is_device_cpu(torch_dev):
        vram_set_state = VRAMState.DISABLED
    else:
        vram_set_state = vram_state

274
275
276
    if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
        model_size = model.model_size()
        current_free_mem = get_free_memory(torch_dev)
277
        lowvram_model_memory = int(max(256 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
278
        if model_size > (current_free_mem - minimum_inference_memory()): #only switch to lowvram if really necessary
279
280
            vram_set_state = VRAMState.LOW_VRAM

281
    current_loaded_model = model
282

283
    if vram_set_state == VRAMState.DISABLED:
Yurii Mazurevich's avatar
Yurii Mazurevich committed
284
        pass
285
    elif vram_set_state == VRAMState.NORMAL_VRAM or vram_set_state == VRAMState.HIGH_VRAM or vram_set_state == VRAMState.SHARED:
286
        model_accelerated = False
287
        real_model.to(torch_dev)
288
    else:
289
        if vram_set_state == VRAMState.NO_VRAM:
290
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "256MiB", "cpu": "16GiB"})
291
292
        elif vram_set_state == VRAMState.LOW_VRAM:
            device_map = accelerate.infer_auto_device_map(real_model, max_memory={0: "{}MiB".format(lowvram_model_memory // (1024 * 1024)), "cpu": "16GiB"})
comfyanonymous's avatar
comfyanonymous committed
293

294
        accelerate.dispatch_model(real_model, device_map=device_map, main_device=torch_dev)
295
        model_accelerated = True
296
    return current_loaded_model
297

298
def load_controlnet_gpu(control_models):
comfyanonymous's avatar
comfyanonymous committed
299
    global current_gpu_controlnets
300
    global vram_state
301
    if vram_state == VRAMState.DISABLED:
302
        return
303

304
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
305
306
307
        for m in control_models:
            if hasattr(m, 'set_lowvram'):
                m.set_lowvram(True)
308
309
310
        #don't load controlnets like this if low vram because they will be loaded right before running and unloaded right after
        return

311
312
313
314
    models = []
    for m in control_models:
        models += m.get_models()

comfyanonymous's avatar
comfyanonymous committed
315
316
317
318
    for m in current_gpu_controlnets:
        if m not in models:
            m.cpu()

319
    device = get_torch_device()
comfyanonymous's avatar
comfyanonymous committed
320
321
    current_gpu_controlnets = []
    for m in models:
322
        current_gpu_controlnets.append(m.to(device))
comfyanonymous's avatar
comfyanonymous committed
323

324

325
326
def load_if_low_vram(model):
    global vram_state
327
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
comfyanonymous's avatar
comfyanonymous committed
328
        return model.to(get_torch_device())
329
330
331
332
    return model

def unload_if_low_vram(model):
    global vram_state
333
    if vram_state == VRAMState.LOW_VRAM or vram_state == VRAMState.NO_VRAM:
334
335
336
        return model.cpu()
    return model

337
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
338
    if vram_state == VRAMState.HIGH_VRAM:
339
340
341
342
        return get_torch_device()
    else:
        return torch.device("cpu")

343
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
344
    if args.gpu_only:
345
346
347
348
        return get_torch_device()
    else:
        return torch.device("cpu")

349
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
350
    if args.gpu_only:
351
        return get_torch_device()
352
353
354
355
356
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
        if torch.get_num_threads() < 8: #leaving the text encoder on the CPU is faster than shifting it if the CPU is fast enough.
            return get_torch_device()
        else:
            return torch.device("cpu")
357
358
359
    else:
        return torch.device("cpu")

360
361
362
363
def vae_device():
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
364
    if args.gpu_only:
365
366
367
368
        return get_torch_device()
    else:
        return torch.device("cpu")

369
370
371
372
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
373

374

375
def xformers_enabled():
376
377
    global xpu_available
    global directml_enabled
378
379
    global cpu_state
    if cpu_state != CPUState.GPU:
380
        return False
381
382
383
384
    if xpu_available:
        return False
    if directml_enabled:
        return False
385
    return XFORMERS_IS_AVAILABLE
386

387
388
389
390
391

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
392
393

    return XFORMERS_ENABLED_VAE
394

395
def pytorch_attention_enabled():
396
    global ENABLE_PYTORCH_ATTENTION
397
398
    return ENABLE_PYTORCH_ATTENTION

399
400
401
402
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
403
        if is_nvidia(): #pytorch flash attention only works on Nvidia
404
405
406
            return True
    return False

407
def get_free_memory(dev=None, torch_free_too=False):
408
    global xpu_available
409
    global directml_enabled
410
    if dev is None:
411
        dev = get_torch_device()
412

Yurii Mazurevich's avatar
Yurii Mazurevich committed
413
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
414
415
416
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
417
418
419
420
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
        elif xpu_available:
421
422
423
424
425
426
427
428
429
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - torch.xpu.memory_allocated(dev)
            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
430
431
432
433
434

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
435
436
437

def maximum_batch_area():
    global vram_state
438
    if vram_state == VRAMState.NO_VRAM:
439
440
441
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
442
    if xformers_enabled() or pytorch_attention_flash_attention():
443
        #TODO: this needs to be tweaked
444
        area = 20 * memory_free
445
446
447
    else:
        #TODO: this formula is because AMD sucks and has memory management issues which might be fixed in the future
        area = ((memory_free - 1024) * 0.9) / (0.6)
448
    return int(max(area, 0))
449
450

def cpu_mode():
451
452
    global cpu_state
    return cpu_state == CPUState.CPU
453

Yurii Mazurevich's avatar
Yurii Mazurevich committed
454
def mps_mode():
455
456
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
457

458
459
def is_device_cpu(device):
    if hasattr(device, 'type'):
comfyanonymous's avatar
comfyanonymous committed
460
461
462
463
464
465
466
        if (device.type == 'cpu'):
            return True
    return False

def is_device_mps(device):
    if hasattr(device, 'type'):
        if (device.type == 'mps'):
467
468
469
            return True
    return False

470
def should_use_fp16(device=None, model_params=0):
471
    global xpu_available
472
473
    global directml_enabled

474
475
476
    if FORCE_FP16:
        return True

477
    if device is not None: #TODO
comfyanonymous's avatar
comfyanonymous committed
478
        if is_device_cpu(device) or is_device_mps(device):
479
            return False
480

481
482
483
    if FORCE_FP32:
        return False

484
485
486
    if directml_enabled:
        return False

487
    if cpu_mode() or mps_mode() or xpu_available:
488
489
490
491
492
        return False #TODO ?

    if torch.cuda.is_bf16_supported():
        return True

comfyanonymous's avatar
comfyanonymous committed
493
    props = torch.cuda.get_device_properties("cuda")
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
    if props.major < 6:
        return False

    fp16_works = False
    #FP16 is confirmed working on a 1080 (GP104) but it's a bit slower than FP32 so it should only be enabled
    #when the model doesn't actually fit on the card
    #TODO: actually test if GP106 and others have the same type of behavior
    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050"]
    for x in nvidia_10_series:
        if x in props.name.lower():
            fp16_works = True

    if fp16_works:
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
        if model_params * 4 > free_model_memory:
            return True

511
512
513
    if props.major < 7:
        return False

514
    #FP16 is just broken on these cards
515
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600"]
516
517
518
519
520
521
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

522
523
def soft_empty_cache():
    global xpu_available
524
525
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
526
527
        torch.mps.empty_cache()
    elif xpu_available:
528
529
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
530
        if is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
531
532
533
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
#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()