try to fix #142

* retry to unload already merged model as intended * use send_to_trash() for removed model
wkpark · Jul 5, 2024 · 32d5947 · 32d5947
1 parent ccdda9a
commit 32d5947
Showing 1 changed file with 15 additions and 10 deletions.
diff --git a/scripts/model_mixer.py b/scripts/model_mixer.py
@@ -4464,17 +4464,22 @@ def fake_checkpoint(checkpoint_info, metadata, model_name, sha256, fake=True):
             sd_models.unload_model_weights()
             #sd_models.model_data.__init__()
 
+          if shared.sd_model is not None:
+            if partial_update:
+              if getattr(sd_models.model_data, "loaded_sd_models", None) is not None:
+                if len(sd_models.model_data.loaded_sd_models) > shared.opts.sd_checkpoints_limit > 1:
+                    try:
+                        print(" - unload current merged model from loaded_sd_models...")
+                        sd_models.model_data.loaded_sd_models.remove(shared.sd_model)
+                    except ValueError:
+                        print(f" - Failed to unload current merged model...")
+                    else:
+                        print(" - trash unloaded model...")
+                        sd_models.send_model_to_trash(shared.sd_model)
 
-          if partial_update:
-            if getattr(sd_models.model_data, "loaded_sd_models", None) is not None:
-                try:
-                    print(" - unload current merged model from loaded_sd_models...")
-                    sd_models.model_data.loaded_sd_models.remove(shared.sd_model)
-                except ValueError:
-                    print(f" - Failed to unload current merged model...")
-                    pass
-
-                sd_models.send_model_to_trash(shared.sd_model)
+              sd_models.model_data.sd_model = None
+              shared.sd_model = None
+              devices.torch_gc()
 
           if sd_models.model_data.sd_model:
             send_model_to_cpu(sd_models.model_data.sd_model)