model_management.py 25.2 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

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

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

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

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

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

302
        if is_intel_xpu() and not args.disable_ipex_optimize:
303
            self.real_model = torch.xpu.optimize(self.real_model.eval(), inplace=True, auto_kernel_selection=True, graph_mode=True)
304

comfyanonymous's avatar
comfyanonymous committed
305
        return self.real_model
306

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

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

comfyanonymous's avatar
comfyanonymous committed
314
315
316
317
318
319
320
321
322
323
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:
comfyanonymous's avatar
comfyanonymous committed
324
        logging.debug("unload clone {}".format(i))
comfyanonymous's avatar
comfyanonymous committed
325
326
327
328
329
        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
330
331
332
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
comfyanonymous's avatar
comfyanonymous committed
333
334
335
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
comfyanonymous's avatar
comfyanonymous committed
336
337
338
                m = current_loaded_models.pop(i)
                m.model_unload()
                del m
comfyanonymous's avatar
comfyanonymous committed
339
340
341
342
                unloaded_model = True

    if unloaded_model:
        soft_empty_cache()
343
344
345
346
347
    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
348
349

def load_models_gpu(models, memory_required=0):
350
351
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
352
353
354
355
356
357
358
359
360
361
362
363
364
    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:
365
            if hasattr(x, "model"):
comfyanonymous's avatar
comfyanonymous committed
366
                logging.info(f"Requested to load {x.model.__class__.__name__}")
comfyanonymous's avatar
comfyanonymous committed
367
368
369
370
371
372
373
            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)
374
375
        return

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

comfyanonymous's avatar
comfyanonymous committed
378
379
380
381
    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
382

comfyanonymous's avatar
comfyanonymous committed
383
384
385
    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
386

comfyanonymous's avatar
comfyanonymous committed
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)
398
            lowvram_model_memory = int(max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
comfyanonymous's avatar
comfyanonymous committed
399
400
401
402
            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
403

comfyanonymous's avatar
comfyanonymous committed
404
        if vram_set_state == VRAMState.NO_VRAM:
405
            lowvram_model_memory = 64 * 1024 * 1024
406

comfyanonymous's avatar
comfyanonymous committed
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
        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
425

426
427
428
429
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
430
431
432
433
434
435
436
    elif dtype == torch.float32:
        dtype_size = 4
    else:
        try:
            dtype_size = dtype.itemsize
        except: #Old pytorch doesn't have .itemsize
            pass
437
438
    return dtype_size

439
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
440
    if vram_state == VRAMState.HIGH_VRAM:
441
442
443
444
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
445
446
447
448
449
450
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")
451
452
453
    if DISABLE_SMART_MEMORY:
        return cpu_dev

454
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
455
456
457
458
459
460
461
462

    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
463
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
464
465
    if args.bf16_unet:
        return torch.bfloat16
466
467
    if args.fp16_unet:
        return torch.float16
468
469
470
471
    if args.fp8_e4m3fn_unet:
        return torch.float8_e4m3fn
    if args.fp8_e5m2_unet:
        return torch.float8_e5m2
472
    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
473
474
        if torch.float16 in supported_dtypes:
            return torch.float16
475
    if should_use_bf16(device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
476
477
        if torch.bfloat16 in supported_dtypes:
            return torch.bfloat16
478
479
    return torch.float32

480
# None means no manual cast
comfyanonymous's avatar
comfyanonymous committed
481
def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
482
483
484
    if weight_dtype == torch.float32:
        return None

comfyanonymous's avatar
comfyanonymous committed
485
    fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)
486
487
488
    if fp16_supported and weight_dtype == torch.float16:
        return None

comfyanonymous's avatar
comfyanonymous committed
489
490
491
492
493
    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:
494
        return torch.float16
comfyanonymous's avatar
comfyanonymous committed
495
496
497

    elif bf16_supported and torch.bfloat16 in supported_dtypes:
        return torch.bfloat16
498
499
500
    else:
        return torch.float32

501
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
502
    if args.gpu_only:
503
504
505
506
        return get_torch_device()
    else:
        return torch.device("cpu")

507
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
508
    if args.gpu_only:
509
        return get_torch_device()
510
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
511
512
        if is_intel_xpu():
            return torch.device("cpu")
513
        if should_use_fp16(prioritize_performance=False):
514
515
516
            return get_torch_device()
        else:
            return torch.device("cpu")
517
518
519
    else:
        return torch.device("cpu")

520
521
522
523
524
525
526
527
528
529
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

530
531
532
    if is_device_cpu(device):
        return torch.float16

533
534
    return torch.float16

535

536
537
538
539
540
541
def intermediate_device():
    if args.gpu_only:
        return get_torch_device()
    else:
        return torch.device("cpu")

542
def vae_device():
543
544
    if args.cpu_vae:
        return torch.device("cpu")
545
546
547
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
548
    if args.gpu_only:
549
550
551
552
        return get_torch_device()
    else:
        return torch.device("cpu")

553
def vae_dtype():
554
555
    global VAE_DTYPE
    return VAE_DTYPE
556

557
558
559
560
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
561

562
563
564
def supports_dtype(device, dtype): #TODO
    if dtype == torch.float32:
        return True
565
    if is_device_cpu(device):
566
567
568
569
570
571
572
        return False
    if dtype == torch.float16:
        return True
    if dtype == torch.bfloat16:
        return True
    return False

573
574
575
576
577
def device_supports_non_blocking(device):
    if is_device_mps(device):
        return False #pytorch bug? mps doesn't support non blocking
    return True

578
579
580
581
582
583
584
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
585
586
        elif is_intel_xpu():
            device_supports_cast = True
587

588
    non_blocking = device_supports_non_blocking(device)
comfyanonymous's avatar
comfyanonymous committed
589

590
591
592
    if device_supports_cast:
        if copy:
            if tensor.device == device:
comfyanonymous's avatar
comfyanonymous committed
593
594
                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)
595
        else:
comfyanonymous's avatar
comfyanonymous committed
596
            return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
597
    else:
comfyanonymous's avatar
comfyanonymous committed
598
        return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
599

600
def xformers_enabled():
601
    global directml_enabled
602
603
    global cpu_state
    if cpu_state != CPUState.GPU:
604
        return False
605
    if is_intel_xpu():
606
607
608
        return False
    if directml_enabled:
        return False
609
    return XFORMERS_IS_AVAILABLE
610

611
612
613
614
615

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
616
617

    return XFORMERS_ENABLED_VAE
618

619
def pytorch_attention_enabled():
620
    global ENABLE_PYTORCH_ATTENTION
621
622
    return ENABLE_PYTORCH_ATTENTION

623
624
625
626
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
627
        if is_nvidia(): #pytorch flash attention only works on Nvidia
628
629
630
            return True
    return False

631
def get_free_memory(dev=None, torch_free_too=False):
632
    global directml_enabled
633
    if dev is None:
634
        dev = get_torch_device()
635

Yurii Mazurevich's avatar
Yurii Mazurevich committed
636
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
637
638
639
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
640
641
642
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
643
        elif is_intel_xpu():
644
645
646
647
648
            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
649
            mem_free_total = torch.xpu.get_device_properties(dev).total_memory - mem_allocated
650
651
652
653
654
655
656
        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
657
658
659
660
661

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

663
def cpu_mode():
664
665
    global cpu_state
    return cpu_state == CPUState.CPU
666

Yurii Mazurevich's avatar
Yurii Mazurevich committed
667
def mps_mode():
668
669
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
670

671
def is_device_type(device, type):
672
    if hasattr(device, 'type'):
673
        if (device.type == type):
comfyanonymous's avatar
comfyanonymous committed
674
675
676
            return True
    return False

677
678
679
def is_device_cpu(device):
    return is_device_type(device, 'cpu')

comfyanonymous's avatar
comfyanonymous committed
680
def is_device_mps(device):
681
682
683
684
    return is_device_type(device, 'mps')

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

686
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
687
688
    global directml_enabled

689
690
691
692
    if device is not None:
        if is_device_cpu(device):
            return False

693
694
695
    if FORCE_FP16:
        return True

696
    if device is not None:
697
        if is_device_mps(device):
698
            return True
699

700
701
702
    if FORCE_FP32:
        return False

703
704
705
    if directml_enabled:
        return False

706
707
708
709
710
    if mps_mode():
        return True

    if cpu_mode():
        return False
711

712
    if is_intel_xpu():
comfyanonymous's avatar
comfyanonymous committed
713
714
        return True

715
    if torch.version.hip:
716
717
        return True

comfyanonymous's avatar
comfyanonymous committed
718
    props = torch.cuda.get_device_properties("cuda")
719
720
721
    if props.major >= 8:
        return True

722
723
724
725
726
727
728
    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
729
    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"]
730
731
732
733
    for x in nvidia_10_series:
        if x in props.name.lower():
            fp16_works = True

734
    if fp16_works or manual_cast:
735
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
736
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
737
738
            return True

739
740
741
    if props.major < 7:
        return False

742
    #FP16 is just broken on these cards
743
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"]
744
745
746
747
748
749
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

750
751
752
753
754
755
756
757
758
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

759
760
761
    if FORCE_FP32:
        return False

762
763
764
765
766
767
    if directml_enabled:
        return False

    if cpu_mode() or mps_mode():
        return False

comfyanonymous's avatar
comfyanonymous committed
768
769
770
771
772
773
774
775
776
777
    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

778
779
780
781
782
783
784
    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
785
786
    return False

787
def soft_empty_cache(force=False):
788
789
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
790
        torch.mps.empty_cache()
791
    elif is_intel_xpu():
792
793
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
794
        if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
795
796
797
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

798
799
800
801
def unload_all_models():
    free_memory(1e30, get_torch_device())


802
def resolve_lowvram_weight(weight, model, key): #TODO: remove
comfyanonymous's avatar
comfyanonymous committed
803
804
    return weight

805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
#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()