From 6b7f170b3e8dee1db52972fe5a7ce1f31be84962 Mon Sep 17 00:00:00 2001 From: Rattus Date: Wed, 11 Mar 2026 18:57:37 +1000 Subject: [PATCH] implement touch accounting Implement accounting of touching mmapped tensors. --- comfy/model_management.py | 3 +++ comfy/utils.py | 4 +++- 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 813b927be..98ffcc1f1 100644 --- a/comfy/model_management.py +++ b/comfy/model_management.py @@ -1227,6 +1227,9 @@ def cast_to_gathered(tensors, r, non_blocking=False, stream=None): continue if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view): continue + storage = tensor.untyped_storage() + if hasattr(storage, "_comfy_tensor_mmap_touched"): + storage._comfy_tensor_mmap_touched = True dest_view.copy_(tensor, non_blocking=non_blocking) diff --git a/comfy/utils.py b/comfy/utils.py index c30366f8e..795350830 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -104,9 +104,11 @@ def load_safetensors(ckpt): #We are working with read-only RAM by design warnings.filterwarnings("ignore", message="The given buffer is not writable") tensor = torch.frombuffer(mv[start:end], dtype=_TYPES[info["dtype"]]).view(info["shape"]) - setattr(tensor.untyped_storage(), + storage = tensor.untyped_storage() + setattr(storage, "_comfy_tensor_file_slice", comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start)) + setattr(storage, "_comfy_tensor_mmap_touched", False) sd[name] = tensor return sd, header.get("__metadata__", {}),