Merge a5acc660b2 into e14f3b6610

Merge branch 'master' into unetloader_remain_mmap
enable only when mmap is enabled
2026-01-10 06:10:50 +08:00 · 2026-01-06 11:35:32 +02:00 · 2025-12-18 16:20:06 +08:00 · 2025-12-15 16:39:54 +08:00 · 2025-12-12 17:06:06 +08:00
1 changed files with 4 additions and 1 deletions
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@ -306,7 +306,10 @@ class BaseModel(torch.nn.Module):
                to_load[k[len(unet_prefix):]] = sd.pop(k)

        to_load = self.model_config.process_unet_state_dict(to_load)
-        m, u = self.diffusion_model.load_state_dict(to_load, strict=False)
+        # assign=True will reuse the tensor storage in state dict, this will avoid copy and saving CPU memory
+        # when loading large models with mmap. 
+        delay_copy_with_assign = utils.MMAP_TORCH_FILES or not utils.DISABLE_MMAP
+        m, u = self.diffusion_model.load_state_dict(to_load, strict=False, assign=delay_copy_with_assign)
        if len(m) > 0:
            logging.warning("unet missing: {}".format(m))
Author	SHA1	Message	Date
Xiaoyu Xu	d66e5dd864	Merge `a5acc660b2` into `e14f3b6610`	2026-01-06 11:35:32 +02:00
Xiaoyu Xu	a5acc660b2	Merge branch 'master' into unetloader_remain_mmap	2025-12-18 16:20:06 +08:00
strint	5f10b70d9c	enable only when mmap is enabled	2025-12-15 16:39:54 +08:00
strint	afe05d68ee	load diffusion model from state dict with assign	2025-12-12 17:06:06 +08:00