From e306633f20d2209b83670b0477bff421684b9cf8 Mon Sep 17 00:00:00 2001 From: psychedelicious <4822129+psychedelicious@users.noreply.github.com> Date: Tue, 29 Nov 2022 08:08:28 +1100 Subject: [PATCH 1/2] Converts ESRGAN image input to RGB - Also adds typing for image input. - Partially resolves #1604 --- ldm/invoke/restoration/realesrgan.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/ldm/invoke/restoration/realesrgan.py b/ldm/invoke/restoration/realesrgan.py index 97c2de33fea..0a45eec69e0 100644 --- a/ldm/invoke/restoration/realesrgan.py +++ b/ldm/invoke/restoration/realesrgan.py @@ -5,7 +5,7 @@ from ldm.invoke.globals import Globals from PIL import Image - +from PIL.Image import Image as ImageType class ESRGAN(): def __init__(self, bg_tile_size=400) -> None: @@ -41,7 +41,7 @@ def load_esrgan_bg_upsampler(self): return bg_upsampler - def process(self, image, strength: float, seed: str = None, upsampler_scale: int = 2): + def process(self, image: ImageType, strength: float, seed: str = None, upsampler_scale: int = 2): with warnings.catch_warnings(): warnings.filterwarnings('ignore', category=DeprecationWarning) warnings.filterwarnings('ignore', category=UserWarning) @@ -62,7 +62,9 @@ def process(self, image, strength: float, seed: str = None, upsampler_scale: int print( f'>> Real-ESRGAN Upscaling seed:{seed} : scale:{upsampler_scale}x' ) - + # ESRGAN outputs images with partial transparency if given RGBA images; convert to RGB + image = image.convert("RGB") + # REALSRGAN expects a BGR np array; make array and flip channels bgr_image_array = np.array(image, dtype=np.uint8)[...,::-1] From 81e0c9c9c3fe2a5d16f918ddc023633ae1b962d1 Mon Sep 17 00:00:00 2001 From: Kyle Schouviller Date: Mon, 28 Nov 2022 21:17:30 -0800 Subject: [PATCH 2/2] ensure there are unmasked pixels before color matching --- ldm/invoke/generator/base.py | 21 ++++++++++++--------- 1 file changed, 12 insertions(+), 9 deletions(-) diff --git a/ldm/invoke/generator/base.py b/ldm/invoke/generator/base.py index a8ef62822d0..ce199fefb15 100644 --- a/ldm/invoke/generator/base.py +++ b/ldm/invoke/generator/base.py @@ -141,15 +141,18 @@ def repaste_and_color_correct(self, result: Image.Image, init_image: Image.Image np_init_rgb_pixels_masked = init_rgb_pixels[mask_pixels, :] np_image_masked = np_image[mask_pixels, :] - init_means = np_init_rgb_pixels_masked.mean(axis=0) - init_std = np_init_rgb_pixels_masked.std(axis=0) - gen_means = np_image_masked.mean(axis=0) - gen_std = np_image_masked.std(axis=0) - - # Color correct - np_matched_result = np_image.copy() - np_matched_result[:,:,:] = (((np_matched_result[:,:,:].astype(np.float32) - gen_means[None,None,:]) / gen_std[None,None,:]) * init_std[None,None,:] + init_means[None,None,:]).clip(0, 255).astype(np.uint8) - matched_result = Image.fromarray(np_matched_result, mode='RGB') + if np_init_rgb_pixels_masked.size > 0: + init_means = np_init_rgb_pixels_masked.mean(axis=0) + init_std = np_init_rgb_pixels_masked.std(axis=0) + gen_means = np_image_masked.mean(axis=0) + gen_std = np_image_masked.std(axis=0) + + # Color correct + np_matched_result = np_image.copy() + np_matched_result[:,:,:] = (((np_matched_result[:,:,:].astype(np.float32) - gen_means[None,None,:]) / gen_std[None,None,:]) * init_std[None,None,:] + init_means[None,None,:]).clip(0, 255).astype(np.uint8) + matched_result = Image.fromarray(np_matched_result, mode='RGB') + else: + matched_result = Image.fromarray(np_image, mode='RGB') # Blur the mask out (into init image) by specified amount if mask_blur_radius > 0: