feat(api): make chain pipeline work without a source image

ssube · Jul 1, 2023 · af416c2 · af416c2
1 parent fd3e65e
commit af416c2
Show file tree

Hide file tree

Showing 2 changed files with 30 additions and 15 deletions.
diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py
@@ -99,27 +99,36 @@ def __call__(
             callback = ChainProgress.from_progress(callback)
 
         start = monotonic()
-        logger.info(
-            "running pipeline on source image with dimensions %sx%s",
-            source.width,
-            source.height,
-        )
         image = source
 
+        if source is not None:
+            logger.info(
+                "running pipeline on source image with dimensions %sx%s",
+                source.width,
+                source.height,
+            )
+        else:
+            logger.info("running pipeline without source image")
+
         for stage_pipe, stage_params, stage_kwargs in self.stages:
             name = stage_params.name or stage_pipe.__name__
             kwargs = stage_kwargs or {}
             kwargs = {**pipeline_kwargs, **kwargs}
 
-            logger.debug(
-                "running stage %s on image with dimensions %sx%s, %s",
-                name,
-                image.width,
-                image.height,
-                kwargs.keys(),
-            )
+            if image is not None:
+                logger.debug(
+                    "running stage %s on source image with dimensions %sx%s, %s",
+                    name,
+                    image.width,
+                    image.height,
+                    kwargs.keys(),
+                )
+            else:
+                logger.debug(
+                    "running stage %s without source image, %s", name, kwargs.keys()
+                )
 
-            if (
+            if image is not None and (
                 image.width > stage_params.tile_size
                 or image.height > stage_params.tile_size
             ):

diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py
@@ -3,7 +3,13 @@
 
 from PIL import Image
 
-from ..chain import blend_img2img, blend_mask, upscale_highres, upscale_outpaint
+from ..chain import (
+    blend_img2img,
+    blend_mask,
+    source_txt2img,
+    upscale_highres,
+    upscale_outpaint,
+)
 from ..chain.base import ChainPipeline
 from ..output import save_image
 from ..params import (
@@ -36,7 +42,7 @@ def run_txt2img_pipeline(
     # prepare the chain pipeline and first stage
     chain = ChainPipeline()
     stage = StageParams()
-    chain.append((blend_img2img, stage, None))
+    chain.append((source_txt2img, stage, None))
 
     # apply upscaling and correction, before highres
     first_upscale, after_upscale = split_upscale(upscale)