Unload all models if there's an OOM error.

This commit is contained in:
comfyanonymous 2024-08-06 03:30:28 -04:00
parent c14ac98fed
commit de17a9755e
1 changed files with 5 additions and 0 deletions

View File

@ -188,6 +188,11 @@ def recursive_execute(server, prompt, outputs, current_item, extra_data, execute
"current_inputs": input_data_formatted, "current_inputs": input_data_formatted,
"current_outputs": output_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) return (False, error_details, ex)
executed.add(unique_id) executed.add(unique_id)