From 812ceb49e7b8c316458a90cf49a60255e84ab044 Mon Sep 17 00:00:00 2001 From: radna0 Date: Sun, 17 Nov 2024 13:29:07 +0000 Subject: [PATCH] Add XLA/TPU support with sync callback for diffusion models --- README.md | 19 ++ comfy/cli_args.py | 3 + comfy/model_management.py | 505 ++++++++++++++++++++++++++++++-------- latent_preview.py | 83 +++++-- 4 files changed, 493 insertions(+), 117 deletions(-) diff --git a/README.md b/README.md index 7476b5788..e2e9a6409 100644 --- a/README.md +++ b/README.md @@ -200,6 +200,25 @@ You can install ComfyUI in Apple Mac silicon (M1 or M2) with any recent macOS ve ```pip install torch-directml``` Then you can launch ComfyUI with: ```python main.py --directml``` +#### TPU/XLA Devices +Users with TPU/XLA devices can install the PyTorch XLA stable build with the following command: + +```pip install torch~=2.5.0 torch_xla[tpu]~=2.5.0 -f https://storage.googleapis.com/libtpu-releases/index.html``` + +This is the command to install the nightly 2.6.0 which might have some performance improvements: + +``` +pip3 install --pre torch torchvision --index-url https://download.pytorch.org/whl/nightly/cpu +pip install 'torch_xla[tpu] @ https://storage.googleapis.com/pytorch-xla-releases/wheels/tpuvm/torch_xla-2.6.0.dev-cp310-cp310-linux_x86_64.whl' -f https://storage.googleapis.com/libtpu-releases/index.html +``` + + +To get memory info for TPU devices, install the [tpu-info](https://github.com/AI-Hypercomputer/cloud-accelerator-diagnostics/tree/main/tpu_info) package with the following command: + +```pip install tpu-info``` + + + # Running ```python main.py``` diff --git a/comfy/cli_args.py b/comfy/cli_args.py index 20b9f4749..7524a33d5 100644 --- a/comfy/cli_args.py +++ b/comfy/cli_args.py @@ -138,6 +138,9 @@ parser.add_argument("--multi-user", action="store_true", help="Enables per-user parser.add_argument("--verbose", default='INFO', const='DEBUG', nargs="?", choices=['DEBUG', 'INFO', 'WARNING', 'ERROR', 'CRITICAL'], help='Set the logging level') +parser.add_argument("--xla", action="store_true", help="To use the XLA devices for everything.") +# parser.add_argument("--xla_fsdp", action="store_true", help="To use the XLA devices with FSDP for everything.") + # The default built-in provider hosted under web/ DEFAULT_VERSION_STRING = "comfyanonymous/ComfyUI@latest" diff --git a/comfy/model_management.py b/comfy/model_management.py index fd493aff0..d57739525 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -23,19 +23,29 @@ from comfy.cli_args import args import torch import sys import platform +import importlib.util + +if sys.version_info < (3, 8): + import importlib_metadata +else: + import importlib.metadata as importlib_metadata + class VRAMState(Enum): - 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 + 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 LOW_VRAM = 2 NORMAL_VRAM = 3 HIGH_VRAM = 4 - SHARED = 5 #No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both. + SHARED = 5 # No dedicated vram: memory shared between CPU and GPU but models still need to be moved between both. + class CPUState(Enum): GPU = 0 CPU = 1 MPS = 2 + XLA = 3 + # Determine VRAM State vram_state = VRAMState.NORMAL_VRAM @@ -48,7 +58,10 @@ xpu_available = False torch_version = "" try: torch_version = torch.version.__version__ - xpu_available = (int(torch_version[0]) < 2 or (int(torch_version[0]) == 2 and int(torch_version[2]) <= 4)) and torch.xpu.is_available() + xpu_available = ( + int(torch_version[0]) < 2 + or (int(torch_version[0]) == 2 and int(torch_version[2]) <= 4) + ) and torch.xpu.is_available() except: pass @@ -60,22 +73,43 @@ if args.deterministic: directml_enabled = False if args.directml is not None: import torch_directml + directml_enabled = True device_index = args.directml if device_index < 0: directml_device = torch_directml.device() else: directml_device = torch_directml.device(device_index) - logging.info("Using directml with device: {}".format(torch_directml.device_name(device_index))) + logging.info( + "Using directml with device: {}".format( + torch_directml.device_name(device_index) + ) + ) # torch_directml.disable_tiled_resources(True) - lowvram_available = False #TODO: need to find a way to get free memory in directml before this can be enabled by default. + lowvram_available = False # TODO: need to find a way to get free memory in directml before this can be enabled by default. try: import intel_extension_for_pytorch as ipex + _ = torch.xpu.device_count() xpu_available = torch.xpu.is_available() except: - xpu_available = xpu_available or (hasattr(torch, "xpu") and torch.xpu.is_available()) + xpu_available = xpu_available or ( + hasattr(torch, "xpu") and torch.xpu.is_available() + ) + +XLA_AVAILABLE = False +if args.xla or args.xla_fsdp: + try: + XLA_AVAILABLE = importlib.util.find_spec("torch_xla") is not None + xla_version = importlib_metadata.version("torch_xla") + logging.info(f"PyTorch XLA version {xla_version} available.") + import torch_xla as xla + import torch_xla.core.xla_model as xm + + except ImportError: + raise ImportError("PyTorch XLA not installed.") + try: if torch.backends.mps.is_available(): @@ -87,6 +121,18 @@ except: if args.cpu: cpu_state = CPUState.CPU +if (args.xla or args.xla_fsdp) and XLA_AVAILABLE: + cpu_state = CPUState.XLA + + +def get_xla_memory_info(dev): + # xm.get_memory_info(dev) only has bytes_limit and bytes_used + mem_reserved = xm.get_memory_info(dev)["bytes_used"] + # minus 1GB for safety + mem_total = xm.get_memory_info(dev)["bytes_limit"] + return (mem_reserved, mem_total) + + def is_intel_xpu(): global cpu_state global xpu_available @@ -95,6 +141,7 @@ def is_intel_xpu(): return True return False + def get_torch_device(): global directml_enabled global cpu_state @@ -105,32 +152,37 @@ def get_torch_device(): return torch.device("mps") if cpu_state == CPUState.CPU: return torch.device("cpu") + if cpu_state == CPUState.XLA: + return xla.device() else: if is_intel_xpu(): return torch.device("xpu", torch.xpu.current_device()) 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'): + 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 = 1024 * 1024 * 1024 # TODO mem_total_torch = mem_total elif is_intel_xpu(): stats = torch.xpu.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] + mem_reserved = stats["reserved_bytes.all.current"] mem_total_torch = mem_reserved mem_total = torch.xpu.get_device_properties(dev).total_memory + elif cpu_state == CPUState.XLA: + mem_total_torch, mem_total = get_xla_memory_info(dev) else: stats = torch.cuda.memory_stats(dev) - mem_reserved = stats['reserved_bytes.all.current'] + 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 @@ -140,9 +192,12 @@ def get_total_memory(dev=None, torch_total_too=False): else: return mem_total + total_vram = get_total_memory(get_torch_device()) / (1024 * 1024) total_ram = psutil.virtual_memory().total / (1024 * 1024) -logging.info("Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram)) +logging.info( + "Total VRAM {:0.0f} MB, total RAM {:0.0f} MB".format(total_vram, total_ram) +) try: logging.info("pytorch version: {}".format(torch_version)) @@ -162,6 +217,7 @@ else: try: import xformers import xformers.ops + XFORMERS_IS_AVAILABLE = True try: XFORMERS_IS_AVAILABLE = xformers._has_cpp_library @@ -171,14 +227,19 @@ else: XFORMERS_VERSION = xformers.version.__version__ logging.info("xformers version: {}".format(XFORMERS_VERSION)) if XFORMERS_VERSION.startswith("0.0.18"): - 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") + 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" + ) XFORMERS_ENABLED_VAE = False except: pass except: XFORMERS_IS_AVAILABLE = False + def is_nvidia(): global cpu_state if cpu_state == CPUState.GPU: @@ -186,6 +247,7 @@ def is_nvidia(): return True return False + ENABLE_PYTORCH_ATTENTION = False if args.use_pytorch_cross_attention: ENABLE_PYTORCH_ATTENTION = True @@ -196,12 +258,23 @@ VAE_DTYPES = [torch.float32] try: if is_nvidia(): if int(torch_version[0]) >= 2: - if ENABLE_PYTORCH_ATTENTION == False and args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + if ( + ENABLE_PYTORCH_ATTENTION == False + and args.use_split_cross_attention == False + and args.use_quad_cross_attention == False + ): ENABLE_PYTORCH_ATTENTION = True - if torch.cuda.is_bf16_supported() and torch.cuda.get_device_properties(torch.cuda.current_device()).major >= 8: + if ( + torch.cuda.is_bf16_supported() + and torch.cuda.get_device_properties(torch.cuda.current_device()).major + >= 8 + ): VAE_DTYPES = [torch.bfloat16] + VAE_DTYPES if is_intel_xpu(): - if args.use_split_cross_attention == False and args.use_quad_cross_attention == False: + if ( + args.use_split_cross_attention == False + and args.use_quad_cross_attention == False + ): ENABLE_PYTORCH_ATTENTION = True except: pass @@ -241,9 +314,10 @@ if lowvram_available: vram_state = set_vram_to -if cpu_state != CPUState.GPU: +if cpu_state != CPUState.GPU and cpu_state != CPUState.XLA: vram_state = VRAMState.DISABLED + if cpu_state == CPUState.MPS: vram_state = VRAMState.SHARED @@ -254,14 +328,17 @@ DISABLE_SMART_MEMORY = args.disable_smart_memory if DISABLE_SMART_MEMORY: logging.info("Disabling smart memory management") + def get_torch_device_name(device): - if hasattr(device, 'type'): + if hasattr(device, "type"): if device.type == "cuda": try: allocator_backend = torch.cuda.get_allocator_backend() except: allocator_backend = "" - return "{} {} : {}".format(device, torch.cuda.get_device_name(device), allocator_backend) + return "{} {} : {}".format( + device, torch.cuda.get_device_name(device), allocator_backend + ) else: return "{}".format(device.type) elif is_intel_xpu(): @@ -269,6 +346,7 @@ def get_torch_device_name(device): else: return "CUDA {}: {}".format(device, torch.cuda.get_device_name(device)) + try: logging.info("Device: {}".format(get_torch_device_name(get_torch_device()))) except: @@ -277,6 +355,7 @@ except: current_loaded_models = [] + def module_size(module): module_mem = 0 sd = module.state_dict() @@ -285,6 +364,7 @@ def module_size(module): module_mem += t.nelement() * t.element_size() return module_mem + class LoadedModel: def __init__(self, model): self.model = model @@ -320,15 +400,30 @@ class LoadedModel: self.model_use_more_vram(use_more_vram) else: try: - self.real_model = self.model.patch_model(device_to=patch_model_to, lowvram_model_memory=lowvram_model_memory, load_weights=load_weights, force_patch_weights=force_patch_weights) + self.real_model = self.model.patch_model( + device_to=patch_model_to, + lowvram_model_memory=lowvram_model_memory, + load_weights=load_weights, + force_patch_weights=force_patch_weights, + ) except Exception as e: self.model.unpatch_model(self.model.offload_device) self.model_unload() raise e - if is_intel_xpu() and not args.disable_ipex_optimize and 'ipex' in globals() and self.real_model is not None: + if ( + is_intel_xpu() + and not args.disable_ipex_optimize + and "ipex" in globals() + and self.real_model is not None + ): with torch.no_grad(): - self.real_model = ipex.optimize(self.real_model.eval(), inplace=True, graph_mode=True, concat_linear=True) + self.real_model = ipex.optimize( + self.real_model.eval(), + inplace=True, + graph_mode=True, + concat_linear=True, + ) self.weights_loaded = True return self.real_model @@ -341,10 +436,14 @@ class LoadedModel: def model_unload(self, memory_to_free=None, unpatch_weights=True): if memory_to_free is not None: if memory_to_free < self.model.loaded_size(): - freed = self.model.partially_unload(self.model.offload_device, memory_to_free) + freed = self.model.partially_unload( + self.model.offload_device, memory_to_free + ) if freed >= memory_to_free: return False - self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights) + self.model.unpatch_model( + self.model.offload_device, unpatch_weights=unpatch_weights + ) self.model.model_patches_to(self.model.offload_device) self.weights_loaded = self.weights_loaded and not unpatch_weights self.real_model = None @@ -356,6 +455,7 @@ class LoadedModel: def __eq__(self, other): return self.model is other.model + def use_more_memory(extra_memory, loaded_models, device): for m in loaded_models: if m.device == device: @@ -363,6 +463,7 @@ def use_more_memory(extra_memory, loaded_models, device): if extra_memory <= 0: break + def offloaded_memory(loaded_models, device): offloaded_mem = 0 for m in loaded_models: @@ -370,22 +471,32 @@ def offloaded_memory(loaded_models, device): offloaded_mem += m.model_offloaded_memory() return offloaded_mem + WINDOWS = any(platform.win32_ver()) EXTRA_RESERVED_VRAM = 400 * 1024 * 1024 if WINDOWS: - EXTRA_RESERVED_VRAM = 600 * 1024 * 1024 #Windows is higher because of the shared vram issue + EXTRA_RESERVED_VRAM = ( + 600 * 1024 * 1024 + ) # Windows is higher because of the shared vram issue if args.reserve_vram is not None: EXTRA_RESERVED_VRAM = args.reserve_vram * 1024 * 1024 * 1024 - logging.debug("Reserving {}MB vram for other applications.".format(EXTRA_RESERVED_VRAM / (1024 * 1024))) + logging.debug( + "Reserving {}MB vram for other applications.".format( + EXTRA_RESERVED_VRAM / (1024 * 1024) + ) + ) + def extra_reserved_memory(): return EXTRA_RESERVED_VRAM + def minimum_inference_memory(): return (1024 * 1024 * 1024) * 0.8 + extra_reserved_memory() + def unload_model_clones(model, unload_weights_only=True, force_unload=True): to_unload = [] for i in range(len(current_loaded_models)): @@ -417,16 +528,24 @@ def unload_model_clones(model, unload_weights_only=True, force_unload=True): return unload_weight + def free_memory(memory_required, device, keep_loaded=[]): unloaded_model = [] can_unload = [] unloaded_models = [] - for i in range(len(current_loaded_models) -1, -1, -1): + 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: - can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i)) + can_unload.append( + ( + -shift_model.model_offloaded_memory(), + sys.getrefcount(shift_model.model), + shift_model.model_memory(), + i, + ) + ) shift_model.currently_used = False for x in sorted(can_unload): @@ -437,7 +556,9 @@ def free_memory(memory_required, device, keep_loaded=[]): if free_mem > memory_required: break memory_to_free = memory_required - free_mem - logging.debug(f"Unloading {current_loaded_models[i].model.model.__class__.__name__}") + logging.debug( + f"Unloading {current_loaded_models[i].model.model.__class__.__name__}" + ) if current_loaded_models[i].model_unload(memory_to_free): unloaded_model.append(i) @@ -448,12 +569,21 @@ def free_memory(memory_required, device, keep_loaded=[]): soft_empty_cache() else: if vram_state != VRAMState.HIGH_VRAM: - mem_free_total, mem_free_torch = get_free_memory(device, torch_free_too=True) + 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() return unloaded_models -def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimum_memory_required=None, force_full_load=False): + +def load_models_gpu( + models, + memory_required=0, + force_patch_weights=False, + minimum_memory_required=None, + force_full_load=False, +): global vram_state inference_memory = minimum_inference_memory() @@ -461,7 +591,9 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu if minimum_memory_required is None: minimum_memory_required = extra_mem else: - minimum_memory_required = max(inference_memory, minimum_memory_required + extra_reserved_memory()) + minimum_memory_required = max( + inference_memory, minimum_memory_required + extra_reserved_memory() + ) models = set(models) @@ -478,8 +610,12 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu 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) + 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: loaded.currently_used = True @@ -494,35 +630,57 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu devs = set(map(lambda a: a.device, models_already_loaded)) for d in devs: if d != torch.device("cpu"): - free_memory(extra_mem + offloaded_memory(models_already_loaded, d), d, models_already_loaded) + free_memory( + extra_mem + offloaded_memory(models_already_loaded, d), + d, + models_already_loaded, + ) free_mem = get_free_memory(d) if free_mem < minimum_memory_required: - logging.info("Unloading models for lowram load.") #TODO: partial model unloading when this case happens, also handle the opposite case where models can be unlowvramed. + logging.info( + "Unloading models for lowram load." + ) # TODO: partial model unloading when this case happens, also handle the opposite case where models can be unlowvramed. models_to_load = free_memory(minimum_memory_required, d) logging.info("{} models unloaded.".format(len(models_to_load))) else: - use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d) + use_more_memory( + free_mem - minimum_memory_required, models_already_loaded, d + ) if len(models_to_load) == 0: return - logging.info(f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}") + logging.info( + f"Loading {len(models_to_load)} new model{'s' if len(models_to_load) > 1 else ''}" + ) total_memory_required = {} for loaded_model in models_to_load: - unload_model_clones(loaded_model.model, unload_weights_only=True, force_unload=False) #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) + unload_model_clones( + loaded_model.model, unload_weights_only=True, force_unload=False + ) # 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) for loaded_model in models_already_loaded: - total_memory_required[loaded_model.device] = total_memory_required.get(loaded_model.device, 0) + loaded_model.model_memory_required(loaded_model.device) + total_memory_required[loaded_model.device] = total_memory_required.get( + loaded_model.device, 0 + ) + loaded_model.model_memory_required(loaded_model.device) for loaded_model in models_to_load: - 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 + 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 for device in total_memory_required: if device != torch.device("cpu"): - free_memory(total_memory_required[device] * 1.1 + extra_mem, device, models_already_loaded) + free_memory( + total_memory_required[device] * 1.1 + extra_mem, + device, + models_already_loaded, + ) for loaded_model in models_to_load: model = loaded_model.model @@ -532,32 +690,52 @@ def load_models_gpu(models, memory_required=0, force_patch_weights=False, minimu 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) and not force_full_load: + if ( + lowvram_available + and ( + vram_set_state == VRAMState.LOW_VRAM + or vram_set_state == VRAMState.NORMAL_VRAM + ) + and not force_full_load + ): model_size = loaded_model.model_memory_required(torch_dev) current_free_mem = get_free_memory(torch_dev) - lowvram_model_memory = max(64 * (1024 * 1024), (current_free_mem - minimum_memory_required), min(current_free_mem * 0.4, current_free_mem - minimum_inference_memory())) - if model_size <= lowvram_model_memory: #only switch to lowvram if really necessary + lowvram_model_memory = max( + 64 * (1024 * 1024), + (current_free_mem - minimum_memory_required), + min( + current_free_mem * 0.4, + current_free_mem - minimum_inference_memory(), + ), + ) + if ( + model_size <= lowvram_model_memory + ): # only switch to lowvram if really necessary lowvram_model_memory = 0 if vram_set_state == VRAMState.NO_VRAM: lowvram_model_memory = 64 * 1024 * 1024 - cur_loaded_model = loaded_model.model_load(lowvram_model_memory, force_patch_weights=force_patch_weights) + cur_loaded_model = loaded_model.model_load( + lowvram_model_memory, force_patch_weights=force_patch_weights + ) current_loaded_models.insert(0, loaded_model) - devs = set(map(lambda a: a.device, models_already_loaded)) for d in devs: if d != torch.device("cpu"): free_mem = get_free_memory(d) if free_mem > minimum_memory_required: - use_more_memory(free_mem - minimum_memory_required, models_already_loaded, d) + use_more_memory( + free_mem - minimum_memory_required, models_already_loaded, d + ) return def load_model_gpu(model): return load_models_gpu([model]) + def loaded_models(only_currently_used=False): output = [] for m in current_loaded_models: @@ -568,16 +746,19 @@ def loaded_models(only_currently_used=False): output.append(m.model) return output + def cleanup_models(keep_clone_weights_loaded=False): to_delete = [] for i in range(len(current_loaded_models)): - #TODO: very fragile function needs improvement + # TODO: very fragile function needs improvement num_refs = sys.getrefcount(current_loaded_models[i].model) if num_refs <= 2: 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 + # 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 for i in to_delete: @@ -585,6 +766,7 @@ def cleanup_models(keep_clone_weights_loaded=False): x.model_unload() del x + def dtype_size(dtype): dtype_size = 4 if dtype == torch.float16 or dtype == torch.bfloat16: @@ -594,16 +776,18 @@ def dtype_size(dtype): else: try: dtype_size = dtype.itemsize - except: #Old pytorch doesn't have .itemsize + except: # Old pytorch doesn't have .itemsize pass return dtype_size + def unet_offload_device(): if vram_state == VRAMState.HIGH_VRAM: return get_torch_device() else: return torch.device("cpu") + def unet_inital_load_device(parameters, dtype): torch_dev = get_torch_device() if vram_state == VRAMState.HIGH_VRAM: @@ -622,10 +806,16 @@ def unet_inital_load_device(parameters, dtype): else: return cpu_dev -def maximum_vram_for_weights(device=None): - return (get_total_memory(device) * 0.88 - minimum_inference_memory()) -def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): +def maximum_vram_for_weights(device=None): + return get_total_memory(device) * 0.88 - minimum_inference_memory() + + +def unet_dtype( + device=None, + model_params=0, + supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], +): if model_params < 0: model_params = 1000000000000000000000 if args.bf16_unet: @@ -647,7 +837,9 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor pass if fp8_dtype is not None: - if supports_fp8_compute(device): #if fp8 compute is supported the casting is most likely not expensive + if supports_fp8_compute( + device + ): # if fp8 compute is supported the casting is most likely not expensive return fp8_dtype free_model_memory = maximum_vram_for_weights(device) @@ -655,7 +847,9 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor return fp8_dtype for dt in supported_dtypes: - if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params): + if dt == torch.float16 and should_use_fp16( + device=device, model_params=model_params + ): if torch.float16 in supported_dtypes: return torch.float16 if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params): @@ -663,17 +857,26 @@ def unet_dtype(device=None, model_params=0, supported_dtypes=[torch.float16, tor return torch.bfloat16 for dt in supported_dtypes: - if dt == torch.float16 and should_use_fp16(device=device, model_params=model_params, manual_cast=True): + if dt == torch.float16 and should_use_fp16( + device=device, model_params=model_params, manual_cast=True + ): if torch.float16 in supported_dtypes: return torch.float16 - if dt == torch.bfloat16 and should_use_bf16(device, model_params=model_params, manual_cast=True): + if dt == torch.bfloat16 and should_use_bf16( + device, model_params=model_params, manual_cast=True + ): if torch.bfloat16 in supported_dtypes: return torch.bfloat16 return torch.float32 + # None means no manual cast -def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.float16, torch.bfloat16, torch.float32]): +def unet_manual_cast( + weight_dtype, + inference_device, + supported_dtypes=[torch.float16, torch.bfloat16, torch.float32], +): if weight_dtype == torch.float32: return None @@ -694,12 +897,14 @@ def unet_manual_cast(weight_dtype, inference_device, supported_dtypes=[torch.flo return torch.float32 + def text_encoder_offload_device(): if args.gpu_only: return get_torch_device() else: return torch.device("cpu") + def text_encoder_device(): if args.gpu_only: return get_torch_device() @@ -711,6 +916,7 @@ def text_encoder_device(): else: return torch.device("cpu") + def text_encoder_initial_device(load_device, offload_device, model_size=0): if load_device == offload_device or model_size <= 1024 * 1024 * 1024: return offload_device @@ -725,6 +931,7 @@ def text_encoder_initial_device(load_device, offload_device, model_size=0): else: return offload_device + def text_encoder_dtype(device=None): if args.fp8_e4m3fn_text_enc: return torch.float8_e4m3fn @@ -747,17 +954,20 @@ def intermediate_device(): else: return torch.device("cpu") + def vae_device(): if args.cpu_vae: return torch.device("cpu") return get_torch_device() + def vae_offload_device(): if args.gpu_only: return get_torch_device() else: return torch.device("cpu") + def vae_dtype(device=None, allowed_dtypes=[]): global VAE_DTYPES if args.fp16_vae: @@ -775,12 +985,14 @@ def vae_dtype(device=None, allowed_dtypes=[]): return VAE_DTYPES[0] + def get_autocast_device(dev): - if hasattr(dev, 'type'): + if hasattr(dev, "type"): return dev.type return "cuda" -def supports_dtype(device, dtype): #TODO + +def supports_dtype(device, dtype): # TODO if dtype == torch.float32: return True if is_device_cpu(device): @@ -791,12 +1003,13 @@ def supports_dtype(device, dtype): #TODO return True return False -def supports_cast(device, dtype): #TODO + +def supports_cast(device, dtype): # TODO if dtype == torch.float32: return True if dtype == torch.float16: return True - if directml_enabled: #TODO: test this + if directml_enabled: # TODO: test this return False if dtype == torch.bfloat16: return True @@ -808,6 +1021,7 @@ def supports_cast(device, dtype): #TODO return True return False + def pick_weight_dtype(dtype, fallback_dtype, device=None): if dtype is None: dtype = fallback_dtype @@ -819,30 +1033,36 @@ def pick_weight_dtype(dtype, fallback_dtype, device=None): return dtype + def device_supports_non_blocking(device): if is_device_mps(device): - return False #pytorch bug? mps doesn't support non blocking + return False # pytorch bug? mps doesn't support non blocking if is_intel_xpu(): return False - if args.deterministic: #TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews) + if ( + args.deterministic + ): # TODO: figure out why deterministic breaks non blocking from gpu to cpu (previews) return False if directml_enabled: return False return True + def device_should_use_non_blocking(device): if not device_supports_non_blocking(device): return False return False # return True #TODO: figure out why this causes memory issues on Nvidia and possibly others + def force_channels_last(): if args.force_channels_last: return True - #TODO + # TODO return False + def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False): if device is None or weight.device == device: if not copy: @@ -854,9 +1074,12 @@ def cast_to(weight, dtype=None, device=None, non_blocking=False, copy=False): r.copy_(weight, non_blocking=non_blocking) return r + def cast_to_device(tensor, device, dtype, copy=False): non_blocking = device_supports_non_blocking(device) - return cast_to(tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy) + return cast_to( + tensor, dtype=dtype, device=device, non_blocking=non_blocking, copy=copy + ) def xformers_enabled(): @@ -878,25 +1101,30 @@ def xformers_enabled_vae(): return XFORMERS_ENABLED_VAE + def pytorch_attention_enabled(): global ENABLE_PYTORCH_ATTENTION return ENABLE_PYTORCH_ATTENTION + def pytorch_attention_flash_attention(): global ENABLE_PYTORCH_ATTENTION if ENABLE_PYTORCH_ATTENTION: - #TODO: more reliable way of checking for flash attention? - if is_nvidia(): #pytorch flash attention only works on Nvidia + # TODO: more reliable way of checking for flash attention? + if is_nvidia(): # pytorch flash attention only works on Nvidia return True if is_intel_xpu(): return True return False + def force_upcast_attention_dtype(): upcast = args.force_upcast_attention try: macos_version = tuple(int(n) for n in platform.mac_ver()[0].split(".")) - if (14, 5) <= macos_version <= (15, 2): # black image bug on recent versions of macOS + if ( + (14, 5) <= macos_version <= (15, 2) + ): # black image bug on recent versions of macOS upcast = True except: pass @@ -905,29 +1133,36 @@ def force_upcast_attention_dtype(): else: return None + def get_free_memory(dev=None, torch_free_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'): + if hasattr(dev, "type") and (dev.type == "cpu" or dev.type == "mps"): mem_free_total = psutil.virtual_memory().available mem_free_torch = mem_free_total else: if directml_enabled: - mem_free_total = 1024 * 1024 * 1024 #TODO + mem_free_total = 1024 * 1024 * 1024 # TODO mem_free_torch = mem_free_total elif is_intel_xpu(): stats = torch.xpu.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] + mem_active = stats["active_bytes.all.current"] + mem_reserved = stats["reserved_bytes.all.current"] mem_free_torch = mem_reserved - mem_active - mem_free_xpu = torch.xpu.get_device_properties(dev).total_memory - mem_reserved + mem_free_xpu = ( + torch.xpu.get_device_properties(dev).total_memory - mem_reserved + ) mem_free_total = mem_free_xpu + mem_free_torch + elif cpu_state == CPUState.XLA: + mem_reserved, mem_total = get_xla_memory_info(dev) + mem_free_total = mem_total - mem_reserved + mem_free_torch = mem_free_total else: stats = torch.cuda.memory_stats(dev) - mem_active = stats['active_bytes.all.current'] - mem_reserved = stats['reserved_bytes.all.current'] + 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 @@ -937,30 +1172,44 @@ def get_free_memory(dev=None, torch_free_too=False): else: return mem_free_total + def cpu_mode(): global cpu_state return cpu_state == CPUState.CPU + +def xla_mode(): + global cpu_state + return cpu_state == CPUState.XLA + + def mps_mode(): global cpu_state return cpu_state == CPUState.MPS + def is_device_type(device, type): - if hasattr(device, 'type'): - if (device.type == type): + if hasattr(device, "type"): + if device.type == type: return True return False + def is_device_cpu(device): - return is_device_type(device, 'cpu') + return is_device_type(device, "cpu") + def is_device_mps(device): - return is_device_type(device, 'mps') + return is_device_type(device, "mps") + def is_device_cuda(device): - return is_device_type(device, 'cuda') + return is_device_type(device, "cuda") -def should_use_fp16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): + +def should_use_fp16( + device=None, model_params=0, prioritize_performance=True, manual_cast=False +): global directml_enabled if device is not None: @@ -983,6 +1232,9 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if mps_mode(): return True + if xla_mode(): + return True + if cpu_mode(): return False @@ -999,14 +1251,31 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if props.major < 6: return False - #FP16 is confirmed working on a 1080 (GP104) and on latest pytorch actually seems faster than fp32 - nvidia_10_series = ["1080", "1070", "titan x", "p3000", "p3200", "p4000", "p4200", "p5000", "p5200", "p6000", "1060", "1050", "p40", "p100", "p6", "p4"] + # FP16 is confirmed working on a 1080 (GP104) and on latest pytorch actually seems faster than fp32 + nvidia_10_series = [ + "1080", + "1070", + "titan x", + "p3000", + "p3200", + "p4000", + "p4200", + "p5000", + "p5200", + "p6000", + "1060", + "1050", + "p40", + "p100", + "p6", + "p4", + ] for x in nvidia_10_series: if x in props.name.lower(): if WINDOWS or manual_cast: return True else: - return False #weird linux behavior where fp32 is faster + return False # weird linux behavior where fp32 is faster if manual_cast: free_model_memory = maximum_vram_for_weights(device) @@ -1016,17 +1285,33 @@ def should_use_fp16(device=None, model_params=0, prioritize_performance=True, ma if props.major < 7: return False - #FP16 is just broken on these cards - nvidia_16_series = ["1660", "1650", "1630", "T500", "T550", "T600", "MX550", "MX450", "CMP 30HX", "T2000", "T1000", "T1200"] + # FP16 is just broken on these cards + nvidia_16_series = [ + "1660", + "1650", + "1630", + "T500", + "T550", + "T600", + "MX550", + "MX450", + "CMP 30HX", + "T2000", + "T1000", + "T1200", + ] for x in nvidia_16_series: if x in props.name: return False return True -def should_use_bf16(device=None, model_params=0, prioritize_performance=True, manual_cast=False): + +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 + if is_device_cpu(device): # TODO ? bf16 works on CPU but is extremely slow return False if device is not None: @@ -1042,6 +1327,9 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma if mps_mode(): return True + if xla_mode(): + return True + if cpu_mode(): return False @@ -1061,7 +1349,11 @@ def should_use_bf16(device=None, model_params=0, prioritize_performance=True, ma return False + def supports_fp8_compute(device=None): + if XLA_AVAILABLE: + return True + if not is_nvidia(): return False @@ -1073,15 +1365,18 @@ def supports_fp8_compute(device=None): if props.minor < 9: return False - if int(torch_version[0]) < 2 or (int(torch_version[0]) == 2 and int(torch_version[2]) < 3): + if int(torch_version[0]) < 2 or ( + int(torch_version[0]) == 2 and int(torch_version[2]) < 3 + ): return False if WINDOWS: - if (int(torch_version[0]) == 2 and int(torch_version[2]) < 4): + if int(torch_version[0]) == 2 and int(torch_version[2]) < 4: return False return True + def soft_empty_cache(force=False): global cpu_state if cpu_state == CPUState.MPS: @@ -1089,39 +1384,51 @@ def soft_empty_cache(force=False): elif is_intel_xpu(): torch.xpu.empty_cache() elif torch.cuda.is_available(): - if force or is_nvidia(): #This seems to make things worse on ROCm so I only do it for cuda + if ( + force or is_nvidia() + ): # This seems to make things worse on ROCm so I only do it for cuda torch.cuda.empty_cache() torch.cuda.ipc_collect() + def unload_all_models(): free_memory(1e30, get_torch_device()) -def resolve_lowvram_weight(weight, model, key): #TODO: remove - print("WARNING: The comfy.model_management.resolve_lowvram_weight function will be removed soon, please stop using it.") +def resolve_lowvram_weight(weight, model, key): # TODO: remove + print( + "WARNING: The comfy.model_management.resolve_lowvram_weight function will be removed soon, please stop using it." + ) return weight -#TODO: might be cleaner to put this somewhere else + +# 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 diff --git a/latent_preview.py b/latent_preview.py index d60e68d55..cff05d97c 100644 --- a/latent_preview.py +++ b/latent_preview.py @@ -11,12 +11,22 @@ import logging MAX_PREVIEW_RESOLUTION = args.preview_size -def preview_to_image(latent_image): - latents_ubyte = (((latent_image + 1.0) / 2.0).clamp(0, 1) # change scale from -1..1 to 0..1 - .mul(0xFF) # to 0..255 - ).to(device="cpu", dtype=torch.uint8, non_blocking=comfy.model_management.device_supports_non_blocking(latent_image.device)) - return Image.fromarray(latents_ubyte.numpy()) +def preview_to_image(latent_image): + latents_ubyte = ( + ((latent_image + 1.0) / 2.0) + .clamp(0, 1) # change scale from -1..1 to 0..1 + .mul(0xFF) # to 0..255 + ).to( + device="cpu", + dtype=torch.uint8, + non_blocking=comfy.model_management.device_supports_non_blocking( + latent_image.device + ), + ) + + return Image.fromarray(latents_ubyte.numpy()) + class LatentPreviewer: def decode_latent_to_preview(self, x0): @@ -26,6 +36,7 @@ class LatentPreviewer: preview_image = self.decode_latent_to_preview(x0) return ("JPEG", preview_image, MAX_PREVIEW_RESOLUTION) + class TAESDPreviewerImpl(LatentPreviewer): def __init__(self, taesd): self.taesd = taesd @@ -37,22 +48,34 @@ class TAESDPreviewerImpl(LatentPreviewer): class Latent2RGBPreviewer(LatentPreviewer): def __init__(self, latent_rgb_factors, latent_rgb_factors_bias=None): - self.latent_rgb_factors = torch.tensor(latent_rgb_factors, device="cpu").transpose(0, 1) + self.latent_rgb_factors = torch.tensor( + latent_rgb_factors, device="cpu" + ).transpose(0, 1) self.latent_rgb_factors_bias = None if latent_rgb_factors_bias is not None: - self.latent_rgb_factors_bias = torch.tensor(latent_rgb_factors_bias, device="cpu") + self.latent_rgb_factors_bias = torch.tensor( + latent_rgb_factors_bias, device="cpu" + ) def decode_latent_to_preview(self, x0): - self.latent_rgb_factors = self.latent_rgb_factors.to(dtype=x0.dtype, device=x0.device) + self.latent_rgb_factors = self.latent_rgb_factors.to( + dtype=x0.dtype, device=x0.device + ) if self.latent_rgb_factors_bias is not None: - self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to(dtype=x0.dtype, device=x0.device) + self.latent_rgb_factors_bias = self.latent_rgb_factors_bias.to( + dtype=x0.dtype, device=x0.device + ) if x0.ndim == 5: x0 = x0[0, :, 0] else: x0 = x0[0] - latent_image = torch.nn.functional.linear(x0.movedim(0, -1), self.latent_rgb_factors, bias=self.latent_rgb_factors_bias) + latent_image = torch.nn.functional.linear( + x0.movedim(0, -1), + self.latent_rgb_factors, + bias=self.latent_rgb_factors_bias, + ) # latent_image = x0[0].permute(1, 2, 0) @ self.latent_rgb_factors return preview_to_image(latent_image) @@ -66,27 +89,44 @@ def get_previewer(device, latent_format): taesd_decoder_path = None if latent_format.taesd_decoder_name is not None: taesd_decoder_path = next( - (fn for fn in folder_paths.get_filename_list("vae_approx") - if fn.startswith(latent_format.taesd_decoder_name)), - "" + ( + fn + for fn in folder_paths.get_filename_list("vae_approx") + if fn.startswith(latent_format.taesd_decoder_name) + ), + "", + ) + taesd_decoder_path = folder_paths.get_full_path( + "vae_approx", taesd_decoder_path ) - taesd_decoder_path = folder_paths.get_full_path("vae_approx", taesd_decoder_path) if method == LatentPreviewMethod.Auto: method = LatentPreviewMethod.Latent2RGB if method == LatentPreviewMethod.TAESD: if taesd_decoder_path: - taesd = TAESD(None, taesd_decoder_path, latent_channels=latent_format.latent_channels).to(device) + taesd = TAESD( + None, + taesd_decoder_path, + latent_channels=latent_format.latent_channels, + ).to(device) previewer = TAESDPreviewerImpl(taesd) else: - logging.warning("Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format(latent_format.taesd_decoder_name)) + logging.warning( + "Warning: TAESD previews enabled, but could not find models/vae_approx/{}".format( + latent_format.taesd_decoder_name + ) + ) if previewer is None: if latent_format.latent_rgb_factors is not None: - previewer = Latent2RGBPreviewer(latent_format.latent_rgb_factors, latent_format.latent_rgb_factors_bias) + previewer = Latent2RGBPreviewer( + latent_format.latent_rgb_factors, + latent_format.latent_rgb_factors_bias, + ) return previewer + def prepare_callback(model, steps, x0_output_dict=None): preview_format = "JPEG" if preview_format not in ["JPEG", "PNG"]: @@ -95,6 +135,7 @@ def prepare_callback(model, steps, x0_output_dict=None): previewer = get_previewer(model.load_device, model.model.latent_format) pbar = comfy.utils.ProgressBar(steps) + def callback(step, x0, x, total_steps): if x0_output_dict is not None: x0_output_dict["x0"] = x0 @@ -103,5 +144,11 @@ def prepare_callback(model, steps, x0_output_dict=None): if previewer: preview_bytes = previewer.decode_latent_to_preview_image(preview_format, x0) pbar.update_absolute(step + 1, total_steps, preview_bytes) - return callback + if args.xla or args.xla_fsdp: + import torch_xla as xla + import torch_xla.core.xla_model as xm + logging.info(f"Step {step + 1}/{total_steps}") + xla.sync() + + return callback