From 112d5a8876899b44c3d01f26ec1f797b0fefad23 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Thu, 4 Jan 2024 09:01:19 -0600
Subject: [PATCH] start passing metadata through stages

---
 api/onnx_web/chain/result.py          |  3 ++-
 api/onnx_web/chain/upscale_resrgan.py |  2 +-
 api/onnx_web/server/api.py            |  2 +-
 gui/src/types/api-v2.ts               | 37 +++++++++++++++++++--------
 4 files changed, 30 insertions(+), 14 deletions(-)

diff --git a/api/onnx_web/chain/result.py b/api/onnx_web/chain/result.py
index ce51badc..541fb3e7 100644
--- a/api/onnx_web/chain/result.py
+++ b/api/onnx_web/chain/result.py
@@ -205,6 +205,7 @@ class StageResult:
         self,
         arrays: Optional[List[np.ndarray]] = None,
         images: Optional[List[Image.Image]] = None,
+        metadata: Optional[List[ImageMetadata]] = None,
         source: Optional[Any] = None,
     ) -> None:
         if sum([arrays is not None, images is not None, source is not None]) > 1:
@@ -215,7 +216,7 @@ class StageResult:
         self.arrays = arrays
         self.images = images
         self.source = source
-        self.metadata = []
+        self.metadata = metadata or []
 
     def __len__(self) -> int:
         if self.arrays is not None:
diff --git a/api/onnx_web/chain/upscale_resrgan.py b/api/onnx_web/chain/upscale_resrgan.py
index 33338b43..f3d48195 100644
--- a/api/onnx_web/chain/upscale_resrgan.py
+++ b/api/onnx_web/chain/upscale_resrgan.py
@@ -117,4 +117,4 @@ class UpscaleRealESRGANStage(BaseStage):
             logger.info("final output image size: %s", output.shape)
             outputs.append(output)
 
-        return StageResult(arrays=outputs)
+        return StageResult(arrays=outputs, metadata=sources.metadata)
diff --git a/api/onnx_web/server/api.py b/api/onnx_web/server/api.py
index bc8cb5e4..266b35dc 100644
--- a/api/onnx_web/server/api.py
+++ b/api/onnx_web/server/api.py
@@ -129,7 +129,7 @@ def image_reply(
     }
 
     if outputs is not None:
-        if metadata is None:
+        if metadata is not None:
             logger.error("metadata is required with outputs")
             return error_reply("metadata is required with outputs")
 
diff --git a/gui/src/types/api-v2.ts b/gui/src/types/api-v2.ts
index 578c1294..9d25f920 100644
--- a/gui/src/types/api-v2.ts
+++ b/gui/src/types/api-v2.ts
@@ -70,8 +70,12 @@ export interface CancelledJobResponse extends BaseJobResponse {
   reason?: string;
 }
 
-export interface UnknownJobResponse extends BaseJobResponse {
-  status: JobStatus.UNKNOWN;
+/**
+ * Failed image job with error information.
+ */
+export interface FailedJobResponse extends BaseJobResponse {
+  status: JobStatus.FAILED;
+  error: string;
 }
 
 /**
@@ -86,14 +90,6 @@ export interface RunningJobResponse extends BaseJobResponse {
   status: JobStatus.RUNNING;
 }
 
-/**
- * Failed image job with error information.
- */
-export interface FailedJobResponse extends BaseJobResponse {
-  status: JobStatus.FAILED;
-  error: string;
-}
-
 /**
  * Successful txt2img image job with output keys and metadata.
  */
@@ -145,9 +141,25 @@ export interface SuccessBlendJobResponse extends BaseJobResponse {
 export interface SuccessChainJobResponse extends BaseJobResponse {
   status: JobStatus.SUCCESS;
   outputs: Array<string>;
-  metadata: Array<ImageMetadata<BaseImgParams, JobType>>; // TODO: could be all kinds
+  metadata: Array<
+  ImageMetadata<Txt2ImgParams, JobType.TXT2IMG>
+  | ImageMetadata<Img2ImgParams, JobType.IMG2IMG>
+  | ImageMetadata<InpaintParams, JobType.INPAINT>
+  | ImageMetadata<BaseImgParams, JobType.UPSCALE>
+  | ImageMetadata<BaseImgParams, JobType.BLEND>
+  >;
 }
 
+/**
+ * Unknown job type with no additional information.
+ */
+export interface UnknownJobResponse extends BaseJobResponse {
+  status: JobStatus.UNKNOWN;
+}
+
+/**
+ * All successful job types.
+ */
 export type SuccessJobResponse
   = SuccessTxt2ImgJobResponse
   | SuccessImg2ImgJobResponse
@@ -156,6 +168,9 @@ export type SuccessJobResponse
   | SuccessBlendJobResponse
   | SuccessChainJobResponse;
 
+/**
+ * All job types.
+ */
 export type JobResponse
   = CancelledJobResponse
   | PendingJobResponse