fix(api): remove size limit on inpainting stage

ssube · Apr 29, 2023 · f782f39 · f782f39
1 parent 0666d81
commit f782f39
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 6 deletions.
diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py
@@ -56,7 +56,7 @@ def upscale_outpaint(
 
     draw_mask = ImageDraw.Draw(stage_mask)
     full_size = Size(*full_dims)
-    full_latents = get_latents_from_seed(params.seed, full_size)
+    full_latents = get_latents_from_seed(params.seed, full_size.latent_size())
 
     if is_debug():
         save_image(server, "last-source.png", source)

diff --git a/api/onnx_web/chain/utils.py b/api/onnx_web/chain/utils.py
@@ -21,6 +21,19 @@ def __call__(self, image: Image.Image, dims: Tuple[int, int, int]) -> Image.Imag
         pass
 
 
+def complete_tile(
+    source: Image.Image,
+    tile: int,
+) -> Image.Image:
+    if source.width < tile or source.height < tile:
+        full_source = Image.new(source.mode, (tile, tile))
+        full_source.paste(source)
+        return full_source
+
+    return source
+
+
+
 def process_tile_grid(
     source: Image.Image,
     tile: int,
@@ -29,10 +42,10 @@ def process_tile_grid(
     **kwargs,
 ) -> Image.Image:
     width, height = source.size
-    image = Image.new("RGB", (width * scale, height * scale))
+    image = Image.new(source.mode, (width * scale, height * scale))
 
-    tiles_x = width // tile
-    tiles_y = height // tile
+    tiles_x = ceil(width / tile)
+    tiles_y = ceil(height / tile)
     total = tiles_x * tiles_y
 
     for y in range(tiles_y):
@@ -41,7 +54,9 @@ def process_tile_grid(
             left = x * tile
             top = y * tile
             logger.debug("processing tile %s of %s, %s.%s", idx + 1, total, y, x)
+
             tile_image = source.crop((left, top, left + tile, top + tile))
+            tile_image = complete_tile(tile_image)
 
             for filter in filters:
                 tile_image = filter(tile_image, (left, top, tile))
@@ -74,6 +89,7 @@ def process_tile_spiral(
         logger.debug("processing tile %s of %s, %sx%s", counter, len(tiles), left, top)
 
         tile_image = image.crop((left, top, left + tile, top + tile))
+        tile_image = complete_tile(tile_image)
 
         for filter in filters:
             tile_image = filter(tile_image, (left, top, tile))

diff --git a/api/onnx_web/convert/diffusion/diffusers.py b/api/onnx_web/convert/diffusion/diffusers.py
@@ -22,7 +22,6 @@
     OnnxStableDiffusionPipeline,
     StableDiffusionControlNetPipeline,
     StableDiffusionInstructPix2PixPipeline,
-    StableDiffusionPanoramaPipeline,
     StableDiffusionPipeline,
     StableDiffusionUpscalePipeline,
 )
@@ -42,7 +41,7 @@
     "img2img": StableDiffusionPipeline,
     "inpaint": StableDiffusionPipeline,
     "lpw": StableDiffusionPipeline,
-    "panorama": StableDiffusionPanoramaPipeline,
+    "panorama": StableDiffusionPipeline,
     "pix2pix": StableDiffusionInstructPix2PixPipeline,
     "txt2img": StableDiffusionPipeline,
     "upscale": StableDiffusionUpscalePipeline,

diff --git a/api/onnx_web/params.py b/api/onnx_web/params.py
@@ -1,5 +1,6 @@
 from enum import IntEnum
 from logging import getLogger
+from math import ceil
 from typing import Any, Dict, List, Literal, Optional, Tuple, Union
 
 from .models.meta import NetworkModel
@@ -78,6 +79,12 @@ def add_border(self, border: Border):
             border.top + self.height + border.bottom,
         )
 
+    def latent_size(self):
+        return Size(
+            ceil(self.width / 8),
+            ceil(self.height / 8),
+        )
+
     def tojson(self) -> Dict[str, int]:
         return {
             "height": self.height,