model_management.py 21.1 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
63
def get_torch_device():
    global xpu_available
    global directml_enabled
64
    global cpu_state
65
66
67
    if directml_enabled:
        global directml_device
        return directml_device
68
    if cpu_state == CPUState.MPS:
69
        return torch.device("mps")
70
    if cpu_state == CPUState.CPU:
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
        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:
92
93
            stats = torch.xpu.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
94
            mem_total = torch.xpu.get_device_properties(dev).total_memory
95
            mem_total_torch = mem_reserved
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
        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

116
117
118
119
120
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

121
122
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
123
124
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
125
126
127
128
else:
    try:
        import xformers
        import xformers.ops
129
        XFORMERS_IS_AVAILABLE = True
130
131
132
133
134
135
136
137
138
139
140
        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
141
    except:
142
        XFORMERS_IS_AVAILABLE = False
143

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

150
ENABLE_PYTORCH_ATTENTION = args.use_pytorch_cross_attention
151
VAE_DTYPE = torch.float32
152

153
154
155
156
157
158

try:
    if is_nvidia():
        torch_version = torch.version.__version__
        if int(torch_version[0]) >= 2:
            if ENABLE_PYTORCH_ATTENTION == False and XFORMERS_IS_AVAILABLE == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
159
                ENABLE_PYTORCH_ATTENTION = True
160
161
162
163
164
165
166
167
168
169
170
171
            if torch.cuda.is_bf16_supported():
                VAE_DTYPE = torch.bfloat16
except:
    pass

if args.fp16_vae:
    VAE_DTYPE = torch.float16
elif args.bf16_vae:
    VAE_DTYPE = torch.bfloat16
elif args.fp32_vae:
    VAE_DTYPE = torch.float32

172

173
if ENABLE_PYTORCH_ATTENTION:
174
175
176
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
177
    XFORMERS_IS_AVAILABLE = False
178

179
180
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
181
    lowvram_available = True
182
183
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
184
elif args.highvram or args.gpu_only:
185
    vram_state = VRAMState.HIGH_VRAM
186

187
FORCE_FP32 = False
188
FORCE_FP16 = False
189
190
191
192
if args.force_fp32:
    print("Forcing FP32, if this improves things please report it.")
    FORCE_FP32 = True

193
194
195
196
if args.force_fp16:
    print("Forcing FP16.")
    FORCE_FP16 = True

197
if lowvram_available:
198
199
    try:
        import accelerate
200
201
        if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
            vram_state = set_vram_to
202
203
204
    except Exception as e:
        import traceback
        print(traceback.format_exc())
205
206
        print("ERROR: LOW VRAM MODE NEEDS accelerate.")
        lowvram_available = False
207

208

209
210
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
211

212
213
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
214

215
print(f"Set vram state to: {vram_state.name}")
216

217
218
219
220
DISABLE_SMART_MEMORY = args.disable_smart_memory

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

222
def get_torch_device_name(device):
223
    global xpu_available
224
    if hasattr(device, 'type'):
225
        if device.type == "cuda":
226
227
228
229
230
            try:
                allocator_backend = torch.cuda.get_allocator_backend()
            except:
                allocator_backend = ""
            return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)
231
232
        else:
            return "{}".format(device.type)
233
234
    elif xpu_available:
        return "{} {}".format(device, torch.xpu.get_device_name(device))
235
236
    else:
        return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))
237
238

try:
239
    print("Device:", get_torch_device_name(get_torch_device()))
240
241
242
except:
    print("Could not pick default device.")

243
print("VAE dtype:", VAE_DTYPE)
244

comfyanonymous's avatar
comfyanonymous committed
245
current_loaded_models = []
246

comfyanonymous's avatar
comfyanonymous committed
247
248
249
250
251
class LoadedModel:
    def __init__(self, model):
        self.model = model
        self.model_accelerated = False
        self.device = model.load_device
252

comfyanonymous's avatar
comfyanonymous committed
253
254
    def model_memory(self):
        return self.model.model_size()
255

comfyanonymous's avatar
comfyanonymous committed
256
257
258
259
260
    def model_memory_required(self, device):
        if device == self.model.current_device:
            return 0
        else:
            return self.model_memory()
261

comfyanonymous's avatar
comfyanonymous committed
262
    def model_load(self, lowvram_model_memory=0):
263
        global xpu_available
comfyanonymous's avatar
comfyanonymous committed
264
265
266
        patch_model_to = None
        if lowvram_model_memory == 0:
            patch_model_to = self.device
267

comfyanonymous's avatar
comfyanonymous committed
268
269
        self.model.model_patches_to(self.device)
        self.model.model_patches_to(self.model.model_dtype())
270

comfyanonymous's avatar
comfyanonymous committed
271
272
273
274
275
276
        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
277

comfyanonymous's avatar
comfyanonymous committed
278
279
280
281
282
        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
283

284
        if xpu_available and not args.disable_ipex_optimize:
285
            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
286

comfyanonymous's avatar
comfyanonymous committed
287
        return self.real_model
288

comfyanonymous's avatar
comfyanonymous committed
289
290
291
292
    def model_unload(self):
        if self.model_accelerated:
            accelerate.hooks.remove_hook_from_submodules(self.real_model)
            self.model_accelerated = False
293

comfyanonymous's avatar
comfyanonymous committed
294
295
        self.model.unpatch_model(self.model.offload_device)
        self.model.model_patches_to(self.model.offload_device)
296

comfyanonymous's avatar
comfyanonymous committed
297
298
    def __eq__(self, other):
        return self.model is other.model
comfyanonymous's avatar
comfyanonymous committed
299

comfyanonymous's avatar
comfyanonymous committed
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
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
316
317
318
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
comfyanonymous's avatar
comfyanonymous committed
319
320
321
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
comfyanonymous's avatar
comfyanonymous committed
322
323
324
                m = current_loaded_models.pop(i)
                m.model_unload()
                del m
comfyanonymous's avatar
comfyanonymous committed
325
326
327
328
329
330
331
                unloaded_model = True

    if unloaded_model:
        soft_empty_cache()


def load_models_gpu(models, memory_required=0):
332
333
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
    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:
            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)
354
355
        return

comfyanonymous's avatar
comfyanonymous committed
356
    print("loading new")
357

comfyanonymous's avatar
comfyanonymous committed
358
359
360
361
    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
362

comfyanonymous's avatar
comfyanonymous committed
363
364
365
    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
366

comfyanonymous's avatar
comfyanonymous committed
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
    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
383

comfyanonymous's avatar
comfyanonymous committed
384
385
        if vram_set_state == VRAMState.NO_VRAM:
            lowvram_model_memory = 256 * 1024 * 1024
386

comfyanonymous's avatar
comfyanonymous committed
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
        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)):
        print(sys.getrefcount(current_loaded_models[i].model))
        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
406

407
408
409
410
411
412
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
    return dtype_size

413
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
414
    if vram_state == VRAMState.HIGH_VRAM:
415
416
417
418
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
419
420
421
422
423
424
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")
425
426
427
    if DISABLE_SMART_MEMORY:
        return cpu_dev

428
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
429
430
431
432
433
434
435
436

    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

437
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
438
    if args.gpu_only:
439
440
441
442
        return get_torch_device()
    else:
        return torch.device("cpu")

443
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
444
    if args.gpu_only:
445
        return get_torch_device()
446
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
447
        if should_use_fp16(prioritize_performance=False):
448
449
450
            return get_torch_device()
        else:
            return torch.device("cpu")
451
452
453
    else:
        return torch.device("cpu")

454
455
456
457
def vae_device():
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
458
    if args.gpu_only:
459
460
461
462
        return get_torch_device()
    else:
        return torch.device("cpu")

463
def vae_dtype():
464
465
    global VAE_DTYPE
    return VAE_DTYPE
466

467
468
469
470
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
471

472

473
def xformers_enabled():
474
475
    global xpu_available
    global directml_enabled
476
477
    global cpu_state
    if cpu_state != CPUState.GPU:
478
        return False
479
480
481
482
    if xpu_available:
        return False
    if directml_enabled:
        return False
483
    return XFORMERS_IS_AVAILABLE
484

485
486
487
488
489

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
490
491

    return XFORMERS_ENABLED_VAE
492

493
def pytorch_attention_enabled():
494
    global ENABLE_PYTORCH_ATTENTION
495
496
    return ENABLE_PYTORCH_ATTENTION

497
498
499
500
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
501
        if is_nvidia(): #pytorch flash attention only works on Nvidia
502
503
504
            return True
    return False

505
def get_free_memory(dev=None, torch_free_too=False):
506
    global xpu_available
507
    global directml_enabled
508
    if dev is None:
509
        dev = get_torch_device()
510

Yurii Mazurevich's avatar
Yurii Mazurevich committed
511
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
512
513
514
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
515
516
517
518
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
        elif xpu_available:
519
520
521
522
523
            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
524
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
525
526
527
528
529
530
531
        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
532
533
534
535
536

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

comfyanonymous's avatar
comfyanonymous committed
538
539
540
541
542
543
544
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)

545
546
def maximum_batch_area():
    global vram_state
547
    if vram_state == VRAMState.NO_VRAM:
548
549
550
        return 0

    memory_free = get_free_memory() / (1024 * 1024)
551
    if xformers_enabled() or pytorch_attention_flash_attention():
552
        #TODO: this needs to be tweaked
553
        area = 20 * memory_free
554
555
556
    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)
557
    return int(max(area, 0))
558
559

def cpu_mode():
560
561
    global cpu_state
    return cpu_state == CPUState.CPU
562

Yurii Mazurevich's avatar
Yurii Mazurevich committed
563
def mps_mode():
564
565
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
566

567
568
def is_device_cpu(device):
    if hasattr(device, 'type'):
comfyanonymous's avatar
comfyanonymous committed
569
570
571
572
573
574
575
        if (device.type == 'cpu'):
            return True
    return False

def is_device_mps(device):
    if hasattr(device, 'type'):
        if (device.type == 'mps'):
576
577
578
            return True
    return False

579
def should_use_fp16(device=None, model_params=0, prioritize_performance=True):
580
    global xpu_available
581
582
    global directml_enabled

583
584
585
586
    if device is not None:
        if is_device_cpu(device):
            return False

587
588
589
    if FORCE_FP16:
        return True

590
    if device is not None: #TODO
591
        if is_device_mps(device):
592
            return False
593

594
595
596
    if FORCE_FP32:
        return False

597
598
599
    if directml_enabled:
        return False

600
    if cpu_mode() or mps_mode():
601
602
        return False #TODO ?

comfyanonymous's avatar
comfyanonymous committed
603
604
605
606
    if xpu_available:
        return True

    if torch.cuda.is_bf16_supported():
607
608
        return True

comfyanonymous's avatar
comfyanonymous committed
609
    props = torch.cuda.get_device_properties("cuda")
610
611
612
613
614
615
616
617
618
619
620
621
622
623
    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())
624
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
625
626
            return True

627
628
629
    if props.major < 7:
        return False

630
    #FP16 is just broken on these cards
631
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX"]
632
633
634
635
636
637
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

638
639
def soft_empty_cache():
    global xpu_available
640
641
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
642
643
        torch.mps.empty_cache()
    elif xpu_available:
644
645
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
646
        if is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
647
648
649
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

comfyanonymous's avatar
comfyanonymous committed
650
651
652
653
654
655
656
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

657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
#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()