fred-dev/comfy_ui_ali
0
1"""2 This file is part of ComfyUI.3 Copyright (C) 2024 Comfy4 5 This program is free software: you can redistribute it and/or modify6 it under the terms of the GNU General Public License as published by7 the Free Software Foundation, either version 3 of the License, or8 (at your option) any later version.9 10 This program is distributed in the hope that it will be useful,11 but WITHOUT ANY WARRANTY; without even the implied warranty of12 MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the13 GNU General Public License for more details.14 15 You should have received a copy of the GNU General Public License16 along with this program. If not, see <https://www.gnu.org/licenses/>.17"""18 19import psutil20import logging21from enum import Enum22from comfy.cli_args import args, PerformanceFeature23import torch24import sys25import platform26import weakref27import gc28 29class VRAMState(Enum):30 DISABLED = 0 #No vram present: no need to move models to vram31 NO_VRAM = 1 #Very low vram: enable all the options to save vram32 LOW_VRAM = 233 NORMAL_VRAM = 334 HIGH_VRAM = 435 SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both.36 37class CPUState(Enum):38 GPU = 039 CPU = 140 MPS = 241 42# Determine VRAM State43vram_state = VRAMState.NORMAL_VRAM44set_vram_to = VRAMState.NORMAL_VRAM45cpu_state = CPUState.GPU46 47total_vram = 048 49xpu_available = False50torch_version = ""51try:52 torch_version = torch.version.__version__53 temp = torch_version.split(".")54 torch_version_numeric = (int(temp[0]), int(temp[1]))55 xpu_available = (torch_version_numeric[0] < 2 or (torch_version_numeric[0] == 2 and torch_version_numeric[1] <= 4)) and torch.xpu.is_available()56except:57 pass58 59lowvram_available = True60if args.deterministic:61 logging.info("Using deterministic algorithms for pytorch")62 torch.use_deterministic_algorithms(True, warn_only=True)63 64directml_enabled = False65if args.directml is not None:66 import torch_directml67 directml_enabled = True68 device_index = args.directml69 if device_index < 0:70 directml_device = torch_directml.device()71 else:72 directml_device = torch_directml.device(device_index)73 logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index)))74 # torch_directml.disable_tiled_resources(True)75 lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default.76 77try:78 import intel_extension_for_pytorch as ipex79 _ = torch.xpu.device_count()80 xpu_available = xpu_available or torch.xpu.is_available()81except:82 xpu_available = xpu_available or (hasattr(torch, "xpu") and torch.xpu.is_available())83 84try:85 if torch.backends.mps.is_available():86 cpu_state = CPUState.MPS87 import torch.mps88except:89 pass90 91try:92 import torch_npu # noqa: F40193 _ = torch.npu.device_count()94 npu_available = torch.npu.is_available()95except:96 npu_available = False97 98try:99 import torch_mlu # noqa: F401100 _ = torch.mlu.device_count()101 mlu_available = torch.mlu.is_available()102except:103 mlu_available = False104 105if args.cpu:106 cpu_state = CPUState.CPU107 108def is_intel_xpu():109 global cpu_state110 global xpu_available111 if cpu_state == CPUState.GPU:112 if xpu_available:113 return True114 return False115 116def is_ascend_npu():117 global npu_available118 if npu_available:119 return True120 return False121 122def is_mlu():123 global mlu_available124 if mlu_available:125 return True126 return False127 128def get_torch_device():129 global directml_enabled130 global cpu_state131 if directml_enabled:132 global directml_device133 return directml_device134 if cpu_state == CPUState.MPS:135 return torch.device("mps")136 if cpu_state == CPUState.CPU:137 return torch.device("cpu")138 else:139 if is_intel_xpu():140 return torch.device("xpu", torch.xpu.current_device())141 elif is_ascend_npu():142 return torch.device("npu", torch.npu.current_device())143 elif is_mlu():144 return torch.device("mlu", torch.mlu.current_device())145 else:146 return torch.device(torch.cuda.current_device())147 148def get_total_memory(dev=None, torch_total_too=False):149 global directml_enabled150 if dev is None:151 dev = get_torch_device()152 153 if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):154 mem_total = psutil.virtual_memory().total155 mem_total_torch = mem_total156 else:157 if directml_enabled:158 mem_total = 1024 * 1024 * 1024 #TODO159 mem_total_torch = mem_total160 elif is_intel_xpu():161 stats = torch.xpu.memory_stats(dev)162 mem_reserved = stats['reserved_bytes.all.current']163 mem_total_torch = mem_reserved164 mem_total = torch.xpu.get_device_properties(dev).total_memory165 elif is_ascend_npu():166 stats = torch.npu.memory_stats(dev)167 mem_reserved = stats['reserved_bytes.all.current']168 _, mem_total_npu = torch.npu.mem_get_info(dev)169 mem_total_torch = mem_reserved170 mem_total = mem_total_npu171 elif is_mlu():172 stats = torch.mlu.memory_stats(dev)173 mem_reserved = stats['reserved_bytes.all.current']174 _, mem_total_mlu = torch.mlu.mem_get_info(dev)175 mem_total_torch = mem_reserved176 mem_total = mem_total_mlu177 else:178 stats = torch.cuda.memory_stats(dev)179 mem_reserved = stats['reserved_bytes.all.current']180 _, mem_total_cuda = torch.cuda.mem_get_info(dev)181 mem_total_torch = mem_reserved182 mem_total = mem_total_cuda183 184 if torch_total_too:185 return (mem_total, mem_total_torch)186 else:187 return mem_total188 189def mac_version():190 try:191 return tuple(int(n) for n in platform.mac_ver()[0].split("."))192 except:193 return None194 195total_vram = get_total_memory(get_torch_device()) / (1024 * 1024)196total_ram = psutil.virtual_memory().total / (1024 * 1024)197logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram))198 199try:200 logging.info("pytorch version: {}".format(torch_version))201 mac_ver = mac_version()202 if mac_ver is not None:203 logging.info("Mac Version {}".format(mac_ver))204except:205 pass206 207try:208 OOM_EXCEPTION = torch.cuda.OutOfMemoryError209except:210 OOM_EXCEPTION = Exception211 212XFORMERS_VERSION = ""213XFORMERS_ENABLED_VAE = True214if args.disable_xformers:215 XFORMERS_IS_AVAILABLE = False216else:217 try:218 import xformers219 import xformers.ops220 XFORMERS_IS_AVAILABLE = True221 try:222 XFORMERS_IS_AVAILABLE = xformers._has_cpp_library223 except:224 pass225 try:226 XFORMERS_VERSION = xformers.version.__version__227 logging.info("xformers version: {}".format(XFORMERS_VERSION))228 if XFORMERS_VERSION.startswith("0.0.18"):229 logging.warning("\nWARNING: This version of xformers has a major bug where you will get black images when generating high resolution images.")230 logging.warning("Please downgrade or upgrade xformers to a different version.\n")231 XFORMERS_ENABLED_VAE = False232 except:233 pass234 except:235 XFORMERS_IS_AVAILABLE = False236 237def is_nvidia():238 global cpu_state239 if cpu_state == CPUState.GPU:240 if torch.version.cuda:241 return True242 return False243 244def is_amd():245 global cpu_state246 if cpu_state == CPUState.GPU:247 if torch.version.hip:248 return True249 return False250 251MIN_WEIGHT_MEMORY_RATIO = 0.4252if is_nvidia():253 MIN_WEIGHT_MEMORY_RATIO = 0.0254 255ENABLE_PYTORCH_ATTENTION = False256if args.use_pytorch_cross_attention:257 ENABLE_PYTORCH_ATTENTION = True258 XFORMERS_IS_AVAILABLE = False259 260try:261 if is_nvidia():262 if torch_version_numeric[0] >= 2:263 if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False:264 ENABLE_PYTORCH_ATTENTION = True265 if is_intel_xpu() or is_ascend_npu() or is_mlu():266 if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:267 ENABLE_PYTORCH_ATTENTION = True268except:269 pass270 271 272try:273 if is_amd():274 arch = torch.cuda.get_device_properties(get_torch_device()).gcnArchName275 logging.info("AMD arch: {}".format(arch))276 if args.use_split_cross_attention == False and args.use_quad_cross_attention == False:277 if torch_version_numeric[0] >= 2 and torch_version_numeric[1] >= 7: # works on 2.6 but doesn't actually seem to improve much278 if any((a in arch) for a in ["gfx1100", "gfx1101"]): # TODO: more arches279 ENABLE_PYTORCH_ATTENTION = True280except:281 pass282 283 284if ENABLE_PYTORCH_ATTENTION:285 torch.backends.cuda.enable_math_sdp(True)286 torch.backends.cuda.enable_flash_sdp(True)287 torch.backends.cuda.enable_mem_efficient_sdp(True)288 289 290PRIORITIZE_FP16 = False # TODO: remove and replace with something that shows exactly which dtype is faster than the other291try:292 if is_nvidia() and PerformanceFeature.Fp16Accumulation in args.fast:293 torch.backends.cuda.matmul.allow_fp16_accumulation = True294 PRIORITIZE_FP16 = True # TODO: limit to cards where it actually boosts performance295 logging.info("Enabled fp16 accumulation.")296except:297 pass298 299try:300 if torch_version_numeric[0] == 2 and torch_version_numeric[1] >= 5:301 torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True)302except:303 logging.warning("Warning, could not set allow_fp16_bf16_reduction_math_sdp")304 305if args.lowvram:306 set_vram_to = VRAMState.LOW_VRAM307 lowvram_available = True308elif args.novram:309 set_vram_to = VRAMState.NO_VRAM310elif args.highvram or args.gpu_only:311 vram_state = VRAMState.HIGH_VRAM312 313FORCE_FP32 = False314if args.force_fp32:315 logging.info("Forcing FP32, if this improves things please report it.")316 FORCE_FP32 = True317 318if lowvram_available:319 if set_vram_to in (VRAMState.LOW_VRAM, VRAMState.NO_VRAM):320 vram_state = set_vram_to321 322 323if cpu_state != CPUState.GPU:324 vram_state = VRAMState.DISABLED325 326if cpu_state == CPUState.MPS:327 vram_state = VRAMState.SHARED328 329logging.info(f"Set vram state to: {vram_state.name}")330 331DISABLE_SMART_MEMORY = args.disable_smart_memory332 333if DISABLE_SMART_MEMORY:334 logging.info("Disabling smart memory management")335 336def get_torch_device_name(device):337 if hasattr(device, 'type'):338 if device.type == "cuda":339 try:340 allocator_backend = torch.cuda.get_allocator_backend()341 except:342 allocator_backend = ""343 return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend)344 else:345 return "{}".format(device.type)346 elif is_intel_xpu():347 return "{} {}".format(device, torch.xpu.get_device_name(device))348 elif is_ascend_npu():349 return "{} {}".format(device, torch.npu.get_device_name(device))350 elif is_mlu():351 return "{} {}".format(device, torch.mlu.get_device_name(device))352 else:353 return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device))354 355try:356 logging.info("Device: {}".format(get_torch_device_name(get_torch_device())))357except:358 logging.warning("Could not pick default device.")359 360 361current_loaded_models = []362 363def module_size(module):364 module_mem = 0365 sd = module.state_dict()366 for k in sd:367 t = sd[k]368 module_mem += t.nelement() * t.element_size()369 return module_mem370 371class LoadedModel:372 def __init__(self, model):373 self._set_model(model)374 self.device = model.load_device375 self.real_model = None376 self.currently_used = True377 self.model_finalizer = None378 self._patcher_finalizer = None379 380 def _set_model(self, model):381 self._model = weakref.ref(model)382 if model.parent is not None:383 self._parent_model = weakref.ref(model.parent)384 self._patcher_finalizer = weakref.finalize(model, self._switch_parent)385 386 def _switch_parent(self):387 model = self._parent_model()388 if model is not None:389 self._set_model(model)390 391 @property392 def model(self):393 return self._model()394 395 def model_memory(self):396 return self.model.model_size()397 398 def model_loaded_memory(self):399 return self.model.loaded_size()400 401 def model_offloaded_memory(self):402 return self.model.model_size() - self.model.loaded_size()403 404 def model_memory_required(self, device):405 if device == self.model.current_loaded_device():406 return self.model_offloaded_memory()407 else:408 return self.model_memory()409 410 def model_load(self, lowvram_model_memory=0, force_patch_weights=False):411 self.model.model_patches_to(self.device)412 self.model.model_patches_to(self.model.model_dtype())413 414 # if self.model.loaded_size() > 0:415 use_more_vram = lowvram_model_memory416 if use_more_vram == 0:417 use_more_vram = 1e32418 self.model_use_more_vram(use_more_vram, force_patch_weights=force_patch_weights)419 real_model = self.model.model420 421 if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and real_model is not None:422 with torch.no_grad():423 real_model = ipex.optimize(real_model.eval(), inplace=True, graph_mode=True, concat_linear=True)424 425 self.real_model = weakref.ref(real_model)426 self.model_finalizer = weakref.finalize(real_model, cleanup_models)427 return real_model428 429 def should_reload_model(self, force_patch_weights=False):430 if force_patch_weights and self.model.lowvram_patch_counter() > 0:431 return True432 return False433 434 def model_unload(self, memory_to_free=None, unpatch_weights=True):435 if memory_to_free is not None:436 if memory_to_free < self.model.loaded_size():437 freed = self.model.partially_unload(self.model.offload_device, memory_to_free)438 if freed >= memory_to_free:439 return False440 self.model.detach(unpatch_weights)441 self.model_finalizer.detach()442 self.model_finalizer = None443 self.real_model = None444 return True445 446 def model_use_more_vram(self, extra_memory, force_patch_weights=False):447 return self.model.partially_load(self.device, extra_memory, force_patch_weights=force_patch_weights)448 449 def __eq__(self, other):450 return self.model is other.model451 452 def __del__(self):453 if self._patcher_finalizer is not None:454 self._patcher_finalizer.detach()455 456 def is_dead(self):457 return self.real_model() is not None and self.model is None458 459 460def use_more_memory(extra_memory, loaded_models, device):461 for m in loaded_models:462 if m.device == device:463 extra_memory -= m.model_use_more_vram(extra_memory)464 if extra_memory <= 0:465 break466 467def offloaded_memory(loaded_models, device):468 offloaded_mem = 0469 for m in loaded_models:470 if m.device == device:471 offloaded_mem += m.model_offloaded_memory()472 return offloaded_mem473 474WINDOWS = any(platform.win32_ver())475 476EXTRA_RESERVED_VRAM = 400 * 1024 * 1024477if WINDOWS:478 EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue479 480if args.reserve_vram is not None:481 EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024482 logging.debug("Reserving {}MB vram for other applications.".format(EXTRA_RESERVED_VRAM / (1024 * 1024)))483 484def extra_reserved_memory():485 return EXTRA_RESERVED_VRAM486 487def minimum_inference_memory():488 return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory()489 490def free_memory(memory_required, device, keep_loaded=[]):491 cleanup_models_gc()492 unloaded_model = []493 can_unload = []494 unloaded_models = []495 496 for i in range(len(current_loaded_models) -1, -1, -1):497 shift_model = current_loaded_models[i]498 if shift_model.device == device:499 if shift_model not in keep_loaded and not shift_model.is_dead():500 can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i))501 shift_model.currently_used = False502 503 for x in sorted(can_unload):504 i = x[-1]505 memory_to_free = None506 if not DISABLE_SMART_MEMORY:507 free_mem = get_free_memory(device)508 if free_mem > memory_required:509 break510 memory_to_free = memory_required - free_mem511 logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}")512 if current_loaded_models[i].model_unload(memory_to_free):513 unloaded_model.append(i)514 515 for i in sorted(unloaded_model, reverse=True):516 unloaded_models.append(current_loaded_models.pop(i))517 518 if len(unloaded_model) > 0:519 soft_empty_cache()520 else:521 if vram_state != VRAMState.HIGH_VRAM:522 mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True)523 if mem_free_torch > mem_free_total * 0.25:524 soft_empty_cache()525 return unloaded_models526 527def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False):528 cleanup_models_gc()529 global vram_state530 531 inference_memory = minimum_inference_memory()532 extra_mem = max(inference_memory, memory_required + extra_reserved_memory())533 if minimum_memory_required is None:534 minimum_memory_required = extra_mem535 else:536 minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory())537 538 models = set(models)539 540 models_to_load = []541 542 for x in models:543 loaded_model = LoadedModel(x)544 try:545 loaded_model_index = current_loaded_models.index(loaded_model)546 except:547 loaded_model_index = None548 549 if loaded_model_index is not None:550 loaded = current_loaded_models[loaded_model_index]551 loaded.currently_used = True552 models_to_load.append(loaded)553 else:554 if hasattr(x, "model"):555 logging.info(f"Requested to load {x.model.__class__.__name__}")556 models_to_load.append(loaded_model)557 558 for loaded_model in models_to_load:559 to_unload = []560 for i in range(len(current_loaded_models)):561 if loaded_model.model.is_clone(current_loaded_models[i].model):562 to_unload = [i] + to_unload563 for i in to_unload:564 current_loaded_models.pop(i).model.detach(unpatch_all=False)565 566 total_memory_required = {}567 for loaded_model in models_to_load:568 total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device)569 570 for device in total_memory_required:571 if device != torch.device("cpu"):572 free_memory(total_memory_required[device] * 1.1 + extra_mem, device)573 574 for device in total_memory_required:575 if device != torch.device("cpu"):576 free_mem = get_free_memory(device)577 if free_mem < minimum_memory_required:578 models_l = free_memory(minimum_memory_required, device)579 logging.info("{} models unloaded.".format(len(models_l)))580 581 for loaded_model in models_to_load:582 model = loaded_model.model583 torch_dev = model.load_device584 if is_device_cpu(torch_dev):585 vram_set_state = VRAMState.DISABLED586 else:587 vram_set_state = vram_state588 lowvram_model_memory = 0589 if lowvram_available and (vram_set_state == VRAMState.LOW_VRAM or vram_set_state == VRAMState.NORMAL_VRAM) and not force_full_load:590 loaded_memory = loaded_model.model_loaded_memory()591 current_free_mem = get_free_memory(torch_dev) + loaded_memory592 593 lowvram_model_memory = max(128 * 1024 * 1024, (current_free_mem - minimum_memory_required), min(current_free_mem * MIN_WEIGHT_MEMORY_RATIO, current_free_mem - minimum_inference_memory()))594 lowvram_model_memory = max(0.1, lowvram_model_memory - loaded_memory)595 596 if vram_set_state == VRAMState.NO_VRAM:597 lowvram_model_memory = 0.1598 599 loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights)600 current_loaded_models.insert(0, loaded_model)601 return602 603def load_model_gpu(model):604 return load_models_gpu([model])605 606def loaded_models(only_currently_used=False):607 output = []608 for m in current_loaded_models:609 if only_currently_used:610 if not m.currently_used:611 continue612 613 output.append(m.model)614 return output615 616 617def cleanup_models_gc():618 do_gc = False619 for i in range(len(current_loaded_models)):620 cur = current_loaded_models[i]621 if cur.is_dead():622 logging.info("Potential memory leak detected with model {}, doing a full garbage collect, for maximum performance avoid circular references in the model code.".format(cur.real_model().__class__.__name__))623 do_gc = True624 break625 626 if do_gc:627 gc.collect()628 soft_empty_cache()629 630 for i in range(len(current_loaded_models)):631 cur = current_loaded_models[i]632 if cur.is_dead():633 logging.warning("WARNING, memory leak with model {}. Please make sure it is not being referenced from somewhere.".format(cur.real_model().__class__.__name__))634 635 636 637def cleanup_models():638 to_delete = []639 for i in range(len(current_loaded_models)):640 if current_loaded_models[i].real_model() is None:641 to_delete = [i] + to_delete642 643 for i in to_delete:644 x = current_loaded_models.pop(i)645 del x646 647def dtype_size(dtype):648 dtype_size = 4649 if dtype == torch.float16 or dtype == torch.bfloat16:650 dtype_size = 2651 elif dtype == torch.float32:652 dtype_size = 4653 else:654 try:655 dtype_size = dtype.itemsize656 except: #Old pytorch doesn't have .itemsize657 pass658 return dtype_size659 660def unet_offload_device():661 if vram_state == VRAMState.HIGH_VRAM:662 return get_torch_device()663 else:664 return torch.device("cpu")665 666def unet_inital_load_device(parameters, dtype):667 torch_dev = get_torch_device()668 if vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.SHARED:669 return torch_dev670 671 cpu_dev = torch.device("cpu")672 if DISABLE_SMART_MEMORY:673 return cpu_dev674 675 model_size = dtype_size(dtype) * parameters676 677 mem_dev = get_free_memory(torch_dev)678 mem_cpu = get_free_memory(cpu_dev)679 if mem_dev > mem_cpu and model_size < mem_dev:680 return torch_dev681 else:682 return cpu_dev683 684def maximum_vram_for_weights(device=None):685 return (get_total_memory(device) * 0.88 - minimum_inference_memory())686 687def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], weight_dtype=None):688 if model_params < 0:689 model_params = 1000000000000000000000690 if args.fp32_unet:691 return torch.float32692 if args.fp64_unet:693 return torch.float64694 if args.bf16_unet:695 return torch.bfloat16696 if args.fp16_unet:697 return torch.float16698 if args.fp8_e4m3fn_unet:699 return torch.float8_e4m3fn700 if args.fp8_e5m2_unet:701 return torch.float8_e5m2702 703 fp8_dtype = None704 try:705 if weight_dtype in [torch.float8_e4m3fn, torch.float8_e5m2]:706 fp8_dtype = weight_dtype707 except:708 pass709 710 if fp8_dtype is not None:711 if supports_fp8_compute(device): #if fp8 compute is supported the casting is most likely not expensive712 return fp8_dtype713 714 free_model_memory = maximum_vram_for_weights(device)715 if model_params * 2 > free_model_memory:716 return fp8_dtype717 718 if PRIORITIZE_FP16 or weight_dtype == torch.float16:719 if torch.float16 in supported_dtypes and should_use_fp16(device=device, model_params=model_params):720 return torch.float16721 722 for dt in supported_dtypes:723 if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params):724 if torch.float16 in supported_dtypes:725 return torch.float16726 if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params):727 if torch.bfloat16 in supported_dtypes:728 return torch.bfloat16729 730 for dt in supported_dtypes:731 if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True):732 if torch.float16 in supported_dtypes:733 return torch.float16734 if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True):735 if torch.bfloat16 in supported_dtypes:736 return torch.bfloat16737 738 return torch.float32739 740# None means no manual cast741def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]):742 if weight_dtype == torch.float32 or weight_dtype == torch.float64:743 return None744 745 fp16_supported = should_use_fp16(inference_device, prioritize_performance=False)746 if fp16_supported and weight_dtype == torch.float16:747 return None748 749 bf16_supported = should_use_bf16(inference_device)750 if bf16_supported and weight_dtype == torch.bfloat16:751 return None752 753 fp16_supported = should_use_fp16(inference_device, prioritize_performance=True)754 if PRIORITIZE_FP16 and fp16_supported and torch.float16 in supported_dtypes:755 return torch.float16756 757 for dt in supported_dtypes:758 if dt == torch.float16 and fp16_supported:759 return torch.float16760 if dt == torch.bfloat16 and bf16_supported:761 return torch.bfloat16762 763 return torch.float32764 765def text_encoder_offload_device():766 if args.gpu_only:767 return get_torch_device()768 else:769 return torch.device("cpu")770 771def text_encoder_device():772 if args.gpu_only:773 return get_torch_device()774 elif vram_state == VRAMState.HIGH_VRAM or vram_state == VRAMState.NORMAL_VRAM:775 if should_use_fp16(prioritize_performance=False):776 return get_torch_device()777 else:778 return torch.device("cpu")779 else:780 return torch.device("cpu")781 782def text_encoder_initial_device(load_device, offload_device, model_size=0):783 if load_device == offload_device or model_size <= 1024 * 1024 * 1024:784 return offload_device785 786 if is_device_mps(load_device):787 return load_device788 789 mem_l = get_free_memory(load_device)790 mem_o = get_free_memory(offload_device)791 if mem_l > (mem_o * 0.5) and model_size * 1.2 < mem_l:792 return load_device793 else:794 return offload_device795 796def text_encoder_dtype(device=None):797 if args.fp8_e4m3fn_text_enc:798 return torch.float8_e4m3fn799 elif args.fp8_e5m2_text_enc:800 return torch.float8_e5m2801 elif args.fp16_text_enc:802 return torch.float16803 elif args.fp32_text_enc:804 return torch.float32805 806 if is_device_cpu(device):807 return torch.float16808 809 return torch.float16810 811 812def intermediate_device():813 if args.gpu_only:814 return get_torch_device()815 else:816 return torch.device("cpu")817 818def vae_device():819 if args.cpu_vae:820 return torch.device("cpu")821 return get_torch_device()822 823def vae_offload_device():824 if args.gpu_only:825 return get_torch_device()826 else:827 return torch.device("cpu")828 829def vae_dtype(device=None, allowed_dtypes=[]):830 if args.fp16_vae:831 return torch.float16832 elif args.bf16_vae:833 return torch.bfloat16834 elif args.fp32_vae:835 return torch.float32836 837 for d in allowed_dtypes:838 if d == torch.float16 and should_use_fp16(device):839 return d840 841 # NOTE: bfloat16 seems to work on AMD for the VAE but is extremely slow in some cases compared to fp32842 if d == torch.bfloat16 and (not is_amd()) and should_use_bf16(device):843 return d844 845 return torch.float32846 847def get_autocast_device(dev):848 if hasattr(dev, 'type'):849 return dev.type850 return "cuda"851 852def supports_dtype(device, dtype): #TODO853 if dtype == torch.float32:854 return True855 if is_device_cpu(device):856 return False857 if dtype == torch.float16:858 return True859 if dtype == torch.bfloat16:860 return True861 return False862 863def supports_cast(device, dtype): #TODO864 if dtype == torch.float32:865 return True866 if dtype == torch.float16:867 return True868 if directml_enabled: #TODO: test this869 return False870 if dtype == torch.bfloat16:871 return True872 if is_device_mps(device):873 return False874 if dtype == torch.float8_e4m3fn:875 return True876 if dtype == torch.float8_e5m2:877 return True878 return False879 880def pick_weight_dtype(dtype, fallback_dtype, device=None):881 if dtype is None:882 dtype = fallback_dtype883 elif dtype_size(dtype) > dtype_size(fallback_dtype):884 dtype = fallback_dtype885 886 if not supports_cast(device, dtype):887 dtype = fallback_dtype888 889 return dtype890 891def device_supports_non_blocking(device):892 if is_device_mps(device):893 return False #pytorch bug? mps doesn't support non blocking894 if is_intel_xpu():895 return False896 if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews)897 return False898 if directml_enabled:899 return False900 return True901 902def device_should_use_non_blocking(device):903 if not device_supports_non_blocking(device):904 return False905 return False906 # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others907 908def force_channels_last():909 if args.force_channels_last:910 return True911 912 #TODO913 return False914 915def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False):916 if device is None or weight.device == device:917 if not copy:918 if dtype is None or weight.dtype == dtype:919 return weight920 return weight.to(dtype=dtype, copy=copy)921 922 r = torch.empty_like(weight, dtype=dtype, device=device)923 r.copy_(weight, non_blocking=non_blocking)924 return r925 926def cast_to_device(tensor, device, dtype, copy=False):927 non_blocking = device_supports_non_blocking(device)928 return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy)929 930def sage_attention_enabled():931 return args.use_sage_attention932 933def flash_attention_enabled():934 return args.use_flash_attention935 936def xformers_enabled():937 global directml_enabled938 global cpu_state939 if cpu_state != CPUState.GPU:940 return False941 if is_intel_xpu():942 return False943 if is_ascend_npu():944 return False945 if is_mlu():946 return False947 if directml_enabled:948 return False949 return XFORMERS_IS_AVAILABLE950 951 952def xformers_enabled_vae():953 enabled = xformers_enabled()954 if not enabled:955 return False956 957 return XFORMERS_ENABLED_VAE958 959def pytorch_attention_enabled():960 global ENABLE_PYTORCH_ATTENTION961 return ENABLE_PYTORCH_ATTENTION962 963def pytorch_attention_enabled_vae():964 if is_amd():965 return False # enabling pytorch attention on AMD currently causes crash when doing high res966 return pytorch_attention_enabled()967 968def pytorch_attention_flash_attention():969 global ENABLE_PYTORCH_ATTENTION970 if ENABLE_PYTORCH_ATTENTION:971 #TODO: more reliable way of checking for flash attention?972 if is_nvidia(): #pytorch flash attention only works on Nvidia973 return True974 if is_intel_xpu():975 return True976 if is_ascend_npu():977 return True978 if is_mlu():979 return True980 if is_amd():981 return True #if you have pytorch attention enabled on AMD it probably supports at least mem efficient attention982 return False983 984def force_upcast_attention_dtype():985 upcast = args.force_upcast_attention986 987 macos_version = mac_version()988 if macos_version is not None and ((14, 5) <= macos_version < (16,)): # black image bug on recent versions of macOS989 upcast = True990 991 if upcast:992 return {torch.float16: torch.float32}993 else:994 return None995 996def get_free_memory(dev=None, torch_free_too=False):997 global directml_enabled998 if dev is None:999 dev = get_torch_device()1000 1001 if hasattr(dev, 'type') and (dev.type == 'cpu' or dev.type == 'mps'):1002 mem_free_total = psutil.virtual_memory().available1003 mem_free_torch = mem_free_total1004 else:1005 if directml_enabled:1006 mem_free_total = 1024 * 1024 * 1024 #TODO1007 mem_free_torch = mem_free_total1008 elif is_intel_xpu():1009 stats = torch.xpu.memory_stats(dev)1010 mem_active = stats['active_bytes.all.current']1011 mem_reserved = stats['reserved_bytes.all.current']1012 mem_free_torch = mem_reserved - mem_active1013 mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved1014 mem_free_total = mem_free_xpu + mem_free_torch1015 elif is_ascend_npu():1016 stats = torch.npu.memory_stats(dev)1017 mem_active = stats['active_bytes.all.current']1018 mem_reserved = stats['reserved_bytes.all.current']1019 mem_free_npu, _ = torch.npu.mem_get_info(dev)1020 mem_free_torch = mem_reserved - mem_active1021 mem_free_total = mem_free_npu + mem_free_torch1022 elif is_mlu():1023 stats = torch.mlu.memory_stats(dev)1024 mem_active = stats['active_bytes.all.current']1025 mem_reserved = stats['reserved_bytes.all.current']1026 mem_free_mlu, _ = torch.mlu.mem_get_info(dev)1027 mem_free_torch = mem_reserved - mem_active1028 mem_free_total = mem_free_mlu + mem_free_torch1029 else:1030 stats = torch.cuda.memory_stats(dev)1031 mem_active = stats['active_bytes.all.current']1032 mem_reserved = stats['reserved_bytes.all.current']1033 mem_free_cuda, _ = torch.cuda.mem_get_info(dev)1034 mem_free_torch = mem_reserved - mem_active1035 mem_free_total = mem_free_cuda + mem_free_torch1036 1037 if torch_free_too:1038 return (mem_free_total, mem_free_torch)1039 else:1040 return mem_free_total1041 1042def cpu_mode():1043 global cpu_state1044 return cpu_state == CPUState.CPU1045 1046def mps_mode():1047 global cpu_state1048 return cpu_state == CPUState.MPS1049 1050def is_device_type(device, type):1051 if hasattr(device, 'type'):1052 if (device.type == type):1053 return True1054 return False1055 1056def is_device_cpu(device):1057 return is_device_type(device, 'cpu')1058 1059def is_device_mps(device):1060 return is_device_type(device, 'mps')1061 1062def is_device_cuda(device):1063 return is_device_type(device, 'cuda')1064 1065def is_directml_enabled():1066 global directml_enabled1067 if directml_enabled:1068 return True1069 1070 return False1071 1072def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):1073 if device is not None:1074 if is_device_cpu(device):1075 return False1076 1077 if args.force_fp16:1078 return True1079 1080 if FORCE_FP32:1081 return False1082 1083 if is_directml_enabled():1084 return True1085 1086 if (device is not None and is_device_mps(device)) or mps_mode():1087 return True1088 1089 if cpu_mode():1090 return False1091 1092 if is_intel_xpu():1093 return True1094 1095 if is_ascend_npu():1096 return True1097 1098 if is_mlu():1099 return True1100 1101 if torch.version.hip:1102 return True1103 1104 props = torch.cuda.get_device_properties(device)1105 if props.major >= 8:1106 return True1107 1108 if props.major < 6:1109 return False1110 1111 #FP16 is confirmed working on a 1080 (GP104) and on latest pytorch actually seems faster than fp321112 nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"]1113 for x in nvidia_10_series:1114 if x in props.name.lower():1115 if WINDOWS or manual_cast:1116 return True1117 else:1118 return False #weird linux behavior where fp32 is faster1119 1120 if manual_cast:1121 free_model_memory = maximum_vram_for_weights(device)1122 if (not prioritize_performance) or model_params * 4 > free_model_memory:1123 return True1124 1125 if props.major < 7:1126 return False1127 1128 #FP16 is just broken on these cards1129 nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"]1130 for x in nvidia_16_series:1131 if x in props.name:1132 return False1133 1134 return True1135 1136def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False):1137 if device is not None:1138 if is_device_cpu(device): #TODO ? bf16 works on CPU but is extremely slow1139 return False1140 1141 if FORCE_FP32:1142 return False1143 1144 if directml_enabled:1145 return False1146 1147 if (device is not None and is_device_mps(device)) or mps_mode():1148 if mac_version() < (14,):1149 return False1150 return True1151 1152 if cpu_mode():1153 return False1154 1155 if is_intel_xpu():1156 return True1157 1158 if is_ascend_npu():1159 return True1160 1161 if is_amd():1162 arch = torch.cuda.get_device_properties(device).gcnArchName1163 if any((a in arch) for a in ["gfx1030", "gfx1031", "gfx1010", "gfx1011", "gfx1012", "gfx906", "gfx900", "gfx803"]): # RDNA2 and older don't support bf161164 if manual_cast:1165 return True1166 return False1167 1168 props = torch.cuda.get_device_properties(device)1169 1170 if is_mlu():1171 if props.major > 3:1172 return True1173 1174 if props.major >= 8:1175 return True1176 1177 bf16_works = torch.cuda.is_bf16_supported()1178 1179 if bf16_works and manual_cast:1180 free_model_memory = maximum_vram_for_weights(device)1181 if (not prioritize_performance) or model_params * 4 > free_model_memory:1182 return True1183 1184 return False1185 1186def supports_fp8_compute(device=None):1187 if not is_nvidia():1188 return False1189 1190 props = torch.cuda.get_device_properties(device)1191 if props.major >= 9:1192 return True1193 if props.major < 8:1194 return False1195 if props.minor < 9:1196 return False1197 1198 if torch_version_numeric[0] < 2 or (torch_version_numeric[0] == 2 and torch_version_numeric[1] < 3):1199 return False1200