From 75a9b1bbd92f6f4e5fdddd0e9a74e39905679d3a Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Mon, 12 Sep 2022 18:00:53 +0300
Subject: [PATCH] Noisy image previews #257

---
 modules/sd_samplers.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index 6b7979e2..95d24299 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -58,12 +58,14 @@ def p_sample_ddim_hook(sampler_wrapper, x_dec, cond, ts, *args, **kwargs):
         img_orig = sampler_wrapper.sampler.model.q_sample(sampler_wrapper.init_latent, ts)
         x_dec = img_orig * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec
 
-        store_latent(sampler_wrapper.init_latent * sampler_wrapper.mask + sampler_wrapper.nmask * x_dec)
+    res = sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs)
 
+    if sampler_wrapper.mask is not None:
+        store_latent(sampler_wrapper.init_latent * sampler_wrapper.mask + sampler_wrapper.nmask * res[1])
     else:
-        store_latent(x_dec)
+        store_latent(res[1])
 
-    return sampler_wrapper.orig_p_sample_ddim(x_dec, cond, ts, *args, **kwargs)
+    return res
 
 
 def extended_tdqm(sequence, *args, desc=None, **kwargs):