From ce6cf08684435ee705787bb8964d578426abb4e1 Mon Sep 17 00:00:00 2001
From: Sean Sube <seansube@gmail.com>
Date: Sat, 28 Jan 2023 08:44:24 -0600
Subject: [PATCH] fix outpaint and txt2img stages

---
 api/onnx_web/chain/generate_txt2img.py | 33 ++++++++++++--------------
 api/onnx_web/chain/upscale_outpaint.py |  4 ++--
 2 files changed, 17 insertions(+), 20 deletions(-)

diff --git a/api/onnx_web/chain/generate_txt2img.py b/api/onnx_web/chain/generate_txt2img.py
index 953e3c22..0f6d2774 100644
--- a/api/onnx_web/chain/generate_txt2img.py
+++ b/api/onnx_web/chain/generate_txt2img.py
@@ -44,26 +44,23 @@ def generate_txt2img(
     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)
+    pipe = load_pipeline(OnnxStableDiffusionPipeline,
+                            params.model, params.provider, params.scheduler)
 
-        latents = get_latents_from_seed(params.seed, size)
-        rng = np.random.RandomState(params.seed)
+    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])
+    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,
+    )
+    output = result.images[0]
 
     print('final output image size', output.size)
     return output
diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py
index 4a386378..828c814a 100644
--- a/api/onnx_web/chain/upscale_outpaint.py
+++ b/api/onnx_web/chain/upscale_outpaint.py
@@ -65,7 +65,7 @@ def upscale_outpaint(
         mask_image.save(base_join(ctx.output_path, 'last-mask.png'))
         noise_image.save(base_join(ctx.output_path, 'last-noise.png'))
 
-    def outpaint():
+    def outpaint(image: Image.Image):
         pipe = load_pipeline(OnnxStableDiffusionInpaintPipeline,
                              params.model, params.provider, params.scheduler)
 
@@ -77,7 +77,7 @@ def upscale_outpaint(
             generator=rng,
             guidance_scale=params.cfg,
             height=size.height,
-            image=source_image,
+            image=image,
             latents=latents,
             mask_image=mask_image,
             negative_prompt=params.negative_prompt,