execution_hook.py

# -*- coding: utf-8 -*-
"""
在 ComfyUI 执行任意工作流前检查：若非“仅 llama 工作流”，则先结束 llama-server，让出 CUDA 显存。
"""
import logging
from .llama_controller import is_llama_only_workflow, stop_llama_server


def _on_before_execute(prompt):
    """执行前回调：若 prompt 不是纯 llama 工作流，则停止 llama-server。"""
    if not is_llama_only_workflow(prompt):
        logging.info("[ComfyUI-LlamaServer] Non-llama workflow detected, stopping llama-server to free GPU.")
        stop_llama_server()


def install_execution_hook():
    """
    尝试挂载到 ComfyUI 的 PromptExecutor.execute_async，在真正执行前调用 _on_before_execute(prompt)。
    若当前 ComfyUI 版本结构不同则跳过，仅打日志。
    """
    try:
        import execution  # type: ignore  # ComfyUI core module, resolved at runtime
        executor_class = getattr(execution, "PromptExecutor", None)
        if executor_class is None:
            logging.warning("[ComfyUI-LlamaServer] execution.PromptExecutor not found, skip execution hook.")
            return
        original = getattr(executor_class, "execute_async", None)
        if not callable(original):
            logging.warning("[ComfyUI-LlamaServer] PromptExecutor.execute_async not found, skip execution hook.")
            return

        async def patched_execute_async(self, prompt, prompt_id, extra_data=None, execute_outputs=None):
            extra_data = extra_data or {}
            execute_outputs = execute_outputs or []
            _on_before_execute(prompt)
            return await original(self, prompt, prompt_id, extra_data, execute_outputs)

        executor_class.execute_async = patched_execute_async
        logging.info("[ComfyUI-LlamaServer] Execution hook installed: stop llama-server before non-llama workflows.")
    except Exception as e:
        logging.warning("[ComfyUI-LlamaServer] Could not install execution hook: %s", e)