feat(api): use per-tile latents for outpainting

api/onnx_web/chain/upscale_outpaint.py

@@ -7,6 +7,7 @@ from typing import Callable, Tuple
 
 from ..diffusion.load import (
     get_latents_from_seed,
+    get_tile_latents,
     load_pipeline,
 )
 from ..image import (
@@ -56,7 +57,7 @@ def upscale_outpaint(
         # if no mask was provided, keep the full source image
         mask_image = Image.new('RGB', source_image.size, 'black')
 
-    source_image, mask_image, noise_image, _full_dims = expand_image(
+    source_image, mask_image, noise_image, full_dims = expand_image(
         source_image,
         mask_image,
         border,
@@ -65,6 +66,8 @@ def upscale_outpaint(
         mask_filter=mask_filter)
 
     draw_mask = ImageDraw.Draw(mask_image)
+    full_size = Size(*full_dims)
+    full_latents = get_latents_from_seed(params.seed, full_size)
 
     if is_debug():
         source_image.save(base_join(ctx.output_path, 'last-source.png'))
@@ -85,7 +88,8 @@ def upscale_outpaint(
         pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline,
                              model, params.provider, params.scheduler)
 
-        latents = get_latents_from_seed(params.seed, size)
+        # TODO: take a subset of the full latents
+        latents = get_tile_latents(full_latents, dims, tile)
         rng = np.random.RandomState(params.seed)
 
         result = pipe(

api/onnx_web/diffusion/load.py

@@ -2,7 +2,7 @@ from diffusers import (
     DiffusionPipeline,
 )
 from logging import getLogger
-from typing import Any, Optional
+from typing import Any, Optional, Tuple
 
 from ..params import (
     Size,
@@ -31,6 +31,18 @@ def get_latents_from_seed(seed: int, size: Size) -> np.ndarray:
     return image_latents
 
 
+def get_tile_latents(full_latents: np.ndarray, dims: Tuple[int, int, int]) -> np.ndarray:
+    x, y, tile = dims
+    x = x // 8
+    y = y // 8
+
+    t = tile // 8
+    xt = x + t
+    yt = y + t
+
+    return full_latents[:,:,x:xt,y:yt]
+
+
 def load_pipeline(pipeline: DiffusionPipeline, model: str, provider: str, scheduler: Any, device: Optional[str] = None):
     global last_pipeline_instance
     global last_pipeline_scheduler