Custom node for ComfyUI to make interrupting sampling more responsive

comfyui_fast_terminate.py

# Place in custom_nodes as comfyui_fast_terminate.py
# Will add a ModelPatchFastTerminate node.

from comfy.model_management import throw_exception_if_processing_interrupted

class ModelPatchFastTerminateNode:
  RETURN_TYPES = ("MODEL",)
  FUNCTION = "go"
  CATEGORY = "hacks"
  DESCRIPTION = "Patches a model to check if processing is interrupted at the start of every block. May make interrupting generations more responsive. Only works for models that support patching like SD15 and SDXL."

  @classmethod
  def INPUT_TYPES(cls):
    return { "required": { "model": ("MODEL",) } }

  @classmethod
  def go(cls, model):
    m = model.clone()

    def input_block_patch(h, _transformer_options):
      throw_exception_if_processing_interrupted()
      return h

    def output_block_patch(h, hsp, _transformer_options):
      throw_exception_if_processing_interrupted()
      return h, hsp

    m.set_model_input_block_patch(input_block_patch)
    m.set_model_output_block_patch(output_block_patch)

    return (m,)

NODE_CLASS_MAPPINGS = {
  "ModelPatchFastTerminate": ModelPatchFastTerminateNode,
}