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

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

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

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

27
total_vram = 0
28

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

32
if args.deterministic:
comfyanonymous's avatar
comfyanonymous committed
33
    logging.info("Using deterministic algorithms for pytorch")
34
35
    torch.use_deterministic_algorithms(True, warn_only=True)

36
directml_enabled = False
37
if args.directml is not None:
38
39
    import torch_directml
    directml_enabled = True
40
41
42
43
44
    device_index = args.directml
    if device_index < 0:
        directml_device = torch_directml.device()
    else:
        directml_device = torch_directml.device(device_index)
comfyanonymous's avatar
comfyanonymous committed
45
    logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index)))
46
    # torch_directml.disable_tiled_resources(True)
47
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
48

49
try:
50
51
52
    import intel_extension_for_pytorch as ipex
    if torch.xpu.is_available():
        xpu_available = True
53
54
55
except:
    pass

56
57
58
try:
    if torch.backends.mps.is_available():
        cpu_state = CPUState.MPS
KarryCharon's avatar
KarryCharon committed
59
        import torch.mps
60
61
62
63
64
65
except:
    pass

if args.cpu:
    cpu_state = CPUState.CPU

66
67
def is_intel_xpu():
    global cpu_state
68
    global xpu_available
69
70
71
72
73
74
    if cpu_state == CPUState.GPU:
        if xpu_available:
            return True
    return False

def get_torch_device():
75
    global directml_enabled
76
    global cpu_state
77
78
79
    if directml_enabled:
        global directml_device
        return directml_device
80
    if cpu_state == CPUState.MPS:
81
        return torch.device("mps")
82
    if cpu_state == CPUState.CPU:
83
84
        return torch.device("cpu")
    else:
85
        if is_intel_xpu():
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
            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
102
        elif is_intel_xpu():
103
104
            stats = torch.xpu.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
105
            mem_total = torch.xpu.get_device_properties(dev).total_memory
106
            mem_total_torch = mem_reserved
107
108
109
110
111
112
113
114
115
116
117
118
119
120
        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)
comfyanonymous's avatar
comfyanonymous committed
121
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
122
123
if not args.normalvram and not args.cpu:
    if lowvram_available and total_vram <= 4096:
124
        logging.warning("Trying to enable lowvram mode because your GPU seems to have 4GB or less. If you don't want this use: --normalvram")
125
126
        set_vram_to = VRAMState.LOW_VRAM

127
128
129
130
131
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

132
133
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
134
135
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
136
137
138
139
else:
    try:
        import xformers
        import xformers.ops
140
        XFORMERS_IS_AVAILABLE = True
141
142
143
144
        try:
            XFORMERS_IS_AVAILABLE = xformers._has_cpp_library
        except:
            pass
145
146
        try:
            XFORMERS_VERSION = xformers.version.__version__
comfyanonymous's avatar
comfyanonymous committed
147
            logging.info("xformers version: {}".format(XFORMERS_VERSION))
148
            if XFORMERS_VERSION.startswith("0.0.18"):
149
150
                logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")
                logging.warning("Please downgrade or upgrade xformers to a different version.\n")
151
152
153
                XFORMERS_ENABLED_VAE = False
        except:
            pass
154
    except:
155
        XFORMERS_IS_AVAILABLE = False
156

157
158
159
160
161
def is_nvidia():
    global cpu_state
    if cpu_state == CPUState.GPU:
        if torch.version.cuda:
            return True
162
    return False
163

164
165
166
167
168
ENABLE_PYTORCH_ATTENTION = False
if args.use_pytorch_cross_attention:
    ENABLE_PYTORCH_ATTENTION = True
    XFORMERS_IS_AVAILABLE = False

169
VAE_DTYPE = torch.float32
170

171
172
173
174
try:
    if is_nvidia():
        torch_version = torch.version.__version__
        if int(torch_version[0]) >= 2:
175
            if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
176
                ENABLE_PYTORCH_ATTENTION = True
177
            if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8:
178
                VAE_DTYPE = torch.bfloat16
179
180
181
    if is_intel_xpu():
        if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:
            ENABLE_PYTORCH_ATTENTION = True
182
183
184
except:
    pass

185
186
187
if is_intel_xpu():
    VAE_DTYPE = torch.bfloat16

188
189
190
if args.cpu_vae:
    VAE_DTYPE = torch.float32

191
192
193
194
195
196
197
if args.fp16_vae:
    VAE_DTYPE = torch.float16
elif args.bf16_vae:
    VAE_DTYPE = torch.bfloat16
elif args.fp32_vae:
    VAE_DTYPE = torch.float32

198

199
if ENABLE_PYTORCH_ATTENTION:
200
201
202
    torch.backends.cuda.enable_math_sdp(True)
    torch.backends.cuda.enable_flash_sdp(True)
    torch.backends.cuda.enable_mem_efficient_sdp(True)
203

204
205
if args.lowvram:
    set_vram_to = VRAMState.LOW_VRAM
206
    lowvram_available = True
207
208
elif args.novram:
    set_vram_to = VRAMState.NO_VRAM
209
elif args.highvram or args.gpu_only:
210
    vram_state = VRAMState.HIGH_VRAM
211

212
FORCE_FP32 = False
213
FORCE_FP16 = False
214
if args.force_fp32:
comfyanonymous's avatar
comfyanonymous committed
215
    logging.info("Forcing FP32, if this improves things please report it.")
216
217
    FORCE_FP32 = True

218
if args.force_fp16:
comfyanonymous's avatar
comfyanonymous committed
219
    logging.info("Forcing FP16.")
220
221
    FORCE_FP16 = True

222
if lowvram_available:
223
224
    if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):
        vram_state = set_vram_to
225

226

227
228
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
229

230
231
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
232

comfyanonymous's avatar
comfyanonymous committed
233
logging.info(f"Set vram state to: {vram_state.name}")
234

235
236
237
DISABLE_SMART_MEMORY = args.disable_smart_memory

if DISABLE_SMART_MEMORY:
comfyanonymous's avatar
comfyanonymous committed
238
    logging.info("Disabling smart memory management")
239

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

try:
comfyanonymous's avatar
comfyanonymous committed
256
    logging.info("Device: {}".format(get_torch_device_name(get_torch_device())))
257
except:
258
    logging.warning("Could not pick default device.")
259

comfyanonymous's avatar
comfyanonymous committed
260
logging.info("VAE dtype: {}".format(VAE_DTYPE))
261

comfyanonymous's avatar
comfyanonymous committed
262
current_loaded_models = []
263

264
265
266
267
268
269
270
271
def module_size(module):
    module_mem = 0
    sd = module.state_dict()
    for k in sd:
        t = sd[k]
        module_mem += t.nelement() * t.element_size()
    return module_mem

comfyanonymous's avatar
comfyanonymous committed
272
273
274
275
class LoadedModel:
    def __init__(self, model):
        self.model = model
        self.device = model.load_device
276
        self.weights_loaded = False
277

comfyanonymous's avatar
comfyanonymous committed
278
279
    def model_memory(self):
        return self.model.model_size()
280

comfyanonymous's avatar
comfyanonymous committed
281
282
283
284
285
    def model_memory_required(self, device):
        if device == self.model.current_device:
            return 0
        else:
            return self.model_memory()
286

comfyanonymous's avatar
comfyanonymous committed
287
    def model_load(self, lowvram_model_memory=0):
288
        patch_model_to = self.device
289

comfyanonymous's avatar
comfyanonymous committed
290
291
        self.model.model_patches_to(self.device)
        self.model.model_patches_to(self.model.model_dtype())
292

293
294
        load_weights = not self.weights_loaded

comfyanonymous's avatar
comfyanonymous committed
295
        try:
296
            if lowvram_model_memory > 0 and load_weights:
297
298
                self.real_model = self.model.patch_model_lowvram(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory)
            else:
299
                self.real_model = self.model.patch_model(device_to=patch_model_to, patch_weights=load_weights)
comfyanonymous's avatar
comfyanonymous committed
300
301
302
303
        except Exception as e:
            self.model.unpatch_model(self.model.offload_device)
            self.model_unload()
            raise e
304

305
        if is_intel_xpu() and not args.disable_ipex_optimize:
306
            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
307

308
        self.weights_loaded = True
comfyanonymous's avatar
comfyanonymous committed
309
        return self.real_model
310

311
312
    def model_unload(self, unpatch_weights=True):
        self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
comfyanonymous's avatar
comfyanonymous committed
313
        self.model.model_patches_to(self.model.offload_device)
314
        self.weights_loaded = self.weights_loaded and not unpatch_weights
315

comfyanonymous's avatar
comfyanonymous committed
316
317
    def __eq__(self, other):
        return self.model is other.model
comfyanonymous's avatar
comfyanonymous committed
318

comfyanonymous's avatar
comfyanonymous committed
319
320
321
def minimum_inference_memory():
    return (1024 * 1024 * 1024)

322
323
324
def unload_model_clones(loaded_model, unload_weights_only=True):
    model = loaded_model.model

comfyanonymous's avatar
comfyanonymous committed
325
326
327
328
329
    to_unload = []
    for i in range(len(current_loaded_models)):
        if model.is_clone(current_loaded_models[i].model):
            to_unload = [i] + to_unload

330
331
332
333
    if len(to_unload) == 0:
        return

    same_weights = 0
comfyanonymous's avatar
comfyanonymous committed
334
    for i in to_unload:
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
        if model.clone_has_same_weights(current_loaded_models[i].model):
            same_weights += 1

    if same_weights == len(to_unload):
        unload_weight = False
    else:
        unload_weight = True

    if unload_weights_only and unload_weight == False:
        return

    for i in to_unload:
        logging.debug("unload clone {} {}".format(i, unload_weight))
        current_loaded_models.pop(i).model_unload(unpatch_weights=unload_weight)

    loaded_model.weights_loaded = not unload_weight
comfyanonymous's avatar
comfyanonymous committed
351
352
353
354

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
355
356
357
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
comfyanonymous's avatar
comfyanonymous committed
358
359
360
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
comfyanonymous's avatar
comfyanonymous committed
361
362
363
                m = current_loaded_models.pop(i)
                m.model_unload()
                del m
comfyanonymous's avatar
comfyanonymous committed
364
365
366
367
                unloaded_model = True

    if unloaded_model:
        soft_empty_cache()
368
369
370
371
372
    else:
        if vram_state != VRAMState.HIGH_VRAM:
            mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)
            if mem_free_torch > mem_free_total * 0.25:
                soft_empty_cache()
comfyanonymous's avatar
comfyanonymous committed
373
374

def load_models_gpu(models, memory_required=0):
375
376
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
377
378
379
380
381
382
383
384
385
386
387
388
389
    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:
390
            if hasattr(x, "model"):
comfyanonymous's avatar
comfyanonymous committed
391
                logging.info(f"Requested to load {x.model.__class__.__name__}")
comfyanonymous's avatar
comfyanonymous committed
392
393
394
395
396
397
398
            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)
399
400
        return

comfyanonymous's avatar
comfyanonymous committed
401
    logging.info(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}")
402

comfyanonymous's avatar
comfyanonymous committed
403
404
    total_memory_required = {}
    for loaded_model in models_to_load:
405
        unload_model_clones(loaded_model, unload_weights_only=True) #unload clones where the weights are different
comfyanonymous's avatar
comfyanonymous committed
406
        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
407

comfyanonymous's avatar
comfyanonymous committed
408
409
410
    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
411

412
413
414
    for loaded_model in models_to_load:
        unload_model_clones(loaded_model, unload_weights_only=False) #unload the rest of the clones where the weights can stay loaded

comfyanonymous's avatar
comfyanonymous committed
415
416
417
418
419
420
421
422
423
424
425
    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)
426
            lowvram_model_memory = int(max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
comfyanonymous's avatar
comfyanonymous committed
427
428
429
430
            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
431

comfyanonymous's avatar
comfyanonymous committed
432
        if vram_set_state == VRAMState.NO_VRAM:
433
            lowvram_model_memory = 64 * 1024 * 1024
434

comfyanonymous's avatar
comfyanonymous committed
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
        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
453

454
455
456
457
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
458
459
460
461
462
463
464
    elif dtype == torch.float32:
        dtype_size = 4
    else:
        try:
            dtype_size = dtype.itemsize
        except: #Old pytorch doesn't have .itemsize
            pass
465
466
    return dtype_size

467
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
468
    if vram_state == VRAMState.HIGH_VRAM:
469
470
471
472
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
473
474
475
476
477
478
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")
479
480
481
    if DISABLE_SMART_MEMORY:
        return cpu_dev

482
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
483
484
485
486
487
488
489
490

    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

comfyanonymous's avatar
comfyanonymous committed
491
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
492
493
    if args.bf16_unet:
        return torch.bfloat16
494
495
    if args.fp16_unet:
        return torch.float16
496
497
498
499
    if args.fp8_e4m3fn_unet:
        return torch.float8_e4m3fn
    if args.fp8_e5m2_unet:
        return torch.float8_e5m2
500
    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
501
502
        if torch.float16 in supported_dtypes:
            return torch.float16
503
    if should_use_bf16(device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
504
505
        if torch.bfloat16 in supported_dtypes:
            return torch.bfloat16
506
507
    return torch.float32

508
# None means no manual cast
comfyanonymous's avatar
comfyanonymous committed
509
def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
510
511
512
    if weight_dtype == torch.float32:
        return None

comfyanonymous's avatar
comfyanonymous committed
513
    fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)
514
515
516
    if fp16_supported and weight_dtype == torch.float16:
        return None

comfyanonymous's avatar
comfyanonymous committed
517
518
519
520
521
    bf16_supported = should_use_bf16(inference_device)
    if bf16_supported and weight_dtype == torch.bfloat16:
        return None

    if fp16_supported and torch.float16 in supported_dtypes:
522
        return torch.float16
comfyanonymous's avatar
comfyanonymous committed
523
524
525

    elif bf16_supported and torch.bfloat16 in supported_dtypes:
        return torch.bfloat16
526
527
528
    else:
        return torch.float32

529
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
530
    if args.gpu_only:
531
532
533
534
        return get_torch_device()
    else:
        return torch.device("cpu")

535
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
536
    if args.gpu_only:
537
        return get_torch_device()
538
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
539
540
        if is_intel_xpu():
            return torch.device("cpu")
541
        if should_use_fp16(prioritize_performance=False):
542
543
544
            return get_torch_device()
        else:
            return torch.device("cpu")
545
546
547
    else:
        return torch.device("cpu")

548
549
550
551
552
553
554
555
556
557
def text_encoder_dtype(device=None):
    if args.fp8_e4m3fn_text_enc:
        return torch.float8_e4m3fn
    elif args.fp8_e5m2_text_enc:
        return torch.float8_e5m2
    elif args.fp16_text_enc:
        return torch.float16
    elif args.fp32_text_enc:
        return torch.float32

558
559
560
    if is_device_cpu(device):
        return torch.float16

561
562
    return torch.float16

563

564
565
566
567
568
569
def intermediate_device():
    if args.gpu_only:
        return get_torch_device()
    else:
        return torch.device("cpu")

570
def vae_device():
571
572
    if args.cpu_vae:
        return torch.device("cpu")
573
574
575
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
576
    if args.gpu_only:
577
578
579
580
        return get_torch_device()
    else:
        return torch.device("cpu")

581
def vae_dtype():
582
583
    global VAE_DTYPE
    return VAE_DTYPE
584

585
586
587
588
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
589

590
591
592
def supports_dtype(device, dtype): #TODO
    if dtype == torch.float32:
        return True
593
    if is_device_cpu(device):
594
595
596
597
598
599
600
        return False
    if dtype == torch.float16:
        return True
    if dtype == torch.bfloat16:
        return True
    return False

601
602
603
604
605
def device_supports_non_blocking(device):
    if is_device_mps(device):
        return False #pytorch bug? mps doesn't support non blocking
    return True

606
607
608
609
610
611
612
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
613
614
        elif is_intel_xpu():
            device_supports_cast = True
615

616
    non_blocking = device_supports_non_blocking(device)
comfyanonymous's avatar
comfyanonymous committed
617

618
619
620
    if device_supports_cast:
        if copy:
            if tensor.device == device:
comfyanonymous's avatar
comfyanonymous committed
621
622
                return tensor.to(dtype, copy=copy, non_blocking=non_blocking)
            return tensor.to(device, copy=copy, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
623
        else:
comfyanonymous's avatar
comfyanonymous committed
624
            return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
625
    else:
comfyanonymous's avatar
comfyanonymous committed
626
        return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
627

628
def xformers_enabled():
629
    global directml_enabled
630
631
    global cpu_state
    if cpu_state != CPUState.GPU:
632
        return False
633
    if is_intel_xpu():
634
635
636
        return False
    if directml_enabled:
        return False
637
    return XFORMERS_IS_AVAILABLE
638

639
640
641
642
643

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
644
645

    return XFORMERS_ENABLED_VAE
646

647
def pytorch_attention_enabled():
648
    global ENABLE_PYTORCH_ATTENTION
649
650
    return ENABLE_PYTORCH_ATTENTION

651
652
653
654
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
655
        if is_nvidia(): #pytorch flash attention only works on Nvidia
656
657
658
            return True
    return False

659
def get_free_memory(dev=None, torch_free_too=False):
660
    global directml_enabled
661
    if dev is None:
662
        dev = get_torch_device()
663

Yurii Mazurevich's avatar
Yurii Mazurevich committed
664
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
665
666
667
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
668
669
670
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
671
        elif is_intel_xpu():
672
673
674
675
676
            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
677
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
678
679
680
681
682
683
684
        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
685
686
687
688
689

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

691
def cpu_mode():
692
693
    global cpu_state
    return cpu_state == CPUState.CPU
694

Yurii Mazurevich's avatar
Yurii Mazurevich committed
695
def mps_mode():
696
697
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
698

699
def is_device_type(device, type):
700
    if hasattr(device, 'type'):
701
        if (device.type == type):
comfyanonymous's avatar
comfyanonymous committed
702
703
704
            return True
    return False

705
706
707
def is_device_cpu(device):
    return is_device_type(device, 'cpu')

comfyanonymous's avatar
comfyanonymous committed
708
def is_device_mps(device):
709
710
711
712
    return is_device_type(device, 'mps')

def is_device_cuda(device):
    return is_device_type(device, 'cuda')
713

714
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
715
716
    global directml_enabled

717
718
719
720
    if device is not None:
        if is_device_cpu(device):
            return False

721
722
723
    if FORCE_FP16:
        return True

724
    if device is not None:
725
        if is_device_mps(device):
726
            return True
727

728
729
730
    if FORCE_FP32:
        return False

731
732
733
    if directml_enabled:
        return False

734
735
736
737
738
    if mps_mode():
        return True

    if cpu_mode():
        return False
739

740
    if is_intel_xpu():
comfyanonymous's avatar
comfyanonymous committed
741
742
        return True

743
    if torch.version.hip:
744
745
        return True

comfyanonymous's avatar
comfyanonymous committed
746
    props = torch.cuda.get_device_properties("cuda")
747
748
749
    if props.major >= 8:
        return True

750
751
752
753
754
755
756
    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
757
    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"]
758
759
760
761
    for x in nvidia_10_series:
        if x in props.name.lower():
            fp16_works = True

762
    if fp16_works or manual_cast:
763
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
764
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
765
766
            return True

767
768
769
    if props.major < 7:
        return False

770
    #FP16 is just broken on these cards
771
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"]
772
773
774
775
776
777
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

778
779
780
781
782
783
784
785
786
def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
    if device is not None:
        if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow
            return False

    if device is not None: #TODO not sure about mps bf16 support
        if is_device_mps(device):
            return False

787
788
789
    if FORCE_FP32:
        return False

790
791
792
793
794
795
    if directml_enabled:
        return False

    if cpu_mode() or mps_mode():
        return False

comfyanonymous's avatar
comfyanonymous committed
796
797
798
799
800
801
802
803
804
805
    if is_intel_xpu():
        return True

    if device is None:
        device = torch.device("cuda")

    props = torch.cuda.get_device_properties(device)
    if props.major >= 8:
        return True

806
807
808
809
810
811
812
    bf16_works = torch.cuda.is_bf16_supported()

    if bf16_works or manual_cast:
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
            return True

comfyanonymous's avatar
comfyanonymous committed
813
814
    return False

815
def soft_empty_cache(force=False):
816
817
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
818
        torch.mps.empty_cache()
819
    elif is_intel_xpu():
820
821
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
822
        if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
823
824
825
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

826
827
828
829
def unload_all_models():
    free_memory(1e30, get_torch_device())


830
def resolve_lowvram_weight(weight, model, key): #TODO: remove
comfyanonymous's avatar
comfyanonymous committed
831
832
    return weight

833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
#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()