wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-06-20 23:09:27 +08:00
Code Issues Actions 28 Packages Projects Releases Wiki Activity

implement touch accounting

Browse Source 
Implement accounting of touching mmapped tensors.
This commit is contained in:
Rattus 2026-03-11 18:57:37 +10:00
parent 34647e8345
commit 6b7f170b3e
2 changed files with 6 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/model_management.py
View File

@ -1227,6 +1227,9 @@ def cast_to_gathered(tensors, r, non_blocking=False, stream=None):
continue continue
if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view): if comfy.memory_management.read_tensor_file_slice_into(tensor, dest_view):
continue 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) dest_view.copy_(tensor, non_blocking=non_blocking)

4
comfy/utils.py
View File

@ -104,9 +104,11 @@ def load_safetensors(ckpt):
#We are working with read-only RAM by design #We are working with read-only RAM by design
warnings.filterwarnings("ignore", message="The given buffer is not writable") warnings.filterwarnings("ignore", message="The given buffer is not writable")
tensor = torch.frombuffer(mv[start:end], dtype=_TYPES[info["dtype"]]).view(info["shape"]) 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_tensor_file_slice",
comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start)) comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start))
setattr(storage, "_comfy_tensor_mmap_touched", False)
sd[name] = tensor sd[name] = tensor
return sd, header.get("__metadata__", {}), return sd, header.get("__metadata__", {}),