feat(api): synthesize a mask for outpaint stages

api/onnx_web/chain/__init__.py

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

api/onnx_web/chain/generate_txt2img.py

@@ -0,0 +1,69 @@
+from diffusers import (
+    OnnxStableDiffusionPipeline,
+)
+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 generate_txt2img(
+    ctx: ServerContext,
+    stage: StageParams,
+    params: ImageParams,
+    source_image: Image.Image,
+    *,
+    size: Size,
+) -> Image:
+    print('generating image using txt2img', params.prompt)
+
+    if source_image is not None:
+        print('a source image was passed to a txt2img stage, but will be discarded')
+
+    def txt2img():
+        pipe = load_pipeline(OnnxStableDiffusionPipeline,
+                             params.model, params.provider, params.scheduler)
+
+        latents = get_latents_from_seed(params.seed, size)
+        rng = np.random.RandomState(params.seed)
+
+        result = pipe(
+            params.prompt,
+            height=size.height,
+            width=size.width,
+            generator=rng,
+            guidance_scale=params.cfg,
+            latents=latents,
+            negative_prompt=params.negative_prompt,
+            num_inference_steps=params.steps,
+        )
+        return result.images[0]
+
+    output = process_tiles(output, 512, 1, [txt2img])
+
+    print('final output image size', output.size)
+    return output

api/onnx_web/chain/upscale_outpaint.py

@@ -1,7 +1,7 @@
 from diffusers import (
     OnnxStableDiffusionInpaintPipeline,
 )
-from PIL import Image
+from PIL import Image, ImageDraw
 from typing import Callable
 
 from ..diffusion import (
@@ -38,37 +38,40 @@ def upscale_outpaint(
     source_image: Image.Image,
     *,
     expand: Border,
-    mask_image: Image.Image,
+    mask_image: Image.Image = None,
     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)
+    if mask_image is None:
+        mask_image = Image.new('RGB', source_image.size, fill_color)
+        draw = ImageDraw.Draw(mask_image)
+        draw.rectangle((expand.left, expand.top, expand.left +
+                       source_image.width, expand.top + source_image.height), fill='black')
+
+    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)
+    size = Size(*full_dims)
+
+    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'))
 
     def outpaint():
         pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline,
-                         params.model, params.provider, params.scheduler)
+                             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,

api/onnx_web/serve.py

@@ -27,6 +27,7 @@
     upscale_outpaint,
     upscale_resrgan,
     upscale_stable_diffusion,
+    ChainPipeline,
 )
 from .diffusion import (
     run_img2img_pipeline,
@@ -51,6 +52,7 @@
     Border,
     ImageParams,
     Size,
+    StageParams,
     UpscaleParams,
 )
 from .utils import (
@@ -538,6 +540,19 @@ def upscale():
 @app.route('/api/chain', methods=['POST'])
 def chain():
     print('TODO: run chain pipeline')
+
+    params, size = pipeline_from_request()
+
+    example = ChainPipeline(stages=[
+        (source_txt2img, StageParams(), None),
+        (upscale_outpaint, StageParams(outscale=4), {
+            'expand': Border(256, 256, 256, 256),
+        }),
+    ])
+
+    output = make_output_name('chain', params, size)
+    executor.submit_stored(output, example, context, params, None)
+
     # parse body as json, list of stages
     # build and run chain pipeline
     return jsonify({})