From 39f84c4a49f48060cd5e3d0b04f6718712bac41e Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Tue, 31 Jan 2023 08:16:57 -0600
Subject: [PATCH] fix(api): convert images to ndarray before passing to GFPGAN

---
 api/onnx_web/chain/correct_gfpgan.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/api/onnx_web/chain/correct_gfpgan.py b/api/onnx_web/chain/correct_gfpgan.py
index 07bc318ea..d7117859a 100644
--- a/api/onnx_web/chain/correct_gfpgan.py
+++ b/api/onnx_web/chain/correct_gfpgan.py
@@ -17,6 +17,8 @@
     load_resrgan,
 )
 
+import numpy as np
+
 logger = getLogger(__name__)
 
 
@@ -56,7 +58,7 @@ def correct_gfpgan(
     ctx: ServerContext,
     _stage: StageParams,
     _params: ImageParams,
-    image: Image.Image,
+    source_image: Image.Image,
     *,
     upscale: UpscaleParams,
     upsampler: Optional[RealESRGANer] = None,
@@ -69,7 +71,9 @@ def correct_gfpgan(
     logger.info('correcting faces with GFPGAN model: %s', upscale.correction_model)
     gfpgan = load_gfpgan(ctx, upscale, upsampler=upsampler)
 
+    output = np.array(source_image)
     _, _, output = gfpgan.enhance(
         image, has_aligned=False, only_center_face=False, paste_back=True, weight=upscale.face_strength)
+    output = Image.fromarray(output, 'RGB')
 
     return output