From de17a9755ecb8419cb167fe8504791df5b07246f Mon Sep 17 00:00:00 2001
From: comfyanonymous <comfyanonymous@protonmail.com>
Date: Tue, 6 Aug 2024 03:30:28 -0400
Subject: [PATCH] Unload all models if there's an OOM error.

---
 execution.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/execution.py b/execution.py
index 0a2e62e7..d207e1b9 100644
--- a/execution.py
+++ b/execution.py
@@ -188,6 +188,11 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
             "current_inputs": input_data_formatted,
             "current_outputs": output_data_formatted
         }
+
+        if isinstance(ex, comfy.model_management.OOM_EXCEPTION):
+            logging.error("Got an OOM, unloading all loaded models.")
+            comfy.model_management.unload_all_models()
+
         return (False, error_details, ex)
 
     executed.add(unique_id)