From 0b64633584e95e94ea2023bbe5640250f9e23984 Mon Sep 17 00:00:00 2001
From: AUTOMATIC1111 <16777216c@gmail.com>
Date: Fri, 28 Jun 2024 09:23:41 +0300
Subject: [PATCH] fix img2img

---
 modules/models/sd3/sd3_model.py   | 3 +++
 modules/processing.py             | 6 +++---
 modules/sd_samplers_kdiffusion.py | 5 ++++-
 3 files changed, 10 insertions(+), 4 deletions(-)

diff --git a/modules/models/sd3/sd3_model.py b/modules/models/sd3/sd3_model.py
index f8a4c96cc..c17fd97e9 100644
--- a/modules/models/sd3/sd3_model.py
+++ b/modules/models/sd3/sd3_model.py
@@ -73,3 +73,6 @@ class SD3Inferencer(torch.nn.Module):
             (self, 'cond_stage_model'),
             (self, 'model'),
         ]
+
+    def add_noise_to_latent(self, x, noise, amount):
+        return x * (1 - amount) + noise * amount
diff --git a/modules/processing.py b/modules/processing.py
index d32a1811e..c3ce975ea 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -1737,10 +1737,10 @@ class StableDiffusionProcessingImg2Img(StableDiffusionProcessing):
             latmask = latmask[0]
             if self.mask_round:
                 latmask = np.around(latmask)
-            latmask = np.tile(latmask[None], (4, 1, 1))
+            latmask = np.tile(latmask[None], (self.init_latent.shape[1], 1, 1))
 
-            self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(self.sd_model.dtype)
-            self.nmask = torch.asarray(latmask).to(shared.device).type(self.sd_model.dtype)
+            self.mask = torch.asarray(1.0 - latmask).to(shared.device).type(devices.dtype)
+            self.nmask = torch.asarray(latmask).to(shared.device).type(devices.dtype)
 
             # this needs to be fixed to be done in sample() using actual seeds for batches
             if self.inpainting_fill == 2:
diff --git a/modules/sd_samplers_kdiffusion.py b/modules/sd_samplers_kdiffusion.py
index cede0760a..8398299f8 100644
--- a/modules/sd_samplers_kdiffusion.py
+++ b/modules/sd_samplers_kdiffusion.py
@@ -133,7 +133,10 @@ class KDiffusionSampler(sd_samplers_common.Sampler):
         sigmas = self.get_sigmas(p, steps)
         sigma_sched = sigmas[steps - t_enc - 1:]
 
-        xi = x + noise * sigma_sched[0]
+        if hasattr(shared.sd_model, 'add_noise_to_latent'):
+            xi = shared.sd_model.add_noise_to_latent(x, noise, sigma_sched[0])
+        else:
+            xi = x + noise * sigma_sched[0]
 
         if opts.img2img_extra_noise > 0:
             p.extra_generation_params["Extra noise"] = opts.img2img_extra_noise