Prioritize freeing partially offloaded models first.

comfyanonymous · Sep 4, 2024 · c742737 · c742737
1 parent 22d1241
commit c742737
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/comfy/model_management.py b/comfy/model_management.py
@@ -426,7 +426,7 @@ def free_memory(memory_required, device, keep_loaded=[]):
         shift_model = current_loaded_models[i]
         if shift_model.device == device:
             if shift_model not in keep_loaded:
-                can_unload.append((sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
+                can_unload.append((-shift_model.model_offloaded_memory(), sys.getrefcount(shift_model.model), shift_model.model_memory(), i))
                 shift_model.currently_used = False
 
     for x in sorted(can_unload):