From 1f6ab7dbfb4c1e214a33d9b000ad5682a357791e Mon Sep 17 00:00:00 2001
From: bigcat88 <bigcat88@icloud.com>
Date: Tue, 24 Dec 2024 17:41:49 +0200
Subject: [PATCH] support for "unload_models" flag when creating a task

Signed-off-by: bigcat88 <bigcat88@icloud.com>
---
 main.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/main.py b/main.py
index ccc99fdc..591d813b 100644
--- a/main.py
+++ b/main.py
@@ -166,6 +166,14 @@ def prompt_worker(q, server_instance):
         queue_item = q.get(timeout=timeout)
         if queue_item is not None:
             item, item_id = queue_item
+
+            if item[3].get("unload_models"):
+                # For those cases where the flag is set, to clear memory before execution
+                comfy.model_management.unload_all_models()
+                gc.collect()
+                comfy.model_management.soft_empty_cache()
+                last_gc_collect = time.perf_counter()
+
             execution_start_time = time.perf_counter()
             prompt_id = item[1]
             server_instance.last_prompt_id = prompt_id