Skip to content

GitLab

Commit a352c021 authored by blepping's avatar blepping
Browse files

Allow custom samplers to request discard penultimate sigma

parent fd73b5ee
Hide whitespace changes
Inline Side-by-side
Showing with 2 additions and 1 deletion
comfy/samplers.py
View file @ a352c021
...@@ -652,6 +652,7 @@ def sampler_object(name): ...@@ -652,6 +652,7 @@ def sampler_object(name):
class KSampler: class KSampler:
SCHEDULERS = SCHEDULER_NAMES SCHEDULERS = SCHEDULER_NAMES
SAMPLERS = SAMPLER_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={}): def __init__(self, model, steps, device, sampler=None, scheduler=None, denoise=None, model_options={}):
self.model = model self.model = model
...@@ -670,7 +671,7 @@ class KSampler: ...@@ -670,7 +671,7 @@ class KSampler:
sigmas = None sigmas = None
discard_penultimate_sigma = False 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 steps += 1
discard_penultimate_sigma = True discard_penultimate_sigma = True
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment