Corrects color channels in face restoration; Fixes #1167

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

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