fix(api): use correct part of the gfpgan output

ssube · Dec 31, 2023 · de8ab6f · de8ab6f
1 parent c6718a3
commit de8ab6f
Showing 1 changed file with 5 additions and 5 deletions.
diff --git a/api/onnx_web/chain/correct_gfpgan.py b/api/onnx_web/chain/correct_gfpgan.py
@@ -73,15 +73,15 @@ def run(
         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)