feat(api): add inpaint as a chain stage

api/onnx_web/chain/__init__.py

@@ -7,6 +7,9 @@ from .base import (
 from .correct_gfpgan import (
   correct_gfpgan,
 )
+from .upscale_outpaint import (
+  upscale_outpaint,
+)
 from .upscale_resrgan import (
   upscale_resrgan,
 )

api/onnx_web/chain/base.py

@@ -1,6 +1,6 @@
 from PIL import Image
 from os import path
-from typing import Any, Callable, List, Optional, Protocol, Tuple
+from typing import Any, List, Optional, Protocol, Tuple
 
 from ..params import (
     ImageParams,
@@ -9,6 +9,9 @@ from ..params import (
 from ..utils import (
     ServerContext,
 )
+from .utils import (
+    process_tiles,
+)
 
 
 class StageCallback(Protocol):
@@ -26,35 +29,6 @@ class StageCallback(Protocol):
 PipelineStage = Tuple[StageCallback, StageParams, Optional[dict]]
 
 
-def process_tiles(
-    source: Image,
-    tile: int,
-    scale: int,
-    filters: List[Callable],
-) -> Image:
-    width, height = source.size
-    image = Image.new('RGB', (width * scale, height * scale))
-
-    tiles_x = width // tile
-    tiles_y = height // tile
-    total = tiles_x * tiles_y
-
-    for y in range(tiles_y):
-        for x in range(tiles_x):
-            idx = (y * tiles_x) + x
-            left = x * tile
-            top = y * tile
-            print('processing tile %s of %s, %s.%s' % (idx, total, y, x))
-            tile_image = source.crop((left, top, left + tile, top + tile))
-
-            for filter in filters:
-                tile_image = filter(tile_image)
-
-            image.paste(tile_image, (left * scale, top * scale))
-
-    return image
-
-
 class ChainPipeline:
     '''
     Run many stages in series, passing the image results from each to the next, and processing

api/onnx_web/chain/upscale_outpaint.py

@@ -0,0 +1,89 @@
+from diffusers import (
+    OnnxStableDiffusionInpaintPipeline,
+)
+from PIL import Image
+from typing import Callable
+
+from ..diffusion import (
+    get_latents_from_seed,
+    load_pipeline,
+)
+from ..image import (
+    expand_image,
+    mask_filter_none,
+    noise_source_histogram,
+)
+from ..params import (
+    Border,
+    ImageParams,
+    Size,
+    StageParams,
+)
+from ..utils import (
+    base_join,
+    is_debug,
+    ServerContext,
+)
+from .utils import (
+    process_tiles,
+)
+
+import numpy as np
+
+
+def upscale_outpaint(
+    ctx: ServerContext,
+    stage: StageParams,
+    params: ImageParams,
+    source_image: Image.Image,
+    *,
+    expand: Border,
+    mask_image: Image.Image,
+    fill_color: str = 'white',
+    mask_filter: Callable = mask_filter_none,
+    noise_source: Callable = noise_source_histogram,
+) -> Image:
+    print('upscaling image by expanding borders', expand)
+
+    output = expand_image(source_image, mask_image, expand)
+    size = Size(*output.size)
+
+    def outpaint():
+        pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline,
+                         params.model, params.provider, params.scheduler)
+
+        latents = get_latents_from_seed(params.seed, size)
+        rng = np.random.RandomState(params.seed)
+
+        print('applying mask filter and generating noise source')
+        source_image, mask_image, noise_image, _full_dims = expand_image(
+            source_image,
+            mask_image,
+            expand,
+            fill=fill_color,
+            noise_source=noise_source,
+            mask_filter=mask_filter)
+
+        if is_debug():
+            source_image.save(base_join(ctx.output_path, 'last-source.png'))
+            mask_image.save(base_join(ctx.output_path, 'last-mask.png'))
+            noise_image.save(base_join(ctx.output_path, 'last-noise.png'))
+
+        result = pipe(
+            params.prompt,
+            generator=rng,
+            guidance_scale=params.cfg,
+            height=size.height,
+            image=source_image,
+            latents=latents,
+            mask_image=mask_image,
+            negative_prompt=params.negative_prompt,
+            num_inference_steps=params.steps,
+            width=size.width,
+        )
+        return result.images[0]
+
+    output = process_tiles(output, 256, 4, [outpaint])
+
+    print('final output image size', output.size)
+    return output

api/onnx_web/chain/utils.py

@@ -0,0 +1,31 @@
+from PIL import Image
+from typing import Callable, List
+
+
+def process_tiles(
+    source: Image.Image,
+    tile: int,
+    scale: int,
+    filters: List[Callable],
+) -> Image:
+    width, height = source.size
+    image = Image.new('RGB', (width * scale, height * scale))
+
+    tiles_x = width // tile
+    tiles_y = height // tile
+    total = tiles_x * tiles_y
+
+    for y in range(tiles_y):
+        for x in range(tiles_x):
+            idx = (y * tiles_x) + x
+            left = x * tile
+            top = y * tile
+            print('processing tile %s of %s, %s.%s' % (idx, total, y, x))
+            tile_image = source.crop((left, top, left + tile, top + tile))
+
+            for filter in filters:
+                tile_image = filter(tile_image)
+
+            image.paste(tile_image, (left * scale, top * scale))
+
+    return image

api/onnx_web/serve.py

@@ -22,6 +22,12 @@ from onnxruntime import get_available_providers
 from os import makedirs, path
 from typing import Tuple
 
+from .chain import (
+    correct_gfpgan,
+    upscale_outpaint,
+    upscale_resrgan,
+    upscale_stable_diffusion,
+)
 from .diffusion import (
     run_img2img_pipeline,
     run_inpaint_pipeline,
@@ -47,11 +53,6 @@ from .params import (
     Size,
     UpscaleParams,
 )
-from .upscale import (
-    correct_gfpgan,
-    upscale_resrgan,
-    upscale_stable_diffusion,
-)
 from .utils import (
     is_debug,
     get_and_clamp_float,
@@ -109,6 +110,7 @@ mask_filters = {
 }
 chain_stages = {
     'correction-gfpgan': correct_gfpgan,
+    'upscaling-outpaint': upscale_outpaint,
     'upscaling-resrgan': upscale_resrgan,
     'upscaling-stable-diffusion': upscale_stable_diffusion,
 }