From a352c021eca1d39c3fc128a6ee65c0527f5c6f3c Mon Sep 17 00:00:00 2001
From: blepping <blepping@tuta.io>
Date: Thu, 8 Feb 2024 02:24:23 -0700
Subject: [PATCH] Allow custom samplers to request discard penultimate sigma

---
 comfy/samplers.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/comfy/samplers.py b/comfy/samplers.py
index f4c3e268f..5dd72f3fa 100644
--- a/comfy/samplers.py
+++ b/comfy/samplers.py
@@ -652,6 +652,7 @@ def sampler_object(name):
 class KSampler:
     SCHEDULERS = SCHEDULER_NAMES
     SAMPLERS = SAMPLER_NAMES
+    DISCARD_PENULTIMATE_SIGMA_SAMPLERS = set(('dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2'))
 
     def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
         self.model = model
@@ -670,7 +671,7 @@ class KSampler:
         sigmas = None
 
         discard_penultimate_sigma = False
-        if self.sampler in ['dpm_2', 'dpm_2_ancestral', 'uni_pc', 'uni_pc_bh2']:
+        if self.sampler in self.DISCARD_PENULTIMATE_SIGMA_SAMPLERS:
             steps += 1
             discard_penultimate_sigma = True