如何终止Hugging Face Pipeline运行？自定义StoppingCriteria无效求助

我尝试用自定义StoppingCriteria终止Hugging Face Diffusion Pipeline的运行，但没生效。我在return flag行设了断点，想调试时修改flag值来终止，但没用。我的代码如下：

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()

问题原因与修复方案 #

核心问题 #

你踩了两个关键坑：一是误用了transformers库的StoppingCriteriaList，但Diffusers的Pipeline需要用自身的StoppingCriteria体系；二是自定义停止条件必须是继承自diffusers.StoppingCriteria的类，不能是普通函数。另外全局变量flag在Pipeline的运行流程里，无法实时读取到你的调试修改。

修复后的代码 #

import torch
from diffusers import DiffusionPipeline, StoppingCriteria, 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"

class CustomStoppingCriteria(StoppingCriteria):
    def __init__(self):
        self.flag = False

    def __call__(self, step: int, timestep: int, latents: torch.FloatTensor) -> bool:
        # 在这里设置断点，调试时把self.flag改成True就能终止生成
        return self.flag

# 实例化自定义停止条件
stop_criteria = CustomStoppingCriteria()
# 传入Pipeline运行
image = pipe(prompt, stopping_criteria=StoppingCriteriaList([stop_criteria])).images[0]
image.show()

关键说明 #

• Diffusers的停止条件要求必须继承diffusers.StoppingCriteria，并且__call__方法的参数是step（当前步数）、timestep（当前时间步）、latents（隐变量），和transformers的参数完全不同，之前参数不对自然不会触发停止逻辑。
• 用类的成员变量self.flag代替全局变量，调试时直接修改实例的flag属性，能确保停止条件函数实时读取到最新值。
• 导入diffusers的StoppingCriteriaList来包装自定义条件，和Pipeline的要求完全匹配。

内容的提问来源于stack exchange，提问作者Intolighter