Allow custom samplers to request discard penultimate sigma

comfyanonymous · Feb 8, 2024 · a352c02 · a352c02
1 parent fd73b5e
commit a352c02
Showing 1 changed file with 2 additions and 1 deletion.
diff --git 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 @@ def calculate_sigmas(self, steps):
         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