Add LCMScheduler and tweak settings for SamplerLCMAlternative

__init__.py

@@ -1,3 +1,3 @@
-from sampler_lcm_alt import NODE_CLASS_MAPPINGS
+from .sampler_lcm_alt import NODE_CLASS_MAPPINGS
 
 __all__ = ['NODE_CLASS_MAPPINGS']

sampler_lcm_alt.py

@@ -4,20 +4,25 @@ from tqdm.auto import trange, tqdm
 import torch
 
 @torch.no_grad()
-def sample_lcm_alt(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, sigma_limit=0.115):
+def sample_lcm_alt(model, x, sigmas, extra_args=None, callback=None, disable=None, noise_sampler=None, euler_steps=-3, ancestral=0.0):
     extra_args = {} if extra_args is None else extra_args
     noise_sampler = default_noise_sampler(x) if noise_sampler is None else noise_sampler
     s_in = x.new_ones([x.shape[0]])
-    sigma_min = float(model.inner_model.inner_model.model_sampling.sigma_min)
-    sigma_max = float(model.inner_model.inner_model.model_sampling.sigma_max)
-    sigma_limit = sigma_min + sigma_limit * (sigma_max - sigma_min)
+    euler_limit = euler_steps%(len(sigmas)-1)
     for i in trange(len(sigmas) - 1, disable=disable):
         denoised = model(x, sigmas[i] * s_in, **extra_args)
         if callback is not None:
             callback({'x': x, 'i': i, 'sigma': sigmas[i], 'sigma_hat': sigmas[i], 'denoised': denoised})
 
-        if sigmas[i+1] > sigma_limit:
-            noise = (x - denoised) / sigmas[i]
+        if i < euler_limit:
+            if ancestral < 1.0:
+                removed_noise = (x - denoised) / sigmas[i]
+            if ancestral > 0.0:
+                noise = noise_sampler(sigmas[i], sigmas[i + 1])
+                if ancestral < 1.0:
+                    noise = (ancestral**0.5) * noise + ((1.0 - ancestral)**0.5) * removed_noise
+            elif ancestral == 0.0:
+                noise = removed_noise
         elif sigmas[i + 1] > 0:
             noise = noise_sampler(sigmas[i], sigmas[i + 1])
         else:
@@ -27,23 +32,42 @@ def sample_lcm_alt(model, x, sigmas, extra_args=None, callback=None, disable=Non
             x += sigmas[i + 1] * noise
     return x
 
+class LCMScheduler:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {"required":
+                    {"model": ("MODEL",),
+                     "steps": ("INT", {"default": 8, "min": 1, "max": 10000}),
+                      }
+               }
+    RETURN_TYPES = ("SIGMAS",)
+    CATEGORY = "sampling/custom_sampling/schedulers"
+
+    FUNCTION = "get_sigmas"
+
+    def get_sigmas(self, model, scheduler, steps, step_to_multiply, multiplier):
+        sigmas = comfy.samplers.calculate_sigmas_scheduler(model.model, "sgm_uniform", steps).cpu()
+        return (sigmas, )
+
 class SamplerLCMAlternative:
     @classmethod
     def INPUT_TYPES(s):
         return {"required":
-                    {"sigma_limit": ("FLOAT", {"default": 0.115, "min": 0.0, "max": 1.0, "step":0.001, "round": False}),
-                      }
+                    {"euler_steps": ("INT", {"default": 0, "min": -10000, "max": 10000}),
+                     "ancestral": ("FLOAT", {"default": 0, "min": 0, "max": 1.0, "step": 0.01, "round": False}),
+                    }
                }
     RETURN_TYPES = ("SAMPLER",)
     CATEGORY = "sampling/custom_sampling/samplers"
 
     FUNCTION = "get_sampler"
 
-    def get_sampler(self, sigma_limit):
-        sampler = comfy.samplers.KSAMPLER(sample_lcm_alt, extra_options={"sigma_limit": sigma_limit})
+    def get_sampler(self, euler_steps, ancestral):
+        sampler = comfy.samplers.KSAMPLER(sample_lcm_alt, extra_options={"euler_steps": euler_steps, "ancestral": ancestral})
         return (sampler, )
 
 
 NODE_CLASS_MAPPINGS = {
+    "LCMAlternativeScheduler": LCMAlternativeScheduler,
     "SamplerLCMAlternative": SamplerLCMAlternative,
 }