Fix bug when model cannot be partially unloaded.

comfyanonymous · Aug 13, 2024 · 39fb74c · 39fb74c
1 parent 74e124f
commit 39fb74c
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -338,8 +338,9 @@ def should_reload_model(self, force_patch_weights=False):
  def model_unload(self, memory_to_free=None, unpatch_weights=True):
  if memory_to_free is not None:
  if memory_to_free < self.model.loaded_size():
- self.model.partially_unload(self.model.offload_device, memory_to_free)
- return False
+ freed = self.model.partially_unload(self.model.offload_device, memory_to_free)
+ if freed >= memory_to_free:
+ return False
  self.model.unpatch_model(self.model.offload_device, unpatch_weights=unpatch_weights)
  self.model.model_patches_to(self.model.offload_device)
  self.weights_loaded = self.weights_loaded and not unpatch_weights