From 8781f7fda0f3e8523b8b4d17a9d4913b4d93432b Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sun, 2 Jul 2023 18:25:15 -0500
Subject: [PATCH] fix(api): make tiling compatible with source stages

---
 api/onnx_web/chain/tile.py                          | 13 +++++++++----
 .../test-refs/img2img-sd-v1-5-256-pumpkin-0.png     |  4 ++--
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/api/onnx_web/chain/tile.py b/api/onnx_web/chain/tile.py
index 1e3bfec0..1f80ef57 100644
--- a/api/onnx_web/chain/tile.py
+++ b/api/onnx_web/chain/tile.py
@@ -29,6 +29,9 @@ def complete_tile(
     source: Image.Image,
     tile: int,
 ) -> Image.Image:
+    if source is None:
+        return source
+
     if source.width < tile or source.height < tile:
         full_source = Image.new(source.mode, (tile, tile))
         full_source.paste(source)
@@ -158,7 +161,7 @@ def process_tile_grid(
     overlap: float = 0.0,
     **kwargs,
 ) -> Image.Image:
-    width, height = source.size
+    width, height = kwargs.get("size", source.size if source else None)
 
     adj_tile = int(float(tile) * (1.0 - overlap))
     tiles_x = ceil(width / adj_tile)
@@ -182,7 +185,9 @@ def process_tile_grid(
             top = y * adj_tile
             logger.info("processing tile %s of %s, %s.%s", idx + 1, total, y, x)
 
-            tile_image = source.crop((left, top, left + tile, top + tile))
+            tile_image = (
+                source.crop((left, top, left + tile, top + tile)) if source else None
+            )
             tile_image = complete_tile(tile_image, tile)
 
             for filter in filters:
@@ -204,7 +209,7 @@ def process_tile_spiral(
     if scale != 1:
         raise ValueError("unsupported scale")
 
-    width, height = source.size
+    width, height = kwargs.get("size", source.size if source else None)
 
     # spiral uses the previous run and needs a scratch texture for 3x memory
     image = Image.new("RGB", (width * scale, height * scale))
@@ -222,7 +227,7 @@ def process_tile_spiral(
             "processing tile %s of %s, %sx%s", counter, len(tile_coords), left, top
         )
 
-        tile_image = image.crop((left, top, left + tile, top + tile))
+        tile_image = image.crop((left, top, left + tile, top + tile)) if image else None
         tile_image = complete_tile(tile_image, tile)
 
         for filter in filters:
diff --git a/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png b/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
index b61dab55..28a2a046 100644
--- a/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
+++ b/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:18ab09b41c89f20cae4028770cd169e68a36245822c69ec7ce19ce622921b1bf
-size 215130
+oid sha256:294bc6c6aaa08dc5d5b3a831d5bbc092880fbc82b3c3437f8459f991cee3d03e
+size 133168