1 files changed, 11 insertions, 5 deletions
diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index fc0c94b4..dff89c09 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -4,7 +4,6 @@ import torch
 import tqdm
 from PIL import Image
 import inspect
-
 import k_diffusion.sampling
 import ldm.models.diffusion.ddim
 import ldm.models.diffusion.plms
@@ -23,6 +22,8 @@ samplers_k_diffusion = [
     ('Heun', 'sample_heun', ['k_heun']),
     ('DPM2', 'sample_dpm_2', ['k_dpm_2']),
     ('DPM2 a', 'sample_dpm_2_ancestral', ['k_dpm_2_a']),
+    ('DPM fast', 'sample_dpm_fast', ['k_dpm_fast']),
+    ('DPM adaptive', 'sample_dpm_adaptive', ['k_dpm_ad']),
 ]
 
 samplers_data_k_diffusion = [
@@ -36,7 +37,7 @@ samplers = [
     SamplerData('DDIM', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.ddim.DDIMSampler, model), []),
     SamplerData('PLMS', lambda model: VanillaStableDiffusionSampler(ldm.models.diffusion.plms.PLMSSampler, model), []),
 ]
-samplers_for_img2img = [x for x in samplers if x.name != 'PLMS']
+samplers_for_img2img = [x for x in samplers if x.name not in ['PLMS', 'DPM fast', 'DPM adaptive']]
 
 sampler_extra_params = {
     'sample_euler': ['s_churn', 's_tmin', 's_tmax', 's_noise'],
@@ -309,8 +310,13 @@ class KDiffusionSampler:
         x = x * sigmas[0]
 
         extra_params_kwargs = self.initialize(p)
-
-        samples = self.func(self.model_wrap_cfg, x, sigmas, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs)
-
+        if 'sigma_min' in inspect.signature(self.func).parameters:
+            extra_params_kwargs['sigma_min'] = self.model_wrap.sigmas[0].item()
+            extra_params_kwargs['sigma_max'] = self.model_wrap.sigmas[-1].item()
+            if 'n' in inspect.signature(self.func).parameters:
+                extra_params_kwargs['n'] = steps
+        else:
+            extra_params_kwargs['sigmas'] = sigmas
+        samples = self.func(self.model_wrap_cfg, x, extra_args={'cond': conditioning, 'uncond': unconditional_conditioning, 'cond_scale': p.cfg_scale}, disable=False, callback=self.callback_state, **extra_params_kwargs)
         return samples