use tile size when completing partial latents

ssube · Jul 10, 2023 · 95cad90 · 95cad90
1 parent 8acfaaa
commit 95cad90
Show file tree

Hide file tree

Showing 3 changed files with 34 additions and 23 deletions.
diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py
@@ -146,7 +146,9 @@ def __call__(
                     )
 
                     def stage_tile(
-                        source_tile: Image.Image, tile_mask: Image.Image, dims: Tuple[int, int, int]
+                        source_tile: Image.Image,
+                        tile_mask: Image.Image,
+                        dims: Tuple[int, int, int],
                     ) -> Image.Image:
                         output_tile = stage_pipe.run(
                             job,

diff --git a/api/onnx_web/chain/source_txt2img.py b/api/onnx_web/chain/source_txt2img.py
@@ -6,7 +6,12 @@
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import encode_prompt, get_latents_from_seed, get_tile_latents, parse_prompt
+from ..diffusers.utils import (
+    encode_prompt,
+    get_latents_from_seed,
+    get_tile_latents,
+    parse_prompt,
+)
 from ..params import ImageParams, Size, SizeChart, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
@@ -50,21 +55,21 @@ def run(
 
         tile_size = params.tiles
 
+        if max(size) > tile_size:
+            latent_size = Size(tile_size, tile_size)
+            pipe_width = pipe_height = tile_size
+        else:
+            latent_size = Size(size.width, size.height)
+            pipe_width = size.width
+            pipe_height = size.height
+
         # generate new latents or slice existing
         if latents is None:
-            if max(size) > tile_size:
-                latent_size = Size(tile_size, tile_size)
-                pipe_width = pipe_height = tile_size
-            else:
-                latent_size = Size(size.width, size.height)
-                pipe_width = size.width
-                pipe_height = size.height
-
             # generate new latents
             latents = get_latents_from_seed(params.seed, latent_size, params.batch)
         else:
             # slice existing latents
-            latents = get_tile_latents(latents, dims, size)
+            latents = get_tile_latents(latents, dims, Size(tile_size, tile_size))
             pipe_width = pipe_height = tile_size
 
         pipe_type = params.get_valid_pipeline("txt2img")

diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py
@@ -6,7 +6,12 @@
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import encode_prompt, get_latents_from_seed, get_tile_latents, parse_prompt
+from ..diffusers.utils import (
+    encode_prompt,
+    get_latents_from_seed,
+    get_tile_latents,
+    parse_prompt,
+)
 from ..image import mask_filter_none, noise_source_histogram
 from ..output import save_image
 from ..params import Border, ImageParams, Size, SizeChart, StageParams
@@ -68,23 +73,22 @@ def run(
 
             size = Size(*source.size)
             tile_size = params.tiles
+            if max(size) > tile_size:
+                latent_size = Size(tile_size, tile_size)
+                pipe_width = pipe_height = tile_size
+            else:
+                latent_size = Size(size.width, size.height)
+                pipe_width = size.width
+                pipe_height = size.height
 
             # generate new latents or slice existing
             if latents is None:
-                if max(size) > tile_size:
-                    latent_size = Size(tile_size, tile_size)
-                    pipe_width = pipe_height = tile_size
-                else:
-                    latent_size = Size(size.width, size.height)
-                    pipe_width = size.width
-                    pipe_height = size.height
-
                 # generate new latents
                 latents = get_latents_from_seed(params.seed, latent_size, params.batch)
             else:
-                # slice existing latents
-                latents = get_tile_latents(latents, dims, size)
-                pipe_width, pipe_height, _tile_size = dims
+                # slice existing latents and make sure there is a complete tile
+                latents = get_tile_latents(latents, dims, Size(tile_size, tile_size))
+                pipe_width = pipe_height = tile_size
 
             if params.lpw():
                 logger.debug("using LPW pipeline for inpaint")