diff --git a/comfy/model_management.py b/comfy/model_management.py index 1ada985881a..fe1e32c3347 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -49,8 +49,8 @@ def get_mmap_mem_threshold_gb(): logging.debug(f"MMAP_MEM_THRESHOLD_GB: {mmap_mem_threshold_gb}") return mmap_mem_threshold_gb -def get_free_disk(): - return psutil.disk_usage("/").free +def get_free_disk(dir: str = "/"): + return psutil.disk_usage(dir).free class VRAMState(Enum): DISABLED = 0 #No vram present: no need to move models to vram @@ -829,7 +829,7 @@ def model_unload(self, memory_to_free=None, unpatch_weights=True): logging.debug(f"before unload, available_memory of offload device {self.model.offload_device}: {available_memory/(1024*1024*1024)} GB") mmap_mem_threshold = get_mmap_mem_threshold_gb() * 1024 * 1024 * 1024 # this is reserved memory for other system usage - if min(memory_to_free, model_loaded_size) > available_memory - mmap_mem_threshold or memory_to_free < model_loaded_size: + if memory_to_free < model_loaded_size: partially_unload = True else: partially_unload = False diff --git a/comfy/model_patcher.py b/comfy/model_patcher.py index 3a37c5549f4..25cbbe01552 100644 --- a/comfy/model_patcher.py +++ b/comfy/model_patcher.py @@ -32,6 +32,7 @@ import weakref import gc import tqdm +import mmap import comfy.float import comfy.hooks @@ -47,10 +48,13 @@ from comfy.patcher_extension import CallbacksMP, PatcherInjection, WrappersMP from comfy.model_management import get_free_memory, get_mmap_mem_threshold_gb, get_free_disk -def need_mmap() -> bool: +_USE_GDS_OFFLOAD = bool(os.environ.get("USE_GDS_OFFLOAD", "False").lower() in ("true", "1", "yes")) +_CUDA_GDS_AVAILABLE = hasattr(torch, "cuda") and hasattr(torch.cuda, "gds") and hasattr(torch.cuda.gds, "GdsFile") + +def need_mmap(offload_size: int = 0) -> bool: free_cpu_mem = get_free_memory(torch.device("cpu")) mmap_mem_threshold_gb = get_mmap_mem_threshold_gb() - if free_cpu_mem < mmap_mem_threshold_gb * 1024 * 1024 * 1024: + if free_cpu_mem - offload_size < mmap_mem_threshold_gb * 1024 * 1024 * 1024: logging.debug(f"Enabling mmap, current free cpu memory {free_cpu_mem/(1024*1024*1024)} GB < {mmap_mem_threshold_gb} GB") return True return False @@ -61,32 +65,56 @@ def to_mmap(t: torch.Tensor, filename: Optional[str] = None) -> torch.Tensor: """ # Create temporary file if filename is None: - temp_file = tempfile.mkstemp(suffix='.pt', prefix='comfy_mmap_')[1] + fd, temp_file = tempfile.mkstemp(suffix='.bin', prefix='comfy_mmap_') + os.close(fd) else: temp_file = filename - # Save tensor to file - cpu_tensor = t.cpu() - torch.save(cpu_tensor, temp_file) - - # If we created a CPU copy from other device, delete it to free memory - if not t.device.type == 'cpu': - del cpu_tensor - gc.collect() + use_gds = ( + _USE_GDS_OFFLOAD + and _CUDA_GDS_AVAILABLE + and t.is_contiguous() + and t.storage_offset() == 0 + and t.untyped_storage().nbytes() == t.numel() * t.element_size() + and t.is_cuda + ) + if use_gds: + file = torch.cuda.gds.GdsFile(temp_file, os.O_CREAT | os.O_RDWR) + file.save_storage(t.untyped_storage(), offset=0) + t_type = t.dtype + t_shape = t.shape + num = t.numel() * t.element_size() + del file + + with open(temp_file, "rb") as fo: + mm = mmap.mmap(fo.fileno(), length=num, access=mmap.ACCESS_COPY) + mmap_tensor = torch.frombuffer(mm, dtype=t_type).reshape(t_shape).cpu() + mmap_tensor._mmap = mm + else: + cpu_tensor = t.cpu() + torch.save(cpu_tensor, temp_file) - # Load with mmap - this doesn't load all data into RAM - mmap_tensor = torch.load(temp_file, map_location='cpu', mmap=True, weights_only=False) + # If we created a CPU copy from other device, delete it to free memory + if not t.device.type == 'cpu': + del cpu_tensor + gc.collect() + + # Load with mmap - this doesn't load all data into RAM + mmap_tensor = torch.load(temp_file, map_location='cpu', mmap=True, weights_only=False) # Register cleanup callback - will be called when tensor is garbage collected - def _cleanup(): + def _cleanup(temp_file, mmap_file): try: + if mmap_file is not None: + mmap_file.close() if os.path.exists(temp_file): os.remove(temp_file) logging.debug(f"Cleaned up mmap file: {temp_file}") except Exception: pass - weakref.finalize(mmap_tensor, _cleanup) + mmap_file = getattr(mmap_tensor, "_mmap", None) + weakref.finalize(mmap_tensor.untyped_storage(), _cleanup, temp_file, mmap_file) return mmap_tensor @@ -109,12 +137,18 @@ def model_to_mmap(model: torch.nn.Module): The same model with all tensors converted to memory-mapped format """ free_cpu_mem = get_free_memory(torch.device("cpu")) + free_disk_mem = get_free_disk(dir=tempfile.gettempdir()) + model_mem = comfy.model_management.module_size(model) + if model_mem > free_disk_mem: + logging.error(f"Not enough free disk memory to convert model to mmap. Model size: {model_mem/(1024*1024*1024)} GB, free disk memory: {free_disk_mem/(1024*1024*1024)} GB") + raise ValueError("Not enough free disk memory to convert model to mmap") logging.debug(f"Converting model {model.__class__.__name__} to mmap, current free cpu memory: {free_cpu_mem/(1024*1024*1024)} GB") def convert_fn(t): if isinstance(t, QuantizedTensor): - logging.debug(f"QuantizedTensor detected, mmap skipped, tensor meta info: size {t.size()}, dtype {t.dtype}, device {t.device}, is_contiguous {t.is_contiguous()}") - return t + inner_tensor_names, quant_ctx = t.__tensor_flatten__() + inner_tensors = {name: to_mmap(getattr(t, name)) for name in inner_tensor_names} + return QuantizedTensor.__tensor_unflatten__(inner_tensors, quant_ctx, t.size(), t.stride()) elif isinstance(t, torch.nn.Parameter): new_tensor = to_mmap(t.detach()) return torch.nn.Parameter(new_tensor, requires_grad=t.requires_grad) @@ -1240,9 +1274,13 @@ def unpatch_model(self, device_to=None, unpatch_weights=True): if device_to is not None: - if need_mmap(): + if need_mmap(offload_size=self.loaded_size()): # offload to mmap - model_to_mmap(self.model) + try: + model_to_mmap(self.model) + except Exception as e: + logging.warning(f"Error occurred while offloading model to mmap: {e}, fall back to normal offload") + self.model.to(device_to) else: self.model.to(device_to) self.model.device = device_to @@ -1305,12 +1343,13 @@ def partially_unload(self, device_to, memory_to_free=0, force_patch_weights=Fals bias_key = "{}.bias".format(n) if move_weight: cast_weight = self.force_cast_weights - if need_mmap(): - if get_free_disk() < module_mem: - logging.warning(f"Not enough disk space to offload {n} to mmap, current free disk space {get_free_disk()/(1024*1024*1024)} GB < {module_mem/(1024*1024*1024)} GB") - break - # offload to mmap - model_to_mmap(m) + if need_mmap(offload_size=module_mem): + try: + # offload to mmap + model_to_mmap(m) + except Exception as e: + logging.warning(f"Error occurred while offloading {n} to mmap: {e}, fall back to normal offload") + m.to(device_to) else: m.to(device_to) module_mem += move_weight_functions(m, device_to) diff --git a/comfy/ops.py b/comfy/ops.py index ff64aad5924..699f77d9c47 100644 --- a/comfy/ops.py +++ b/comfy/ops.py @@ -1101,7 +1101,11 @@ def _quantized_apply(module, fn, recurse=True): p = fn(param) if (not torch.is_inference_mode_enabled()) and p.is_inference(): p = p.clone() - module.register_parameter(key, torch.nn.Parameter(p, requires_grad=False)) + new_param = torch.nn.Parameter(p, requires_grad=False) + if isinstance(p, QuantizedTensor): + new_param._qdata = p._qdata + new_param._params = p._params + module.register_parameter(key, new_param) for key, buf in module._buffers.items(): if buf is not None: module._buffers[key] = fn(buf) diff --git a/main.py b/main.py index b6f7d6cef73..511d750e8b3 100644 --- a/main.py +++ b/main.py @@ -185,7 +185,8 @@ def execute_script(script_path): node_paths = folder_paths.get_folder_paths("custom_nodes") for custom_node_path in node_paths: - possible_modules = os.listdir(custom_node_path) + possible_modules = sorted(os.listdir(custom_node_path)) + logging.debug("!!!possible_modules1!!!:\n" + "\n".join(possible_modules)) node_prestartup_times = [] for possible_module in possible_modules: diff --git a/nodes.py b/nodes.py index fa3a7794955..e840fa51715 100644 --- a/nodes.py +++ b/nodes.py @@ -2353,7 +2353,8 @@ async def init_external_custom_nodes(): node_paths = folder_paths.get_folder_paths("custom_nodes") node_import_times = [] for custom_node_path in node_paths: - possible_modules = os.listdir(os.path.realpath(custom_node_path)) + possible_modules = sorted(os.listdir(os.path.realpath(custom_node_path))) + logging.debug("!!!possible_modules2!!!:\n" + "\n".join(possible_modules)) if "__pycache__" in possible_modules: possible_modules.remove("__pycache__")