diff --git a/ldm/invoke/restoration/codeformer.py b/ldm/invoke/restoration/codeformer.py index 0d13ae0a36..8be2f34245 100644 --- a/ldm/invoke/restoration/codeformer.py +++ b/ldm/invoke/restoration/codeformer.py @@ -3,6 +3,7 @@ import torch import numpy as np import warnings import sys +import cv2 pretrained_model_url = 'https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/codeformer.pth' @@ -40,11 +41,12 @@ class CodeFormerRestoration(): cf.load_state_dict(checkpoint) cf.eval() - image = image.convert('RGB') + # Codeformer expects BGR image data + bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) face_helper = FaceRestoreHelper(upscale_factor=1, use_parse=True, device=device) face_helper.clean_all() - face_helper.read_image(np.array(image, dtype=np.uint8)) + face_helper.read_image(bgrImage) face_helper.get_face_landmarks_5(resize=640, eye_dist_threshold=5) face_helper.align_warp_face() @@ -71,7 +73,8 @@ class CodeFormerRestoration(): restored_img = face_helper.paste_faces_to_input_image() - res = Image.fromarray(restored_img) + # Convert back to RGB for PIL + res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)) if strength < 1.0: # Resize the image to the new image if the sizes have changed diff --git a/ldm/invoke/restoration/gfpgan.py b/ldm/invoke/restoration/gfpgan.py index 473d708961..d45338edc2 100644 --- a/ldm/invoke/restoration/gfpgan.py +++ b/ldm/invoke/restoration/gfpgan.py @@ -3,6 +3,7 @@ import warnings import os import sys import numpy as np +import cv2 from PIL import Image @@ -53,15 +54,18 @@ class GFPGAN(): f'>> Download https://github.com/TencentARC/GFPGAN/releases/download/v1.3.0/GFPGANv1.4.pth to {self.model_path}, \nor change GFPGAN directory with --gfpgan_dir.' ) - image = image.convert('RGB') + # GFPGAN expects BGR image data + bgrImage = cv2.cvtColor(np.array(image), cv2.COLOR_RGB2BGR) _, _, restored_img = self.gfpgan.enhance( - np.array(image, dtype=np.uint8), + bgrImage, has_aligned=False, only_center_face=False, paste_back=True, ) - res = Image.fromarray(restored_img) + + # Convert back to RGB for PIL + res = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB)) if strength < 1.0: # Resize the image to the new image if the sizes have changed