From 49597bfa3e36c78635db4234106611846fbc4117 Mon Sep 17 00:00:00 2001 From: strint Date: Fri, 17 Oct 2025 21:43:49 +0800 Subject: [PATCH] load remains mmap --- comfy/ldm/flux/model.py | 6 +++--- comfy/ldm/modules/diffusionmodules/openaimodel.py | 6 +++--- comfy/model_base.py | 4 ++-- comfy/sd.py | 6 +++--- comfy/utils.py | 2 +- 5 files changed, 12 insertions(+), 12 deletions(-) diff --git a/comfy/ldm/flux/model.py b/comfy/ldm/flux/model.py index 263cdae26..da46ed2ed 100644 --- a/comfy/ldm/flux/model.py +++ b/comfy/ldm/flux/model.py @@ -280,13 +280,13 @@ class Flux(nn.Module): out = out[:, :img_tokens] return rearrange(out, "b (h w) (c ph pw) -> b c (h ph) (w pw)", h=h_len, w=w_len, ph=2, pw=2)[:,:,:h_orig,:w_orig] - def load_state_dict(self, state_dict, strict=True): - import pdb; pdb.set_trace() + def load_state_dict(self, state_dict, strict=True, assign=False): + # import pdb; pdb.set_trace() """Override load_state_dict() to add logging""" logging.info(f"Flux load_state_dict start, strict={strict}, state_dict keys count={len(state_dict)}") # Call parent's load_state_dict method - result = super().load_state_dict(state_dict, strict=strict) + result = super().load_state_dict(state_dict, strict=strict, assign=assign) logging.info(f"Flux load_state_dict end, strict={strict}, state_dict keys count={len(state_dict)}") diff --git a/comfy/ldm/modules/diffusionmodules/openaimodel.py b/comfy/ldm/modules/diffusionmodules/openaimodel.py index e847700c6..2cdf711d4 100644 --- a/comfy/ldm/modules/diffusionmodules/openaimodel.py +++ b/comfy/ldm/modules/diffusionmodules/openaimodel.py @@ -913,13 +913,13 @@ class UNetModel(nn.Module): return self.out(h) - def load_state_dict(self, state_dict, strict=True): - import pdb; pdb.set_trace() + def load_state_dict(self, state_dict, strict=True, assign=False): + # import pdb; pdb.set_trace() """Override load_state_dict() to add logging""" logging.info(f"UNetModel load_state_dict start, strict={strict}, state_dict keys count={len(state_dict)}") # Call parent's load_state_dict method - result = super().load_state_dict(state_dict, strict=strict) + result = super().load_state_dict(state_dict, strict=strict, assign=assign) logging.info(f"UNetModel load_state_dict end, strict={strict}, state_dict keys count={len(state_dict)}") diff --git a/comfy/model_base.py b/comfy/model_base.py index 34dd16037..409e7fb87 100644 --- a/comfy/model_base.py +++ b/comfy/model_base.py @@ -305,9 +305,9 @@ class BaseModel(torch.nn.Module): logging.info(f"load model {self.model_config} weights process end") # TODO(sf): to mmap # diffusion_model is UNetModel - import pdb; pdb.set_trace() + # import pdb; pdb.set_trace() # TODO(sf): here needs to avoid load mmap into cpu mem - m, u = self.diffusion_model.load_state_dict(to_load, strict=False) + m, u = self.diffusion_model.load_state_dict(to_load, strict=False, assign=True) free_cpu_memory = get_free_memory(torch.device("cpu")) logging.info(f"load model {self.model_config} weights end, free cpu memory size {free_cpu_memory/(1024*1024*1024)} GB") if len(m) > 0: diff --git a/comfy/sd.py b/comfy/sd.py index 7005a1b53..a956884fb 100644 --- a/comfy/sd.py +++ b/comfy/sd.py @@ -1338,7 +1338,7 @@ def load_diffusion_model_state_dict(sd, model_options={}): model_config.optimizations["fp8"] = True model = model_config.get_model(new_sd, "") - import pdb; pdb.set_trace() + # import pdb; pdb.set_trace() model = model.to(offload_device) model.load_model_weights(new_sd, "") left_over = sd.keys() @@ -1351,10 +1351,10 @@ def load_diffusion_model(unet_path, model_options={}): # TODO(sf): here load file into mem sd = comfy.utils.load_torch_file(unet_path) logging.info(f"load model start, path {unet_path}") - import pdb; pdb.set_trace() + # import pdb; pdb.set_trace() model = load_diffusion_model_state_dict(sd, model_options=model_options) logging.info(f"load model end, path {unet_path}") - import pdb; pdb.set_trace() + # import pdb; pdb.set_trace() if model is None: logging.error("ERROR UNSUPPORTED DIFFUSION MODEL {}".format(unet_path)) raise RuntimeError("ERROR: Could not detect model type of: {}\n{}".format(unet_path, model_detection_error_hint(unet_path, sd))) diff --git a/comfy/utils.py b/comfy/utils.py index a66402451..4c22f684c 100644 --- a/comfy/utils.py +++ b/comfy/utils.py @@ -102,7 +102,7 @@ def load_torch_file(ckpt, safe_load=False, device=None, return_metadata=False): sd = pl_sd else: sd = pl_sd - import pdb; pdb.set_trace() + # import pdb; pdb.set_trace() return (sd, metadata) if return_metadata else sd def save_torch_file(sd, ckpt, metadata=None):