fix up highres metadata

api/onnx_web/chain/blend_img2img.py

@@ -102,7 +102,9 @@ class BlendImg2ImgStage(BaseStage):
 
             outputs.extend(result.images)
 
-        metadata = [metadata.child(params) for metadata in sources.metadata]
+        metadata = [
+            metadata.child(params, metadata.size) for metadata in sources.metadata
+        ]
         return StageResult(images=outputs, metadata=metadata)
 
     def steps(

api/onnx_web/chain/edit_metadata.py

@@ -0,0 +1,39 @@
+from typing import Any
+
+from ..params import HighresParams, ImageParams, Size, StageParams, UpscaleParams
+from ..server import ServerContext
+from ..worker import WorkerContext
+from .result import StageResult
+
+
+class EditMetadataStage:
+    def run(
+        self,
+        _worker: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        source: StageResult,
+        *,
+        size: Size = None,
+        upscale: UpscaleParams = None,
+        highres: HighresParams = None,
+        note: str = None,
+        **kwargs,
+    ) -> Any:
+        # Modify the source image's metadata using the provided parameters
+        for metadata in source.metadata:
+            if note is not None:
+                metadata.note = note
+
+            if size is not None:
+                metadata.size = size
+
+            if upscale is not None:
+                metadata.upscale = upscale
+
+            if highres is not None:
+                metadata.highres = highres
+
+        # Return the modified source image
+        return source

api/onnx_web/chain/highres.py

@@ -2,6 +2,7 @@ from logging import getLogger
 from typing import Optional
 
 from ..chain.blend_img2img import BlendImg2ImgStage
+from ..chain.edit_metadata import EditMetadataStage
 from ..chain.upscale import stage_upscale_correction
 from ..chain.upscale_simple import UpscaleSimpleStage
 from ..params import HighresParams, ImageParams, StageParams, UpscaleParams
@@ -66,4 +67,11 @@ def stage_highres(
             strength=highres.strength,
         )
 
+    # add highres parameters to the image metadata
+    chain.stage(
+        EditMetadataStage(),
+        stage.with_args(outscale=1),
+        highres=highres,
+    )
+
     return chain

api/onnx_web/chain/result.py

@@ -28,6 +28,7 @@ class NetworkMetadata:
 
 class ImageMetadata:
     ancestors: List["ImageMetadata"]
+    note: str
     params: ImageParams
     size: Size
 
@@ -70,6 +71,7 @@ class ImageMetadata:
         self.loras = loras or []
         self.models = models or []
         self.ancestors = ancestors or []
+        self.note = ""
 
     def child(
         self,

api/onnx_web/chain/stages.py

@@ -9,6 +9,7 @@ from .blend_linear import BlendLinearStage
 from .blend_mask import BlendMaskStage
 from .correct_codeformer import CorrectCodeformerStage
 from .correct_gfpgan import CorrectGFPGANStage
+from .edit_metadata import EditMetadataStage
 from .persist_disk import PersistDiskStage
 from .persist_s3 import PersistS3Stage
 from .reduce_crop import ReduceCropStage
@@ -38,6 +39,7 @@ CHAIN_STAGES = {
     "blend-mask": BlendMaskStage,
     "correct-codeformer": CorrectCodeformerStage,
     "correct-gfpgan": CorrectGFPGANStage,
+    "edit-metadata": EditMetadataStage,
     "persist-disk": PersistDiskStage,
     "persist-s3": PersistS3Stage,
     "reduce-crop": ReduceCropStage,