How to Stop a Hugging Face Pipeline Operation (Diffusers & Transformers)

To stop a Hugging Face pipeline operation, you must use the correct interruption mechanism for the specific library you are using. The transformers pipeline class supports a stopping_criteria parameter, while the diffusers pipeline class requires setting the _interrupt property to True inside a callback function passed to callback_on_step_end. Attempting to use stopping_criteria with a diffusers pipeline will not work because the two libraries have separate interruption implementations.

The Full Answer

Hugging Face provides two major libraries for running machine learning models: transformers (for text, audio, and vision models) and diffusers (for diffusion models used in image generation). Each library has its own pipeline class and its own method for interrupting or stopping a running pipeline operation. Confusing the two is a common mistake that leads to code that appears correct but does nothing.

Understanding the Two Libraries

• transformers pipeline: Used for tasks like text classification, question answering, summarization, and text generation. It supports a stopping_criteria parameter that accepts a StoppingCriteriaList object. This list contains callable objects that are evaluated at each generation step. If any callable returns True, generation stops.
• diffusers pipeline: Used for image generation with diffusion models (e.g., Stable Diffusion, Kandinsky). It does not support stopping_criteria. Instead, it provides a callback_on_step_end parameter that accepts a function called after each denoising step. To interrupt generation, you set pipeline._interrupt = True inside that callback.

The Problem with Using stopping_criteria on a diffusers Pipeline

In the original question, the user attempted to stop a diffusers pipeline (specifically DiffusionPipeline from the diffusers library) by passing a stopping_criteria parameter. This approach fails because DiffusionPipeline does not inspect or use the stopping_criteria argument. The parameter is silently ignored. The debugger breakpoint on return flag never triggers because the custom stopping criteria function is never called.

Here is the incorrect code from the question:

import torch
from diffusers import DiffusionPipeline
from transformers import StoppingCriteriaList

pipe = DiffusionPipeline.from_pretrained(
    "kandinsky-community/kandinsky-3", variant="fp16", torch_dtype=torch.float16
)
pipe.enable_sequential_cpu_offload()
prompt = "Cemetery of abandoned vehicles, many different rusty cars"

flag = False

def custom_stopping_criteria(
    input_ids: torch.LongTensor, score: torch.FloatTensor, **kwargs
) -> bool:
    return flag

stopping_criteria = StoppingCriteriaList([custom_stopping_criteria])

image = pipe(prompt, stopping_criteria=stopping_criteria).images[0]
image.show()

This code will run the full 50 denoising steps (or whatever the default is for the model) without interruption, regardless of the flag variable.

The Correct Approach for diffusers Pipelines

As confirmed in the accepted answer on Stack Overflow, the diffusers library has its own interruption mechanism. You must use the callback_on_step_end parameter to provide a function that will be called at the end of each denoising step. Inside that function, you set pipeline._interrupt = True to stop the generation.

Here is the corrected code, adapted from the accepted answer:

import torch
from diffusers import DiffusionPipeline

# Using a smaller model for demonstration
model_id = "OFA-Sys/small-stable-diffusion-v0"

pipe = DiffusionPipeline.from_pretrained(
    model_id, torch_dtype=torch.float16
)
pipe.enable_sequential_cpu_offload()

prompt = "Cemetery of abandoned vehicles, many different rusty cars"

def interrupt_callback(pipeline, i, t, callback_kwargs):
    pipeline._interrupt = True
    return callback_kwargs

image = pipe(
    prompt,
    callback_on_step_end=interrupt_callback,
)

What this code does:

1. It loads a diffusion pipeline from the diffusers library. The example uses OFA-Sys/small-stable-diffusion-v0 for demonstration, but you can replace it with any model ID (e.g., kandinsky-community/kandinsky-3).
2. It enables sequential CPU offloading to save GPU memory (optional but recommended for large models).
3. It defines a callback function interrupt_callback that receives four arguments:
   • pipeline: the pipeline object itself.
   • i: the current step number (0-indexed).
   • t: the current timestep value. The exact meaning of t is not documented in the source, but it is passed by the pipeline.
   • callback_kwargs: a dictionary of keyword arguments that the pipeline passes to the callback. This must be returned from the callback to maintain state.
4. Inside the callback, it sets pipeline._interrupt = True. This is a special attribute that the pipeline checks after each step. If it is True, the pipeline stops the generation loop immediately.
5. The callback returns callback_kwargs to satisfy the pipeline's expectation.
6. The pipeline is called with callback_on_step_end=interrupt_callback. The generation will stop after the first step because the callback sets _interrupt to True on the first call.

Output from the accepted answer:

1/50 [00:01<01:21, 1.67s/it]
Potential NSFW content was detected in one or more images. A black image will be returned instead. Try again with a different prompt and/or seed.

This output shows that only 1 out of 50 steps was completed before the pipeline was interrupted. The NSFW warning is incidental to the example and not related to the interruption mechanism.

Making the Interruption Conditional

In practice, you will not want to interrupt the pipeline on the very first step. You will want to stop it based on some condition, such as user input, a timeout, or a quality metric. Here is how to implement a conditional interruption:

import torch
from diffusers import DiffusionPipeline
import threading
import time

model_id = "OFA-Sys/small-stable-diffusion-v0"

pipe = DiffusionPipeline.from_pretrained(
    model_id, torch_dtype=torch.float16
)
pipe.enable_sequential_cpu_offload()

prompt = "Cemetery of abandoned vehicles, many different rusty cars"

# A flag that can be set from another thread (e.g., a UI button)
stop_flag = False

def interrupt_callback(pipeline, i, t, callback_kwargs):
    if stop_flag:
        pipeline._interrupt = True
    return callback_kwargs

# Simulate a user pressing stop after 3 seconds
def delayed_stop():
    global stop_flag
    time.sleep(3)
    stop_flag = True

threading.Thread(target=delayed_stop, daemon=True).start()

image = pipe(
    prompt,
    callback_on_step_end=interrupt_callback,
)

This example uses a global stop_flag that can be set from any thread. The callback checks the flag at each step and sets _interrupt when the flag becomes True. The pipeline will stop at the end of the current step.

The Correct Approach for transformers Pipelines

If you are using a transformers pipeline (e.g., for text generation), the stopping_criteria parameter works as expected. Here is a brief example for completeness:

from transformers import pipeline, StoppingCriteriaList, StoppingCriteria
import torch

class StopOnKeyword(StoppingCriteria):
    def __init__(self, keyword):
        self.keyword = keyword

    def __call__(self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs) -> bool:
        # Check if the last generated token matches the keyword
        last_token = input_ids[0, -1].item()
        return last_token == self.keyword

pipe = pipeline("text-generation", model="gpt2")

# This will stop generation when the token for "stop" is generated
# Note: this is a simplified example; you would need the actual token ID
stopping_criteria = StoppingCriteriaList([StopOnKeyword(50256)])  # 50256 is the EOS token for GPT-2

output = pipe("Hello, I am", stopping_criteria=stopping_criteria)
print(output)

This is not directly related to the original question, but it clarifies the correct usage for transformers pipelines.