From de8ab6f1c933242207daee119f6ff1ac8cda4845 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 31 Dec 2023 07:04:35 -0600
Subject: [PATCH] fix(api): use correct part of the gfpgan output

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

diff --git a/api/onnx_web/chain/correct_gfpgan.py b/api/onnx_web/chain/correct_gfpgan.py
index e1db8bcb..f3ce33f3 100644
--- a/api/onnx_web/chain/correct_gfpgan.py
+++ b/api/onnx_web/chain/correct_gfpgan.py
@@ -73,15 +73,15 @@ class CorrectGFPGANStage(BaseStage):
         device = worker.get_device()
         gfpgan = self.load(server, stage, upscale, device)
 
-        outputs = [
-            gfpgan.enhance(
+        outputs = []
+        for source in sources.as_numpy():
+            cropped, restored, result = gfpgan.enhance(
                 source,
                 has_aligned=False,
                 only_center_face=False,
                 paste_back=True,
                 weight=upscale.face_strength,
             )
-            for source in sources.as_numpy()
-        ]
+            outputs.append(result)
 
-        return StageResult(images=outputs)
+        return StageResult.from_arrays(outputs)