From 52e071da2a04acfd19cf8f6e69e006bd59937447 Mon Sep 17 00:00:00 2001
From: rewbs <robin@soal.org>
Date: Thu, 8 Sep 2022 02:35:26 +0000
Subject: [PATCH 1/3] Add color correction to img2img loopback to avoid a
 progressive skew to magenta. Based on codedealer's PR to hlky's repo here:
 https://github.com/sd-webui/stable-diffusion-webui/pull/698/files.

---
 modules/img2img.py | 28 +++++++++++++++++++++++++++-
 requirements.txt   |  1 +
 2 files changed, 28 insertions(+), 1 deletion(-)

diff --git a/modules/img2img.py b/modules/img2img.py
index 3129798d..2c74842d 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -1,4 +1,6 @@
 import math
+import cv2
+import numpy as np
 from PIL import Image
 
 from modules.processing import Processed, StableDiffusionProcessingImg2Img, process_images
@@ -57,8 +59,19 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
 
         state.job_count = n_iter
 
+        do_color_correction = False
+        try:
+            from skimage import exposure
+            do_color_correction = True
+        except:
+            print("Install scikit-image to perform color correction on loopback")
+
+
         for i in range(n_iter):
 
+            if do_color_correction and i == 0:
+                correction_target = cv2.cvtColor(np.asarray(init_img.copy()), cv2.COLOR_RGB2LAB)
+
             p.n_iter = 1
             p.batch_size = 1
             p.do_not_save_grid = True
@@ -69,8 +82,21 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
             if initial_seed is None:
                 initial_seed = processed.seed
                 initial_info = processed.info
+            
+            init_img = processed.images[0]
 
-            p.init_images = [processed.images[0]]
+            if do_color_correction and correction_target is not None:
+                print("Colour correcting input...")
+                init_img = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
+                    cv2.cvtColor(
+                        np.asarray(init_img),
+                        cv2.COLOR_RGB2LAB
+                    ),
+                    correction_target,
+                    channel_axis=2
+                ), cv2.COLOR_LAB2RGB).astype("uint8"))
+
+            p.init_images = [init_img]
             p.seed = processed.seed + 1
             p.denoising_strength = max(p.denoising_strength * 0.95, 0.1)
             history.append(processed.images[0])
diff --git a/requirements.txt b/requirements.txt
index c9e3f2fc..ba1bc281 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,5 +10,6 @@ omegaconf
 pytorch_lightning
 diffusers
 invisible-watermark
+scikit-image
 git+https://github.com/crowsonkb/k-diffusion.git
 git+https://github.com/TencentARC/GFPGAN.git

From 1e7a36fd79612d924f0aca18061f1b3bd947b02a Mon Sep 17 00:00:00 2001
From: rewbs <robin@soal.org>
Date: Thu, 8 Sep 2022 02:53:13 +0000
Subject: [PATCH 2/3] Remove debug print.

---
 modules/img2img.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/modules/img2img.py b/modules/img2img.py
index 2c74842d..52971785 100644
--- a/modules/img2img.py
+++ b/modules/img2img.py
@@ -86,7 +86,6 @@ def img2img(prompt: str, init_img, init_img_with_mask, steps: int, sampler_index
             init_img = processed.images[0]
 
             if do_color_correction and correction_target is not None:
-                print("Colour correcting input...")
                 init_img = Image.fromarray(cv2.cvtColor(exposure.match_histograms(
                     cv2.cvtColor(
                         np.asarray(init_img),

From bc12eddb408c3503717b234e1a8bb635049f4a91 Mon Sep 17 00:00:00 2001
From: rewbs <robin@soal.org>
Date: Thu, 8 Sep 2022 05:57:22 +0000
Subject: [PATCH 3/3] Add scikit-image dependency to requirements_versions.txt
 for windows users.

---
 requirements_versions.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements_versions.txt b/requirements_versions.txt
index 177c6b58..e8a7470c 100644
--- a/requirements_versions.txt
+++ b/requirements_versions.txt
@@ -8,3 +8,4 @@ torch
 transformers==4.19.2
 omegaconf==2.1.1
 pytorch_lightning==1.7.2
+scikit-image==0.19.2