From f7b171f198bc70b0b6d7c3a974e265baae28bf44 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 14 Jan 2024 09:41:54 -0600
Subject: [PATCH] fix highres output size

---
 api/onnx_web/chain/highres.py |  3 ++-
 api/onnx_web/chain/result.py  | 23 +++++++++++++++++------
 2 files changed, 19 insertions(+), 7 deletions(-)

diff --git a/api/onnx_web/chain/highres.py b/api/onnx_web/chain/highres.py
index 7e47ff1f..2318223d 100644
--- a/api/onnx_web/chain/highres.py
+++ b/api/onnx_web/chain/highres.py
@@ -67,11 +67,12 @@ def stage_highres(
             strength=highres.strength,
         )
 
-    # add highres parameters to the image metadata
+    # add highres parameters to the image metadata and clear upscale stage
     chain.stage(
         EditMetadataStage(),
         stage.with_args(outscale=1, tile_size=SizeChart.max),
         highres=highres,
+        upscale=UpscaleParams(upscale.upscale_model, upscale=False),
         replace_params=params,
     )
 
diff --git a/api/onnx_web/chain/result.py b/api/onnx_web/chain/result.py
index 4e357e57..2ef5ff4d 100644
--- a/api/onnx_web/chain/result.py
+++ b/api/onnx_web/chain/result.py
@@ -202,22 +202,20 @@ class ImageMetadata:
             }
         )
 
-        # calculate final output size
-        output_size = self.size
+        # add optional params
         if self.border is not None:
             json["border"] = self.border.tojson()
-            output_size = output_size.add_border(self.border)
 
         if self.highres is not None:
             json["highres"] = self.highres.tojson()
-            output_size = self.highres.resize(output_size)
 
         if self.upscale is not None:
             json["upscale"] = self.upscale.tojson()
-            output_size = self.upscale.resize(output_size)
 
-        json["size"] = output_size.tojson()
+        # calculate final output size
+        json["size"] = self.get_output_size().tojson()
 
+        # hash and add models and networks
         if self.inversions is not None:
             for name, weight in self.inversions:
                 model_hash = self.get_network_hash(server, name, "inversion")[1]
@@ -241,6 +239,19 @@ class ImageMetadata:
 
         return json
 
+    def get_output_size(self) -> Size:
+        output_size = self.size
+        if self.border is not None:
+            output_size = output_size.add_border(self.border)
+
+        if self.highres is not None:
+            output_size = self.highres.resize(output_size)
+
+        if self.upscale is not None:
+            output_size = self.upscale.resize(output_size)
+
+        return output_size
+
     @staticmethod
     def from_exif(input: str) -> "ImageMetadata":
         lines = input.splitlines()