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.

2 files changed, 29 insertions, 2 deletions

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])

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