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

8
class VRAMState(Enum):
9
10
    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
11
12
13
    LOW_VRAM = 2
    NORMAL_VRAM = 3
    HIGH_VRAM = 4
14
    SHARED = 5      #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.
15
16
17
18
19

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

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

26
total_vram = 0
27

28
lowvram_available = True
藍+85CD's avatar
藍+85CD committed
29
xpu_available = False
30

31
directml_enabled = False
32
if args.directml is not None:
33
34
    import torch_directml
    directml_enabled = True
35
36
37
38
39
40
    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))
41
    # torch_directml.disable_tiled_resources(True)
42
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
43

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

51
52
53
try:
    if torch.backends.mps.is_available():
        cpu_state = CPUState.MPS
KarryCharon's avatar
KarryCharon committed
54
        import torch.mps
55
56
57
58
59
60
except:
    pass

if args.cpu:
    cpu_state = CPUState.CPU

61
62
def is_intel_xpu():
    global cpu_state
63
    global xpu_available
64
65
66
67
68
69
    if cpu_state == CPUState.GPU:
        if xpu_available:
            return True
    return False

def get_torch_device():
70
    global directml_enabled
71
    global cpu_state
72
73
74
    if directml_enabled:
        global directml_device
        return directml_device
75
    if cpu_state == CPUState.MPS:
76
        return torch.device("mps")
77
    if cpu_state == CPUState.CPU:
78
79
        return torch.device("cpu")
    else:
80
        if is_intel_xpu():
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
            return torch.device("xpu")
        else:
            return torch.device(torch.cuda.current_device())

def get_total_memory(dev=None, torch_total_too=False):
    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
97
        elif is_intel_xpu():
98
99
            stats = torch.xpu.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
100
            mem_total = torch.xpu.get_device_properties(dev).total_memory
101
            mem_total_torch = mem_reserved
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
        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

122
123
124
125
126
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

127
128
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
129
130
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
131
132
133
134
else:
    try:
        import xformers
        import xformers.ops
135
        XFORMERS_IS_AVAILABLE = True
136
137
138
139
140
141
142
143
144
145
146
        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
147
    except:
148
        XFORMERS_IS_AVAILABLE = False
149

150
151
152
153
154
def is_nvidia():
    global cpu_state
    if cpu_state == CPUState.GPU:
        if torch.version.cuda:
            return True
155
    return False
156

157
158
159
160
161
ENABLE_PYTORCH_ATTENTION = False
if args.use_pytorch_cross_attention:
    ENABLE_PYTORCH_ATTENTION = True
    XFORMERS_IS_AVAILABLE = False

162
VAE_DTYPE = torch.float32
163

164
165
166
167
try:
    if is_nvidia():
        torch_version = torch.version.__version__
        if int(torch_version[0]) >= 2:
168
            if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
169
                ENABLE_PYTORCH_ATTENTION = True
170
171
            if torch.cuda.is_bf16_supported():
                VAE_DTYPE = torch.bfloat16
172
173
174
    if is_intel_xpu():
        if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
            ENABLE_PYTORCH_ATTENTION = True
175
176
177
except:
    pass

178
179
180
if is_intel_xpu():
    VAE_DTYPE = torch.bfloat16

181
182
183
184
185
186
187
if args.fp16_vae:
    VAE_DTYPE = torch.float16
elif args.bf16_vae:
    VAE_DTYPE = torch.bfloat16
elif args.fp32_vae:
    VAE_DTYPE = torch.float32

188

189
if ENABLE_PYTORCH_ATTENTION:
190
191
192
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
193

194
195
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
196
    lowvram_available = True
197
198
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
199
elif args.highvram or args.gpu_only:
200
    vram_state = VRAMState.HIGH_VRAM
201

202
FORCE_FP32 = False
203
FORCE_FP16 = False
204
205
206
207
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

208
209
210
211
if args.force_fp16:
    print("Forcing FP16.")
    FORCE_FP16 = True

212
if lowvram_available:
213
214
    try:
        import accelerate
215
216
        if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
            vram_state = set_vram_to
217
218
219
    except Exception as e:
        import traceback
        print(traceback.format_exc())
220
221
        print("ERROR: LOW VRAM MODE NEEDS accelerate.")
        lowvram_available = False
222

223

224
225
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
226

227
228
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
229

230
print(f"Set vram state to: {vram_state.name}")
231

232
233
234
235
DISABLE_SMART_MEMORY = args.disable_smart_memory

if DISABLE_SMART_MEMORY:
    print("Disabling smart memory management")
236

237
238
def get_torch_device_name(device):
    if hasattr(device, 'type'):
239
        if device.type == "cuda":
240
241
242
243
244
            try:
                allocator_backend = torch.cuda.get_allocator_backend()
            except:
                allocator_backend = ""
            return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
245
246
        else:
            return "{}".format(device.type)
247
    elif is_intel_xpu():
248
        return "{} {}".format(device, torch.xpu.get_device_name(device))
249
250
    else:
        return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
251
252

try:
253
    print("Device:", get_torch_device_name(get_torch_device()))
254
255
256
except:
    print("Could not pick default device.")

257
print("VAE dtype:", VAE_DTYPE)
258

comfyanonymous's avatar
comfyanonymous committed
259
current_loaded_models = []
260

comfyanonymous's avatar
comfyanonymous committed
261
262
263
264
265
class LoadedModel:
    def __init__(self, model):
        self.model = model
        self.model_accelerated = False
        self.device = model.load_device
266

comfyanonymous's avatar
comfyanonymous committed
267
268
    def model_memory(self):
        return self.model.model_size()
269

comfyanonymous's avatar
comfyanonymous committed
270
271
272
273
274
    def model_memory_required(self, device):
        if device == self.model.current_device:
            return 0
        else:
            return self.model_memory()
275

comfyanonymous's avatar
comfyanonymous committed
276
277
278
279
    def model_load(self, lowvram_model_memory=0):
        patch_model_to = None
        if lowvram_model_memory == 0:
            patch_model_to = self.device
280

comfyanonymous's avatar
comfyanonymous committed
281
282
        self.model.model_patches_to(self.device)
        self.model.model_patches_to(self.model.model_dtype())
283

comfyanonymous's avatar
comfyanonymous committed
284
285
286
287
288
289
        try:
            self.real_model = self.model.patch_model(device_to=patch_model_to) #TODO: do something with loras and offloading to CPU
        except Exception as e:
            self.model.unpatch_model(self.model.offload_device)
            self.model_unload()
            raise e
290

comfyanonymous's avatar
comfyanonymous committed
291
292
293
294
295
        if lowvram_model_memory > 0:
            print("loading in lowvram mode", lowvram_model_memory/(1024 * 1024))
            device_map = accelerate.infer_auto_device_map(self.real_model, max_memory={0: "{}MiB".format(lowvram_model_memory // (1024 * 1024)), "cpu": "16GiB"})
            accelerate.dispatch_model(self.real_model, device_map=device_map, main_device=self.device)
            self.model_accelerated = True
296

297
        if is_intel_xpu() and not args.disable_ipex_optimize:
298
            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
299

comfyanonymous's avatar
comfyanonymous committed
300
        return self.real_model
301

comfyanonymous's avatar
comfyanonymous committed
302
303
304
305
    def model_unload(self):
        if self.model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(self.real_model)
            self.model_accelerated = False
306

comfyanonymous's avatar
comfyanonymous committed
307
308
        self.model.unpatch_model(self.model.offload_device)
        self.model.model_patches_to(self.model.offload_device)
309

comfyanonymous's avatar
comfyanonymous committed
310
311
    def __eq__(self, other):
        return self.model is other.model
comfyanonymous's avatar
comfyanonymous committed
312

comfyanonymous's avatar
comfyanonymous committed
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
def minimum_inference_memory():
    return (1024 * 1024 * 1024)

def unload_model_clones(model):
    to_unload = []
    for i in range(len(current_loaded_models)):
        if model.is_clone(current_loaded_models[i].model):
            to_unload = [i] + to_unload

    for i in to_unload:
        print("unload clone", i)
        current_loaded_models.pop(i).model_unload()

def free_memory(memory_required, device, keep_loaded=[]):
    unloaded_model = False
    for i in range(len(current_loaded_models) -1, -1, -1):
comfyanonymous's avatar
comfyanonymous committed
329
330
331
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
comfyanonymous's avatar
comfyanonymous committed
332
333
334
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
comfyanonymous's avatar
comfyanonymous committed
335
336
337
                m = current_loaded_models.pop(i)
                m.model_unload()
                del m
comfyanonymous's avatar
comfyanonymous committed
338
339
340
341
342
343
344
                unloaded_model = True

    if unloaded_model:
        soft_empty_cache()


def load_models_gpu(models, memory_required=0):
345
346
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
347
348
349
350
351
352
353
354
355
356
357
358
359
    inference_memory = minimum_inference_memory()
    extra_mem = max(inference_memory, memory_required)

    models_to_load = []
    models_already_loaded = []
    for x in models:
        loaded_model = LoadedModel(x)

        if loaded_model in current_loaded_models:
            index = current_loaded_models.index(loaded_model)
            current_loaded_models.insert(0, current_loaded_models.pop(index))
            models_already_loaded.append(loaded_model)
        else:
360
361
            if hasattr(x, "model"):
                print(f"Requested to load {x.model.__class__.__name__}")
comfyanonymous's avatar
comfyanonymous committed
362
363
364
365
366
367
368
            models_to_load.append(loaded_model)

    if len(models_to_load) == 0:
        devs = set(map(lambda a: a.device, models_already_loaded))
        for d in devs:
            if d != torch.device("cpu"):
                free_memory(extra_mem, d, models_already_loaded)
369
370
        return

371
    print(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
372

comfyanonymous's avatar
comfyanonymous committed
373
374
375
376
    total_memory_required = {}
    for loaded_model in models_to_load:
        unload_model_clones(loaded_model.model)
        total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)
comfyanonymous's avatar
comfyanonymous committed
377

comfyanonymous's avatar
comfyanonymous committed
378
379
380
    for device in total_memory_required:
        if device != torch.device("cpu"):
            free_memory(total_memory_required[device] * 1.3 + extra_mem, device, models_already_loaded)
comfyanonymous's avatar
comfyanonymous committed
381

comfyanonymous's avatar
comfyanonymous committed
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
    for loaded_model in models_to_load:
        model = loaded_model.model
        torch_dev = model.load_device
        if is_device_cpu(torch_dev):
            vram_set_state = VRAMState.DISABLED
        else:
            vram_set_state = vram_state
        lowvram_model_memory = 0
        if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM):
            model_size = loaded_model.model_memory_required(torch_dev)
            current_free_mem = get_free_memory(torch_dev)
            lowvram_model_memory = int(max(256 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
            if model_size > (current_free_mem - inference_memory): #only switch to lowvram if really necessary
                vram_set_state = VRAMState.LOW_VRAM
            else:
                lowvram_model_memory = 0
398

comfyanonymous's avatar
comfyanonymous committed
399
400
        if vram_set_state == VRAMState.NO_VRAM:
            lowvram_model_memory = 256 * 1024 * 1024
401

comfyanonymous's avatar
comfyanonymous committed
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
        cur_loaded_model = loaded_model.model_load(lowvram_model_memory)
        current_loaded_models.insert(0, loaded_model)
    return


def load_model_gpu(model):
    return load_models_gpu([model])

def cleanup_models():
    to_delete = []
    for i in range(len(current_loaded_models)):
        if sys.getrefcount(current_loaded_models[i].model) <= 2:
            to_delete = [i] + to_delete

    for i in to_delete:
        x = current_loaded_models.pop(i)
        x.model_unload()
        del x
420

421
422
423
424
425
426
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
    return dtype_size

427
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
428
    if vram_state == VRAMState.HIGH_VRAM:
429
430
431
432
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
433
434
435
436
437
438
def unet_inital_load_device(parameters, dtype):
    torch_dev = get_torch_device()
    if vram_state == VRAMState.HIGH_VRAM:
        return torch_dev

    cpu_dev = torch.device("cpu")
439
440
441
    if DISABLE_SMART_MEMORY:
        return cpu_dev

442
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
443
444
445
446
447
448
449
450

    mem_dev = get_free_memory(torch_dev)
    mem_cpu = get_free_memory(cpu_dev)
    if mem_dev > mem_cpu and model_size < mem_dev:
        return torch_dev
    else:
        return cpu_dev

451
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
452
    if args.gpu_only:
453
454
455
456
        return get_torch_device()
    else:
        return torch.device("cpu")

457
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
458
    if args.gpu_only:
459
        return get_torch_device()
460
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
461
462
        if is_intel_xpu():
            return torch.device("cpu")
463
        if should_use_fp16(prioritize_performance=False):
464
465
466
            return get_torch_device()
        else:
            return torch.device("cpu")
467
468
469
    else:
        return torch.device("cpu")

470
471
472
473
def vae_device():
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
474
    if args.gpu_only:
475
476
477
478
        return get_torch_device()
    else:
        return torch.device("cpu")

479
def vae_dtype():
480
481
    global VAE_DTYPE
    return VAE_DTYPE
482

483
484
485
486
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
487

488
489
490
491
492
493
494
def cast_to_device(tensor, device, dtype, copy=False):
    device_supports_cast = False
    if tensor.dtype == torch.float32 or tensor.dtype == torch.float16:
        device_supports_cast = True
    elif tensor.dtype == torch.bfloat16:
        if hasattr(device, 'type') and device.type.startswith("cuda"):
            device_supports_cast = True
495
496
        elif is_intel_xpu():
            device_supports_cast = True
497
498
499
500
501
502
503
504
505
506

    if device_supports_cast:
        if copy:
            if tensor.device == device:
                return tensor.to(dtype, copy=copy)
            return tensor.to(device, copy=copy).to(dtype)
        else:
            return tensor.to(device).to(dtype)
    else:
        return tensor.to(dtype).to(device, copy=copy)
507

508
def xformers_enabled():
509
    global directml_enabled
510
511
    global cpu_state
    if cpu_state != CPUState.GPU:
512
        return False
513
    if is_intel_xpu():
514
515
516
        return False
    if directml_enabled:
        return False
517
    return XFORMERS_IS_AVAILABLE
518

519
520
521
522
523

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
524
525

    return XFORMERS_ENABLED_VAE
526

527
def pytorch_attention_enabled():
528
    global ENABLE_PYTORCH_ATTENTION
529
530
    return ENABLE_PYTORCH_ATTENTION

531
532
533
534
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
535
        if is_nvidia(): #pytorch flash attention only works on Nvidia
536
537
538
            return True
    return False

539
def get_free_memory(dev=None, torch_free_too=False):
540
    global directml_enabled
541
    if dev is None:
542
        dev = get_torch_device()
543

Yurii Mazurevich's avatar
Yurii Mazurevich committed
544
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
545
546
547
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
548
549
550
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
551
        elif is_intel_xpu():
552
553
554
555
556
            stats = torch.xpu.memory_stats(dev)
            mem_active = stats['active_bytes.all.current']
            mem_allocated = stats['allocated_bytes.all.current']
            mem_reserved = stats['reserved_bytes.all.current']
            mem_free_torch = mem_reserved - mem_active
557
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
558
559
560
561
562
563
564
        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
565
566
567
568
569

    if torch_free_too:
        return (mem_free_total, mem_free_torch)
    else:
        return mem_free_total
570

comfyanonymous's avatar
comfyanonymous committed
571
572
573
574
575
576
577
def batch_area_memory(area):
    if xformers_enabled() or pytorch_attention_flash_attention():
        #TODO: these formulas are copied from maximum_batch_area below
        return (area / 20) * (1024 * 1024)
    else:
        return (((area * 0.6) / 0.9) + 1024) * (1024 * 1024)

578
579
def maximum_batch_area():
    global vram_state
580
    if vram_state == VRAMState.NO_VRAM:
581
582
583
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
584
    if xformers_enabled() or pytorch_attention_flash_attention():
585
        #TODO: this needs to be tweaked
586
        area = 20 * memory_free
587
588
589
    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)
590
    return int(max(area, 0))
591
592

def cpu_mode():
593
594
    global cpu_state
    return cpu_state == CPUState.CPU
595

Yurii Mazurevich's avatar
Yurii Mazurevich committed
596
def mps_mode():
597
598
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
599

600
601
def is_device_cpu(device):
    if hasattr(device, 'type'):
comfyanonymous's avatar
comfyanonymous committed
602
603
604
605
606
607
608
        if (device.type == 'cpu'):
            return True
    return False

def is_device_mps(device):
    if hasattr(device, 'type'):
        if (device.type == 'mps'):
609
610
611
            return True
    return False

612
def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
613
614
    global directml_enabled

615
616
617
618
    if device is not None:
        if is_device_cpu(device):
            return False

619
620
621
    if FORCE_FP16:
        return True

622
    if device is not None: #TODO
623
        if is_device_mps(device):
624
            return False
625

626
627
628
    if FORCE_FP32:
        return False

629
630
631
    if directml_enabled:
        return False

632
    if cpu_mode() or mps_mode():
633
634
        return False #TODO ?

635
    if is_intel_xpu():
comfyanonymous's avatar
comfyanonymous committed
636
637
638
        return True

    if torch.cuda.is_bf16_supported():
639
640
        return True

comfyanonymous's avatar
comfyanonymous committed
641
    props = torch.cuda.get_device_properties("cuda")
642
643
644
645
646
647
648
649
650
651
652
653
654
655
    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())
656
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
657
658
            return True

659
660
661
    if props.major < 7:
        return False

662
    #FP16 is just broken on these cards
663
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX"]
664
665
666
667
668
669
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

670
def soft_empty_cache(force=False):
671
672
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
673
        torch.mps.empty_cache()
674
    elif is_intel_xpu():
675
676
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
677
        if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
678
679
680
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

comfyanonymous's avatar
comfyanonymous committed
681
682
683
684
685
686
687
def resolve_lowvram_weight(weight, model, key):
    if weight.device == torch.device("meta"): #lowvram NOTE: this depends on the inner working of the accelerate library so it might break.
        key_split = key.split('.')              # I have no idea why they don't just leave the weight there instead of using the meta device.
        op = comfy.utils.get_attr(model, '.'.join(key_split[:-1]))
        weight = op._hf_hook.weights_map[key_split[-1]]
    return weight

688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
#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()