From 0aa79685035559e513ee54bb95febedf2a6cee39 Mon Sep 17 00:00:00 2001
From: cmdr2 <secondary.cmdr2@gmail.com>
Date: Sun, 11 Dec 2022 19:34:07 +0530
Subject: [PATCH] Move color correction to diffusionkit; Rename color
 correction to 'Preserve color profile'

---
 ui/index.html              |  2 +-
 ui/media/js/main.js        |  6 +++---
 ui/sd_internal/runtime2.py | 18 ++----------------
 3 files changed, 6 insertions(+), 20 deletions(-)

diff --git a/ui/index.html b/ui/index.html
index 30648424..0273b574 100644
--- a/ui/index.html
+++ b/ui/index.html
@@ -213,7 +213,7 @@
                     <div><ul>
                     <li><b class="settings-subheader">Render Settings</b></li>
                     <li class="pl-5"><input id="stream_image_progress" name="stream_image_progress" type="checkbox"> <label for="stream_image_progress">Show a live preview <small>(uses more VRAM, slower images)</small></label></li>
-                    <li id="apply_color_correction_setting" class="pl-5"><input id="apply_color_correction" name="apply_color_correction" type="checkbox"> <label for="apply_color_correction">Apply color correction <small>(helps during inpainting)</small></label></li>
+                    <li id="apply_color_correction_setting" class="pl-5"><input id="apply_color_correction" name="apply_color_correction" type="checkbox"> <label for="apply_color_correction">Preserve color profile <small>(helps during inpainting)</small></label></li>
                     <li class="pl-5"><input id="use_face_correction" name="use_face_correction" type="checkbox"> <label for="use_face_correction">Fix incorrect faces and eyes <small>(uses GFPGAN)</small></label></li>
                     <li class="pl-5">
                         <input id="use_upscale" name="use_upscale" type="checkbox"> <label for="use_upscale">Upscale image by 4x with </label>
diff --git a/ui/media/js/main.js b/ui/media/js/main.js
index e058f9e9..6e00f8c0 100644
--- a/ui/media/js/main.js
+++ b/ui/media/js/main.js
@@ -761,8 +761,8 @@ function createTask(task) {
         taskConfig += `, <b>Hypernetwork:</b> ${task.reqBody.use_hypernetwork_model}`
         taskConfig += `, <b>Hypernetwork Strength:</b> ${task.reqBody.hypernetwork_strength}`
     }
-    if (task.reqBody.apply_color_correction) {
-        taskConfig += `, <b>Color Correction:</b> true`
+    if (task.reqBody.preserve_init_image_color_profile) {
+        taskConfig += `, <b>Preserve Color Profile:</b> true`
     }
 
     let taskEntry = document.createElement('div')
@@ -871,7 +871,7 @@ function getCurrentUserRequest() {
         if (maskSetting.checked) {
             newTask.reqBody.mask = imageInpainter.getImg()
         }
-        newTask.reqBody.apply_color_correction = applyColorCorrectionField.checked
+        newTask.reqBody.preserve_init_image_color_profile = applyColorCorrectionField.checked
         newTask.reqBody.sampler = 'ddim'
     } else {
         newTask.reqBody.sampler = samplerField.value
diff --git a/ui/sd_internal/runtime2.py b/ui/sd_internal/runtime2.py
index 434f489b..0e1efc2c 100644
--- a/ui/sd_internal/runtime2.py
+++ b/ui/sd_internal/runtime2.py
@@ -62,10 +62,8 @@ def reload_models_if_necessary(task_data: TaskData):
         if thread_data.model_paths.get(model_type) != model_path_in_req:
             thread_data.model_paths[model_type] = model_path_in_req
 
-            if thread_data.model_paths[model_type] is not None:
-                model_loader.load_model(thread_data, model_type)
-            else:
-                model_loader.unload_model(thread_data, model_type)
+            action_fn = model_loader.unload_model if thread_data.model_paths[model_type] is None else model_loader.load_model
+            action_fn(thread_data, model_type)
 
 def make_images(req: GenerateImageRequest, task_data: TaskData, data_queue: queue.Queue, task_temp_images: list, step_callback):
     try:
@@ -94,7 +92,6 @@ def _make_images_internal(req: GenerateImageRequest, task_data: TaskData, data_q
     print(metadata)
 
     images, user_stopped = generate_images(req, data_queue, task_temp_images, step_callback, task_data.stream_image_progress)
-    images = apply_color_correction(req, images, user_stopped)
     images = apply_filters(task_data, images, user_stopped, task_data.show_only_filtered_image)
 
     if task_data.save_to_disk_path is not None:
@@ -139,17 +136,6 @@ def generate_images(req: GenerateImageRequest, data_queue: queue.Queue, task_tem
 
     return images, user_stopped
 
-def apply_color_correction(req: GenerateImageRequest, images: list, user_stopped):
-    if user_stopped or req.init_image is None or not req.apply_color_correction:
-        return images
-
-    for i, img_info in enumerate(images):
-        img, seed, filtered = img_info
-        img = image_utils.apply_color_correction(orig_image=req.init_image, image_to_correct=img)
-        images[i] = (img, seed, filtered)
-
-    return images
-
 def apply_filters(task_data: TaskData, images: list, user_stopped, show_only_filtered_image):
     if user_stopped or (task_data.use_face_correction is None and task_data.use_upscale is None):
         return images