Skip to content

Commit

Permalink
feat: expose interrupt_current_processing from comfyui
Browse files Browse the repository at this point in the history
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 Haidra-Org/horde-worker-reGen#262.
  • Loading branch information
tazlin committed Aug 26, 2024
1 parent cdcd32c commit 569c971
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions hordelib/comfy_horde.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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")
Expand Down Expand Up @@ -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

Expand Down Expand Up @@ -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."""

Expand Down

0 comments on commit 569c971

Please sign in to comment.