From be40797e590eb80dc13a4918dae46ef38f5bf803 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 6 Jan 2024 17:40:36 -0600
Subject: [PATCH] use txt2img params for highres output

---
 api/onnx_web/chain/edit_metadata.py | 17 +++++++++++------
 api/onnx_web/chain/highres.py       |  1 +
 2 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/api/onnx_web/chain/edit_metadata.py b/api/onnx_web/chain/edit_metadata.py
index 522ddb81..61aae34c 100644
--- a/api/onnx_web/chain/edit_metadata.py
+++ b/api/onnx_web/chain/edit_metadata.py
@@ -1,12 +1,13 @@
-from typing import Any
+from typing import Any, Optional
 
 from ..params import HighresParams, ImageParams, Size, StageParams, UpscaleParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .base import BaseStage
 from .result import StageResult
 
 
-class EditMetadataStage:
+class EditMetadataStage(BaseStage):
     def run(
         self,
         _worker: WorkerContext,
@@ -15,10 +16,11 @@ class EditMetadataStage:
         _params: ImageParams,
         source: StageResult,
         *,
-        size: Size = None,
-        upscale: UpscaleParams = None,
-        highres: HighresParams = None,
-        note: str = None,
+        size: Optional[Size] = None,
+        upscale: Optional[UpscaleParams] = None,
+        highres: Optional[HighresParams] = None,
+        note: Optional[str] = None,
+        replace_params: Optional[ImageParams] = None,
         **kwargs,
     ) -> Any:
         # Modify the source image's metadata using the provided parameters
@@ -26,6 +28,9 @@ class EditMetadataStage:
             if note is not None:
                 metadata.note = note
 
+            if replace_params is not None:
+                metadata.params = replace_params
+
             if size is not None:
                 metadata.size = size
 
diff --git a/api/onnx_web/chain/highres.py b/api/onnx_web/chain/highres.py
index 5e6fd22e..b962e503 100644
--- a/api/onnx_web/chain/highres.py
+++ b/api/onnx_web/chain/highres.py
@@ -72,6 +72,7 @@ def stage_highres(
         EditMetadataStage(),
         stage.with_args(outscale=1),
         highres=highres,
+        replace_params=params,
     )
 
     return chain