fix: correct sampling when gamma is 0

comfyanonymous · Jun 3, 2024 · 6fd91a8 · 6fd91a8
1 parent 9b8f97c
commit 6fd91a8
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/comfy/k_diffusion/sampling.py b/comfy/k_diffusion/sampling.py
@@ -841,5 +841,7 @@ def sample_tcd(
         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