From d8b4016d6799eca80784276b993cc82894f7f2cb Mon Sep 17 00:00:00 2001
From: Rattus <rattus128@gmail.com>
Date: Wed, 11 Mar 2026 19:48:40 +1000
Subject: [PATCH] utils: use the aimdo mmap to load sft files

---
 comfy/utils.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/comfy/utils.py b/comfy/utils.py
index 795350830..9931fe3b4 100644
--- a/comfy/utils.py
+++ b/comfy/utils.py
@@ -20,6 +20,8 @@
 import torch
 import math
 import struct
+import ctypes
+import os
 import comfy.memory_management
 import safetensors.torch
 import numpy as np
@@ -32,7 +34,6 @@ from einops import rearrange
 from comfy.cli_args import args
 import json
 import time
-import mmap
 import threading
 import warnings
 
@@ -82,12 +83,15 @@ _TYPES = {
 }
 
 def load_safetensors(ckpt):
-    f = open(ckpt, "rb", buffering=0)
-    mapping = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
-    mv = memoryview(mapping)
+    import comfy_aimdo.model_mmap
 
-    header_size = struct.unpack("<Q", mapping[:8])[0]
-    header = json.loads(mapping[8:8+header_size].decode("utf-8"))
+    f = open(ckpt, "rb", buffering=0)
+    model_mmap = comfy_aimdo.model_mmap.ModelMMAP(ckpt)
+    file_size = os.path.getsize(ckpt)
+    mv = memoryview((ctypes.c_uint8 * file_size).from_address(model_mmap.get()))
+
+    header_size = struct.unpack("<Q", mv[:8])[0]
+    header = json.loads(mv[8:8 + header_size].tobytes().decode("utf-8"))
 
     mv = mv[(data_base_offset := 8 + header_size):]
 
@@ -108,6 +112,7 @@ def load_safetensors(ckpt):
                 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_refs", (model_mmap, mv))
                 setattr(storage, "_comfy_tensor_mmap_touched", False)
                 sd[name] = tensor