fix self conditioning in elucidated imagen, thanks to @qiyan98

lucidrains · Jan 14, 2023 · 4cf311f · 4cf311f
1 parent e077796
commit 4cf311f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
diff --git a/imagen_pytorch/elucidated_imagen.py b/imagen_pytorch/elucidated_imagen.py
@@ -483,7 +483,9 @@ def one_unet_sample(
 
                 # second order correction, if not the last timestep
 
-                if sigma_next != 0:
+                has_second_order_correction = sigma_next != 0
+
+                if has_second_order_correction:
                     self_cond = model_output if unet.self_cond else None
 
                     model_output_next = self.preconditioned_network_forward(
@@ -504,7 +506,7 @@ def one_unet_sample(
                     repaint_noise = torch.randn(shape, device = self.device)
                     images = images + (sigma - sigma_next) * repaint_noise
 
-                x_start = model_output  # save model output for self conditioning
+                x_start = model_output if not has_second_order_correction else model_output_next # save model output for self conditioning
 
         images = images.clamp(-1., 1.)
 

diff --git a/imagen_pytorch/version.py b/imagen_pytorch/version.py
@@ -1 +1 @@
-__version__ = '1.18.10'
+__version__ = '1.18.11'