feat(api): start using chain pipelines for all images

api/onnx_web/__init__.py

@@ -17,7 +17,7 @@
     run_upscale_pipeline,
 )
 from .diffusers.stub_scheduler import StubScheduler
-from .diffusers.upscale import run_upscale_correction
+from .diffusers.upscale import append_upscale_correction
 from .image.utils import (
     expand_image,
     valid_image,

api/onnx_web/chain/__init__.py

@@ -1,10 +1,8 @@
 from .base import ChainPipeline, PipelineStage, StageCallback, StageParams
-from .blend_controlnet import blend_controlnet
 from .blend_img2img import blend_img2img
 from .blend_inpaint import blend_inpaint
 from .blend_linear import blend_linear
 from .blend_mask import blend_mask
-from .blend_pix2pix import blend_pix2pix
 from .correct_codeformer import correct_codeformer
 from .correct_gfpgan import correct_gfpgan
 from .persist_disk import persist_disk
@@ -16,18 +14,17 @@
 from .source_txt2img import source_txt2img
 from .source_url import source_url
 from .upscale_bsrgan import upscale_bsrgan
+from .upscale_highres import upscale_highres
 from .upscale_outpaint import upscale_outpaint
 from .upscale_resrgan import upscale_resrgan
 from .upscale_stable_diffusion import upscale_stable_diffusion
 from .upscale_swinir import upscale_swinir
 
 CHAIN_STAGES = {
-    "blend-controlnet": blend_controlnet,
     "blend-img2img": blend_img2img,
     "blend-inpaint": blend_inpaint,
     "blend-linear": blend_linear,
     "blend-mask": blend_mask,
-    "blend-pix2pix": blend_pix2pix,
     "correct-codeformer": correct_codeformer,
     "correct-gfpgan": correct_gfpgan,
     "persist-disk": persist_disk,
@@ -39,6 +36,7 @@
     "source-txt2img": source_txt2img,
     "source-url": source_url,
     "upscale-bsrgan": upscale_bsrgan,
+    "upscale-highres": upscale_highres,
     "upscale-outpaint": upscale_outpaint,
     "upscale-resrgan": upscale_resrgan,
     "upscale-stable-diffusion": upscale_stable_diffusion,

api/onnx_web/chain/base.py

@@ -88,7 +88,7 @@ def __call__(
         job: WorkerContext,
         server: ServerContext,
         params: ImageParams,
-        source: Image.Image,
+        source: Optional[Image.Image] = None,
         callback: Optional[ProgressCallback] = None,
         **pipeline_kwargs
     ) -> Image.Image:

api/onnx_web/chain/blend_img2img.py

@@ -6,6 +6,7 @@
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
+from ..diffusers.utils import encode_prompt, parse_prompt
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
@@ -20,8 +21,9 @@ def blend_img2img(
     params: ImageParams,
     source: Image.Image,
     *,
+    strength: float,
     callback: Optional[ProgressCallback] = None,
-    stage_source: Image.Image,
+    stage_source: Optional[Image.Image] = None,
     **kwargs,
 ) -> Image.Image:
     params = params.with_args(**kwargs)
@@ -30,14 +32,28 @@ def blend_img2img(
         "blending image using img2img, %s steps: %s", params.steps, params.prompt
     )
 
-    pipe_type = "lpw" if params.lpw() else "img2img"
+    prompt_pairs, loras, inversions = parse_prompt(params)
+
+    pipe_type = params.get_valid_pipeline("img2img")
     pipe = load_pipeline(
         server,
         params,
         pipe_type,
         job.get_device(),
-        # TODO: add LoRAs and TIs
+        inversions=inversions,
+        loras=loras,
     )
+
+    pipe_params = {}
+    if pipe_type == "controlnet":
+        pipe_params["controlnet_conditioning_scale"] = strength
+    elif pipe_type == "img2img":
+        pipe_params["strength"] = strength
+    elif pipe_type == "panorama":
+        pipe_params["strength"] = strength
+    elif pipe_type == "pix2pix":
+        pipe_params["image_guidance_scale"] = strength
+
     if params.lpw():
         logger.debug("using LPW pipeline for img2img")
         rng = torch.manual_seed(params.seed)
@@ -50,8 +66,13 @@ def blend_img2img(
             num_inference_steps=params.steps,
             strength=params.strength,
             callback=callback,
+            **pipe_params,
         )
     else:
+        # encode and record alternative prompts outside of LPW
+        prompt_embeds = encode_prompt(pipe, prompt_pairs, params.batch, params.do_cfg())
+        pipe.unet.set_prompts(prompt_embeds)
+
         rng = np.random.RandomState(params.seed)
         result = pipe(
             params.prompt,
@@ -62,6 +83,7 @@ def blend_img2img(
             num_inference_steps=params.steps,
             strength=params.strength,
             callback=callback,
+            **pipe_params,
         )
 
     output = result.images[0]

api/onnx_web/chain/source_txt2img.py

@@ -6,7 +6,7 @@
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import get_latents_from_seed
+from ..diffusers.utils import encode_prompt, get_latents_from_seed, parse_prompt
 from ..params import ImageParams, Size, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
@@ -36,14 +36,17 @@ def source_txt2img(
             "a source image was passed to a txt2img stage, and will be discarded"
         )
 
+    prompt_pairs, loras, inversions = parse_prompt(params)
+
     latents = get_latents_from_seed(params.seed, size)
-    pipe_type = "lpw" if params.lpw() else "txt2img"
+    pipe_type = params.get_valid_pipeline("txt2img")
     pipe = load_pipeline(
         server,
         params,
         pipe_type,
         job.get_device(),
-        # TODO: add LoRAs and TIs
+        inversions=inversions,
+        loras=loras,
     )
 
     if params.lpw():
@@ -61,6 +64,10 @@ def source_txt2img(
             callback=callback,
         )
     else:
+        # encode and record alternative prompts outside of LPW
+        prompt_embeds = encode_prompt(pipe, prompt_pairs, params.batch, params.do_cfg())
+        pipe.unet.set_prompts(prompt_embeds)
+
         rng = np.random.RandomState(params.seed)
         result = pipe(
             params.prompt,