From b58bc7774e0d9f858ba52cd06b597f9f31db3793 Mon Sep 17 00:00:00 2001
From: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>
Date: Tue, 4 Jun 2024 21:03:37 +0200
Subject: [PATCH] fix: correct sampling when gamma is 0 (#3093)

---
 ldm_patched/k_diffusion/sampling.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/ldm_patched/k_diffusion/sampling.py b/ldm_patched/k_diffusion/sampling.py
index d1bc1e4b..ea5540a4 100644
--- a/ldm_patched/k_diffusion/sampling.py
+++ b/ldm_patched/k_diffusion/sampling.py
@@ -832,5 +832,7 @@ def sample_tcd(model, x, sigmas, extra_args=None, callback=None, disable=None, n
         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
\ No newline at end of file