From b98acca503385b04caa668a37eb0902aaf541a72 Mon Sep 17 00:00:00 2001
From: lllyasviel <lyuminzhang@outlook.com>
Date: Wed, 31 Jan 2024 10:17:12 -0800
Subject: [PATCH] move recolor

---
 .../preprocessor_compiled.py                  | 88 +++++++++----------
 .../scripts/preprocessor_recolor.py           | 49 +++++++++++
 2 files changed, 93 insertions(+), 44 deletions(-)
 create mode 100644 extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py

diff --git a/extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor_compiled.py b/extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor_compiled.py
index 4277dbd8..8cfcd231 100644
--- a/extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor_compiled.py
+++ b/extensions-builtin/forge_legacy_preprocessors/legacy_preprocessors/preprocessor_compiled.py
@@ -750,50 +750,50 @@ legacy_preprocessors = {
             "OpenPose"
         ]
     },
-    "recolor_intensity": {
-        "label": "recolor_intensity",
-        "call_function": recolor_intensity,
-        "unload_function": None,
-        "managed_model": None,
-        "model_free": False,
-        "no_control_mode": False,
-        "resolution": None,
-        "slider_1": {
-            "label": "Gamma Correction",
-            "value": 1.0,
-            "minimum": 0.1,
-            "maximum": 2.0,
-            "step": 0.001
-        },
-        "slider_2": None,
-        "slider_3": None,
-        "priority": 0,
-        "tags": [
-            "Recolor"
-        ]
-    },
-    "recolor_luminance": {
-        "label": "recolor_luminance",
-        "call_function": recolor_luminance,
-        "unload_function": None,
-        "managed_model": None,
-        "model_free": False,
-        "no_control_mode": False,
-        "resolution": None,
-        "slider_1": {
-            "label": "Gamma Correction",
-            "value": 1.0,
-            "minimum": 0.1,
-            "maximum": 2.0,
-            "step": 0.001
-        },
-        "slider_2": None,
-        "slider_3": None,
-        "priority": 100,
-        "tags": [
-            "Recolor"
-        ]
-    },
+    # "recolor_intensity": {
+    #     "label": "recolor_intensity",
+    #     "call_function": recolor_intensity,
+    #     "unload_function": None,
+    #     "managed_model": None,
+    #     "model_free": False,
+    #     "no_control_mode": False,
+    #     "resolution": None,
+    #     "slider_1": {
+    #         "label": "Gamma Correction",
+    #         "value": 1.0,
+    #         "minimum": 0.1,
+    #         "maximum": 2.0,
+    #         "step": 0.001
+    #     },
+    #     "slider_2": None,
+    #     "slider_3": None,
+    #     "priority": 0,
+    #     "tags": [
+    #         "Recolor"
+    #     ]
+    # },
+    # "recolor_luminance": {
+    #     "label": "recolor_luminance",
+    #     "call_function": recolor_luminance,
+    #     "unload_function": None,
+    #     "managed_model": None,
+    #     "model_free": False,
+    #     "no_control_mode": False,
+    #     "resolution": None,
+    #     "slider_1": {
+    #         "label": "Gamma Correction",
+    #         "value": 1.0,
+    #         "minimum": 0.1,
+    #         "maximum": 2.0,
+    #         "step": 0.001
+    #     },
+    #     "slider_2": None,
+    #     "slider_3": None,
+    #     "priority": 100,
+    #     "tags": [
+    #         "Recolor"
+    #     ]
+    # },
     # "reference_adain": {
     #     "label": "reference_adain",
     #     "call_function": identity,
diff --git a/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py b/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py
new file mode 100644
index 00000000..ed15dbba
--- /dev/null
+++ b/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py
@@ -0,0 +1,49 @@
+import cv2
+import numpy as np
+
+from modules_forge.supported_preprocessor import Preprocessor, PreprocessorParameter
+from modules_forge.shared import add_supported_preprocessor
+
+
+class PreprocessorRecolor(Preprocessor):
+    def __init__(self, name, use_intensity):
+        super().__init__()
+        self.name = name
+        self.use_intensity = False
+        self.tags = ['Recolor']
+        self.model_filename_filters = ['color', 'recolor', 'grey', 'gray']
+        self.slider_resolution = PreprocessorParameter(visible=False)
+        self.slider_1 = PreprocessorParameter(
+            visible=True,
+            label="Gamma Correction",
+            value=1.0,
+            minimum=0.1,
+            maximum=2.0,
+            step=0.001
+        )
+
+    def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, **kwargs):
+        gamma = slider_1
+
+        if self.use_intensity:
+            result = cv2.cvtColor(input_image, cv2.COLOR_BGR2HSV)
+            result = result[:, :, 2].astype(np.float32) / 255.0
+        else:
+            result = cv2.cvtColor(input_image, cv2.COLOR_BGR2LAB)
+            result = result[:, :, 0].astype(np.float32) / 255.0
+
+        result = result ** gamma
+        result = (result * 255.0).clip(0, 255).astype(np.uint8)
+        result = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB)
+        return result
+
+
+add_supported_preprocessor(PreprocessorRecolor(
+    name="recolor_intensity",
+    use_intensity=True
+))
+
+add_supported_preprocessor(PreprocessorRecolor(
+    name="recolor_luminance",
+    use_intensity=False
+))