From 569c971f1f61f009c3fc3f73b00d9f58a99357b3 Mon Sep 17 00:00:00 2001 From: tazlin Date: Sun, 25 Aug 2024 09:24:47 -0400 Subject: [PATCH] feat: expose `interrupt_current_processing` from comfyui I am hoping to one day implement a graceful way to interrupt generations in comfyui, such as from a user-initiated abort. I suspect that exposing this function will get us a some percent of the way to accomplishing that goal. See also https://github.com/Haidra-Org/horde-worker-reGen/issues/262. --- hordelib/comfy_horde.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/hordelib/comfy_horde.py b/hordelib/comfy_horde.py index 4266e41..bf7b95f 100644 --- a/hordelib/comfy_horde.py +++ b/hordelib/comfy_horde.py @@ -96,6 +96,8 @@ _comfy_is_changed_cache_get: Callable _comfy_model_patcher_load: Callable +_comfy_interrupt_current_processing: types.FunctionType + _canny: types.ModuleType _hed: types.ModuleType _leres: types.ModuleType @@ -158,6 +160,8 @@ def do_comfy_import( global _comfy_free_memory, _comfy_cleanup_models, _comfy_soft_empty_cache global _canny, _hed, _leres, _midas, _mlsd, _openpose, _pidinet, _uniformer + global _comfy_interrupt_current_processing + if disable_smart_memory: logger.info("Disabling smart memory") sys.argv.append("--disable-smart-memory") @@ -206,6 +210,7 @@ def do_comfy_import( from comfy.model_management import free_memory as _comfy_free_memory from comfy.model_management import cleanup_models as _comfy_cleanup_models from comfy.model_management import soft_empty_cache as _comfy_soft_empty_cache + from comfy.model_management import interrupt_current_processing as _comfy_interrupt_current_processing from comfy.utils import load_torch_file as _comfy_load_torch_file from comfy_extras.chainner_models import model_loading as _comfy_model_loading @@ -379,6 +384,11 @@ def log_free_ram(): ) +def interrupt_comfyui_processing(): + logger.warning("Interrupting comfyui processing") + _comfy_interrupt_current_processing() + + class Comfy_Horde: """Handles horde-specific behavior against ComfyUI."""