Update model_loader_utils.py

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)