wip: add tcd sampler

adapted code from https://github.com/comfyanonymous/ComfyUI/pull/3370
TODO: check if virtual scheduler tcd is needed for using sampling_base ModelSamplingDiscreteDistilled or if it's better to use sgm_uniform directly without patching

ldm_patched/modules/model_sampling.py

@@ -50,17 +50,17 @@ class ModelSamplingDiscrete(torch.nn.Module):
         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)
+        self.set_alphas_cumprod(alphas_cumprod.float())
 
     def set_sigmas(self, sigmas):
         self.register_buffer('sigmas', sigmas)
         self.register_buffer('log_sigmas', sigmas.log())
 
+    def set_alphas_cumprod(self, alphas_cumprod):
+        self.register_buffer("alphas_cumprod", alphas_cumprod.float())
+
     @property
     def sigma_min(self):
         return self.sigmas[0]

ldm_patched/modules/samplers.py

@@ -523,7 +523,7 @@ class UNIPCBH2(Sampler):
 
 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={}):