From 1926a9b7c97a0dc3be77211f2bb132cd64f151f0 Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Thu, 2 Jan 2025 20:37:24 +0800
Subject: [PATCH] Update model_loader_utils.py

---
 model_loader_utils.py | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/model_loader_utils.py b/model_loader_utils.py
index 0127a7b..2b48a25 100644
--- a/model_loader_utils.py
+++ b/model_loader_utils.py
@@ -886,6 +886,8 @@ def msdiffusion_main(image_1, image_2, prompts_dual, width, height, steps, seed,
         controlnet.load_state_dict(cn_state_dict, strict=False)
         controlnet.to(torch.float16)
         pipe = StableDiffusionXLControlNetPipeline.from_pipe(pipe, controlnet=controlnet)
+        del cn_state_dict
+        torch.cuda.empty_cache()
     
     if lora:
         if lora in lora_lightning_list:
@@ -902,7 +904,6 @@ def msdiffusion_main(image_1, image_2, prompts_dual, width, height, steps, seed,
     if device != "mps":
         pipe.enable_model_cpu_offload()
         
-    
     torch.cuda.empty_cache()
     # 预加载 ms
     photomaker_local_path = os.path.join(photomaker_dir, "ms_adapter.bin")
@@ -919,6 +920,15 @@ def msdiffusion_main(image_1, image_2, prompts_dual, width, height, steps, seed,
     image_processor = CLIPImageProcessor()
     image_encoder_type = "clip"
     image_encoder = clip_load(clip_vision)
+    from comfy.model_management import cleanup_models
+    try:
+        cleanup_models()
+    except:
+        try:
+            cleanup_models(keep_clone_weights_loaded=False)
+        except:
+            gc.collect()
+            torch.cuda.empty_cache()
     use_repo = False
     config_path = os.path.join(cur_path, "config", "config.json")
     image_encoder_config = OmegaConf.load(config_path)