feat: add tcd sampler and discrete distilled tcd scheduler based on sgm_uniform (same as lcm) (#2907)

2026-08-16 13:13:16 +02:00 · 2024-05-19 13:08:33 +02:00
parent 3bae73e23e
commit 2e2e8f851a
8 changed files with 65 additions and 13 deletions
@@ -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