feat: add tcd sampler

comfy/k_diffusion/sampling.py

@@ -808,3 +808,40 @@ def sample_heunpp2(model, x, sigmas, extra_args=None, callback=None, disable=Non
             d_prime = w1 * d + w2 * d_2 + w3 * d_3
             x = x + d_prime * dt
     return x
+
+@torch.no_grad()
+def sample_tcd(
+    model,
+    x,
+    sigmas,
+    extra_args=None,
+    callback=None,
+    disable=None,
+    noise_sampler=None,
+    eta=0.3,
+):
+    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]])
+
+    model_sampling = model.inner_model.model_patcher.get_model_object("model_sampling")
+    timesteps_s = torch.floor((1 - eta) * model_sampling.timestep(sigmas)).to(dtype=torch.long).detach().cpu()
+    timesteps_s[-1] = 0
+    alpha_prod_s = model_sampling.alphas_cumprod[timesteps_s]
+    beta_prod_s = 1 - alpha_prod_s
+    for i in trange(len(sigmas) - 1, disable=disable):
+        denoised = model(x, sigmas[i] * s_in, **extra_args)  # predicted_original_sample
+        eps = (x - denoised) / sigmas[i]
+        denoised = alpha_prod_s[i + 1].sqrt() * denoised + beta_prod_s[i + 1].sqrt() * eps
+
+        if callback is not None:
+            callback({"x": x, "i": i, "sigma": sigmas[i], "sigma_hat": sigmas[i], "denoised": denoised})
+
+        x = denoised
+        if eta > 0 and sigmas[i + 1] > 0:
+            noise = noise_sampler(sigmas[i], sigmas[i + 1])
+            x = x / alpha_prod_s[i+1].sqrt() + noise * (sigmas[i+1]**2 + 1 - 1/alpha_prod_s[i+1]).sqrt()
+        else:
+            x *= torch.sqrt(1.0 + sigmas[i + 1] ** 2)
+
+    return x

comfy/model_sampling.py

@@ -64,12 +64,10 @@ def _register_schedule(self, given_betas=None, beta_schedule="linear", timesteps
         self.linear_start = linear_start
         self.linear_end = linear_end
 
-        # self.register_buffer('betas', torch.tensor(betas, dtype=torch.float32))
-        # self.register_buffer('alphas_cumprod', torch.tensor(alphas_cumprod, dtype=torch.float32))
-        # self.register_buffer('alphas_cumprod_prev', torch.tensor(alphas_cumprod_prev, dtype=torch.float32))
-
         sigmas = ((1 - alphas_cumprod) / alphas_cumprod) ** 0.5
         self.set_sigmas(sigmas)
+        # register alphas_cumprod for some sampler, such as tcd.
+        self.register_buffer("alphas_cumprod", alphas_cumprod.float())
 
     def set_sigmas(self, sigmas):
         self.register_buffer('sigmas', sigmas.float())

comfy/samplers.py

@@ -506,7 +506,7 @@ def max_denoise(self, model_wrap, sigmas):
 
 KSAMPLER_NAMES = ["euler", "euler_ancestral", "heun", "heunpp2","dpm_2", "dpm_2_ancestral",
                   "lms", "dpm_fast", "dpm_adaptive", "dpmpp_2s_ancestral", "dpmpp_sde", "dpmpp_sde_gpu",
-                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm"]
+                  "dpmpp_2m", "dpmpp_2m_sde", "dpmpp_2m_sde_gpu", "dpmpp_3m_sde", "dpmpp_3m_sde_gpu", "ddpm", "lcm", "tcd"]
 
 class KSAMPLER(Sampler):
     def __init__(self, sampler_function, extra_options={}, inpaint_options={}):

comfy_extras/nodes_custom_sampler.py

@@ -315,6 +315,23 @@ def get_sampler(self, order, rtol, atol, h_init, pcoeff, icoeff, dcoeff, accept_
                                                               "s_noise":s_noise })
         return (sampler, )
 
+class SamplerTCD:
+    @classmethod
+    def INPUT_TYPES(s):
+        return {
+            "required": {
+                "eta": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}),
+            }
+        }
+    RETURN_TYPES = ("SAMPLER",)
+    CATEGORY = "sampling/custom_sampling/samplers"
+
+    FUNCTION = "get_sampler"
+
+    def get_sampler(self, eta=0.3):
+        sampler = comfy.samplers.ksampler("tcd", {"eta": eta})
+        return (sampler, )
+
 class Noise_EmptyNoise:
     def __init__(self):
         self.seed = 0
@@ -599,6 +616,7 @@ def add_noise(self, model, noise, sigmas, latent_image):
     "SamplerDPMPP_2M_SDE": SamplerDPMPP_2M_SDE,
     "SamplerDPMPP_SDE": SamplerDPMPP_SDE,
     "SamplerDPMAdaptative": SamplerDPMAdaptative,
+    "SamplerTCD": SamplerTCD,
     "SplitSigmas": SplitSigmas,
     "FlipSigmas": FlipSigmas,
 

comfy_extras/nodes_model_advanced.py

@@ -73,7 +73,7 @@ class ModelSamplingDiscrete:
     @classmethod
     def INPUT_TYPES(s):
         return {"required": { "model": ("MODEL",),
-                              "sampling": (["eps", "v_prediction", "lcm", "x0"],),
+                              "sampling": (["eps", "v_prediction", "lcm", "x0", "tcd"],),
                               "zsnr": ("BOOLEAN", {"default": False}),
                               }}
 
@@ -95,6 +95,9 @@ def patch(self, model, sampling, zsnr):
             sampling_base = ModelSamplingDiscreteDistilled
         elif sampling == "x0":
             sampling_type = X0
+        elif sampling == "tcd":
+            sampling_type = comfy.model_sampling.EPS
+            sampling_base= ModelSamplingDiscreteDistilled
 
         class ModelSamplingAdvanced(sampling_base, sampling_type):
             pass