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

modules/async_worker.py

@@ -773,19 +773,19 @@ def worker():
         final_sampler_name = sampler_name
         final_scheduler_name = scheduler_name
 
-        if scheduler_name == 'lcm':
+        if scheduler_name in ['lcm', 'tcd']:
             final_scheduler_name = 'sgm_uniform'
             if pipeline.final_unet is not None:
                 pipeline.final_unet = core.opModelSamplingDiscrete.patch(
                     pipeline.final_unet,
-                    sampling='lcm',
+                    sampling=scheduler_name,
                     zsnr=False)[0]
             if pipeline.final_refiner_unet is not None:
                 pipeline.final_refiner_unet = core.opModelSamplingDiscrete.patch(
                     pipeline.final_refiner_unet,
-                    sampling='lcm',
+                    sampling=scheduler_name,
                     zsnr=False)[0]
-            print('Using lcm scheduler.')
+            print(f'Using {scheduler_name} scheduler.')
 
         async_task.yields.append(['preview', (13, 'Moving model to GPU ...', None)])
 

modules/flags.py

@@ -34,7 +34,8 @@ KSAMPLER = {
     "dpmpp_3m_sde": "",
     "dpmpp_3m_sde_gpu": "",
     "ddpm": "",
-    "lcm": "LCM"
+    "lcm": "LCM",
+    "tcd": "TCD"
 }
 
 SAMPLER_EXTRA = {
@@ -47,7 +48,7 @@ SAMPLERS = KSAMPLER | SAMPLER_EXTRA
 
 KSAMPLER_NAMES = list(KSAMPLER.keys())
 
-SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo"]
+SCHEDULER_NAMES = ["normal", "karras", "exponential", "sgm_uniform", "simple", "ddim_uniform", "lcm", "turbo", "tcd"]
 SAMPLER_NAMES = KSAMPLER_NAMES + list(SAMPLER_EXTRA.keys())
 
 sampler_list = SAMPLER_NAMES

modules/patch_precision.py

@@ -51,6 +51,8 @@ def patched_register_schedule(self, given_betas=None, beta_schedule="linear", ti
     self.linear_end = linear_end
     sigmas = torch.tensor(((1 - alphas_cumprod) / alphas_cumprod) ** 0.5, dtype=torch.float32)
     self.set_sigmas(sigmas)
+    alphas_cumprod = torch.tensor(alphas_cumprod, dtype=torch.float32)
+    self.set_alphas_cumprod(alphas_cumprod)
     return