From 39fb74c5bd13a1dccf4d7293a2f7a755d9f43cbd Mon Sep 17 00:00:00 2001 From: comfyanonymous Date: Tue, 13 Aug 2024 03:57:55 -0400 Subject: [PATCH] Fix bug when model cannot be partially unloaded. --- comfy/model_management.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/comfy/model_management.py b/comfy/model_management.py index 152a76f37d2..a6996709b6f 100644 --- 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