From b312e78c96ee7858970a49abb677539b3cd66564 Mon Sep 17 00:00:00 2001
From: psychedelicious <4822129+psychedelicious@users.noreply.github.com>
Date: Thu, 20 Oct 2022 10:12:20 +0800
Subject: [PATCH] Corrects color channels in face restoration; Fixes #1167

---
 ldm/invoke/restoration/codeformer.py |  9 ++++++---
 ldm/invoke/restoration/gfpgan.py     | 10 +++++++---
 2 files changed, 13 insertions(+), 6 deletions(-)

diff --git a/ldm/invoke/restoration/codeformer.py b/ldm/invoke/restoration/codeformer.py
index 0d13ae0a36b..8be2f342458 100644
--- a/ldm/invoke/restoration/codeformer.py
+++ b/ldm/invoke/restoration/codeformer.py
@@ -3,6 +3,7 @@
 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 @@ def process(self, image, strength, device, seed=None, fidelity=0.75):
             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 @@ def process(self, image, strength, device, seed=None, fidelity=0.75):
 
             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 473d708961b..d45338edc23 100644
--- a/ldm/invoke/restoration/gfpgan.py
+++ b/ldm/invoke/restoration/gfpgan.py
@@ -3,6 +3,7 @@
 import os
 import sys
 import numpy as np
+import cv2
 
 from PIL import Image
 
@@ -53,15 +54,18 @@ def process(self, image, strength: float, seed: str = None):
                 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