From 0d1439734ccd5cc070b45e5b1a91b017e0314686 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Mon, 8 Jan 2024 22:54:41 -0600
Subject: [PATCH] fix output stage

---
 api/onnx_web/chain/persist_disk.py | 19 ++++---------------
 1 file changed, 4 insertions(+), 15 deletions(-)

diff --git a/api/onnx_web/chain/persist_disk.py b/api/onnx_web/chain/persist_disk.py
index 61d39eae..825fabf5 100644
--- a/api/onnx_web/chain/persist_disk.py
+++ b/api/onnx_web/chain/persist_disk.py
@@ -3,7 +3,7 @@ from typing import List, Optional
 
 from PIL import Image
 
-from ..output import save_image
+from ..output import save_image, save_result
 from ..params import ImageParams, Size, SizeChart, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
@@ -18,28 +18,17 @@ class PersistDiskStage(BaseStage):
 
     def run(
         self,
-        _worker: WorkerContext,
+        worker: WorkerContext,
         server: ServerContext,
         _stage: StageParams,
         params: ImageParams,
         sources: StageResult,
         *,
-        output: List[str],
-        size: Optional[Size] = None,
         stage_source: Optional[Image.Image] = None,
         **kwargs,
     ) -> StageResult:
-        logger.info("persisting %s images to disk: %s", len(sources), output)
+        logger.info("persisting %s images to disk", len(sources))
 
-        for name, source, metadata in zip(
-            output, sources.as_images(), sources.metadata
-        ):
-            dest = save_image(
-                server,
-                name,
-                source,
-                metadata=metadata,
-            )
-            logger.info("saved image to %s", dest)
+        save_result(server, sources, worker.job)
 
         return sources