model_management.py 27.8 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
import platform
9

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

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

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

28
total_vram = 0
29

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

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

37
directml_enabled = False
38
if args.directml is not None:
39
40
    import torch_directml
    directml_enabled = True
41
42
43
44
45
    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
46
    logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index)))
47
    # torch_directml.disable_tiled_resources(True)
48
    lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.
49

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

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

if args.cpu:
    cpu_state = CPUState.CPU

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

def get_torch_device():
76
    global directml_enabled
77
    global cpu_state
78
79
80
    if directml_enabled:
        global directml_device
        return directml_device
81
    if cpu_state == CPUState.MPS:
82
        return torch.device("mps")
83
    if cpu_state == CPUState.CPU:
84
85
        return torch.device("cpu")
    else:
86
        if is_intel_xpu():
87
            return torch.device("xpu", torch.xpu.current_device())
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
        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
103
        elif is_intel_xpu():
104
105
106
            stats = torch.xpu.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
            mem_total_torch = mem_reserved
107
            mem_total = torch.xpu.get_device_properties(dev).total_memory
108
109
110
111
112
113
114
115
116
117
118
119
120
121
        else:
            stats = torch.cuda.memory_stats(dev)
            mem_reserved = stats['reserved_bytes.all.current']
            _, mem_total_cuda = torch.cuda.mem_get_info(dev)
            mem_total_torch = mem_reserved
            mem_total = mem_total_cuda

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

total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)
total_ram = psutil.virtual_memory().total / (1024 * 1024)
comfyanonymous's avatar
comfyanonymous committed
122
logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))
123

comfyanonymous's avatar
comfyanonymous committed
124
125
126
127
128
try:
    logging.info("pytorch version: {}".format(torch.version.__version__))
except:
    pass

129
130
131
132
133
try:
    OOM_EXCEPTION = torch.cuda.OutOfMemoryError
except:
    OOM_EXCEPTION = Exception

134
135
XFORMERS_VERSION = ""
XFORMERS_ENABLED_VAE = True
136
137
if args.disable_xformers:
    XFORMERS_IS_AVAILABLE = False
138
139
140
141
else:
    try:
        import xformers
        import xformers.ops
142
        XFORMERS_IS_AVAILABLE = True
143
144
145
146
        try:
            XFORMERS_IS_AVAILABLE = xformers._has_cpp_library
        except:
            pass
147
148
        try:
            XFORMERS_VERSION = xformers.version.__version__
comfyanonymous's avatar
comfyanonymous committed
149
            logging.info("xformers version: {}".format(XFORMERS_VERSION))
150
            if XFORMERS_VERSION.startswith("0.0.18"):
151
152
                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")
153
154
155
                XFORMERS_ENABLED_VAE = False
        except:
            pass
156
    except:
157
        XFORMERS_IS_AVAILABLE = False
158

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

166
167
168
169
170
ENABLE_PYTORCH_ATTENTION = False
if args.use_pytorch_cross_attention:
    ENABLE_PYTORCH_ATTENTION = True
    XFORMERS_IS_AVAILABLE = False

171
VAE_DTYPE = torch.float32
172

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

187
188
189
if is_intel_xpu():
    VAE_DTYPE = torch.bfloat16

190
191
192
if args.cpu_vae:
    VAE_DTYPE = torch.float32

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

200

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

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

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

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

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

228

229
230
if cpu_state != CPUState.GPU:
    vram_state = VRAMState.DISABLED
231

232
233
if cpu_state == CPUState.MPS:
    vram_state = VRAMState.SHARED
234

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

237
238
239
DISABLE_SMART_MEMORY = args.disable_smart_memory

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

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

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

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

comfyanonymous's avatar
comfyanonymous committed
264
current_loaded_models = []
265

266
267
268
269
270
271
272
273
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
274
275
276
277
class LoadedModel:
    def __init__(self, model):
        self.model = model
        self.device = model.load_device
278
        self.weights_loaded = False
279
        self.real_model = None
280

comfyanonymous's avatar
comfyanonymous committed
281
282
    def model_memory(self):
        return self.model.model_size()
283

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

290
    def model_load(self, lowvram_model_memory=0, force_patch_weights=False):
291
        patch_model_to = self.device
292

comfyanonymous's avatar
comfyanonymous committed
293
294
        self.model.model_patches_to(self.device)
        self.model.model_patches_to(self.model.model_dtype())
295

296
297
        load_weights = not self.weights_loaded

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

308
        if is_intel_xpu() and not args.disable_ipex_optimize:
309
            self.real_model = ipex.optimize(self.real_model.eval(), graph_mode=True, concat_linear=True)
310

311
        self.weights_loaded = True
comfyanonymous's avatar
comfyanonymous committed
312
        return self.real_model
313

314
315
316
317
318
    def should_reload_model(self, force_patch_weights=False):
        if force_patch_weights and self.model.lowvram_patch_counter > 0:
            return True
        return False

319
320
    def model_unload(self, unpatch_weights=True):
        self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
comfyanonymous's avatar
comfyanonymous committed
321
        self.model.model_patches_to(self.model.offload_device)
322
        self.weights_loaded = self.weights_loaded and not unpatch_weights
323
        self.real_model = None
324

comfyanonymous's avatar
comfyanonymous committed
325
326
    def __eq__(self, other):
        return self.model is other.model
comfyanonymous's avatar
comfyanonymous committed
327

comfyanonymous's avatar
comfyanonymous committed
328
329
330
def minimum_inference_memory():
    return (1024 * 1024 * 1024)

331
def unload_model_clones(model, unload_weights_only=True, force_unload=True):
comfyanonymous's avatar
comfyanonymous committed
332
333
334
335
336
    to_unload = []
    for i in range(len(current_loaded_models)):
        if model.is_clone(current_loaded_models[i].model):
            to_unload = [i] + to_unload

337
    if len(to_unload) == 0:
338
        return True
339
340

    same_weights = 0
comfyanonymous's avatar
comfyanonymous committed
341
    for i in to_unload:
342
343
344
345
346
347
348
349
        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

350
351
352
    if not force_unload:
        if unload_weights_only and unload_weight == False:
            return None
353
354
355
356
357

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

358
    return unload_weight
comfyanonymous's avatar
comfyanonymous committed
359
360

def free_memory(memory_required, device, keep_loaded=[]):
361
362
363
    unloaded_model = []
    can_unload = []

comfyanonymous's avatar
comfyanonymous committed
364
365
366
367
    for i in range(len(current_loaded_models) -1, -1, -1):
        shift_model = current_loaded_models[i]
        if shift_model.device == device:
            if shift_model not in keep_loaded:
368
369
370
371
372
373
374
375
376
377
378
379
                can_unload.append((sys.getrefcount(shift_model.model), shift_model.model_memory(), i))

    for x in sorted(can_unload):
        i = x[-1]
        if not DISABLE_SMART_MEMORY:
            if get_free_memory(device) > memory_required:
                break
        current_loaded_models[i].model_unload()
        unloaded_model.append(i)

    for i in sorted(unloaded_model, reverse=True):
        current_loaded_models.pop(i)
comfyanonymous's avatar
comfyanonymous committed
380

381
    if len(unloaded_model) > 0:
comfyanonymous's avatar
comfyanonymous committed
382
        soft_empty_cache()
383
384
385
386
387
    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
388

389
def load_models_gpu(models, memory_required=0, force_patch_weights=False):
390
391
    global vram_state

comfyanonymous's avatar
comfyanonymous committed
392
393
394
    inference_memory = minimum_inference_memory()
    extra_mem = max(inference_memory, memory_required)

395
396
    models = set(models)

comfyanonymous's avatar
comfyanonymous committed
397
398
399
400
    models_to_load = []
    models_already_loaded = []
    for x in models:
        loaded_model = LoadedModel(x)
401
        loaded = None
comfyanonymous's avatar
comfyanonymous committed
402

403
404
405
406
407
408
409
410
411
412
413
414
415
416
        try:
            loaded_model_index = current_loaded_models.index(loaded_model)
        except:
            loaded_model_index = None

        if loaded_model_index is not None:
            loaded = current_loaded_models[loaded_model_index]
            if loaded.should_reload_model(force_patch_weights=force_patch_weights): #TODO: cleanup this model reload logic
                current_loaded_models.pop(loaded_model_index).model_unload(unpatch_weights=True)
                loaded = None
            else:
                models_already_loaded.append(loaded)

        if loaded is None:
417
            if hasattr(x, "model"):
comfyanonymous's avatar
comfyanonymous committed
418
                logging.info(f"Requested to load {x.model.__class__.__name__}")
comfyanonymous's avatar
comfyanonymous committed
419
420
421
422
423
424
425
            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)
426
427
        return

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

comfyanonymous's avatar
comfyanonymous committed
430
431
    total_memory_required = {}
    for loaded_model in models_to_load:
432
433
        if unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) == True:#unload clones where the weights are different
            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
434

comfyanonymous's avatar
comfyanonymous committed
435
436
437
    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
438

439
    for loaded_model in models_to_load:
440
441
442
        weights_unloaded = unload_model_clones(loaded_model.model, unload_weights_only=False, force_unload=False) #unload the rest of the clones where the weights can stay loaded
        if weights_unloaded is not None:
            loaded_model.weights_loaded = not weights_unloaded
443

comfyanonymous's avatar
comfyanonymous committed
444
445
446
447
448
449
450
451
452
453
454
    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)
455
            lowvram_model_memory = int(max(64 * (1024 * 1024), (current_free_mem - 1024 * (1024 * 1024)) / 1.3 ))
456
            if model_size <= (current_free_mem - inference_memory): #only switch to lowvram if really necessary
comfyanonymous's avatar
comfyanonymous committed
457
                lowvram_model_memory = 0
458

comfyanonymous's avatar
comfyanonymous committed
459
        if vram_set_state == VRAMState.NO_VRAM:
460
            lowvram_model_memory = 64 * 1024 * 1024
461

462
        cur_loaded_model = loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)
comfyanonymous's avatar
comfyanonymous committed
463
464
465
466
467
468
469
        current_loaded_models.insert(0, loaded_model)
    return


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

470
def cleanup_models(keep_clone_weights_loaded=False):
comfyanonymous's avatar
comfyanonymous committed
471
472
473
    to_delete = []
    for i in range(len(current_loaded_models)):
        if sys.getrefcount(current_loaded_models[i].model) <= 2:
474
475
476
477
478
            if not keep_clone_weights_loaded:
                to_delete = [i] + to_delete
            #TODO: find a less fragile way to do this.
            elif sys.getrefcount(current_loaded_models[i].real_model) <= 3: #references from .real_model + the .model
                to_delete = [i] + to_delete
comfyanonymous's avatar
comfyanonymous committed
479
480
481
482
483

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

485
486
487
488
def dtype_size(dtype):
    dtype_size = 4
    if dtype == torch.float16 or dtype == torch.bfloat16:
        dtype_size = 2
489
490
491
492
493
494
495
    elif dtype == torch.float32:
        dtype_size = 4
    else:
        try:
            dtype_size = dtype.itemsize
        except: #Old pytorch doesn't have .itemsize
            pass
496
497
    return dtype_size

498
def unet_offload_device():
comfyanonymous's avatar
comfyanonymous committed
499
    if vram_state == VRAMState.HIGH_VRAM:
500
501
502
503
        return get_torch_device()
    else:
        return torch.device("cpu")

comfyanonymous's avatar
comfyanonymous committed
504
505
506
507
508
509
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")
510
511
512
    if DISABLE_SMART_MEMORY:
        return cpu_dev

513
    model_size = dtype_size(dtype) * parameters
comfyanonymous's avatar
comfyanonymous committed
514
515
516
517
518
519
520
521

    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
522
def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
523
524
    if args.bf16_unet:
        return torch.bfloat16
525
526
    if args.fp16_unet:
        return torch.float16
527
528
529
530
    if args.fp8_e4m3fn_unet:
        return torch.float8_e4m3fn
    if args.fp8_e5m2_unet:
        return torch.float8_e5m2
531
    if should_use_fp16(device=device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
532
533
        if torch.float16 in supported_dtypes:
            return torch.float16
534
    if should_use_bf16(device, model_params=model_params, manual_cast=True):
comfyanonymous's avatar
comfyanonymous committed
535
536
        if torch.bfloat16 in supported_dtypes:
            return torch.bfloat16
537
538
    return torch.float32

539
# None means no manual cast
comfyanonymous's avatar
comfyanonymous committed
540
def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):
541
542
543
    if weight_dtype == torch.float32:
        return None

comfyanonymous's avatar
comfyanonymous committed
544
    fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)
545
546
547
    if fp16_supported and weight_dtype == torch.float16:
        return None

comfyanonymous's avatar
comfyanonymous committed
548
549
550
551
552
    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:
553
        return torch.float16
comfyanonymous's avatar
comfyanonymous committed
554
555
556

    elif bf16_supported and torch.bfloat16 in supported_dtypes:
        return torch.bfloat16
557
558
559
    else:
        return torch.float32

560
def text_encoder_offload_device():
comfyanonymous's avatar
comfyanonymous committed
561
    if args.gpu_only:
562
563
564
565
        return get_torch_device()
    else:
        return torch.device("cpu")

566
def text_encoder_device():
comfyanonymous's avatar
comfyanonymous committed
567
    if args.gpu_only:
568
        return get_torch_device()
569
    elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:
570
        if should_use_fp16(prioritize_performance=False):
571
572
573
            return get_torch_device()
        else:
            return torch.device("cpu")
574
575
576
    else:
        return torch.device("cpu")

577
578
579
580
581
582
583
584
585
586
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

587
588
589
    if is_device_cpu(device):
        return torch.float16

590
591
    return torch.float16

592

593
594
595
596
597
598
def intermediate_device():
    if args.gpu_only:
        return get_torch_device()
    else:
        return torch.device("cpu")

599
def vae_device():
600
601
    if args.cpu_vae:
        return torch.device("cpu")
602
603
604
    return get_torch_device()

def vae_offload_device():
comfyanonymous's avatar
comfyanonymous committed
605
    if args.gpu_only:
606
607
608
609
        return get_torch_device()
    else:
        return torch.device("cpu")

610
def vae_dtype():
611
612
    global VAE_DTYPE
    return VAE_DTYPE
613

614
615
616
617
def get_autocast_device(dev):
    if hasattr(dev, 'type'):
        return dev.type
    return "cuda"
618

619
620
621
def supports_dtype(device, dtype): #TODO
    if dtype == torch.float32:
        return True
622
    if is_device_cpu(device):
623
624
625
626
627
628
629
        return False
    if dtype == torch.float16:
        return True
    if dtype == torch.bfloat16:
        return True
    return False

630
631
632
def device_supports_non_blocking(device):
    if is_device_mps(device):
        return False #pytorch bug? mps doesn't support non blocking
633
634
    return False
    # return True #TODO: figure out why this causes issues
635

636
637
638
639
640
641
642
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
643
644
        elif is_intel_xpu():
            device_supports_cast = True
645

646
    non_blocking = device_supports_non_blocking(device)
comfyanonymous's avatar
comfyanonymous committed
647

648
649
650
    if device_supports_cast:
        if copy:
            if tensor.device == device:
comfyanonymous's avatar
comfyanonymous committed
651
652
                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)
653
        else:
comfyanonymous's avatar
comfyanonymous committed
654
            return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
655
    else:
comfyanonymous's avatar
comfyanonymous committed
656
        return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
657

658
def xformers_enabled():
659
    global directml_enabled
660
661
    global cpu_state
    if cpu_state != CPUState.GPU:
662
        return False
663
    if is_intel_xpu():
664
665
666
        return False
    if directml_enabled:
        return False
667
    return XFORMERS_IS_AVAILABLE
668

669
670
671
672
673

def xformers_enabled_vae():
    enabled = xformers_enabled()
    if not enabled:
        return False
674
675

    return XFORMERS_ENABLED_VAE
676

677
def pytorch_attention_enabled():
678
    global ENABLE_PYTORCH_ATTENTION
679
680
    return ENABLE_PYTORCH_ATTENTION

681
682
683
684
def pytorch_attention_flash_attention():
    global ENABLE_PYTORCH_ATTENTION
    if ENABLE_PYTORCH_ATTENTION:
        #TODO: more reliable way of checking for flash attention?
685
        if is_nvidia(): #pytorch flash attention only works on Nvidia
686
687
688
            return True
    return False

689
690
691
692
693
694
695
696
697
698
699
700
def force_upcast_attention_dtype():
    upcast = args.force_upcast_attention
    try:
        if platform.mac_ver()[0] in ['14.5']: #black image bug on OSX Sonoma 14.5
            upcast = True
    except:
        pass
    if upcast:
        return torch.float32
    else:
        return None

701
def get_free_memory(dev=None, torch_free_too=False):
702
    global directml_enabled
703
    if dev is None:
704
        dev = get_torch_device()
705

Yurii Mazurevich's avatar
Yurii Mazurevich committed
706
    if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):
707
708
709
        mem_free_total = psutil.virtual_memory().available
        mem_free_torch = mem_free_total
    else:
710
711
712
        if directml_enabled:
            mem_free_total = 1024 * 1024 * 1024 #TODO
            mem_free_torch = mem_free_total
713
        elif is_intel_xpu():
714
715
716
717
            stats = torch.xpu.memory_stats(dev)
            mem_active = stats['active_bytes.all.current']
            mem_reserved = stats['reserved_bytes.all.current']
            mem_free_torch = mem_reserved - mem_active
718
719
            mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved
            mem_free_total = mem_free_xpu + mem_free_torch
720
721
722
723
724
725
726
        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
727
728
729
730
731

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

733
def cpu_mode():
734
735
    global cpu_state
    return cpu_state == CPUState.CPU
736

Yurii Mazurevich's avatar
Yurii Mazurevich committed
737
def mps_mode():
738
739
    global cpu_state
    return cpu_state == CPUState.MPS
Yurii Mazurevich's avatar
Yurii Mazurevich committed
740

741
def is_device_type(device, type):
742
    if hasattr(device, 'type'):
743
        if (device.type == type):
comfyanonymous's avatar
comfyanonymous committed
744
745
746
            return True
    return False

747
748
749
def is_device_cpu(device):
    return is_device_type(device, 'cpu')

comfyanonymous's avatar
comfyanonymous committed
750
def is_device_mps(device):
751
752
753
754
    return is_device_type(device, 'mps')

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

756
def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):
757
758
    global directml_enabled

759
760
761
762
    if device is not None:
        if is_device_cpu(device):
            return False

763
764
765
    if FORCE_FP16:
        return True

766
    if device is not None:
767
        if is_device_mps(device):
768
            return True
769

770
771
772
    if FORCE_FP32:
        return False

773
774
775
    if directml_enabled:
        return False

776
777
778
779
780
    if mps_mode():
        return True

    if cpu_mode():
        return False
781

782
    if is_intel_xpu():
comfyanonymous's avatar
comfyanonymous committed
783
784
        return True

785
    if torch.version.hip:
786
787
        return True

comfyanonymous's avatar
comfyanonymous committed
788
    props = torch.cuda.get_device_properties("cuda")
789
790
791
    if props.major >= 8:
        return True

792
793
794
795
796
797
798
    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
799
    nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"]
800
801
802
803
    for x in nvidia_10_series:
        if x in props.name.lower():
            fp16_works = True

804
    if fp16_works or manual_cast:
805
        free_model_memory = (get_free_memory() * 0.9 - minimum_inference_memory())
806
        if (not prioritize_performance) or model_params * 4 > free_model_memory:
807
808
            return True

809
810
811
    if props.major < 7:
        return False

812
    #FP16 is just broken on these cards
813
    nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"]
814
815
816
817
818
819
    for x in nvidia_16_series:
        if x in props.name:
            return False

    return True

820
821
822
823
824
825
826
827
828
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

829
830
831
    if FORCE_FP32:
        return False

832
833
834
835
836
837
    if directml_enabled:
        return False

    if cpu_mode() or mps_mode():
        return False

comfyanonymous's avatar
comfyanonymous committed
838
839
840
841
842
843
844
845
846
847
    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

848
849
850
851
852
853
854
    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
855
856
    return False

857
def soft_empty_cache(force=False):
858
859
    global cpu_state
    if cpu_state == CPUState.MPS:
comfyanonymous's avatar
comfyanonymous committed
860
        torch.mps.empty_cache()
861
    elif is_intel_xpu():
862
863
        torch.xpu.empty_cache()
    elif torch.cuda.is_available():
864
        if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda
865
866
867
            torch.cuda.empty_cache()
            torch.cuda.ipc_collect()

868
869
870
871
def unload_all_models():
    free_memory(1e30, get_torch_device())


872
def resolve_lowvram_weight(weight, model, key): #TODO: remove
comfyanonymous's avatar
comfyanonymous committed
873
874
    return weight

875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
#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()