From 663af7c1c8bb5aa8ccd6bb1634b51f7ed7a32740 Mon Sep 17 00:00:00 2001
From: Brian Drupieski <brian.drupieski@gmail.com>
Date: Sat, 17 Sep 2022 13:57:01 -0400
Subject: [PATCH 1/4] fix batch color correction

---
 scripts/batch.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/scripts/batch.py b/scripts/batch.py
index 1af4a7b..fc9e6b5 100644
--- a/scripts/batch.py
+++ b/scripts/batch.py
@@ -6,7 +6,7 @@ import traceback
 import modules.scripts as scripts
 import gradio as gr
 
-from modules.processing import Processed, process_images
+from modules.processing import Processed, process_images, setup_color_correction
 from PIL import Image
 from modules.shared import opts, cmd_opts, state
 
@@ -51,6 +51,10 @@ class Script(scripts.Script):
 
             state.job = f"{batch_no} out of {batch_count}: {batch_images[0][1]}"
             p.init_images = [x[0] for x in batch_images]
+
+            if opts.img2img_color_correction:
+                p.color_corrections = [setup_color_correction(i) for i in p.init_images]
+
             proc = process_images(p)
             for image, (_, path) in zip(proc.images, batch_images):
                 filename = os.path.basename(path)

From 7bc5739fe27809802ed24d1833b0ef443ec64aa9 Mon Sep 17 00:00:00 2001
From: Brian Drupieski <brian.drupieski@gmail.com>
Date: Sat, 17 Sep 2022 18:18:30 -0400
Subject: [PATCH 2/4] set loopback color corrections on each iteration

---
 modules/processing.py | 2 ++
 scripts/batch.py      | 5 +----
 scripts/loopback.py   | 4 ++--
 3 files changed, 5 insertions(+), 6 deletions(-)

diff --git a/modules/processing.py b/modules/processing.py
index 6a99d38..a60dec2 100644
--- a/modules/processing.py
+++ b/modules/processing.py
@@ -339,6 +339,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
 
             state.nextjob()
 
+        p.color_corrections = None
+
         unwanted_grid_because_of_img_count = len(output_images) < 2 and opts.grid_only_if_multiple
         if (opts.return_grid or opts.grid_save) and not p.do_not_save_grid and not unwanted_grid_because_of_img_count:
             grid = images.image_grid(output_images, p.batch_size)
diff --git a/scripts/batch.py b/scripts/batch.py
index fc9e6b5..a117fcb 100644
--- a/scripts/batch.py
+++ b/scripts/batch.py
@@ -6,7 +6,7 @@ import traceback
 import modules.scripts as scripts
 import gradio as gr
 
-from modules.processing import Processed, process_images, setup_color_correction
+from modules.processing import Processed, process_images
 from PIL import Image
 from modules.shared import opts, cmd_opts, state
 
@@ -52,9 +52,6 @@ class Script(scripts.Script):
             state.job = f"{batch_no} out of {batch_count}: {batch_images[0][1]}"
             p.init_images = [x[0] for x in batch_images]
 
-            if opts.img2img_color_correction:
-                p.color_corrections = [setup_color_correction(i) for i in p.init_images]
-
             proc = process_images(p)
             for image, (_, path) in zip(proc.images, batch_images):
                 filename = os.path.basename(path)
diff --git a/scripts/loopback.py b/scripts/loopback.py
index 8aca61f..2a3f4d6 100644
--- a/scripts/loopback.py
+++ b/scripts/loopback.py
@@ -40,8 +40,7 @@ class Script(scripts.Script):
         all_images = []
         state.job_count = loops * batch_count
 
-        if opts.img2img_color_correction:
-            p.color_corrections = [processing.setup_color_correction(p.init_images[0])]
+        initial_color_corrections = [processing.setup_color_correction(p.init_images[0])]
 
         for n in range(batch_count):
             history = []
@@ -50,6 +49,7 @@ class Script(scripts.Script):
                 p.n_iter = 1
                 p.batch_size = 1
                 p.do_not_save_grid = True
+                p.color_corrections = initial_color_corrections
 
                 state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}"
 

From cf70528bceba6f636a2c96ceedc985f662877da3 Mon Sep 17 00:00:00 2001
From: Brian Drupieski <brian.drupieski@gmail.com>
Date: Sat, 17 Sep 2022 18:19:21 -0400
Subject: [PATCH 3/4] remove whitespace

---
 scripts/batch.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/scripts/batch.py b/scripts/batch.py
index a117fcb..1af4a7b 100644
--- a/scripts/batch.py
+++ b/scripts/batch.py
@@ -51,7 +51,6 @@ class Script(scripts.Script):
 
             state.job = f"{batch_no} out of {batch_count}: {batch_images[0][1]}"
             p.init_images = [x[0] for x in batch_images]
-
             proc = process_images(p)
             for image, (_, path) in zip(proc.images, batch_images):
                 filename = os.path.basename(path)

From 2dbc9c5903cd2788f4ef21fe8ebfdd8e62f642cb Mon Sep 17 00:00:00 2001
From: Brian Drupieski <brian.drupieski@gmail.com>
Date: Sat, 17 Sep 2022 18:20:43 -0400
Subject: [PATCH 4/4] only set loopback color corrections if corrections
 enabled

---
 scripts/loopback.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/scripts/loopback.py b/scripts/loopback.py
index 2a3f4d6..e90b58d 100644
--- a/scripts/loopback.py
+++ b/scripts/loopback.py
@@ -49,7 +49,9 @@ class Script(scripts.Script):
                 p.n_iter = 1
                 p.batch_size = 1
                 p.do_not_save_grid = True
-                p.color_corrections = initial_color_corrections
+
+                if opts.img2img_color_correction:
+                    p.color_corrections = initial_color_corrections
 
                 state.job = f"Iteration {i + 1}/{loops}, batch {n + 1}/{batch_count}"