From 52e071da2a04acfd19cf8f6e69e006bd59937447 Mon Sep 17 00:00:00 2001
From: rewbs <robin@soal.org>
Date: Thu, 8 Sep 2022 02:35:26 +0000
Subject: 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 | 30 ++++++++++++++++++++++++++++--
 1 file changed, 28 insertions(+), 2 deletions(-)

(limited to 'modules/img2img.py')

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
-
-            p.init_images = [processed.images[0]]
+            
+            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),
+                        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])
-- 
cgit v1.2.3


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

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

(limited to 'modules/img2img.py')

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),
-- 
cgit v1.2.3