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
This commit is contained in:
Manuel Schmid
2024-05-12 19:42:28 +02:00
parent 6308fb8b54
commit 77acf8126a
8 changed files with 65 additions and 13 deletions
+4 -4
View File
@@ -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]