mmap 行为优化

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

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,62 @@ 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:
+        gds = None
+        try:
+            gds = torch.cuda.gds.GdsFile(temp_file, os.O_CREAT | os.O_RDWR)
+            gds.save_storage(t.untyped_storage(), offset=0)
+            t_type = t.dtype
+            t_shape = t.shape
+            num = t.numel() * t.element_size()
+
+            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
+        except Exception as e:
+            logging.warning(f"GDS offload failed, falling back to CPU mmap: {e}")
+            use_gds = False
+        finally:
+            del gds
+    if not use_gds:
+        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 +143,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 +1280,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 +1349,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)

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)

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.info("!!!possible_modules1!!!:\n" + "\n".join(possible_modules))
         node_prestartup_times = []
 
         for possible_module in possible_modules:

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.info("!!!possible_modules2!!!:\n" + "\n".join(possible_modules))
         if "__pycache__" in possible_modules:
             possible_modules.remove("__pycache__")