From 248231d82b0743716c0f54569ab9aaabbe0ae28e Mon Sep 17 00:00:00 2001 From: lllyasviel Date: Wed, 31 Jan 2024 10:26:24 -0800 Subject: [PATCH] Update preprocessor_recolor.py --- .../scripts/preprocessor_recolor.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py b/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py index ed15dbba..019c7c26 100644 --- a/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py +++ b/extensions-builtin/forge_preprocessor_recolor/scripts/preprocessor_recolor.py @@ -21,6 +21,7 @@ class PreprocessorRecolor(Preprocessor): maximum=2.0, step=0.001 ) + self.current_cond = None def __call__(self, input_image, resolution, slider_1=None, slider_2=None, slider_3=None, **kwargs): gamma = slider_1 @@ -37,6 +38,23 @@ class PreprocessorRecolor(Preprocessor): result = cv2.cvtColor(result, cv2.COLOR_GRAY2RGB) return result + def process_before_every_sampling(self, process, cond, *args, **kwargs): + self.current_cond = cond + return + + def process_after_every_sampling(self, process, params, *args, **kwargs): + a1111_batch_result = args[0] + new_results = [] + + for img in a1111_batch_result.images: + new_mean = self.current_cond[0].mean(dim=0, keepdim=True) + img = img - img.mean(dim=0, keepdim=True) + new_mean + img = img.clip(0, 1) + new_results.append(img) + + a1111_batch_result.images = new_results + return + add_supported_preprocessor(PreprocessorRecolor( name="recolor_intensity",