From c1f1bf2b36425d7615b97e370aeb2ed29ac6055e Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 19 Nov 2023 21:05:21 -0600
Subject: [PATCH] various stack fixes

---
 api/onnx_web/chain/pipeline.py | 4 ++--
 api/onnx_web/chain/tile.py     | 7 ++++---
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/api/onnx_web/chain/pipeline.py b/api/onnx_web/chain/pipeline.py
index 50122d7d..359e2bfb 100644
--- a/api/onnx_web/chain/pipeline.py
+++ b/api/onnx_web/chain/pipeline.py
@@ -181,7 +181,7 @@ class ChainPipeline:
                                 server,
                                 stage_params,
                                 per_stage_params,
-                                StageResult(images=[source_tile]),
+                                StageResult(images=source_tile),
                                 tile_mask=tile_mask,
                                 callback=callback,
                                 dims=dims,
@@ -192,7 +192,7 @@ class ChainPipeline:
                                 for j, image in enumerate(tile_result.as_image()):
                                     save_image(server, f"last-tile-{j}.png", image)
 
-                            return tile_result.as_image()
+                            return tile_result
                         except Exception:
                             worker.retries = worker.retries - 1
                             logger.exception(
diff --git a/api/onnx_web/chain/tile.py b/api/onnx_web/chain/tile.py
index aae71af1..a2c3da9a 100644
--- a/api/onnx_web/chain/tile.py
+++ b/api/onnx_web/chain/tile.py
@@ -27,7 +27,7 @@ class TileCallback(Protocol):
 
     def __call__(
         self, image: Image.Image, dims: Tuple[int, int, int]
-    ) -> List[Image.Image]:
+    ) -> StageResult:
         """
         Run this stage against a single tile.
         """
@@ -287,7 +287,7 @@ def process_tile_stack(
                 bottom_margin,
             )
             tile_stack = add_margin(
-                stack,
+                stack.as_image(),
                 left,
                 top,
                 right,
@@ -322,7 +322,7 @@ def process_tile_stack(
         for image_filter in filters:
             tile_stack = image_filter(tile_stack, tile_mask, (left, top, tile))
 
-        tiles.append((left, top, tile_stack))
+        tiles.append((left, top, tile_stack.as_image()))
 
     lefts, tops, stacks = list(zip(*tiles))
     coords = list(zip(lefts, tops))
@@ -331,6 +331,7 @@ def process_tile_stack(
     result = []
     for stack in stacks:
         stack_tiles = zip(coords, stack)
+        stack_tiles = [(left, top, tile) for (left, top), tile in stack_tiles]
         result.append(blend_tiles(stack_tiles, scale, width, height, tile, overlap))
 
     return result