zlei6/ComfyUI
1
0
Fork 0
mirror of https://github.com/Comfy-Org/ComfyUI.git synced 2026-08-25 02:20:18 +08:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #14052 from rattus128/prs/worksplit-t-load-fix

Browse Source 
fixup threaded loader with worksplit multi-gpu
This commit is contained in:
Jedrzej Kosinski
2026-05-22 16:36:33 -07:00
committed by GitHub
parent 4d3d68e473 7a18f9affb
commit cb83c41db7
3 changed files with 21 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
comfy/utils.py
View File

@@ -86,6 +86,7 @@ def load_safetensors(ckpt):
import comfy_aimdo.model_mmap
f = open(ckpt, "rb", buffering=0)
file_lock = threading.Lock()
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()))
@@ -111,7 +112,7 @@ def load_safetensors(ckpt):
storage = tensor.untyped_storage()
setattr(storage,
"_comfy_tensor_file_slice",
comfy.memory_management.TensorFileSlice(f, threading.get_ident(), data_base_offset + start, end - start))
comfy.memory_management.TensorFileSlice(f, file_lock, data_base_offset + start, end - start))
setattr(storage, "_comfy_tensor_mmap_refs", (model_mmap, mv))
sd[name] = tensor