From b44ddcb44398fbe922fd7515f66d8b0c2344bc54 Mon Sep 17 00:00:00 2001
From: AUTOMATIC <16777216c@gmail.com>
Date: Fri, 16 Sep 2022 08:51:21 +0300
Subject: [PATCH]  Prompt editing only applies to images in first batch of
 desired batch size when batch count > 1 #535

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

diff --git a/modules/sd_samplers.py b/modules/sd_samplers.py
index df3a6fe8..5d95bfe0 100644
--- a/modules/sd_samplers.py
+++ b/modules/sd_samplers.py
@@ -115,6 +115,7 @@ class VanillaStableDiffusionSampler:
         self.mask = p.mask
         self.nmask = p.nmask
         self.init_latent = p.init_latent
+        self.step = 0
 
         samples = self.sampler.decode(x1, conditioning, t_enc, unconditional_guidance_scale=p.cfg_scale, unconditional_conditioning=unconditional_conditioning)
 
@@ -127,6 +128,7 @@ class VanillaStableDiffusionSampler:
         self.mask = None
         self.nmask = None
         self.init_latent = None
+        self.step = 0
 
         # existing code fails with cetin step counts, like 9
         try:
@@ -206,6 +208,7 @@ class KDiffusionSampler:
         self.model_wrap_cfg.mask = p.mask
         self.model_wrap_cfg.nmask = p.nmask
         self.model_wrap_cfg.init_latent = p.init_latent
+        self.model_wrap.step = 0
 
         if hasattr(k_diffusion.sampling, 'trange'):
             k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)
@@ -216,6 +219,8 @@ class KDiffusionSampler:
         sigmas = self.model_wrap.get_sigmas(p.steps)
         x = x * sigmas[0]
 
+        self.model_wrap_cfg.step = 0
+
         if hasattr(k_diffusion.sampling, 'trange'):
             k_diffusion.sampling.trange = lambda *args, **kwargs: extended_trange(*args, **kwargs)