From 6a32c06f0646805456ac2744ece289962a2379c6 Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Sat, 23 Mar 2024 17:27:10 -0400
Subject: [PATCH] Move cleanup_models to improve performance.

---
 execution.py | 1 -
 main.py      | 1 +
 2 files changed, 1 insertion(+), 1 deletion(-)

diff --git a/execution.py b/execution.py
index 5f8ea835..1b8f606a 100644
--- a/execution.py
+++ b/execution.py
@@ -368,7 +368,6 @@ class PromptExecutor:
                     d = self.outputs_ui.pop(x)
                     del d
 
-            comfy.model_management.cleanup_models()
             self.add_message("execution_cached",
                           { "nodes": list(current_outputs) , "prompt_id": prompt_id},
                           broadcast=False)
diff --git a/main.py b/main.py
index 3dee72e3..b3a3ebea 100644
--- a/main.py
+++ b/main.py
@@ -139,6 +139,7 @@ def prompt_worker(q, server):
         if need_gc:
             current_time = time.perf_counter()
             if (current_time - last_gc_collect) > gc_collect_interval:
+                comfy.model_management.cleanup_models()
                 gc.collect()
                 comfy.model_management.soft_empty_cache()
                 last_gc_collect = current_time