From 25eb1e3d9072a060a253abce6889e7fef440dfe1 Mon Sep 17 00:00:00 2001 From: Robin Fernandes Date: Thu, 22 Sep 2022 20:54:50 +1000 Subject: [PATCH 1/2] Add option to save before color correction. Add suffix param when saving files, used for special saves without color correction and face restoration. --- modules/images.py | 4 ++-- modules/processing.py | 4 +++- modules/shared.py | 1 + 3 files changed, 6 insertions(+), 3 deletions(-) diff --git a/modules/images.py b/modules/images.py index d1707263..f65f4553 100644 --- a/modules/images.py +++ b/modules/images.py @@ -311,7 +311,7 @@ def get_next_sequence_number(path, basename): return result + 1 -def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None): +def save_image(image, path, basename, seed=None, prompt=None, extension='png', info=None, short_filename=False, no_prompt=False, grid=False, pnginfo_section_name='parameters', p=None, existing_info=None, forced_filename=None, suffix=""): if short_filename or prompt is None or seed is None: file_decoration = "" elif opts.save_to_dirs: @@ -322,7 +322,7 @@ def save_image(image, path, basename, seed=None, prompt=None, extension='png', i if file_decoration != "": file_decoration = "-" + file_decoration.lower() - file_decoration = apply_filename_pattern(file_decoration, p, seed, prompt) + file_decoration = apply_filename_pattern(file_decoration, p, seed, prompt) + suffix if extension == 'png' and opts.enable_pnginfo and info is not None: pnginfo = PngImagePlugin.PngInfo() diff --git a/modules/processing.py b/modules/processing.py index f44c3f26..49fed0d8 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -357,7 +357,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: if p.restore_faces: if opts.save and not p.do_not_save_samples and opts.save_images_before_face_restoration: - images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p) + images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-face-restoration") devices.torch_gc() @@ -366,6 +366,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed: image = Image.fromarray(x_sample) if p.color_corrections is not None and i < len(p.color_corrections): + if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction: + images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction") image = apply_color_correction(p.color_corrections[i], image) if p.overlay_images is not None and i < len(p.overlay_images): diff --git a/modules/shared.py b/modules/shared.py index 841a9877..0b57685e 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -139,6 +139,7 @@ class Options: "enable_pnginfo": OptionInfo(True, "Save text information about generation parameters as chunks to png files"), "add_model_hash_to_info": OptionInfo(False, "Add model hash to generation information"), "img2img_color_correction": OptionInfo(False, "Apply color correction to img2img results to match original colors."), + "save_images_before_color_correction": OptionInfo(False, "Save a copy of image before applying color correction to img2img results"), "img2img_fix_steps": OptionInfo(False, "With img2img, do exactly the amount of steps the slider specifies (normally you'd do less with less denoising)."), "enable_quantization": OptionInfo(False, "Enable quantization in K samplers for sharper and cleaner results. This may change existing seeds. Requires restart to apply."), "font": OptionInfo("", "Font for image grids that have text"), From d26d89377bfb47779daec02e888c3c2c931be1f0 Mon Sep 17 00:00:00 2001 From: Robin Fernandes Date: Fri, 23 Sep 2022 00:57:42 +0000 Subject: [PATCH 2/2] Remove unnecessary duplication --- modules/processing.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/modules/processing.py b/modules/processing.py index 49fed0d8..d27d86e9 100644 --- a/modules/processing.py +++ b/modules/processing.py @@ -20,6 +20,7 @@ import modules.shared as shared import modules.face_restoration import modules.images as images import modules.styles +import logging # some of those options should not be changed at all because they would break the model, so I removed them from options. @@ -28,11 +29,13 @@ opt_f = 8 def setup_color_correction(image): + logging.info("Calibrating color correction.") correction_target = cv2.cvtColor(np.asarray(image.copy()), cv2.COLOR_RGB2LAB) return correction_target def apply_color_correction(correction, image): + logging.info("Applying color correction.") image = Image.fromarray(cv2.cvtColor(exposure.match_histograms( cv2.cvtColor( np.asarray(image), @@ -367,7 +370,7 @@ def process_images(p: StableDiffusionProcessing) -> Processed: if p.color_corrections is not None and i < len(p.color_corrections): if opts.save and not p.do_not_save_samples and opts.save_images_before_color_correction: - images.save_image(Image.fromarray(x_sample), p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction") + images.save_image(image, p.outpath_samples, "", seeds[i], prompts[i], opts.samples_format, info=infotext(n, i), p=p, suffix="-before-color-correction") image = apply_color_correction(p.color_corrections[i], image) if p.overlay_images is not None and i < len(p.overlay_images):