Repository: ssube/onnx-web (language: python)

Issue:
im2img mode



Discussion / hints:
This is mostly implemented, but needs custom image controls.

While img2img is working on CPU, Nvidia (with #6), and some Linux AMD chips, it is failing on my discrete AMD card and WIn10: #20
Fixed #20, so img2img is working well on all three platforms. Just needs different image controls.

Fix patch:
diff --git a/api/onnx_web/__init__.py b/api/onnx_web/__init__.py
index 208d2ffa9..27bd382ec 100644
--- a/api/onnx_web/__init__.py
+++ b/api/onnx_web/__init__.py
@@ -17,10 +17,9 @@
     run_upscale_pipeline,
 )
 from .diffusers.stub_scheduler import StubScheduler
-from .diffusers.upscale import run_upscale_correction
+from .chain.upscale import stage_upscale_correction
 from .image.utils import (
     expand_image,
-    valid_image,
 )
 from .image.mask_filter import (
     mask_filter_gaussian_multiply,
diff --git a/api/onnx_web/chain/__init__.py b/api/onnx_web/chain/__init__.py
index 61e7f5e0c..edda66d21 100644
--- a/api/onnx_web/chain/__init__.py
+++ b/api/onnx_web/chain/__init__.py
@@ -1,46 +1,46 @@
-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
-from .persist_s3 import persist_s3
-from .reduce_crop import reduce_crop
-from .reduce_thumbnail import reduce_thumbnail
-from .source_noise import source_noise
-from .source_s3 import source_s3
-from .source_txt2img import source_txt2img
-from .source_url import source_url
-from .upscale_bsrgan import upscale_bsrgan
-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
+from .base import ChainPipeline, PipelineStage, StageParams
+from .blend_img2img import BlendImg2ImgStage
+from .blend_inpaint import BlendInpaintStage
+from .blend_linear import BlendLinearStage
+from .blend_mask import BlendMaskStage
+from .correct_codeformer import CorrectCodeformerStage
+from .correct_gfpgan import CorrectGFPGANStage
+from .persist_disk import PersistDiskStage
+from .persist_s3 import PersistS3Stage
+from .reduce_crop import ReduceCropStage
+from .reduce_thumbnail import ReduceThumbnailStage
+from .source_noise import SourceNoiseStage
+from .source_s3 import SourceS3Stage
+from .source_txt2img import SourceTxt2ImgStage
+from .source_url import SourceURLStage
+from .upscale_bsrgan import UpscaleBSRGANStage
+from .upscale_highres import UpscaleHighresStage
+from .upscale_outpaint import UpscaleOutpaintStage
+from .upscale_resrgan import UpscaleRealESRGANStage
+from .upscale_simple import UpscaleSimpleStage
+from .upscale_stable_diffusion import UpscaleStableDiffusionStage
+from .upscale_swinir import UpscaleSwinIRStage
 
 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,
-    "persist-s3": persist_s3,
-    "reduce-crop": reduce_crop,
-    "reduce-thumbnail": reduce_thumbnail,
-    "source-noise": source_noise,
-    "source-s3": source_s3,
-    "source-txt2img": source_txt2img,
-    "source-url": source_url,
-    "upscale-bsrgan": upscale_bsrgan,
-    "upscale-outpaint": upscale_outpaint,
-    "upscale-resrgan": upscale_resrgan,
-    "upscale-stable-diffusion": upscale_stable_diffusion,
-    "upscale-swinir": upscale_swinir,
+    "blend-img2img": BlendImg2ImgStage,
+    "blend-inpaint": BlendInpaintStage,
+    "blend-linear": BlendLinearStage,
+    "blend-mask": BlendMaskStage,
+    "correct-codeformer": CorrectCodeformerStage,
+    "correct-gfpgan": CorrectGFPGANStage,
+    "persist-disk": PersistDiskStage,
+    "persist-s3": PersistS3Stage,
+    "reduce-crop": ReduceCropStage,
+    "reduce-thumbnail": ReduceThumbnailStage,
+    "source-noise": SourceNoiseStage,
+    "source-s3": SourceS3Stage,
+    "source-txt2img": SourceTxt2ImgStage,
+    "source-url": SourceURLStage,
+    "upscale-bsrgan": UpscaleBSRGANStage,
+    "upscale-highres": UpscaleHighresStage,
+    "upscale-outpaint": UpscaleOutpaintStage,
+    "upscale-resrgan": UpscaleRealESRGANStage,
+    "upscale-simple": UpscaleSimpleStage,
+    "upscale-stable-diffusion": UpscaleStableDiffusionStage,
+    "upscale-swinir": UpscaleSwinIRStage,
 }
diff --git a/api/onnx_web/chain/base.py b/api/onnx_web/chain/base.py
index cc9fd56b8..0e5acfa18 100644
--- a/api/onnx_web/chain/base.py
+++ b/api/onnx_web/chain/base.py
@@ -1,7 +1,7 @@
 from datetime import timedelta
 from logging import getLogger
 from time import monotonic
-from typing import Any, List, Optional, Protocol, Tuple
+from typing import Any, List, Optional, Tuple
 
 from PIL import Image
 
@@ -10,32 +10,13 @@
 from ..server import ServerContext
 from ..utils import is_debug
 from ..worker import ProgressCallback, WorkerContext
-from .utils import process_tile_order
+from .stage import BaseStage
+from .tile import needs_tile, process_tile_order
 
 logger = getLogger(__name__)
 
 
-class StageCallback(Protocol):
-    """
-    Definition for a stage job function.
-    """
-
-    def __call__(
-        self,
-        job: WorkerContext,
-        server: ServerContext,
-        stage: StageParams,
-        params: ImageParams,
-        source: Image.Image,
-        **kwargs: Any
-    ) -> Image.Image:
-        """
-        Run this stage against a source image.
-        """
-        pass
-
-
-PipelineStage = Tuple[StageCallback, StageParams, Optional[dict]]
+PipelineStage = Tuple[BaseStage, StageParams, Optional[dict]]
 
 
 class ChainProgress:
@@ -76,106 +57,154 @@ def __init__(
         """
         self.stages = list(stages or [])
 
-    def append(self, stage: PipelineStage):
+    def append(self, stage: Optional[PipelineStage]):
         """
         Append an additional stage to this pipeline.
+
+        This requires an already-assembled `PipelineStage`. Use `ChainPipeline.stage` if you want the pipeline to
+        assemble the stage from loose arguments.
         """
         if stage is not None:
             self.stages.append(stage)
 
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        params: ImageParams,
+        sources: List[Image.Image],
+        callback: Optional[ProgressCallback],
+        **kwargs
+    ) -> List[Image.Image]:
+        return self(job, server, params, sources=sources, callback=callback, **kwargs)
+
+    def stage(self, callback: BaseStage, params: StageParams, **kwargs):
+        self.stages.append((callback, params, kwargs))
+        return self
+
     def __call__(
         self,
         job: WorkerContext,
         server: ServerContext,
         params: ImageParams,
-        source: Image.Image,
+        sources: List[Image.Image],
         callback: Optional[ProgressCallback] = None,
         **pipeline_kwargs
-    ) -> Image.Image:
+    ) -> List[Image.Image]:
         """
-        TODO: handle List[Image] inputs and outputs
+        DEPRECATED: use `run` instead
         """
         if callback is not None:
             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 len(sources) > 0:
+            logger.info(
+                "running pipeline on %s source images",
+                len(sources),
+            )
+        else:
+            logger.info("running pipeline without source images")
+
+        stage_sources = sources
         for stage_pipe, stage_params, stage_kwargs in self.stages:
-            name = stage_params.name or stage_pipe.__name__
+            name = stage_params.name or stage_pipe.__class__.__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 len(stage_sources) > 0:
+                logger.debug(
+                    "running stage %s with %s source images, parameters: %s",
+                    name,
+                    len(stage_sources),
+                    kwargs.keys(),
+                )
+            else:
+                logger.debug(
+                    "running stage %s without source images, parameters: %s",
+                    name,
+                    kwargs.keys(),
+                )
+
+            # the stage must be split and tiled if any image is larger than the selected/max tile size
+            must_tile = any(
+                [
+                    needs_tile(
+                        stage_pipe.max_tile,
+                        stage_params.tile_size,
+                        size=kwargs.get("size", None),
+                        source=source,
+                    )
+                    for source in stage_sources
+                ]
             )
 
-            if (
-                image.width > stage_params.tile_size
-                or image.height > stage_params.tile_size
-            ):
-                logger.info(
-                    "image larger than tile size of %s, tiling stage",
-                    stage_params.tile_size,
-                )
+            tile = stage_params.tile_size
+            if stage_pipe.max_tile > 0:
+                tile = min(stage_pipe.max_tile, stage_params.tile_size)
 
-                def stage_tile(tile: Image.Image, _dims) -> Image.Image:
-                    tile = stage_pipe(
-                        job,
-                        server,
-                        stage_params,
-                        params,
+            if must_tile:
+                stage_outputs = []
+                for source in stage_sources:
+                    logger.info(
+                        "image larger than tile size of %s, tiling stage",
                         tile,
-                        callback=callback,
-                        **kwargs
                     )
 
-                    if is_debug():
-                        save_image(server, "last-tile.png", tile)
-
-                    return tile
+                    def stage_tile(source_tile: Image.Image, _dims) -> Image.Image:
+                        output_tile = stage_pipe.run(
+                            job,
+                            server,
+                            stage_params,
+                            params,
+                            [source_tile],
+                            callback=callback,
+                            **kwargs,
+                        )[0]
+
+                        if is_debug():
+                            save_image(server, "last-tile.png", output_tile)
+
+                        return output_tile
+
+                    output = process_tile_order(
+                        stage_params.tile_order,
+                        source,
+                        tile,
+                        stage_params.outscale,
+                        [stage_tile],
+                        **kwargs,
+                    )
+                    stage_outputs.append(output)
 
-                image = process_tile_order(
-                    stage_params.tile_order,
-                    image,
-                    stage_params.tile_size,
-                    stage_params.outscale,
-                    [stage_tile],
-                )
+                stage_sources = stage_outputs
             else:
-                logger.debug("image within tile size, running stage")
-                image = stage_pipe(
+                logger.debug("image within tile size of %s, running stage", tile)
+                stage_sources = stage_pipe.run(
                     job,
                     server,
                     stage_params,
                     params,
-                    image,
+                    stage_sources,
                     callback=callback,
-                    **kwargs
+                    **kwargs,
                 )
 
             logger.debug(
-                "finished stage %s, result size: %sx%s", name, image.width, image.height
+                "finished stage %s with %s results",
+                name,
+                len(stage_sources),
             )
 
             if is_debug():
-                save_image(server, "last-stage.png", image)
+                save_image(server, "last-stage.png", stage_sources[0])
 
         end = monotonic()
         duration = timedelta(seconds=(end - start))
         logger.info(
-            "finished pipeline in %s, result size: %sx%s",
+            "finished pipeline in %s with %s results",
             duration,
-            image.width,
-            image.height,
+            len(stage_sources),
         )
-        return image
+        return stage_sources
diff --git a/api/onnx_web/chain/blend_controlnet.py b/api/onnx_web/chain/blend_controlnet.py
deleted file mode 100644
index c944b767d..000000000
--- a/api/onnx_web/chain/blend_controlnet.py
+++ /dev/null
@@ -1,54 +0,0 @@
-from logging import getLogger
-from typing import Optional
-
-import numpy as np
-from PIL import Image
-
-from ..diffusers.load import load_pipeline
-from ..params import ImageParams, StageParams
-from ..server import ServerContext
-from ..worker import ProgressCallback, WorkerContext
-
-logger = getLogger(__name__)
-
-
-def blend_controlnet(
-    job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    callback: Optional[ProgressCallback] = None,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    source = stage_source or source
-    logger.info(
-        "blending image using ControlNet, %s steps: %s", params.steps, params.prompt
-    )
-
-    pipe = load_pipeline(
-        server,
-        params,
-        "controlnet",
-        job.get_device(),
-    )
-
-    rng = np.random.RandomState(params.seed)
-    result = pipe(
-        params.prompt,
-        generator=rng,
-        guidance_scale=params.cfg,
-        image=source,
-        negative_prompt=params.negative_prompt,
-        num_inference_steps=params.steps,
-        strength=params.strength,  # TODO: ControlNet strength
-        callback=callback,
-    )
-
-    output = result.images[0]
-
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
diff --git a/api/onnx_web/chain/blend_img2img.py b/api/onnx_web/chain/blend_img2img.py
index c825b9e54..4e055105a 100644
--- a/api/onnx_web/chain/blend_img2img.py
+++ b/api/onnx_web/chain/blend_img2img.py
@@ -1,70 +1,99 @@
 from logging import getLogger
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 import torch
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..params import ImageParams, StageParams
+from ..diffusers.utils import encode_prompt, parse_prompt
+from ..params import ImageParams, SizeChart, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def blend_img2img(
-    job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    callback: Optional[ProgressCallback] = None,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    source = stage_source or source
-    logger.info(
-        "blending image using img2img, %s steps: %s", params.steps, params.prompt
-    )
+class BlendImg2ImgStage(BaseStage):
+    max_tile = SizeChart.unlimited
 
-    pipe_type = "lpw" if params.lpw() else "img2img"
-    pipe = load_pipeline(
-        server,
-        params,
-        pipe_type,
-        job.get_device(),
-        # TODO: add LoRAs and TIs
-    )
-    if params.lpw():
-        logger.debug("using LPW pipeline for img2img")
-        rng = torch.manual_seed(params.seed)
-        result = pipe.img2img(
-            params.prompt,
-            generator=rng,
-            guidance_scale=params.cfg,
-            image=source,
-            negative_prompt=params.negative_prompt,
-            num_inference_steps=params.steps,
-            strength=params.strength,
-            callback=callback,
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        strength: float,
+        callback: Optional[ProgressCallback] = None,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        params = params.with_args(**kwargs)
+        logger.info(
+            "blending image using img2img, %s steps: %s", params.steps, params.prompt
         )
-    else:
-        rng = np.random.RandomState(params.seed)
-        result = pipe(
-            params.prompt,
-            generator=rng,
-            guidance_scale=params.cfg,
-            image=source,
-            negative_prompt=params.negative_prompt,
-            num_inference_steps=params.steps,
-            strength=params.strength,
-            callback=callback,
+
+        prompt_pairs, loras, inversions = parse_prompt(params)
+
+        pipe_type = params.get_valid_pipeline("img2img")
+        pipe = load_pipeline(
+            server,
+            params,
+            pipe_type,
+            job.get_device(),
+            inversions=inversions,
+            loras=loras,
         )
 
-    output = result.images[0]
+        pipe_params = {}
+        if pipe_type == "controlnet":
+            pipe_params["controlnet_conditioning_scale"] = strength
+        elif pipe_type == "img2img":
+            pipe_params["strength"] = strength
+        elif pipe_type == "lpw":
+            pipe_params["strength"] = strength
+        elif pipe_type == "panorama":
+            pipe_params["strength"] = strength
+        elif pipe_type == "pix2pix":
+            pipe_params["image_guidance_scale"] = strength
+
+        outputs = []
+        for source in sources:
+            if params.lpw():
+                logger.debug("using LPW pipeline for img2img")
+                rng = torch.manual_seed(params.seed)
+                result = pipe.img2img(
+                    source,
+                    params.prompt,
+                    generator=rng,
+                    guidance_scale=params.cfg,
+                    negative_prompt=params.negative_prompt,
+                    num_inference_steps=params.steps,
+                    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,
+                    generator=rng,
+                    guidance_scale=params.cfg,
+                    image=source,
+                    negative_prompt=params.negative_prompt,
+                    num_inference_steps=params.steps,
+                    callback=callback,
+                    **pipe_params,
+                )
+
+            outputs.extend(result.images)
 
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+        return outputs
diff --git a/api/onnx_web/chain/blend_inpaint.py b/api/onnx_web/chain/blend_inpaint.py
index 3de0e512b..6eef0f117 100644
--- a/api/onnx_web/chain/blend_inpaint.py
+++ b/api/onnx_web/chain/blend_inpaint.py
@@ -1,122 +1,132 @@
 from logging import getLogger
-from typing import Callable, Optional, Tuple
+from typing import Callable, List, Optional, Tuple
 
 import numpy as np
 import torch
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import get_latents_from_seed
+from ..diffusers.utils import get_latents_from_seed, parse_prompt
 from ..image import expand_image, mask_filter_none, noise_source_histogram
 from ..output import save_image
 from ..params import Border, ImageParams, Size, SizeChart, StageParams
 from ..server import ServerContext
 from ..utils import is_debug
 from ..worker import ProgressCallback, WorkerContext
-from .utils import process_tile_order
+from .stage import BaseStage
+from .tile import process_tile_order
 
 logger = getLogger(__name__)
 
 
-def blend_inpaint(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    expand: Border,
-    stage_source: Optional[Image.Image] = None,
-    stage_mask: Optional[Image.Image] = None,
-    fill_color: str = "white",
-    mask_filter: Callable = mask_filter_none,
-    noise_source: Callable = noise_source_histogram,
-    callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    expand = expand.with_args(**kwargs)
-    source = source or stage_source
-    logger.info(
-        "blending image using inpaint, %s steps: %s", params.steps, params.prompt
-    )
+class BlendInpaintStage(BaseStage):
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        expand: Border,
+        stage_source: Optional[Image.Image] = None,
+        stage_mask: Optional[Image.Image] = None,
+        fill_color: str = "white",
+        mask_filter: Callable = mask_filter_none,
+        noise_source: Callable = noise_source_histogram,
+        callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        params = params.with_args(**kwargs)
+        expand = expand.with_args(**kwargs)
+        logger.info(
+            "blending image using inpaint, %s steps: %s", params.steps, params.prompt
+        )
 
-    if stage_mask is None:
-        # if no mask was provided, keep the full source image
-        stage_mask = Image.new("RGB", source.size, "black")
+        _prompt_pairs, loras, inversions = parse_prompt(params)
+        pipe_type = params.get_valid_pipeline("inpaint")
+        pipe = load_pipeline(
+            server,
+            params,
+            pipe_type,
+            job.get_device(),
+            inversions=inversions,
+            loras=loras,
+        )
 
-    source, stage_mask, noise, _full_dims = expand_image(
-        source,
-        stage_mask,
-        expand,
-        fill=fill_color,
-        noise_source=noise_source,
-        mask_filter=mask_filter,
-    )
+        outputs = []
+        for source in sources:
+            if stage_mask is None:
+                # if no mask was provided, keep the full source image
+                stage_mask = Image.new("RGB", source.size, "black")
 
-    if is_debug():
-        save_image(server, "last-source.png", source)
-        save_image(server, "last-mask.png", stage_mask)
-        save_image(server, "last-noise.png", noise)
+            source, stage_mask, noise, _full_dims = expand_image(
+                source,
+                stage_mask,
+                expand,
+                fill=fill_color,
+                noise_source=noise_source,
+                mask_filter=mask_filter,
+            )
 
-    pipe_type = "lpw" if params.lpw() else "inpaint"
-    pipe = load_pipeline(
-        server,
-        params,
-        pipe_type,
-        job.get_device(),
-        # TODO: add LoRAs and TIs
-    )
+            if is_debug():
+                save_image(server, "last-source.png", source)
+                save_image(server, "last-mask.png", stage_mask)
+                save_image(server, "last-noise.png", noise)
 
-    def outpaint(tile_source: Image.Image, dims: Tuple[int, int, int]):
-        left, top, tile = dims
-        size = Size(*tile_source.size)
-        tile_mask = stage_mask.crop((left, top, left + tile, top + tile))
+            def outpaint(tile_source: Image.Image, dims: Tuple[int, int, int]):
+                left, top, tile = dims
+                size = Size(*tile_source.size)
+                tile_mask = stage_mask.crop((left, top, left + tile, top + tile))
 
-        if is_debug():
-            save_image(server, "tile-source.png", tile_source)
-            save_image(server, "tile-mask.png", tile_mask)
+                if is_debug():
+                    save_image(server, "tile-source.png", tile_source)
+                    save_image(server, "tile-mask.png", tile_mask)
 
-        latents = get_latents_from_seed(params.seed, size)
-        if params.lpw():
-            logger.debug("using LPW pipeline for inpaint")
-            rng = torch.manual_seed(params.seed)
-            result = pipe.inpaint(
-                params.prompt,
-                generator=rng,
-                guidance_scale=params.cfg,
-                height=size.height,
-                image=tile_source,
-                latents=latents,
-                mask_image=tile_mask,
-                negative_prompt=params.negative_prompt,
-                num_inference_steps=params.steps,
-                width=size.width,
-                eta=params.eta,
-                callback=callback,
-            )
-        else:
-            rng = np.random.RandomState(params.seed)
-            result = pipe(
-                params.prompt,
-                generator=rng,
-                guidance_scale=params.cfg,
-                height=size.height,
-                image=tile_source,
-                latents=latents,
-                mask_image=stage_mask,
-                negative_prompt=params.negative_prompt,
-                num_inference_steps=params.steps,
-                width=size.width,
-                eta=params.eta,
-                callback=callback,
-            )
+                latents = get_latents_from_seed(params.seed, size)
+                if params.lpw():
+                    logger.debug("using LPW pipeline for inpaint")
+                    rng = torch.manual_seed(params.seed)
+                    result = pipe.inpaint(
+                        params.prompt,
+                        generator=rng,
+                        guidance_scale=params.cfg,
+                        height=size.height,
+                        image=tile_source,
+                        latents=latents,
+                        mask_image=tile_mask,
+                        negative_prompt=params.negative_prompt,
+                        num_inference_steps=params.steps,
+                        width=size.width,
+                        eta=params.eta,
+                        callback=callback,
+                    )
+                else:
+                    rng = np.random.RandomState(params.seed)
+                    result = pipe(
+                        params.prompt,
+                        generator=rng,
+                        guidance_scale=params.cfg,
+                        height=size.height,
+                        image=tile_source,
+                        latents=latents,
+                        mask_image=stage_mask,
+                        negative_prompt=params.negative_prompt,
+                        num_inference_steps=params.steps,
+                        width=size.width,
+                        eta=params.eta,
+                        callback=callback,
+                    )
 
-        return result.images[0]
+                return result.images[0]
 
-    output = process_tile_order(
-        stage.tile_order, source, SizeChart.auto, 1, [outpaint], overlap=params.overlap
-    )
-
-    logger.info("final output image size: %s", output.size)
-    return output
+            outputs.append(
+                process_tile_order(
+                    stage.tile_order,
+                    source,
+                    SizeChart.auto,
+                    1,
+                    [outpaint],
+                    overlap=params.overlap,
+                )
+            )
diff --git a/api/onnx_web/chain/blend_linear.py b/api/onnx_web/chain/blend_linear.py
index d1e6e7336..e56ba2e48 100644
--- a/api/onnx_web/chain/blend_linear.py
+++ b/api/onnx_web/chain/blend_linear.py
@@ -3,27 +3,28 @@
 
 from PIL import Image
 
-from ..image import valid_image
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def blend_linear(
-    _job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    *,
-    alpha: float,
-    sources: Optional[List[Image.Image]] = None,
-    _callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    logger.info("blending image using linear interpolation")
+class BlendLinearStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        alpha: float,
+        stage_source: Optional[Image.Image] = None,
+        _callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        logger.info("blending source images using linear interpolation")
 
-    resized = [valid_image(s) for s in sources]
-
-    return Image.blend(resized[1], resized[0], alpha)
+        return [Image.blend(source, stage_source, alpha) for source in sources]
diff --git a/api/onnx_web/chain/blend_mask.py b/api/onnx_web/chain/blend_mask.py
index 8fe22220f..db956d425 100644
--- a/api/onnx_web/chain/blend_mask.py
+++ b/api/onnx_web/chain/blend_mask.py
@@ -3,40 +3,38 @@
 
 from PIL import Image
 
-from ..image import valid_image
 from ..output import save_image
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..utils import is_debug
 from ..worker import ProgressCallback, WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def blend_mask(
-    _job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    *,
-    sources: Optional[List[Image.Image]] = None,
-    stage_mask: Optional[Image.Image] = None,
-    _callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    logger.info("blending image using mask")
-
-    mult_mask = Image.new("RGBA", stage_mask.size, color="black")
-    mult_mask.alpha_composite(stage_mask)
-    mult_mask = mult_mask.convert("L")
-
-    if is_debug():
-        save_image(server, "last-mask.png", stage_mask)
-        save_image(server, "last-mult-mask.png", mult_mask)
-
-    resized = [
-        valid_image(s, min_dims=mult_mask.size, max_dims=mult_mask.size)
-        for s in sources
-    ]
-
-    return Image.composite(resized[1], resized[0], mult_mask)
+class BlendMaskStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        stage_source: Optional[Image.Image] = None,
+        stage_mask: Optional[Image.Image] = None,
+        _callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        logger.info("blending image using mask")
+
+        mult_mask = Image.new("RGBA", stage_mask.size, color="black")
+        mult_mask.alpha_composite(stage_mask)
+        mult_mask = mult_mask.convert("L")
+
+        if is_debug():
+            save_image(server, "last-mask.png", stage_mask)
+            save_image(server, "last-mult-mask.png", mult_mask)
+
+        return [Image.composite(stage_source, source, mult_mask) for source in sources]
diff --git a/api/onnx_web/chain/blend_pix2pix.py b/api/onnx_web/chain/blend_pix2pix.py
deleted file mode 100644
index a4a2c6143..000000000
--- a/api/onnx_web/chain/blend_pix2pix.py
+++ /dev/null
@@ -1,71 +0,0 @@
-from logging import getLogger
-from typing import Optional
-
-import numpy as np
-import torch
-from PIL import Image
-
-from ..diffusers.load import load_pipeline
-from ..params import ImageParams, StageParams
-from ..server import ServerContext
-from ..worker import ProgressCallback, WorkerContext
-
-logger = getLogger(__name__)
-
-
-def blend_pix2pix(
-    job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    callback: Optional[ProgressCallback] = None,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    source = stage_source or source
-    logger.info(
-        "blending image using instruct pix2pix, %s steps: %s",
-        params.steps,
-        params.prompt,
-    )
-
-    pipe = load_pipeline(
-        server,
-        params,
-        "pix2pix",
-        job.get_device(),
-        # TODO: add LoRAs and TIs
-    )
-    if params.lpw():
-        logger.debug("using LPW pipeline for img2img")
-        rng = torch.manual_seed(params.seed)
-        result = pipe.img2img(
-            params.prompt,
-            generator=rng,
-            guidance_scale=params.cfg,
-            image=source,
-            negative_prompt=params.negative_prompt,
-            num_inference_steps=params.steps,
-            strength=params.strength,
-            callback=callback,
-        )
-    else:
-        rng = np.random.RandomState(params.seed)
-        result = pipe(
-            params.prompt,
-            generator=rng,
-            guidance_scale=params.cfg,
-            image=source,
-            negative_prompt=params.negative_prompt,
-            num_inference_steps=params.steps,
-            strength=params.strength,
-            callback=callback,
-        )
-
-    output = result.images[0]
-
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
diff --git a/api/onnx_web/chain/correct_codeformer.py b/api/onnx_web/chain/correct_codeformer.py
index 60605b437..e86629c13 100644
--- a/api/onnx_web/chain/correct_codeformer.py
+++ b/api/onnx_web/chain/correct_codeformer.py
@@ -1,36 +1,35 @@
 from logging import getLogger
-from typing import Optional
+from typing import List, Optional
 
 from PIL import Image
 
 from ..params import ImageParams, StageParams, UpscaleParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
-device = "cpu"
 
-
-def correct_codeformer(
-    job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    stage_source: Optional[Image.Image] = None,
-    upscale: UpscaleParams,
-    **kwargs,
-) -> Image.Image:
-    # must be within the load function for patch to take effect
-    # TODO: rewrite and remove
-    from codeformer import CodeFormer
-
-    source = stage_source or source
-
-    upscale = upscale.with_args(**kwargs)
-
-    device = job.get_device()
-    pipe = CodeFormer(upscale=upscale.face_outscale).to(device.torch_str())
-    return pipe(source)
+class CorrectCodeformerStage(BaseStage):
+    def run(
+        self,
+        job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        stage_source: Optional[Image.Image] = None,
+        upscale: UpscaleParams,
+        **kwargs,
+    ) -> List[Image.Image]:
+        # must be within the load function for patch to take effect
+        # TODO: rewrite and remove
+        from codeformer import CodeFormer
+
+        upscale = upscale.with_args(**kwargs)
+
+        device = job.get_device()
+        pipe = CodeFormer(upscale=upscale.face_outscale).to(device.torch_str())
+        return [pipe(source) for source in sources]
diff --git a/api/onnx_web/chain/correct_gfpgan.py b/api/onnx_web/chain/correct_gfpgan.py
index 4ef224a58..8dae72ff1 100644
--- a/api/onnx_web/chain/correct_gfpgan.py
+++ b/api/onnx_web/chain/correct_gfpgan.py
@@ -1,84 +1,88 @@
 from logging import getLogger
 from os import path
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 from PIL import Image
 
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import ServerContext
+from ..server import ModelTypes, ServerContext
 from ..utils import run_gc
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def load_gfpgan(
-    server: ServerContext,
-    _stage: StageParams,
-    upscale: UpscaleParams,
-    device: DeviceParams,
-):
-    # must be within the load function for patch to take effect
-    # TODO: rewrite and remove
-    from gfpgan import GFPGANer
-
-    face_path = path.join(server.cache_path, "%s.pth" % (upscale.correction_model))
-    cache_key = (face_path,)
-    cache_pipe = server.cache.get("gfpgan", cache_key)
-
-    if cache_pipe is not None:
-        logger.info("reusing existing GFPGAN pipeline")
-        return cache_pipe
-
-    logger.debug("loading GFPGAN model from %s", face_path)
-
-    # TODO: find a way to pass the ONNX model to underlying architectures
-    gfpgan = GFPGANer(
-        arch="clean",
-        bg_upsampler=None,
-        channel_multiplier=2,
-        device=device.torch_str(),
-        model_path=face_path,
-        upscale=upscale.face_outscale,
-    )
-
-    server.cache.set("gfpgan", cache_key, gfpgan)
-    run_gc([device])
-
-    return gfpgan
-
-
-def correct_gfpgan(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    upscale = upscale.with_args(**kwargs)
-    source = stage_source or source
-
-    if upscale.correction_model is None:
-        logger.warn("no face model given, skipping")
-        return source
-
-    logger.info("correcting faces with GFPGAN model: %s", upscale.correction_model)
-    device = job.get_device()
-    gfpgan = load_gfpgan(server, stage, upscale, device)
-
-    output = np.array(source)
-    _, _, output = gfpgan.enhance(
-        output,
-        has_aligned=False,
-        only_center_face=False,
-        paste_back=True,
-        weight=upscale.face_strength,
-    )
-    output = Image.fromarray(output, "RGB")
-
-    return output
+class CorrectGFPGANStage(BaseStage):
+    def load(
+        self,
+        server: ServerContext,
+        _stage: StageParams,
+        upscale: UpscaleParams,
+        device: DeviceParams,
+    ):
+        # must be within the load function for patch to take effect
+        # TODO: rewrite and remove
+        from gfpgan import GFPGANer
+
+        face_path = path.join(server.cache_path, "%s.pth" % (upscale.correction_model))
+        cache_key = (face_path,)
+        cache_pipe = server.cache.get(ModelTypes.correction, cache_key)
+
+        if cache_pipe is not None:
+            logger.info("reusing existing GFPGAN pipeline")
+            return cache_pipe
+
+        logger.debug("loading GFPGAN model from %s", face_path)
+
+        # TODO: find a way to pass the ONNX model to underlying architectures
+        gfpgan = GFPGANer(
+            arch="clean",
+            bg_upsampler=None,
+            channel_multiplier=2,
+            device=device.torch_str(),
+            model_path=face_path,
+            upscale=upscale.face_outscale,
+        )
+
+        server.cache.set(ModelTypes.correction, cache_key, gfpgan)
+        run_gc([device])
+
+        return gfpgan
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        upscale = upscale.with_args(**kwargs)
+
+        if upscale.correction_model is None:
+            logger.warn("no face model given, skipping")
+            return sources
+
+        logger.info("correcting faces with GFPGAN model: %s", upscale.correction_model)
+        device = job.get_device()
+        gfpgan = self.load(server, stage, upscale, device)
+
+        outputs = []
+        for source in sources:
+            output = np.array(source)
+            _, _, output = gfpgan.enhance(
+                output,
+                has_aligned=False,
+                only_center_face=False,
+                paste_back=True,
+                weight=upscale.face_strength,
+            )
+            outputs.append(Image.fromarray(output, "RGB"))
+
+        return outputs
diff --git a/api/onnx_web/chain/highres.py b/api/onnx_web/chain/highres.py
new file mode 100644
index 000000000..512477484
--- /dev/null
+++ b/api/onnx_web/chain/highres.py
@@ -0,0 +1,64 @@
+from logging import getLogger
+from typing import Optional
+
+from ..chain.base import ChainPipeline
+from ..chain.blend_img2img import BlendImg2ImgStage
+from ..chain.upscale import stage_upscale_correction
+from ..chain.upscale_simple import UpscaleSimpleStage
+from ..params import HighresParams, ImageParams, StageParams, UpscaleParams
+
+logger = getLogger(__name__)
+
+
+def stage_highres(
+    stage: StageParams,
+    params: ImageParams,
+    highres: HighresParams,
+    upscale: UpscaleParams,
+    chain: Optional[ChainPipeline] = None,
+) -> ChainPipeline:
+    logger.info("staging highres pipeline at %s", highres.scale)
+
+    if chain is None:
+        chain = ChainPipeline()
+
+    if not highres.enabled:
+        logger.debug("highres not enabled, skipping")
+        return chain
+
+    if highres.iterations < 1:
+        logger.debug("no highres iterations, skipping")
+        return chain
+
+    for _i in range(highres.iterations):
+        if highres.method == "upscale":
+            logger.debug("using upscaling pipeline for highres")
+            stage_upscale_correction(
+                stage,
+                params,
+                upscale=upscale.with_args(
+                    faces=False,
+                    scale=highres.scale,
+                    outscale=highres.scale,
+                ),
+                chain=chain,
+                overlap=params.overlap,
+            )
+        else:
+            logger.debug("using simple upscaling for highres")
+            chain.stage(
+                UpscaleSimpleStage(),
+                stage,
+                method=highres.method,
+                overlap=params.overlap,
+                upscale=upscale.with_args(scale=highres.scale, outscale=highres.scale),
+            )
+
+        chain.stage(
+            BlendImg2ImgStage(),
+            stage,
+            overlap=params.overlap,
+            strength=highres.strength,
+        )
+
+    return chain
diff --git a/api/onnx_web/chain/persist_disk.py b/api/onnx_web/chain/persist_disk.py
index d5cecdf71..b5d24b965 100644
--- a/api/onnx_web/chain/persist_disk.py
+++ b/api/onnx_web/chain/persist_disk.py
@@ -1,4 +1,5 @@
 from logging import getLogger
+from typing import List
 
 from PIL import Image
 
@@ -6,23 +7,27 @@
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def persist_disk(
-    _job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    output: str,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
+class PersistDiskStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        output: str,
+        stage_source: Image.Image,
+        **kwargs,
+    ) -> List[Image.Image]:
+        for source in sources:
+            # TODO: append index to output name
+            dest = save_image(server, output, source, params=params)
+            logger.info("saved image to %s", dest)
 
-    dest = save_image(server, output, source, params=params)
-    logger.info("saved image to %s", dest)
-    return source
+        return sources
diff --git a/api/onnx_web/chain/persist_s3.py b/api/onnx_web/chain/persist_s3.py
index 3b755f530..a12e39df2 100644
--- a/api/onnx_web/chain/persist_s3.py
+++ b/api/onnx_web/chain/persist_s3.py
@@ -1,6 +1,6 @@
 from io import BytesIO
 from logging import getLogger
-from typing import Optional
+from typing import List, Optional
 
 from boto3 import Session
 from PIL import Image
@@ -8,37 +8,39 @@
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def persist_s3(
-    _job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    output: str,
-    bucket: str,
-    endpoint_url: Optional[str] = None,
-    profile_name: Optional[str] = None,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-
-    session = Session(profile_name=profile_name)
-    s3 = session.client("s3", endpoint_url=endpoint_url)
-
-    data = BytesIO()
-    source.save(data, format=server.image_format)
-    data.seek(0)
-
-    try:
-        s3.upload_fileobj(data, bucket, output)
-        logger.info("saved image to s3://%s/%s", bucket, output)
-    except Exception:
-        logger.exception("error saving image to S3")
-
-    return source
+class PersistS3Stage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        output: str,
+        bucket: str,
+        endpoint_url: Optional[str] = None,
+        profile_name: Optional[str] = None,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        session = Session(profile_name=profile_name)
+        s3 = session.client("s3", endpoint_url=endpoint_url)
+
+        for source in sources:
+            data = BytesIO()
+            source.save(data, format=server.image_format)
+            data.seek(0)
+
+            try:
+                s3.upload_fileobj(data, bucket, output)
+                logger.info("saved image to s3://%s/%s", bucket, output)
+            except Exception:
+                logger.exception("error saving image to S3")
+
+        return sources
diff --git a/api/onnx_web/chain/reduce_crop.py b/api/onnx_web/chain/reduce_crop.py
index 3f2d82db7..c784b1ea9 100644
--- a/api/onnx_web/chain/reduce_crop.py
+++ b/api/onnx_web/chain/reduce_crop.py
@@ -1,29 +1,37 @@
 from logging import getLogger
-from typing import Optional
+from typing import List, Optional
 
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def reduce_crop(
-    _job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    origin: Size,
-    size: Size,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
+class ReduceCropStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        origin: Size,
+        size: Size,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        outputs = []
 
-    image = source.crop((origin.width, origin.height, size.width, size.height))
-    logger.info("created thumbnail with dimensions: %sx%s", image.width, image.height)
-    return image
+        for source in sources:
+            image = source.crop((origin.width, origin.height, size.width, size.height))
+            logger.info(
+                "created thumbnail with dimensions: %sx%s", image.width, image.height
+            )
+            outputs.append(image)
+
+        return outputs
diff --git a/api/onnx_web/chain/reduce_thumbnail.py b/api/onnx_web/chain/reduce_thumbnail.py
index 6df2ed6e1..3ae4f3203 100644
--- a/api/onnx_web/chain/reduce_thumbnail.py
+++ b/api/onnx_web/chain/reduce_thumbnail.py
@@ -1,30 +1,40 @@
 from logging import getLogger
+from typing import List
 
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def reduce_thumbnail(
-    _job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    size: Size,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-    image = source.copy()
-
-    # TODO: should use a call to valid_image
-    image = image.thumbnail((size.width, size.height))
-
-    logger.info("created thumbnail with dimensions: %sx%s", image.width, image.height)
-    return image
+class ReduceThumbnailStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        size: Size,
+        stage_source: Image.Image,
+        **kwargs,
+    ) -> List[Image.Image]:
+        outputs = []
+
+        for source in sources:
+            image = source.copy()
+
+            image = image.thumbnail((size.width, size.height))
+
+            logger.info(
+                "created thumbnail with dimensions: %sx%s", image.width, image.height
+            )
+
+            outputs.append(image)
+
+        return outputs
diff --git a/api/onnx_web/chain/source_noise.py b/api/onnx_web/chain/source_noise.py
index 0092292c7..27b12b03c 100644
--- a/api/onnx_web/chain/source_noise.py
+++ b/api/onnx_web/chain/source_noise.py
@@ -1,34 +1,42 @@
 from logging import getLogger
-from typing import Callable
+from typing import Callable, List
 
 from PIL import Image
 
 from ..params import ImageParams, Size, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def source_noise(
-    _job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    size: Size,
-    noise_source: Callable,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-    logger.info("generating image from noise source")
-
-    if source is not None:
-        logger.warn("a source image was passed to a noise stage, but will be discarded")
-
-    output = noise_source(source, (size.width, size.height), (0, 0))
-
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+class SourceNoiseStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        size: Size,
+        noise_source: Callable,
+        stage_source: Image.Image,
+        **kwargs,
+    ) -> List[Image.Image]:
+        logger.info("generating image from noise source")
+
+        if len(sources) > 0:
+            logger.warn(
+                "source images were passed to a noise stage and will be discarded"
+            )
+
+        outputs = []
+        for source in sources:
+            output = noise_source(source, (size.width, size.height), (0, 0))
+
+            logger.info("final output image size: %sx%s", output.width, output.height)
+            outputs.append(output)
+
+        return outputs
diff --git a/api/onnx_web/chain/source_s3.py b/api/onnx_web/chain/source_s3.py
index 7c4b41584..c10a99d87 100644
--- a/api/onnx_web/chain/source_s3.py
+++ b/api/onnx_web/chain/source_s3.py
@@ -1,6 +1,6 @@
 from io import BytesIO
 from logging import getLogger
-from typing import Optional
+from typing import List, Optional
 
 from boto3 import Session
 from PIL import Image
@@ -8,35 +8,39 @@
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def source_s3(
-    _job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    source_key: str,
-    bucket: str,
-    endpoint_url: Optional[str] = None,
-    profile_name: Optional[str] = None,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-
-    session = Session(profile_name=profile_name)
-    s3 = session.client("s3", endpoint_url=endpoint_url)
-
-    try:
-        logger.info("loading image from s3://%s/%s", bucket, source_key)
-        data = BytesIO()
-        s3.download_fileobj(bucket, source_key, data)
-
-        data.seek(0)
-        return Image.open(data)
-    except Exception:
-        logger.exception("error loading image from S3")
+class SourceS3Stage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        _sources: List[Image.Image],
+        *,
+        source_keys: List[str],
+        bucket: str,
+        endpoint_url: Optional[str] = None,
+        profile_name: Optional[str] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        session = Session(profile_name=profile_name)
+        s3 = session.client("s3", endpoint_url=endpoint_url)
+
+        outputs = []
+        for key in source_keys:
+            try:
+                logger.info("loading image from s3://%s/%s", bucket, key)
+                data = BytesIO()
+                s3.download_fileobj(bucket, key, data)
+
+                data.seek(0)
+                outputs.append(Image.open(data))
+            except Exception:
+                logger.exception("error loading image from S3")
+
+        return outputs
diff --git a/api/onnx_web/chain/source_txt2img.py b/api/onnx_web/chain/source_txt2img.py
index 436bca205..ebc8250f5 100644
--- a/api/onnx_web/chain/source_txt2img.py
+++ b/api/onnx_web/chain/source_txt2img.py
@@ -6,75 +6,90 @@
 from PIL import Image
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import get_latents_from_seed
-from ..params import ImageParams, Size, StageParams
+from ..diffusers.utils import encode_prompt, get_latents_from_seed, parse_prompt
+from ..params import ImageParams, Size, SizeChart, StageParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def source_txt2img(
-    job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    _source: Image.Image,
-    *,
-    size: Size,
-    callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    size = size.with_args(**kwargs)
-    logger.info(
-        "generating image using txt2img, %s steps: %s", params.steps, params.prompt
-    )
+class SourceTxt2ImgStage(BaseStage):
+    max_tile = SizeChart.unlimited
 
-    if "stage_source" in kwargs:
-        logger.warn(
-            "a source image was passed to a txt2img stage, and will be discarded"
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        params: ImageParams,
+        _source: Image.Image,
+        *,
+        size: Size,
+        callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> Image.Image:
+        params = params.with_args(**kwargs)
+        size = size.with_args(**kwargs)
+        logger.info(
+            "generating image using txt2img, %s steps: %s", params.steps, params.prompt
         )
 
-    latents = get_latents_from_seed(params.seed, size)
-    pipe_type = "lpw" if params.lpw() else "txt2img"
-    pipe = load_pipeline(
-        server,
-        params,
-        pipe_type,
-        job.get_device(),
-        # TODO: add LoRAs and TIs
-    )
+        if "stage_source" in kwargs:
+            logger.warn(
+                "a source image was passed to a txt2img stage, and will be discarded"
+            )
 
-    if params.lpw():
-        logger.debug("using LPW pipeline for txt2img")
-        rng = torch.manual_seed(params.seed)
-        result = pipe.text2img(
-            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,
-            callback=callback,
-        )
-    else:
-        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,
-            callback=callback,
+        prompt_pairs, loras, inversions = parse_prompt(params)
+
+        latents = get_latents_from_seed(params.seed, size, params.batch)
+        pipe_type = params.get_valid_pipeline("txt2img")
+        pipe = load_pipeline(
+            server,
+            params,
+            pipe_type,
+            job.get_device(),
+            inversions=inversions,
+            loras=loras,
         )
 
-    output = result.images[0]
+        if params.lpw():
+            logger.debug("using LPW pipeline for txt2img")
+            rng = torch.manual_seed(params.seed)
+            result = pipe.text2img(
+                params.prompt,
+                height=size.height,
+                width=size.width,
+                generator=rng,
+                guidance_scale=params.cfg,
+                latents=latents,
+                negative_prompt=params.negative_prompt,
+                num_images_per_prompt=params.batch,
+                num_inference_steps=params.steps,
+                eta=params.eta,
+                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,
+                height=size.height,
+                width=size.width,
+                generator=rng,
+                guidance_scale=params.cfg,
+                latents=latents,
+                negative_prompt=params.negative_prompt,
+                num_images_per_prompt=params.batch,
+                num_inference_steps=params.steps,
+                eta=params.eta,
+                callback=callback,
+            )
 
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+        return result.images
diff --git a/api/onnx_web/chain/source_url.py b/api/onnx_web/chain/source_url.py
index 3496be953..86c4f67e0 100644
--- a/api/onnx_web/chain/source_url.py
+++ b/api/onnx_web/chain/source_url.py
@@ -1,5 +1,6 @@
 from io import BytesIO
 from logging import getLogger
+from typing import List
 
 import requests
 from PIL import Image
@@ -7,31 +8,37 @@
 from ..params import ImageParams, StageParams
 from ..server import ServerContext
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def source_url(
-    _job: WorkerContext,
-    _server: ServerContext,
-    _stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    source_url: str,
-    stage_source: Image.Image,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-    logger.info("loading image from URL source")
-
-    if source is not None:
-        logger.warn(
-            "a source image was passed to a source stage, and will be discarded"
-        )
-
-    response = requests.get(source_url)
-    output = Image.open(BytesIO(response.content))
-
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+class SourceURLStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        source_urls: List[str],
+        stage_source: Image.Image,
+        **kwargs,
+    ) -> List[Image.Image]:
+        logger.info("loading image from URL source")
+
+        if len(sources) > 0:
+            logger.warn(
+                "a source image was passed to a source stage, and will be discarded"
+            )
+
+        outputs = []
+        for url in source_urls:
+            response = requests.get(url)
+            output = Image.open(BytesIO(response.content))
+
+            logger.info("final output image size: %sx%s", output.width, output.height)
+            outputs.append(output)
+
+        return outputs
diff --git a/api/onnx_web/chain/stage.py b/api/onnx_web/chain/stage.py
new file mode 100644
index 000000000..11c61774b
--- /dev/null
+++ b/api/onnx_web/chain/stage.py
@@ -0,0 +1,31 @@
+from typing import List, Optional
+
+from PIL import Image
+
+from ..params import ImageParams, Size, SizeChart, StageParams
+from ..server.context import ServerContext
+from ..worker.context import WorkerContext
+
+
+class BaseStage:
+    max_tile = SizeChart.auto
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *args,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        raise NotImplementedError()
+
+    def steps(
+        self,
+        _params: ImageParams,
+        size: Size,
+    ) -> int:
+        raise NotImplementedError()
diff --git a/api/onnx_web/chain/utils.py b/api/onnx_web/chain/tile.py
similarity index 88%
rename from api/onnx_web/chain/utils.py
rename to api/onnx_web/chain/tile.py
index f805144f0..e766e7159 100644
--- a/api/onnx_web/chain/utils.py
+++ b/api/onnx_web/chain/tile.py
@@ -1,11 +1,11 @@
 from logging import getLogger
 from math import ceil
-from typing import List, Protocol, Tuple
+from typing import List, Optional, Protocol, Tuple
 
 import numpy as np
 from PIL import Image
 
-from ..params import TileOrder
+from ..params import Size, TileOrder
 
 # from skimage.exposure import match_histograms
 
@@ -29,6 +29,9 @@ def complete_tile(
     source: Image.Image,
     tile: int,
 ) -> Image.Image:
+    if source is None:
+        return source
+
     if source.width < tile or source.height < tile:
         full_source = Image.new(source.mode, (tile, tile))
         full_source.paste(source)
@@ -37,6 +40,23 @@ def complete_tile(
     return source
 
 
+def needs_tile(
+    max_tile: int,
+    stage_tile: int,
+    size: Optional[Size] = None,
+    source: Optional[Image.Image] = None,
+) -> bool:
+    tile = min(max_tile, stage_tile)
+
+    if source is not None:
+        return source.width > tile or source.height > tile
+
+    if size is not None:
+        return size.width > tile or size.height > tile
+
+    return False
+
+
 def get_tile_grads(
     left: int,
     top: int,
@@ -71,7 +91,7 @@ def blend_tiles(
     overlap: float,
 ):
     adj_tile = int(float(tile) * (1.0 - overlap))
-    logger.trace(
+    logger.debug(
         "adjusting tile size from %s to %s based on %s overlap", tile, adj_tile, overlap
     )
 
@@ -92,7 +112,7 @@ def blend_tiles(
 
             # gradient blending
             grad_x, grad_y = get_tile_grads(left, top, adj_tile, width, height)
-            logger.trace("tile gradients: %s, %s, %s", points, grad_x, grad_y)
+            logger.debug("tile gradients: %s, %s, %s", points, grad_x, grad_y)
 
             mult_x = [np.interp(i, points, grad_x) for i in range(tile * scale)]
             mult_y = [np.interp(i, points, grad_y) for i in range(tile * scale)]
@@ -129,6 +149,7 @@ def blend_tiles(
             axis=2,
         )
 
+    logger.trace("mean tiles contributing to each pixel: %s", np.mean(count))
     pixels = np.where(count > 0, value / count, value)
     return Image.fromarray(np.uint8(pixels))
 
@@ -141,12 +162,20 @@ def process_tile_grid(
     overlap: float = 0.0,
     **kwargs,
 ) -> Image.Image:
-    width, height = source.size
+    width, height = kwargs.get("size", source.size if source else None)
 
     adj_tile = int(float(tile) * (1.0 - overlap))
     tiles_x = ceil(width / adj_tile)
     tiles_y = ceil(height / adj_tile)
     total = tiles_x * tiles_y
+    logger.debug(
+        "processing %s tiles (%s x %s) with adjusted size of %s, %s overlap",
+        total,
+        tiles_x,
+        tiles_y,
+        adj_tile,
+        overlap,
+    )
 
     tiles: List[Tuple[int, int, Image.Image]] = []
 
@@ -157,7 +186,9 @@ def process_tile_grid(
             top = y * adj_tile
             logger.info("processing tile %s of %s, %s.%s", idx + 1, total, y, x)
 
-            tile_image = source.crop((left, top, left + tile, top + tile))
+            tile_image = (
+                source.crop((left, top, left + tile, top + tile)) if source else None
+            )
             tile_image = complete_tile(tile_image, tile)
 
             for filter in filters:
@@ -179,7 +210,7 @@ def process_tile_spiral(
     if scale != 1:
         raise ValueError("unsupported scale")
 
-    width, height = source.size
+    width, height = kwargs.get("size", source.size if source else None)
 
     # spiral uses the previous run and needs a scratch texture for 3x memory
     image = Image.new("RGB", (width * scale, height * scale))
@@ -197,7 +228,7 @@ def process_tile_spiral(
             "processing tile %s of %s, %sx%s", counter, len(tile_coords), left, top
         )
 
-        tile_image = image.crop((left, top, left + tile, top + tile))
+        tile_image = image.crop((left, top, left + tile, top + tile)) if image else None
         tile_image = complete_tile(tile_image, tile)
 
         for filter in filters:
diff --git a/api/onnx_web/chain/upscale.py b/api/onnx_web/chain/upscale.py
new file mode 100644
index 000000000..2a299a5e0
--- /dev/null
+++ b/api/onnx_web/chain/upscale.py
@@ -0,0 +1,128 @@
+from logging import getLogger
+from typing import List, Optional, Tuple
+
+from ..params import ImageParams, SizeChart, StageParams, UpscaleParams
+from . import ChainPipeline, PipelineStage
+from .correct_codeformer import CorrectCodeformerStage
+from .correct_gfpgan import CorrectGFPGANStage
+from .upscale_bsrgan import UpscaleBSRGANStage
+from .upscale_resrgan import UpscaleRealESRGANStage
+from .upscale_stable_diffusion import UpscaleStableDiffusionStage
+from .upscale_swinir import UpscaleSwinIRStage
+
+logger = getLogger(__name__)
+
+
+def split_upscale(
+    upscale: UpscaleParams,
+) -> Tuple[Optional[UpscaleParams], UpscaleParams]:
+    if upscale.faces and (
+        upscale.upscale_order == "correction-both"
+        or upscale.upscale_order == "correction-first"
+    ):
+        return (
+            upscale.with_args(
+                scale=1,
+                outscale=1,
+            ),
+            upscale.with_args(
+                upscale_order="correction-last",
+            ),
+        )
+    else:
+        return (
+            None,
+            upscale,
+        )
+
+
+def stage_upscale_correction(
+    stage: StageParams,
+    params: ImageParams,
+    *,
+    upscale: UpscaleParams,
+    chain: Optional[ChainPipeline] = None,
+    pre_stages: Optional[List[PipelineStage]] = None,
+    post_stages: Optional[List[PipelineStage]] = None,
+    **kwargs,
+) -> ChainPipeline:
+    """
+    This is a convenience method for a chain pipeline that will run upscaling and
+    correction, based on the `upscale` params.
+    """
+    logger.info(
+        "staging upscaling and correction pipeline at %s:%s",
+        upscale.scale,
+        upscale.outscale,
+    )
+
+    if chain is None:
+        chain = ChainPipeline()
+
+    if pre_stages is not None:
+        for pre_stage in pre_stages:
+            chain.append(pre_stage)
+
+    upscale_opts = {
+        **kwargs,
+        "upscale": upscale,
+    }
+    upscale_stage: Optional[PipelineStage] = None
+    if upscale.scale > 1:
+        if "bsrgan" in upscale.upscale_model:
+            bsrgan_params = StageParams(
+                tile_size=stage.tile_size,
+                outscale=upscale.outscale,
+            )
+            upscale_stage = (UpscaleBSRGANStage(), bsrgan_params, upscale_opts)
+        elif "esrgan" in upscale.upscale_model:
+            esrgan_params = StageParams(
+                tile_size=stage.tile_size,
+                outscale=upscale.outscale,
+            )
+            upscale_stage = (UpscaleRealESRGANStage(), esrgan_params, upscale_opts)
+        elif "stable-diffusion" in upscale.upscale_model:
+            mini_tile = min(SizeChart.mini, stage.tile_size)
+            sd_params = StageParams(tile_size=mini_tile, outscale=upscale.outscale)
+            upscale_stage = (UpscaleStableDiffusionStage(), sd_params, upscale_opts)
+        elif "swinir" in upscale.upscale_model:
+            swinir_params = StageParams(
+                tile_size=stage.tile_size,
+                outscale=upscale.outscale,
+            )
+            upscale_stage = (UpscaleSwinIRStage(), swinir_params, upscale_opts)
+        else:
+            logger.warn("unknown upscaling model: %s", upscale.upscale_model)
+
+    correct_stage: Optional[PipelineStage] = None
+    if upscale.faces:
+        face_params = StageParams(
+            tile_size=stage.tile_size, outscale=upscale.face_outscale
+        )
+        if upscale.correction_model is None:
+            logger.warn("no correction model set, skipping")
+        elif "codeformer" in upscale.correction_model:
+            correct_stage = (CorrectCodeformerStage(), face_params, upscale_opts)
+        elif "gfpgan" in upscale.correction_model:
+            correct_stage = (CorrectGFPGANStage(), face_params, upscale_opts)
+        else:
+            logger.warn("unknown correction model: %s", upscale.correction_model)
+
+    if upscale.upscale_order == "correction-both":
+        chain.append(correct_stage)
+        chain.append(upscale_stage)
+        chain.append(correct_stage)
+    elif upscale.upscale_order == "correction-first":
+        chain.append(correct_stage)
+        chain.append(upscale_stage)
+    elif upscale.upscale_order == "correction-last":
+        chain.append(upscale_stage)
+        chain.append(correct_stage)
+    else:
+        logger.warn("unknown upscaling order: %s", upscale.upscale_order)
+
+    if post_stages is not None:
+        for post_stage in post_stages:
+            chain.append(post_stage)
+
+    return chain
diff --git a/api/onnx_web/chain/upscale_bsrgan.py b/api/onnx_web/chain/upscale_bsrgan.py
index 37350d090..852d5806b 100644
--- a/api/onnx_web/chain/upscale_bsrgan.py
+++ b/api/onnx_web/chain/upscale_bsrgan.py
@@ -1,118 +1,110 @@
 from logging import getLogger
 from os import path
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 from PIL import Image
 
 from ..models.onnx import OnnxModel
-from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import ServerContext
+from ..params import DeviceParams, ImageParams, Size, StageParams, UpscaleParams
+from ..server import ModelTypes, ServerContext
 from ..utils import run_gc
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def load_bsrgan(
-    server: ServerContext,
-    _stage: StageParams,
-    upscale: UpscaleParams,
-    device: DeviceParams,
-):
-    # must be within the load function for patch to take effect
-    model_path = path.join(server.model_path, "%s.onnx" % (upscale.upscale_model))
-    cache_key = (model_path,)
-    cache_pipe = server.cache.get("bsrgan", cache_key)
-
-    if cache_pipe is not None:
-        logger.debug("reusing existing BSRGAN pipeline")
-        return cache_pipe
-
-    logger.info("loading BSRGAN model from %s", model_path)
-
-    pipe = OnnxModel(
-        server,
-        model_path,
-        provider=device.ort_provider(),
-        sess_options=device.sess_options(),
-    )
-
-    server.cache.set("bsrgan", cache_key, pipe)
-    run_gc([device])
-
-    return pipe
-
-
-def upscale_bsrgan(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    upscale = upscale.with_args(**kwargs)
-    source = stage_source or source
-
-    if upscale.upscale_model is None:
-        logger.warn("no upscaling model given, skipping")
-        return source
-
-    logger.info("upscaling with BSRGAN model: %s", upscale.upscale_model)
-    device = job.get_device()
-    bsrgan = load_bsrgan(server, stage, upscale, device)
-
-    tile_size = (64, 64)
-    tile_x = source.width // tile_size[0]
-    tile_y = source.height // tile_size[1]
-
-    image = np.array(source) / 255.0
-    image = image[:, :, [2, 1, 0]].astype(np.float32).transpose((2, 0, 1))
-    image = np.expand_dims(image, axis=0)
-    logger.trace("BSRGAN input shape: %s", image.shape)
-
-    scale = upscale.outscale
-    dest = np.zeros(
-        (image.shape[0], image.shape[1], image.shape[2] * scale, image.shape[3] * scale)
-    )
-    logger.trace("BSRGAN output shape: %s", dest.shape)
-
-    for x in range(tile_x):
-        for y in range(tile_y):
-            xt = x * tile_size[0]
-            yt = y * tile_size[1]
-
-            ix1 = xt
-            ix2 = xt + tile_size[0]
-            iy1 = yt
-            iy2 = yt + tile_size[1]
-            logger.debug(
-                "running BSRGAN on tile: (%s, %s, %s, %s) -> (%s, %s, %s, %s)",
-                ix1,
-                ix2,
-                iy1,
-                iy2,
-                ix1 * scale,
-                ix2 * scale,
-                iy1 * scale,
-                iy2 * scale,
+class UpscaleBSRGANStage(BaseStage):
+    max_tile = 64
+
+    def load(
+        self,
+        server: ServerContext,
+        _stage: StageParams,
+        upscale: UpscaleParams,
+        device: DeviceParams,
+    ):
+        # must be within the load function for patch to take effect
+        model_path = path.join(server.model_path, "%s.onnx" % (upscale.upscale_model))
+        cache_key = (model_path,)
+        cache_pipe = server.cache.get(ModelTypes.upscaling, cache_key)
+
+        if cache_pipe is not None:
+            logger.debug("reusing existing BSRGAN pipeline")
+            return cache_pipe
+
+        logger.info("loading BSRGAN model from %s", model_path)
+
+        pipe = OnnxModel(
+            server,
+            model_path,
+            provider=device.ort_provider(),
+            sess_options=device.sess_options(),
+        )
+
+        server.cache.set(ModelTypes.upscaling, cache_key, pipe)
+        run_gc([device])
+
+        return pipe
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        upscale = upscale.with_args(**kwargs)
+
+        if upscale.upscale_model is None:
+            logger.warn("no upscaling model given, skipping")
+            return sources
+
+        logger.info("upscaling with BSRGAN model: %s", upscale.upscale_model)
+        device = job.get_device()
+        bsrgan = self.load(server, stage, upscale, device)
+
+        outputs = []
+        for source in sources:
+            image = np.array(source) / 255.0
+            image = image[:, :, [2, 1, 0]].astype(np.float32).transpose((2, 0, 1))
+            image = np.expand_dims(image, axis=0)
+            logger.trace("BSRGAN input shape: %s", image.shape)
+
+            scale = upscale.outscale
+            dest = np.zeros(
+                (
+                    image.shape[0],
+                    image.shape[1],
+                    image.shape[2] * scale,
+                    image.shape[3] * scale,
+                )
             )
+            logger.trace("BSRGAN output shape: %s", dest.shape)
 
-            dest[
-                :,
-                :,
-                ix1 * scale : ix2 * scale,
-                iy1 * scale : iy2 * scale,
-            ] = bsrgan(image[:, :, ix1:ix2, iy1:iy2])
+            dest = bsrgan(image)
 
-    dest = np.clip(np.squeeze(dest, axis=0), 0, 1)
-    dest = dest[[2, 1, 0], :, :].transpose((1, 2, 0))
-    dest = (dest * 255.0).round().astype(np.uint8)
+            dest = np.clip(np.squeeze(dest, axis=0), 0, 1)
+            dest = dest[[2, 1, 0], :, :].transpose((1, 2, 0))
+            dest = (dest * 255.0).round().astype(np.uint8)
 
-    output = Image.fromarray(dest, "RGB")
-    logger.debug("output image size: %s x %s", output.width, output.height)
-    return output
+            output = Image.fromarray(dest, "RGB")
+            logger.debug("output image size: %s x %s", output.width, output.height)
+
+            outputs.append(output)
+
+        return outputs
+
+    def steps(
+        self,
+        params: ImageParams,
+        size: Size,
+    ) -> int:
+        tile = min(params.tiles, self.max_tile)
+        return size.width // tile * size.height // tile
diff --git a/api/onnx_web/chain/upscale_highres.py b/api/onnx_web/chain/upscale_highres.py
new file mode 100644
index 000000000..4ae960141
--- /dev/null
+++ b/api/onnx_web/chain/upscale_highres.py
@@ -0,0 +1,45 @@
+from logging import getLogger
+from typing import List, Optional
+
+from PIL import Image
+
+from ..chain.highres import stage_highres
+from ..params import HighresParams, ImageParams, StageParams, UpscaleParams
+from ..server import ServerContext
+from ..worker import WorkerContext
+from ..worker.context import ProgressCallback
+from .stage import BaseStage
+
+logger = getLogger(__name__)
+
+
+class UpscaleHighresStage(BaseStage):
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *args,
+        highres: HighresParams,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        if highres.scale <= 1:
+            return sources
+
+        chain = stage_highres(stage, params, highres, upscale)
+
+        return [
+            chain(
+                job,
+                server,
+                params,
+                source,
+                callback=callback,
+            )
+            for source in sources
+        ]
diff --git a/api/onnx_web/chain/upscale_outpaint.py b/api/onnx_web/chain/upscale_outpaint.py
index 4db1a3ef4..9276cfbc9 100644
--- a/api/onnx_web/chain/upscale_outpaint.py
+++ b/api/onnx_web/chain/upscale_outpaint.py
@@ -1,155 +1,167 @@
 from logging import getLogger
-from typing import Callable, Optional, Tuple
+from typing import Callable, List, Optional, Tuple
 
 import numpy as np
 import torch
 from PIL import Image, ImageDraw, ImageOps
 
 from ..diffusers.load import load_pipeline
-from ..diffusers.utils import get_latents_from_seed, get_tile_latents
+from ..diffusers.utils import get_latents_from_seed, get_tile_latents, parse_prompt
 from ..image import expand_image, mask_filter_none, noise_source_histogram
 from ..output import save_image
 from ..params import Border, ImageParams, Size, SizeChart, StageParams
 from ..server import ServerContext
 from ..utils import is_debug
 from ..worker import ProgressCallback, WorkerContext
-from .utils import complete_tile, process_tile_grid, process_tile_order
+from .stage import BaseStage
+from .tile import complete_tile, process_tile_grid, process_tile_order
 
 logger = getLogger(__name__)
 
 
-def upscale_outpaint(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    border: Border,
-    stage_source: Optional[Image.Image] = None,
-    stage_mask: Optional[Image.Image] = None,
-    fill_color: str = "white",
-    mask_filter: Callable = mask_filter_none,
-    noise_source: Callable = noise_source_histogram,
-    callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-    logger.info(
-        "upscaling %s x %s image by expanding borders: %s",
-        source.width,
-        source.height,
-        border,
-    )
-
-    margin_x = float(max(border.left, border.right))
-    margin_y = float(max(border.top, border.bottom))
-    overlap = min(margin_x / source.width, margin_y / source.height)
-
-    if stage_mask is None:
-        # if no mask was provided, keep the full source image
-        stage_mask = Image.new("RGB", source.size, "black")
-
-    # masks start as 512x512, resize to cover the source, then trim the extra
-    mask_max = max(source.width, source.height)
-    stage_mask = ImageOps.contain(stage_mask, (mask_max, mask_max))
-    stage_mask = stage_mask.crop((0, 0, source.width, source.height))
-
-    source, stage_mask, noise, full_size = expand_image(
-        source,
-        stage_mask,
-        border,
-        fill=fill_color,
-        noise_source=noise_source,
-        mask_filter=mask_filter,
-    )
-
-    full_latents = get_latents_from_seed(params.seed, Size(*full_size))
-
-    draw_mask = ImageDraw.Draw(stage_mask)
-
-    if is_debug():
-        save_image(server, "last-source.png", source)
-        save_image(server, "last-mask.png", stage_mask)
-        save_image(server, "last-noise.png", noise)
-
-    def outpaint(tile_source: Image.Image, dims: Tuple[int, int, int]):
-        left, top, tile = dims
-        size = Size(*tile_source.size)
-        tile_mask = stage_mask.crop((left, top, left + tile, top + tile))
-        tile_mask = complete_tile(tile_mask, tile)
-
-        if is_debug():
-            save_image(server, "tile-source.png", tile_source)
-            save_image(server, "tile-mask.png", tile_mask)
-
-        latents = get_tile_latents(full_latents, dims, size)
+class UpscaleOutpaintStage(BaseStage):
+    max_tile = SizeChart.unlimited
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        border: Border,
+        stage_source: Optional[Image.Image] = None,
+        stage_mask: Optional[Image.Image] = None,
+        fill_color: str = "white",
+        mask_filter: Callable = mask_filter_none,
+        noise_source: Callable = noise_source_histogram,
+        callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        _prompt_pairs, loras, inversions = parse_prompt(params)
+
         pipe_type = params.get_valid_pipeline("inpaint", params.pipeline)
         pipe = load_pipeline(
             server,
             params,
             pipe_type,
             job.get_device(),
-            # TODO: load LoRAs and TIs
+            inversions=inversions,
+            loras=loras,
         )
-        if params.lpw():
-            logger.debug("using LPW pipeline for inpaint")
-            rng = torch.manual_seed(params.seed)
-            result = pipe.inpaint(
-                tile_source,
-                tile_mask,
-                params.prompt,
-                generator=rng,
-                guidance_scale=params.cfg,
-                height=size.height,
-                latents=latents,
-                negative_prompt=params.negative_prompt,
-                num_inference_steps=params.steps,
-                width=size.width,
-                callback=callback,
-            )
-        else:
-            rng = np.random.RandomState(params.seed)
-            result = pipe(
-                params.prompt,
-                tile_source,
-                tile_mask,
-                height=size.height,
-                width=size.width,
-                num_inference_steps=params.steps,
-                guidance_scale=params.cfg,
-                negative_prompt=params.negative_prompt,
-                generator=rng,
-                latents=latents,
-                callback=callback,
+
+        outputs = []
+        for source in sources:
+            logger.info(
+                "upscaling %s x %s image by expanding borders: %s",
+                source.width,
+                source.height,
+                border,
             )
 
-        # once part of the image has been drawn, keep it
-        draw_mask.rectangle((left, top, left + tile, top + tile), fill="black")
-        return result.images[0]
-
-    if params.pipeline == "panorama":
-        logger.debug("outpainting with one shot panorama, no tiling")
-        return outpaint(source, (0, 0, max(source.width, source.height)))
-    if overlap == 0:
-        logger.debug("outpainting with 0 margin, using grid tiling")
-        output = process_tile_grid(source, SizeChart.auto, 1, [outpaint])
-    elif border.left == border.right and border.top == border.bottom:
-        logger.debug(
-            "outpainting with an even border, using spiral tiling with %s overlap",
-            overlap,
-        )
-        output = process_tile_order(
-            stage.tile_order,
-            source,
-            SizeChart.auto,
-            1,
-            [outpaint],
-            overlap=overlap,
-        )
-    else:
-        logger.debug("outpainting with an uneven border, using grid tiling")
-        output = process_tile_grid(source, SizeChart.auto, 1, [outpaint])
+            margin_x = float(max(border.left, border.right))
+            margin_y = float(max(border.top, border.bottom))
+            overlap = min(margin_x / source.width, margin_y / source.height)
+
+            if stage_mask is None:
+                # if no mask was provided, keep the full source image
+                stage_mask = Image.new("RGB", source.size, "black")
+
+            # masks start as 512x512, resize to cover the source, then trim the extra
+            mask_max = max(source.width, source.height)
+            stage_mask = ImageOps.contain(stage_mask, (mask_max, mask_max))
+            stage_mask = stage_mask.crop((0, 0, source.width, source.height))
+
+            source, stage_mask, noise, full_size = expand_image(
+                source,
+                stage_mask,
+                border,
+                fill=fill_color,
+                noise_source=noise_source,
+                mask_filter=mask_filter,
+            )
 
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+            full_latents = get_latents_from_seed(params.seed, Size(*full_size))
+
+            draw_mask = ImageDraw.Draw(stage_mask)
+
+            if is_debug():
+                save_image(server, "last-source.png", source)
+                save_image(server, "last-mask.png", stage_mask)
+                save_image(server, "last-noise.png", noise)
+
+            def outpaint(tile_source: Image.Image, dims: Tuple[int, int, int]):
+                left, top, tile = dims
+                size = Size(*tile_source.size)
+                tile_mask = stage_mask.crop((left, top, left + tile, top + tile))
+                tile_mask = complete_tile(tile_mask, tile)
+
+                if is_debug():
+                    save_image(server, "tile-source.png", tile_source)
+                    save_image(server, "tile-mask.png", tile_mask)
+
+                latents = get_tile_latents(full_latents, dims, size)
+                if params.lpw():
+                    logger.debug("using LPW pipeline for inpaint")
+                    rng = torch.manual_seed(params.seed)
+                    result = pipe.inpaint(
+                        tile_source,
+                        tile_mask,
+                        params.prompt,
+                        generator=rng,
+                        guidance_scale=params.cfg,
+                        height=size.height,
+                        latents=latents,
+                        negative_prompt=params.negative_prompt,
+                        num_inference_steps=params.steps,
+                        width=size.width,
+                        callback=callback,
+                    )
+                else:
+                    rng = np.random.RandomState(params.seed)
+                    result = pipe(
+                        params.prompt,
+                        tile_source,
+                        tile_mask,
+                        height=size.height,
+                        width=size.width,
+                        num_inference_steps=params.steps,
+                        guidance_scale=params.cfg,
+                        negative_prompt=params.negative_prompt,
+                        generator=rng,
+                        latents=latents,
+                        callback=callback,
+                    )
+
+                # once part of the image has been drawn, keep it
+                draw_mask.rectangle((left, top, left + tile, top + tile), fill="black")
+                return result.images[0]
+
+            if params.pipeline == "panorama":
+                logger.debug("outpainting with one shot panorama, no tiling")
+                output = outpaint(source, (0, 0, max(source.width, source.height)))
+            if overlap == 0:
+                logger.debug("outpainting with 0 margin, using grid tiling")
+                output = process_tile_grid(source, SizeChart.auto, 1, [outpaint])
+            elif border.left == border.right and border.top == border.bottom:
+                logger.debug(
+                    "outpainting with an even border, using spiral tiling with %s overlap",
+                    overlap,
+                )
+                output = process_tile_order(
+                    stage.tile_order,
+                    source,
+                    SizeChart.auto,
+                    1,
+                    [outpaint],
+                    overlap=overlap,
+                )
+            else:
+                logger.debug("outpainting with an uneven border, using grid tiling")
+                output = process_tile_grid(source, SizeChart.auto, 1, [outpaint])
+
+            logger.info("final output image size: %sx%s", output.width, output.height)
+            outputs.append(output)
+
+        return outputs
diff --git a/api/onnx_web/chain/upscale_resrgan.py b/api/onnx_web/chain/upscale_resrgan.py
index f3e68cede..a75cc609b 100644
--- a/api/onnx_web/chain/upscale_resrgan.py
+++ b/api/onnx_web/chain/upscale_resrgan.py
@@ -1,94 +1,101 @@
 from logging import getLogger
 from os import path
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 from PIL import Image
 
 from ..onnx import OnnxRRDBNet
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import ServerContext
+from ..server import ModelTypes, ServerContext
 from ..utils import run_gc
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 TAG_X4_V3 = "real-esrgan-x4-v3"
 
 
-def load_resrgan(
-    server: ServerContext, params: UpscaleParams, device: DeviceParams, tile=0
-):
-    # must be within load function for patches to take effect
-    # TODO: rewrite and remove
-    from realesrgan import RealESRGANer
-
-    model_file = "%s.%s" % (params.upscale_model, params.format)
-    model_path = path.join(server.model_path, model_file)
-
-    cache_key = (model_path, params.format)
-    cache_pipe = server.cache.get("resrgan", cache_key)
-    if cache_pipe is not None:
-        logger.info("reusing existing Real ESRGAN pipeline")
-        return cache_pipe
-
-    if not path.isfile(model_path):
-        raise FileNotFoundError("Real ESRGAN model not found at %s" % model_path)
-
-    # TODO: swap for regular RRDBNet after rewriting wrapper
-    model = OnnxRRDBNet(
-        server,
-        model_file,
-        provider=device.ort_provider(),
-        sess_options=device.sess_options(),
-    )
-
-    dni_weight = None
-    if params.upscale_model == TAG_X4_V3 and params.denoise != 1:
-        wdn_model_path = model_path.replace(TAG_X4_V3, "%s-wdn" % TAG_X4_V3)
-        model_path = [model_path, wdn_model_path]
-        dni_weight = [params.denoise, 1 - params.denoise]
-
-    logger.debug("loading Real ESRGAN upscale model from %s", model_path)
-
-    # TODO: shouldn't need the PTH file
-    model_path_pth = path.join(server.cache_path, ("%s.pth" % params.upscale_model))
-    upsampler = RealESRGANer(
-        scale=params.scale,
-        model_path=model_path_pth,
-        dni_weight=dni_weight,
-        model=model,
-        tile=tile,
-        tile_pad=params.tile_pad,
-        pre_pad=params.pre_pad,
-        half=False,  # TODO: use server optimizations
-    )
-
-    server.cache.set("resrgan", cache_key, upsampler)
-    run_gc([device])
-
-    return upsampler
-
-
-def upscale_resrgan(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    source = stage_source or source
-    logger.info("upscaling image with Real ESRGAN: x%s", upscale.scale)
-
-    output = np.array(source)
-    upsampler = load_resrgan(server, upscale, job.get_device(), tile=stage.tile_size)
-
-    output, _ = upsampler.enhance(output, outscale=upscale.outscale)
-
-    output = Image.fromarray(output, "RGB")
-    logger.info("final output image size: %sx%s", output.width, output.height)
-    return output
+class UpscaleRealESRGANStage(BaseStage):
+    def load(
+        self, server: ServerContext, params: UpscaleParams, device: DeviceParams, tile=0
+    ):
+        # must be within load function for patches to take effect
+        # TODO: rewrite and remove
+        from realesrgan import RealESRGANer
+
+        model_file = "%s.%s" % (params.upscale_model, params.format)
+        model_path = path.join(server.model_path, model_file)
+
+        cache_key = (model_path, params.format)
+        cache_pipe = server.cache.get(ModelTypes.upscaling, cache_key)
+        if cache_pipe is not None:
+            logger.info("reusing existing Real ESRGAN pipeline")
+            return cache_pipe
+
+        if not path.isfile(model_path):
+            raise FileNotFoundError("Real ESRGAN model not found at %s" % model_path)
+
+        # TODO: swap for regular RRDBNet after rewriting wrapper
+        model = OnnxRRDBNet(
+            server,
+            model_file,
+            provider=device.ort_provider(),
+            sess_options=device.sess_options(),
+        )
+
+        dni_weight = None
+        if params.upscale_model == TAG_X4_V3 and params.denoise != 1:
+            wdn_model_path = model_path.replace(TAG_X4_V3, "%s-wdn" % TAG_X4_V3)
+            model_path = [model_path, wdn_model_path]
+            dni_weight = [params.denoise, 1 - params.denoise]
+
+        logger.debug("loading Real ESRGAN upscale model from %s", model_path)
+
+        # TODO: shouldn't need the PTH file
+        model_path_pth = path.join(server.cache_path, ("%s.pth" % params.upscale_model))
+        upsampler = RealESRGANer(
+            scale=params.scale,
+            model_path=model_path_pth,
+            dni_weight=dni_weight,
+            model=model,
+            tile=tile,
+            tile_pad=params.tile_pad,
+            pre_pad=params.pre_pad,
+            half=False,  # TODO: use server optimizations
+        )
+
+        server.cache.set(ModelTypes.upscaling, cache_key, upsampler)
+        run_gc([device])
+
+        return upsampler
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        logger.info("upscaling image with Real ESRGAN: x%s", upscale.scale)
+
+        outputs = []
+        for source in sources:
+            output = np.array(source)
+            upsampler = self.load(
+                server, upscale, job.get_device(), tile=stage.tile_size
+            )
+
+            output, _ = upsampler.enhance(output, outscale=upscale.outscale)
+
+            output = Image.fromarray(output, "RGB")
+            logger.info("final output image size: %sx%s", output.width, output.height)
+            outputs.append(output)
+
+        return outputs
diff --git a/api/onnx_web/chain/upscale_simple.py b/api/onnx_web/chain/upscale_simple.py
new file mode 100644
index 000000000..5e5a7e75f
--- /dev/null
+++ b/api/onnx_web/chain/upscale_simple.py
@@ -0,0 +1,49 @@
+from logging import getLogger
+from typing import List, Optional
+
+from PIL import Image
+
+from ..params import ImageParams, StageParams, UpscaleParams
+from ..server import ServerContext
+from ..worker import WorkerContext
+from .stage import BaseStage
+
+logger = getLogger(__name__)
+
+
+class UpscaleSimpleStage(BaseStage):
+    def run(
+        self,
+        _job: WorkerContext,
+        _server: ServerContext,
+        _stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        method: str,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        if upscale.scale <= 1:
+            logger.debug(
+                "simple upscale stage run with scale of %s, skipping", upscale.scale
+            )
+            return sources
+
+        outputs = []
+        for source in sources:
+            scaled_size = (source.width * upscale.scale, source.height * upscale.scale)
+
+            if method == "bilinear":
+                logger.debug("using bilinear interpolation for highres")
+                source = source.resize(scaled_size, resample=Image.Resampling.BILINEAR)
+            elif method == "lanczos":
+                logger.debug("using Lanczos interpolation for highres")
+                source = source.resize(scaled_size, resample=Image.Resampling.LANCZOS)
+            else:
+                logger.warning("unknown upscaling method: %s", method)
+
+            outputs.append(source)
+
+        return outputs
diff --git a/api/onnx_web/chain/upscale_stable_diffusion.py b/api/onnx_web/chain/upscale_stable_diffusion.py
index f5cd42ee4..f86e2a5c1 100644
--- a/api/onnx_web/chain/upscale_stable_diffusion.py
+++ b/api/onnx_web/chain/upscale_stable_diffusion.py
@@ -1,6 +1,6 @@
 from logging import getLogger
 from os import path
-from typing import Optional
+from typing import List, Optional
 
 import torch
 from PIL import Image
@@ -10,56 +10,63 @@
 from ..params import ImageParams, StageParams, UpscaleParams
 from ..server import ServerContext
 from ..worker import ProgressCallback, WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def upscale_stable_diffusion(
-    job: WorkerContext,
-    server: ServerContext,
-    _stage: StageParams,
-    params: ImageParams,
-    source: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    stage_source: Optional[Image.Image] = None,
-    callback: Optional[ProgressCallback] = None,
-    **kwargs,
-) -> Image.Image:
-    params = params.with_args(**kwargs)
-    upscale = upscale.with_args(**kwargs)
-    source = stage_source or source
-    logger.info(
-        "upscaling with Stable Diffusion, %s steps: %s", params.steps, params.prompt
-    )
+class UpscaleStableDiffusionStage(BaseStage):
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        _stage: StageParams,
+        params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        callback: Optional[ProgressCallback] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        params = params.with_args(**kwargs)
+        upscale = upscale.with_args(**kwargs)
+        logger.info(
+            "upscaling with Stable Diffusion, %s steps: %s", params.steps, params.prompt
+        )
 
-    prompt_pairs, _loras, _inversions = parse_prompt(params)
+        prompt_pairs, _loras, _inversions = parse_prompt(params)
 
-    pipeline = load_pipeline(
-        server,
-        params,
-        "upscale",
-        job.get_device(),
-        model=path.join(server.model_path, upscale.upscale_model),
-    )
-    generator = torch.manual_seed(params.seed)
+        pipeline = load_pipeline(
+            server,
+            params,
+            "upscale",
+            job.get_device(),
+            model=path.join(server.model_path, upscale.upscale_model),
+        )
+        generator = torch.manual_seed(params.seed)
 
-    prompt_embeds = encode_prompt(
-        pipeline,
-        prompt_pairs,
-        num_images_per_prompt=params.batch,
-        do_classifier_free_guidance=params.do_cfg(),
-    )
-    pipeline.unet.set_prompts(prompt_embeds)
+        prompt_embeds = encode_prompt(
+            pipeline,
+            prompt_pairs,
+            num_images_per_prompt=params.batch,
+            do_classifier_free_guidance=params.do_cfg(),
+        )
+        pipeline.unet.set_prompts(prompt_embeds)
 
-    return pipeline(
-        params.prompt,
-        source,
-        generator=generator,
-        guidance_scale=params.cfg,
-        negative_prompt=params.negative_prompt,
-        num_inference_steps=params.steps,
-        eta=params.eta,
-        noise_level=upscale.denoise,
-        callback=callback,
-    ).images[0]
+        outputs = []
+        for source in sources:
+            result = pipeline(
+                params.prompt,
+                source,
+                generator=generator,
+                guidance_scale=params.cfg,
+                negative_prompt=params.negative_prompt,
+                num_inference_steps=params.steps,
+                eta=params.eta,
+                noise_level=upscale.denoise,
+                callback=callback,
+            )
+            outputs.extend(result.images)
+
+        return outputs
diff --git a/api/onnx_web/chain/upscale_swinir.py b/api/onnx_web/chain/upscale_swinir.py
index aaf466cc5..bba808d1e 100644
--- a/api/onnx_web/chain/upscale_swinir.py
+++ b/api/onnx_web/chain/upscale_swinir.py
@@ -1,120 +1,101 @@
 from logging import getLogger
 from os import path
-from typing import Optional
+from typing import List, Optional
 
 import numpy as np
 from PIL import Image
 
 from ..models.onnx import OnnxModel
 from ..params import DeviceParams, ImageParams, StageParams, UpscaleParams
-from ..server import ServerContext
+from ..server import ModelTypes, ServerContext
 from ..utils import run_gc
 from ..worker import WorkerContext
+from .stage import BaseStage
 
 logger = getLogger(__name__)
 
 
-def load_swinir(
-    server: ServerContext,
-    _stage: StageParams,
-    upscale: UpscaleParams,
-    device: DeviceParams,
-):
-    # must be within the load function for patch to take effect
-    model_path = path.join(server.model_path, "%s.onnx" % (upscale.upscale_model))
-    cache_key = (model_path,)
-    cache_pipe = server.cache.get("swinir", cache_key)
-
-    if cache_pipe is not None:
-        logger.info("reusing existing SwinIR pipeline")
-        return cache_pipe
-
-    logger.debug("loading SwinIR model from %s", model_path)
-
-    pipe = OnnxModel(
-        server,
-        model_path,
-        provider=device.ort_provider(),
-        sess_options=device.sess_options(),
-    )
-
-    server.cache.set("swinir", cache_key, pipe)
-    run_gc([device])
-
-    return pipe
-
-
-def upscale_swinir(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    _params: ImageParams,
-    source: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    stage_source: Optional[Image.Image] = None,
-    **kwargs,
-) -> Image.Image:
-    upscale = upscale.with_args(**kwargs)
-    source = stage_source or source
-
-    if upscale.upscale_model is None:
-        logger.warn("no correction model given, skipping")
-        return source
-
-    logger.info("correcting faces with SwinIR model: %s", upscale.upscale_model)
-    device = job.get_device()
-    swinir = load_swinir(server, stage, upscale, device)
-
-    # TODO: add support for other sizes
-    tile_size = (64, 64)
-    tile_x = source.width // tile_size[0]
-    tile_y = source.height // tile_size[1]
-
-    # TODO: add support for grayscale (1-channel) images
-    image = np.array(source) / 255.0
-    image = image[:, :, [2, 1, 0]].astype(np.float32).transpose((2, 0, 1))
-    image = np.expand_dims(image, axis=0)
-    logger.info("SwinIR input shape: %s", image.shape)
-
-    scale = upscale.outscale
-    dest = np.zeros(
-        (image.shape[0], image.shape[1], image.shape[2] * scale, image.shape[3] * scale)
-    )
-    logger.info("SwinIR output shape: %s", dest.shape)
-
-    for x in range(tile_x):
-        for y in range(tile_y):
-            xt = x * tile_size[0]
-            yt = y * tile_size[1]
-
-            ix1 = xt
-            ix2 = xt + tile_size[0]
-            iy1 = yt
-            iy2 = yt + tile_size[1]
-            logger.info(
-                "running SwinIR on tile: (%s, %s, %s, %s) -> (%s, %s, %s, %s)",
-                ix1,
-                ix2,
-                iy1,
-                iy2,
-                ix1 * scale,
-                ix2 * scale,
-                iy1 * scale,
-                iy2 * scale,
+class UpscaleSwinIRStage(BaseStage):
+    max_tile = 64
+
+    def load(
+        self,
+        server: ServerContext,
+        _stage: StageParams,
+        upscale: UpscaleParams,
+        device: DeviceParams,
+    ):
+        # must be within the load function for patch to take effect
+        model_path = path.join(server.model_path, "%s.onnx" % (upscale.upscale_model))
+        cache_key = (model_path,)
+        cache_pipe = server.cache.get(ModelTypes.upscaling, cache_key)
+
+        if cache_pipe is not None:
+            logger.info("reusing existing SwinIR pipeline")
+            return cache_pipe
+
+        logger.debug("loading SwinIR model from %s", model_path)
+
+        pipe = OnnxModel(
+            server,
+            model_path,
+            provider=device.ort_provider(),
+            sess_options=device.sess_options(),
+        )
+
+        server.cache.set(ModelTypes.upscaling, cache_key, pipe)
+        run_gc([device])
+
+        return pipe
+
+    def run(
+        self,
+        job: WorkerContext,
+        server: ServerContext,
+        stage: StageParams,
+        _params: ImageParams,
+        sources: List[Image.Image],
+        *,
+        upscale: UpscaleParams,
+        stage_source: Optional[Image.Image] = None,
+        **kwargs,
+    ) -> List[Image.Image]:
+        upscale = upscale.with_args(**kwargs)
+
+        if upscale.upscale_model is None:
+            logger.warn("no correction model given, skipping")
+            return sources
+
+        logger.info("correcting faces with SwinIR model: %s", upscale.upscale_model)
+        device = job.get_device()
+        swinir = self.load(server, stage, upscale, device)
+
+        outputs = []
+        for source in sources:
+            # TODO: add support for grayscale (1-channel) images
+            image = np.array(source) / 255.0
+            image = image[:, :, [2, 1, 0]].astype(np.float32).transpose((2, 0, 1))
+            image = np.expand_dims(image, axis=0)
+            logger.trace("SwinIR input shape: %s", image.shape)
+
+            scale = upscale.outscale
+            dest = np.zeros(
+                (
+                    image.shape[0],
+                    image.shape[1],
+                    image.shape[2] * scale,
+                    image.shape[3] * scale,
+                )
             )
+            logger.trace("SwinIR output shape: %s", dest.shape)
 
-            dest[
-                :,
-                :,
-                ix1 * scale : ix2 * scale,
-                iy1 * scale : iy2 * scale,
-            ] = swinir(image[:, :, ix1:ix2, iy1:iy2])
+            dest = swinir(image)
+            dest = np.clip(np.squeeze(dest, axis=0), 0, 1)
+            dest = dest[[2, 1, 0], :, :].transpose((1, 2, 0))
+            dest = (dest * 255.0).round().astype(np.uint8)
 
-    dest = np.clip(np.squeeze(dest, axis=0), 0, 1)
-    dest = dest[[2, 1, 0], :, :].transpose((1, 2, 0))
-    dest = (dest * 255.0).round().astype(np.uint8)
+            output = Image.fromarray(dest, "RGB")
+            logger.info("output image size: %s x %s", output.width, output.height)
+            outputs.append(output)
 
-    output = Image.fromarray(dest, "RGB")
-    logger.info("output image size: %s x %s", output.width, output.height)
-    return output
+        return outputs
diff --git a/api/onnx_web/convert/__main__.py b/api/onnx_web/convert/__main__.py
index 7e3315c68..36c974291 100644
--- a/api/onnx_web/convert/__main__.py
+++ b/api/onnx_web/convert/__main__.py
@@ -268,7 +268,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
     model_errors = []
 
     if args.sources and "sources" in models:
-        for model in models.get("sources"):
+        for model in models.get("sources", []):
             model = tuple_to_source(model)
             name = model.get("name")
 
@@ -292,7 +292,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                     model_errors.append(name)
 
     if args.networks and "networks" in models:
-        for network in models.get("networks"):
+        for network in models.get("networks", []):
             name = network["name"]
 
             if name in args.skip:
@@ -316,6 +316,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                             conversion,
                             network,
                             dest,
+                            path.join(conversion.model_path, network_type, name),
                         )
                     if network_type == "inversion" and network_model == "concept":
                         dest, hf = fetch_model(
@@ -342,7 +343,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                     model_errors.append(name)
 
     if args.diffusion and "diffusion" in models:
-        for model in models.get("diffusion"):
+        for model in models.get("diffusion", []):
             model = tuple_to_diffusion(model)
             name = model.get("name")
 
@@ -483,7 +484,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                     model_errors.append(name)
 
     if args.upscaling and "upscaling" in models:
-        for model in models.get("upscaling"):
+        for model in models.get("upscaling", []):
             model = tuple_to_upscaling(model)
             name = model.get("name")
 
@@ -516,7 +517,7 @@ def convert_models(conversion: ConversionContext, args, models: Models):
                     model_errors.append(name)
 
     if args.correction and "correction" in models:
-        for model in models.get("correction"):
+        for model in models.get("correction", []):
             model = tuple_to_correction(model)
             name = model.get("name")
 
diff --git a/api/onnx_web/convert/diffusion/control.py b/api/onnx_web/convert/diffusion/control.py
index ffc9f5c92..7e96801bd 100644
--- a/api/onnx_web/convert/diffusion/control.py
+++ b/api/onnx_web/convert/diffusion/control.py
@@ -1,7 +1,7 @@
 from logging import getLogger
 from os import path
 from pathlib import Path
-from typing import Dict
+from typing import Dict, Optional
 
 import torch
 
@@ -17,16 +17,15 @@ def convert_diffusion_control(
     conversion: ConversionContext,
     model: Dict,
     source: str,
-    model_path: str,
     output_path: str,
-    opset: int,
-    attention_slicing: str,
+    attention_slicing: Optional[str] = None,
 ):
     name = model.get("name")
     source = source or model.get("source")
 
     device = conversion.training_device
     dtype = conversion.torch_dtype()
+    opset = conversion.opset
     logger.debug("using Torch dtype %s for ControlNet", dtype)
 
     output_path = Path(output_path)
@@ -35,7 +34,7 @@ def convert_diffusion_control(
         logger.info("ONNX model already exists, skipping")
         return
 
-    controlnet = ControlNetModel.from_pretrained(model_path, torch_dtype=dtype)
+    controlnet = ControlNetModel.from_pretrained(source, torch_dtype=dtype)
     if attention_slicing is not None:
         logger.info("enabling attention slicing for ControlNet")
         controlnet.set_attention_slice(attention_slicing)
diff --git a/api/onnx_web/convert/diffusion/diffusers.py b/api/onnx_web/convert/diffusion/diffusers.py
index 48bd77979..f096e935e 100644
--- a/api/onnx_web/convert/diffusion/diffusers.py
+++ b/api/onnx_web/convert/diffusion/diffusers.py
@@ -262,7 +262,7 @@ def convert_diffusion_diffusers(
     conversion: ConversionContext,
     model: Dict,
     source: str,
-    format: str,
+    format: Optional[str],
     hf: bool = False,
 ) -> Tuple[bool, str]:
     """
diff --git a/api/onnx_web/diffusers/load.py b/api/onnx_web/diffusers/load.py
index 7b339ee7a..2432e83ad 100644
--- a/api/onnx_web/diffusers/load.py
+++ b/api/onnx_web/diffusers/load.py
@@ -11,7 +11,7 @@
 from ..diffusers.pipelines.upscale import OnnxStableDiffusionUpscalePipeline
 from ..diffusers.utils import expand_prompt
 from ..params import DeviceParams, ImageParams
-from ..server import ServerContext
+from ..server import ModelTypes, ServerContext
 from ..utils import run_gc
 from .patches.unet import UNetWrapper
 from .patches.vae import VAEWrapper
@@ -119,14 +119,14 @@ def load_pipeline(
     scheduler_key = (params.scheduler, model)
     scheduler_type = pipeline_schedulers[params.scheduler]
 
-    cache_pipe = server.cache.get("diffusion", pipe_key)
+    cache_pipe = server.cache.get(ModelTypes.diffusion, pipe_key)
 
     if cache_pipe is not None:
         logger.debug("reusing existing diffusion pipeline")
         pipe = cache_pipe
 
         # update scheduler
-        cache_scheduler = server.cache.get("scheduler", scheduler_key)
+        cache_scheduler = server.cache.get(ModelTypes.scheduler, scheduler_key)
         if cache_scheduler is None:
             logger.debug("loading new diffusion scheduler")
             scheduler = scheduler_type.from_pretrained(
@@ -141,7 +141,7 @@ def load_pipeline(
                 scheduler = scheduler.to(device.torch_str())
 
             pipe.scheduler = scheduler
-            server.cache.set("scheduler", scheduler_key, scheduler)
+            server.cache.set(ModelTypes.scheduler, scheduler_key, scheduler)
             run_gc([device])
 
     else:
@@ -342,8 +342,8 @@ def load_pipeline(
         optimize_pipeline(server, pipe)
         patch_pipeline(server, pipe, pipeline, pipeline_class, params)
 
-        server.cache.set("diffusion", pipe_key, pipe)
-        server.cache.set("scheduler", scheduler_key, components["scheduler"])
+        server.cache.set(ModelTypes.diffusion, pipe_key, pipe)
+        server.cache.set(ModelTypes.scheduler, scheduler_key, components["scheduler"])
 
     if hasattr(pipe, "vae_decoder"):
         pipe.vae_decoder.set_tiled(tiled=params.tiled_vae)
diff --git a/api/onnx_web/diffusers/pipelines/panorama.py b/api/onnx_web/diffusers/pipelines/panorama.py
index a1ba68fb4..ab509be39 100644
--- a/api/onnx_web/diffusers/pipelines/panorama.py
+++ b/api/onnx_web/diffusers/pipelines/panorama.py
@@ -1243,11 +1243,18 @@ def __call__(
         *args,
         **kwargs,
     ):
-        if len(args) > 1 and (
-            isinstance(args[1], np.ndarray) or isinstance(args[1], PIL.Image.Image)
-        ):
-            if len(args) > 2 and (
+        if "image" in kwargs or (
+            len(args) > 1
+            and (
                 isinstance(args[1], np.ndarray) or isinstance(args[1], PIL.Image.Image)
+            )
+        ):
+            if "mask_image" in kwargs or (
+                len(args) > 2
+                and (
+                    isinstance(args[1], np.ndarray)
+                    or isinstance(args[1], PIL.Image.Image)
+                )
             ):
                 logger.debug("running inpaint panorama pipeline")
                 return self.inpaint(*args, **kwargs)
diff --git a/api/onnx_web/diffusers/run.py b/api/onnx_web/diffusers/run.py
index fffb4a30f..d6a7db3e7 100644
--- a/api/onnx_web/diffusers/run.py
+++ b/api/onnx_web/diffusers/run.py
@@ -1,12 +1,18 @@
 from logging import getLogger
-from typing import Any, List, Optional, Tuple
+from typing import Any, List, Optional
 
-import numpy as np
-import torch
 from PIL import Image
 
-from ..chain import blend_mask, upscale_outpaint
-from ..chain.utils import process_tile_order
+from onnx_web.chain.highres import stage_highres
+
+from ..chain import (
+    BlendImg2ImgStage,
+    BlendMaskStage,
+    ChainPipeline,
+    SourceTxt2ImgStage,
+    UpscaleOutpaintStage,
+)
+from ..chain.upscale import split_upscale, stage_upscale_correction
 from ..output import save_image
 from ..params import (
     Border,
@@ -14,213 +20,17 @@
     ImageParams,
     Size,
     StageParams,
-    TileOrder,
     UpscaleParams,
 )
 from ..server import ServerContext
 from ..server.load import get_source_filters
 from ..utils import run_gc, show_system_toast
 from ..worker import WorkerContext
-from ..worker.context import ProgressCallback
-from .load import load_pipeline
-from .upscale import run_upscale_correction
-from .utils import encode_prompt, get_latents_from_seed, parse_prompt
+from .utils import parse_prompt
 
 logger = getLogger(__name__)
 
 
-def run_loopback(
-    job: WorkerContext,
-    server: ServerContext,
-    params: ImageParams,
-    strength: float,
-    image: Image.Image,
-    progress: ProgressCallback,
-    inversions: List[Tuple[str, float]],
-    loras: List[Tuple[str, float]],
-    pipeline: Optional[Any] = None,
-) -> Image.Image:
-    if params.loopback == 0:
-        return image
-
-    # load img2img pipeline once
-    pipe_type = params.get_valid_pipeline("img2img")
-    if pipe_type == "controlnet":
-        logger.debug(
-            "controlnet pipeline cannot be used for loopback, switching to img2img"
-        )
-        pipe_type = "img2img"
-
-    logger.debug("using %s pipeline for loopback", pipe_type)
-
-    pipe = pipeline or load_pipeline(
-        server,
-        params,
-        pipe_type,
-        job.get_device(),
-        inversions=inversions,
-        loras=loras,
-    )
-
-    def loopback_iteration(source: Image.Image):
-        if pipe_type == "lpw":
-            rng = torch.manual_seed(params.seed)
-            result = pipe.img2img(
-                source,
-                params.prompt,
-                generator=rng,
-                guidance_scale=params.cfg,
-                negative_prompt=params.negative_prompt,
-                num_images_per_prompt=1,
-                num_inference_steps=params.steps,
-                strength=strength,
-                eta=params.eta,
-                callback=progress,
-            )
-            return result.images[0]
-        else:
-            rng = np.random.RandomState(params.seed)
-            result = pipe(
-                params.prompt,
-                source,
-                generator=rng,
-                guidance_scale=params.cfg,
-                negative_prompt=params.negative_prompt,
-                num_images_per_prompt=1,
-                num_inference_steps=params.steps,
-                strength=strength,
-                eta=params.eta,
-                callback=progress,
-            )
-            return result.images[0]
-
-    for _i in range(params.loopback):
-        image = loopback_iteration(image)
-
-    return image
-
-
-def run_highres(
-    job: WorkerContext,
-    server: ServerContext,
-    params: ImageParams,
-    size: Size,
-    upscale: UpscaleParams,
-    highres: HighresParams,
-    image: Image.Image,
-    progress: ProgressCallback,
-    inversions: List[Tuple[str, float]],
-    loras: List[Tuple[str, float]],
-    pipeline: Optional[Any] = None,
-) -> Image.Image:
-    if highres.scale <= 1:
-        return image
-
-    if upscale.faces and (
-        upscale.upscale_order == "correction-both"
-        or upscale.upscale_order == "correction-first"
-    ):
-        image = run_upscale_correction(
-            job,
-            server,
-            StageParams(),
-            params,
-            image,
-            upscale=upscale.with_args(
-                scale=1,
-                outscale=1,
-            ),
-            callback=progress,
-        )
-
-    # load img2img pipeline once
-    pipe_type = params.get_valid_pipeline("img2img")
-    logger.debug("using %s pipeline for highres", pipe_type)
-
-    highres_pipe = pipeline or load_pipeline(
-        server,
-        params,
-        pipe_type,
-        job.get_device(),
-        inversions=inversions,
-        loras=loras,
-    )
-
-    def highres_tile(tile: Image.Image, dims):
-        scaled_size = (tile.width * highres.scale, tile.height * highres.scale)
-
-        if highres.method == "bilinear":
-            logger.debug("using bilinear interpolation for highres")
-            tile = tile.resize(scaled_size, resample=Image.Resampling.BILINEAR)
-        elif highres.method == "lanczos":
-            logger.debug("using Lanczos interpolation for highres")
-            tile = tile.resize(scaled_size, resample=Image.Resampling.LANCZOS)
-        else:
-            logger.debug("using upscaling pipeline for highres")
-            tile = run_upscale_correction(
-                job,
-                server,
-                StageParams(),
-                params,
-                tile,
-                upscale=upscale.with_args(
-                    faces=False,
-                    scale=highres.scale,
-                    outscale=highres.scale,
-                ),
-                callback=progress,
-            )
-
-        if pipe_type == "lpw":
-            rng = torch.manual_seed(params.seed)
-            result = highres_pipe.img2img(
-                tile,
-                params.prompt,
-                generator=rng,
-                guidance_scale=params.cfg,
-                negative_prompt=params.negative_prompt,
-                num_images_per_prompt=1,
-                num_inference_steps=highres.steps,
-                strength=highres.strength,
-                eta=params.eta,
-                callback=progress,
-            )
-            return result.images[0]
-        else:
-            rng = np.random.RandomState(params.seed)
-            result = highres_pipe(
-                params.prompt,
-                tile,
-                generator=rng,
-                guidance_scale=params.cfg,
-                negative_prompt=params.negative_prompt,
-                num_images_per_prompt=1,
-                num_inference_steps=highres.steps,
-                strength=highres.strength,
-                eta=params.eta,
-                callback=progress,
-            )
-            return result.images[0]
-
-    logger.info(
-        "running highres fix for %s iterations at %s scale",
-        highres.iterations,
-        highres.scale,
-    )
-
-    for _i in range(highres.iterations):
-        image = process_tile_order(
-            TileOrder.grid,
-            image,
-            size.height // highres.scale,
-            highres.scale,
-            [highres_tile],
-            overlap=params.overlap,
-        )
-
-    return image
-
-
 def run_txt2img_pipeline(
     job: WorkerContext,
     server: ServerContext,
@@ -230,90 +40,51 @@ def run_txt2img_pipeline(
     upscale: UpscaleParams,
     highres: HighresParams,
 ) -> None:
-    latents = get_latents_from_seed(params.seed, size, batch=params.batch)
-    prompt_pairs, loras, inversions = parse_prompt(params)
+    # prepare the chain pipeline and first stage
+    chain = ChainPipeline()
+    stage = StageParams(
+        tile_size=params.tiles,
+    )
+    chain.stage(
+        SourceTxt2ImgStage(),
+        stage,
+        size=size,
+    )
 
-    pipe_type = params.get_valid_pipeline("txt2img")
-    logger.debug("using %s pipeline for txt2img", pipe_type)
+    # apply upscaling and correction, before highres
+    first_upscale, after_upscale = split_upscale(upscale)
+    if first_upscale:
+        stage_upscale_correction(
+            stage,
+            params,
+            upscale=first_upscale,
+            chain=chain,
+        )
 
-    pipe = load_pipeline(
-        server,
+    # apply highres
+    stage_highres(
+        stage,
         params,
-        pipe_type,
-        job.get_device(),
-        inversions=inversions,
-        loras=loras,
+        highres,
+        upscale,
+        chain=chain,
     )
-    progress = job.get_progress_callback()
 
-    if pipe_type == "lpw":
-        rng = torch.manual_seed(params.seed)
-        result = pipe.text2img(
-            params.prompt,
-            height=size.height,
-            width=size.width,
-            generator=rng,
-            guidance_scale=params.cfg,
-            latents=latents,
-            negative_prompt=params.negative_prompt,
-            num_images_per_prompt=params.batch,
-            num_inference_steps=params.steps,
-            eta=params.eta,
-            callback=progress,
-        )
-    else:
-        # encode and record alternative prompts outside of LPW
-        prompt_embeds = encode_prompt(
-            pipe,
-            prompt_pairs,
-            num_images_per_prompt=params.batch,
-            do_classifier_free_guidance=params.do_cfg(),
-        )
-        pipe.unet.set_prompts(prompt_embeds)
-
-        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_images_per_prompt=params.batch,
-            num_inference_steps=params.steps,
-            eta=params.eta,
-            callback=progress,
-        )
+    # apply upscaling and correction, after highres
+    stage_upscale_correction(
+        stage,
+        params,
+        upscale=after_upscale,
+        chain=chain,
+    )
 
-    image_outputs = list(zip(result.images, outputs))
-    del result
-    del pipe
+    # run and save
+    progress = job.get_progress_callback()
+    images = chain(job, server, params, [], callback=progress)
 
-    for image, output in image_outputs:
-        image = run_highres(
-            job,
-            server,
-            params,
-            size,
-            upscale,
-            highres,
-            image,
-            progress,
-            inversions,
-            loras,
-        )
-
-        image = run_upscale_correction(
-            job,
-            server,
-            StageParams(),
-            params,
-            image,
-            upscale=upscale,
-            callback=progress,
-        )
+    _prompt_pairs, loras, inversions = parse_prompt(params, use_input=True)
 
+    for image, output in zip(images, outputs):
         dest = save_image(
             server,
             output,
@@ -326,7 +97,10 @@ def run_txt2img_pipeline(
             loras=loras,
         )
 
+    # clean up
     run_gc([job.get_device()])
+
+    # notify the user
     show_system_toast(f"finished txt2img job: {dest}")
     logger.info("finished txt2img job: %s", dest)
 
@@ -342,110 +116,71 @@ def run_img2img_pipeline(
     strength: float,
     source_filter: Optional[str] = None,
 ) -> None:
-    prompt_pairs, loras, inversions = parse_prompt(params)
-
-    # filter the source image
+    # run filter on the source image
     if source_filter is not None:
         f = get_source_filters().get(source_filter, None)
         if f is not None:
             logger.debug("running source filter: %s", f.__name__)
             source = f(server, source)
 
-    pipe_type = params.get_valid_pipeline("img2img")
-    pipe = load_pipeline(
-        server,
+    # prepare the chain pipeline and first stage
+    chain = ChainPipeline()
+    stage = StageParams(
+        tile_size=params.tiles,
+    )
+    chain.stage(
+        BlendImg2ImgStage(),
+        stage,
+        strength=strength,
+    )
+
+    # apply upscaling and correction, before highres
+    first_upscale, after_upscale = split_upscale(upscale)
+    if first_upscale:
+        stage_upscale_correction(
+            stage,
+            params,
+            upscale=first_upscale,
+            chain=chain,
+        )
+
+    # loopback through multiple img2img iterations
+    for _i in range(params.loopback):
+        chain.stage(
+            BlendImg2ImgStage(),
+            stage,
+            strength=strength,
+        )
+
+    # highres, if selected
+    stage_highres(
+        stage,
         params,
-        pipe_type,
-        job.get_device(),
-        inversions=inversions,
-        loras=loras,
+        highres,
+        upscale,
+        chain=chain,
     )
 
-    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
+    # apply upscaling and correction, after highres
+    stage_upscale_correction(
+        stage,
+        params,
+        upscale=after_upscale,
+        chain=chain,
+    )
 
+    # run and append the filtered source
     progress = job.get_progress_callback()
-    if pipe_type == "lpw":
-        logger.debug("using LPW pipeline for img2img")
-        rng = torch.manual_seed(params.seed)
-        result = pipe.img2img(
-            source,
-            params.prompt,
-            generator=rng,
-            guidance_scale=params.cfg,
-            negative_prompt=params.negative_prompt,
-            num_images_per_prompt=params.batch,
-            num_inference_steps=params.steps,
-            eta=params.eta,
-            callback=progress,
-            **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,
-            source,
-            generator=rng,
-            guidance_scale=params.cfg,
-            negative_prompt=params.negative_prompt,
-            num_images_per_prompt=params.batch,
-            num_inference_steps=params.steps,
-            eta=params.eta,
-            callback=progress,
-            **pipe_params,
-        )
+    images = chain(job, server, params, [source], callback=progress)
 
-    images = result.images
     if source_filter is not None and source_filter != "none":
         images.append(source)
 
-    for image, output in zip(images, outputs):
-        image = run_loopback(
-            job,
-            server,
-            params,
-            strength,
-            image,
-            progress,
-            inversions,
-            loras,
-        )
-
-        size = Size(*source.size)
-        image = run_highres(
-            job,
-            server,
-            params,
-            size,
-            upscale,
-            highres,
-            image,
-            progress,
-            inversions,
-            loras,
-        )
-
-        image = run_upscale_correction(
-            job,
-            server,
-            StageParams(),
-            params,
-            image,
-            upscale=upscale,
-            callback=progress,
-        )
+    # save with metadata
+    _prompt_pairs, loras, inversions = parse_prompt(params, use_input=True)
+    size = Size(*source.size)
 
+    for image, output in zip(images, outputs):
         dest = save_image(
             server,
             output,
@@ -458,7 +193,10 @@ def run_img2img_pipeline(
             loras=loras,
         )
 
+    # clean up
     run_gc([job.get_device()])
+
+    # notify the user
     show_system_toast(f"finished img2img job: {dest}")
     logger.info("finished img2img job: %s", dest)
 
@@ -479,64 +217,71 @@ def run_inpaint_pipeline(
     fill_color: str,
     tile_order: str,
 ) -> None:
-    progress = job.get_progress_callback()
-    stage = StageParams(tile_order=tile_order)
-
-    _prompt_pairs, loras, inversions = parse_prompt(params)
+    logger.debug("building inpaint pipeline")
 
-    logger.debug("applying mask filter and generating noise source")
-    image = upscale_outpaint(
-        job,
-        server,
+    # set up the chain pipeline and base stage
+    chain = ChainPipeline()
+    stage = StageParams(tile_order=tile_order, tile_size=params.tiles)
+    chain.stage(
+        UpscaleOutpaintStage(),
         stage,
-        params,
-        source,
         border=border,
         stage_mask=mask,
         fill_color=fill_color,
         mask_filter=mask_filter,
         noise_source=noise_source,
-        callback=progress,
     )
 
-    image = run_highres(
-        job,
-        server,
+    # apply upscaling and correction, before highres
+    first_upscale, after_upscale = split_upscale(upscale)
+    if first_upscale:
+        stage_upscale_correction(
+            stage,
+            params,
+            upscale=first_upscale,
+            chain=chain,
+        )
+
+    # apply highres
+    stage_highres(
+        stage,
         params,
-        size,
-        upscale,
         highres,
-        image,
-        progress,
-        inversions,
-        loras,
+        upscale,
+        chain=chain,
     )
 
-    image = run_upscale_correction(
-        job,
-        server,
+    # apply upscaling and correction
+    stage_upscale_correction(
         stage,
         params,
-        image,
-        upscale=upscale,
-        callback=progress,
+        upscale=after_upscale,
+        chain=chain,
     )
 
-    dest = save_image(
-        server,
-        outputs[0],
-        image,
-        params,
-        size,
-        upscale=upscale,
-        border=border,
-        inversions=inversions,
-        loras=loras,
-    )
+    # run and save
+    progress = job.get_progress_callback()
+    images = chain(job, server, params, [source], callback=progress)
+
+    _prompt_pairs, loras, inversions = parse_prompt(params, use_input=True)
+    for image, output in zip(images, outputs):
+        dest = save_image(
+            server,
+            output,
+            image,
+            params,
+            size,
+            upscale=upscale,
+            border=border,
+            inversions=inversions,
+            loras=loras,
+        )
 
+    # clean up
     del image
     run_gc([job.get_device()])
 
+    # notify the user
     show_system_toast(f"finished inpaint job: {dest}")
     logger.info("finished inpaint job: %s", dest)
 
@@ -551,34 +296,59 @@ def run_upscale_pipeline(
     highres: HighresParams,
     source: Image.Image,
 ) -> None:
-    progress = job.get_progress_callback()
-    stage = StageParams()
-
-    _prompt_pairs, loras, inversions = parse_prompt(params)
+    # set up the chain pipeline, no base stage for upscaling
+    chain = ChainPipeline()
+    stage = StageParams(tile_size=params.tiles)
+
+    # apply upscaling and correction, before highres
+    first_upscale, after_upscale = split_upscale(upscale)
+    if first_upscale:
+        stage_upscale_correction(
+            stage,
+            params,
+            upscale=first_upscale,
+            chain=chain,
+        )
 
-    image = run_upscale_correction(
-        job, server, stage, params, source, upscale=upscale, callback=progress
+    # apply highres
+    stage_highres(
+        stage,
+        params,
+        highres,
+        upscale,
+        chain=chain,
     )
 
-    # TODO: should this come first?
-    image = run_highres(
-        job,
-        server,
+    # apply upscaling and correction, after highres
+    stage_upscale_correction(
+        stage,
         params,
-        size,
-        upscale,
-        highres,
-        image,
-        progress,
-        inversions,
-        loras,
+        upscale=after_upscale,
+        chain=chain,
     )
 
-    dest = save_image(server, outputs[0], image, params, size, upscale=upscale)
+    # run and save
+    progress = job.get_progress_callback()
+    images = chain(job, server, params, [source], callback=progress)
+
+    _prompt_pairs, loras, inversions = parse_prompt(params, use_input=True)
+    for image, output in zip(images, outputs):
+        dest = save_image(
+            server,
+            output,
+            image,
+            params,
+            size,
+            upscale=upscale,
+            inversions=inversions,
+            loras=loras,
+        )
 
+    # clean up
     del image
     run_gc([job.get_device()])
 
+    # notify the user
     show_system_toast(f"finished upscale job: {dest}")
     logger.info("finished upscale job: %s", dest)
 
@@ -594,28 +364,30 @@ def run_blend_pipeline(
     sources: List[Image.Image],
     mask: Image.Image,
 ) -> None:
-    progress = job.get_progress_callback()
+    # set up the chain pipeline and base stage
+    chain = ChainPipeline()
     stage = StageParams()
+    chain.stage(BlendMaskStage(), stage, stage_source=sources[1], stage_mask=mask)
 
-    image = blend_mask(
-        job,
-        server,
+    # apply upscaling and correction
+    stage_upscale_correction(
         stage,
         params,
-        sources=sources,
-        stage_mask=mask,
-        callback=progress,
+        upscale=upscale,
+        chain=chain,
     )
-    image = image.convert("RGB")
 
-    image = run_upscale_correction(
-        job, server, stage, params, image, upscale=upscale, callback=progress
-    )
+    # run and save
+    progress = job.get_progress_callback()
+    images = chain(job, server, params, sources, callback=progress)
 
-    dest = save_image(server, outputs[0], image, params, size, upscale=upscale)
+    for image, output in zip(images, outputs):
+        dest = save_image(server, output, image, params, size, upscale=upscale)
 
+    # clean up
     del image
     run_gc([job.get_device()])
 
+    # notify the user
     show_system_toast(f"finished blend job: {dest}")
     logger.info("finished blend job: %s", dest)
diff --git a/api/onnx_web/diffusers/upscale.py b/api/onnx_web/diffusers/upscale.py
deleted file mode 100644
index c7d3601ab..000000000
--- a/api/onnx_web/diffusers/upscale.py
+++ /dev/null
@@ -1,114 +0,0 @@
-from logging import getLogger
-from typing import List, Optional
-
-from PIL import Image
-
-from ..chain import (
-    ChainPipeline,
-    PipelineStage,
-    correct_codeformer,
-    correct_gfpgan,
-    upscale_bsrgan,
-    upscale_resrgan,
-    upscale_stable_diffusion,
-    upscale_swinir,
-)
-from ..params import ImageParams, SizeChart, StageParams, UpscaleParams
-from ..server import ServerContext
-from ..worker import ProgressCallback, WorkerContext
-
-logger = getLogger(__name__)
-
-
-def run_upscale_correction(
-    job: WorkerContext,
-    server: ServerContext,
-    stage: StageParams,
-    params: ImageParams,
-    image: Image.Image,
-    *,
-    upscale: UpscaleParams,
-    callback: Optional[ProgressCallback] = None,
-    pre_stages: List[PipelineStage] = None,
-    post_stages: List[PipelineStage] = None,
-) -> Image.Image:
-    """
-    This is a convenience method for a chain pipeline that will run upscaling and
-    correction, based on the `upscale` params.
-    """
-    logger.info(
-        "running upscaling and correction pipeline at %s:%s",
-        upscale.scale,
-        upscale.outscale,
-    )
-
-    chain = ChainPipeline()
-    if pre_stages is not None:
-        for stage, params in pre_stages:
-            chain.append((stage, params))
-
-    upscale_stage = None
-    if upscale.scale > 1:
-        if "bsrgan" in upscale.upscale_model:
-            bsrgan_params = StageParams(
-                tile_size=stage.tile_size,
-                outscale=upscale.outscale,
-            )
-            upscale_stage = (upscale_bsrgan, bsrgan_params, None)
-        elif "esrgan" in upscale.upscale_model:
-            esrgan_params = StageParams(
-                tile_size=stage.tile_size,
-                outscale=upscale.outscale,
-            )
-            upscale_stage = (upscale_resrgan, esrgan_params, None)
-        elif "stable-diffusion" in upscale.upscale_model:
-            mini_tile = min(SizeChart.mini, stage.tile_size)
-            sd_params = StageParams(tile_size=mini_tile, outscale=upscale.outscale)
-            upscale_stage = (upscale_stable_diffusion, sd_params, None)
-        elif "swinir" in upscale.upscale_model:
-            swinir_params = StageParams(
-                tile_size=stage.tile_size,
-                outscale=upscale.outscale,
-            )
-            upscale_stage = (upscale_swinir, swinir_params, None)
-        else:
-            logger.warn("unknown upscaling model: %s", upscale.upscale_model)
-
-    correct_stage = None
-    if upscale.faces:
-        face_params = StageParams(
-            tile_size=stage.tile_size, outscale=upscale.face_outscale
-        )
-        if "codeformer" in upscale.correction_model:
-            correct_stage = (correct_codeformer, face_params, None)
-        elif "gfpgan" in upscale.correction_model:
-            correct_stage = (correct_gfpgan, face_params, None)
-        else:
-            logger.warn("unknown correction model: %s", upscale.correction_model)
-
-    if upscale.upscale_order == "correction-both":
-        chain.append(correct_stage)
-        chain.append(upscale_stage)
-        chain.append(correct_stage)
-    elif upscale.upscale_order == "correction-first":
-        chain.append(correct_stage)
-        chain.append(upscale_stage)
-    elif upscale.upscale_order == "correction-last":
-        chain.append(upscale_stage)
-        chain.append(correct_stage)
-    else:
-        logger.warn("unknown upscaling order: %s", upscale.upscale_order)
-
-    if post_stages is not None:
-        for stage, params in post_stages:
-            chain.append((stage, params))
-
-    return chain(
-        job,
-        server,
-        params,
-        image,
-        prompt=params.prompt,
-        upscale=upscale,
-        callback=callback,
-    )
diff --git a/api/onnx_web/diffusers/utils.py b/api/onnx_web/diffusers/utils.py
index 981b68952..615c78396 100644
--- a/api/onnx_web/diffusers/utils.py
+++ b/api/onnx_web/diffusers/utils.py
@@ -1,7 +1,9 @@
+import random
+from copy import deepcopy
 from logging import getLogger
 from math import ceil
 from re import Pattern, compile
-from typing import List, Optional, Tuple
+from typing import Dict, List, Optional, Tuple
 
 import numpy as np
 import torch
@@ -18,6 +20,8 @@
 CLIP_TOKEN = compile(r"\<clip:([-\w]+):(\d+)\>")
 INVERSION_TOKEN = compile(r"\<inversion:([^:\>]+):(-?[\.|\d]+)\>")
 LORA_TOKEN = compile(r"\<lora:([^:\>]+):(-?[\.|\d]+)\>")
+WILDCARD_TOKEN = compile(r"__([-\w]+)__")
+
 INTERVAL_RANGE = compile(r"(\w+)-{(\d+),(\d+)(?:,(\d+))?}")
 ALTERNATIVE_RANGE = compile(r"\(([^\)]+)\)")
 
@@ -307,14 +311,19 @@ def get_scaled_latents(
 
 def parse_prompt(
     params: ImageParams,
+    use_input: bool = False,
 ) -> Tuple[List[Tuple[str, str]], List[Tuple[str, float]], List[Tuple[str, float]]]:
-    prompt, loras = get_loras_from_prompt(params.input_prompt)
+    prompt, loras = get_loras_from_prompt(
+        params.input_prompt if use_input else params.prompt
+    )
     prompt, inversions = get_inversions_from_prompt(prompt)
     params.prompt = prompt
 
     neg_prompt = None
     if params.input_negative_prompt is not None:
-        neg_prompt, neg_loras = get_loras_from_prompt(params.input_negative_prompt)
+        neg_prompt, neg_loras = get_loras_from_prompt(
+            params.input_negative_prompt if use_input else params.negative_prompt
+        )
         neg_prompt, neg_inversions = get_inversions_from_prompt(neg_prompt)
         params.negative_prompt = neg_prompt
 
@@ -361,3 +370,40 @@ def encode_prompt(
         )
         for prompt, neg_prompt in prompt_pairs
     ]
+
+
+def replace_wildcards(prompt: str, seed: int, wildcards: Dict[str, List[str]]) -> str:
+    next_match = WILDCARD_TOKEN.search(prompt)
+    remaining_prompt = prompt
+
+    # prep a local copy to avoid mutating the main one
+    wildcards = deepcopy(wildcards)
+    random.seed(seed)
+
+    while next_match is not None:
+        logger.debug("found wildcard in prompt: %s", next_match)
+        name, *rest = next_match.groups()
+
+        wildcard = ""
+        if name in wildcards:
+            wildcard = pop_random(wildcards.get(name))
+        else:
+            logger.warning("unknown wildcard: %s", name)
+
+        remaining_prompt = (
+            remaining_prompt[: next_match.start()]
+            + wildcard
+            + remaining_prompt[next_match.end() :]
+        )
+        next_match = WILDCARD_TOKEN.search(remaining_prompt)
+
+    return remaining_prompt
+
+
+def pop_random(list: List[str]) -> str:
+    """
+    From https://stackoverflow.com/a/14088129
+    """
+    i = random.randrange(len(list))
+    list[i], list[-1] = list[-1], list[i]
+    return list.pop()
diff --git a/api/onnx_web/image/__init__.py b/api/onnx_web/image/__init__.py
index 84be6e13a..b0133d41b 100644
--- a/api/onnx_web/image/__init__.py
+++ b/api/onnx_web/image/__init__.py
@@ -1,6 +1,5 @@
 from .utils import (
     expand_image,
-    valid_image,
 )
 from .mask_filter import (
     mask_filter_gaussian_multiply,
diff --git a/api/onnx_web/image/utils.py b/api/onnx_web/image/utils.py
index 643f5595e..e3fe7c43d 100644
--- a/api/onnx_web/image/utils.py
+++ b/api/onnx_web/image/utils.py
@@ -1,6 +1,4 @@
-from typing import Tuple, Union
-
-from PIL import Image, ImageChops, ImageOps
+from PIL import Image, ImageChops
 
 from ..params import Border, Size
 from .mask_filter import mask_filter_none
@@ -31,24 +29,3 @@ def expand_image(
     full_source = Image.composite(full_noise, full_source, full_mask.convert("L"))
 
     return (full_source, full_mask, full_noise, size)
-
-
-def valid_image(
-    image: Image.Image,
-    min_dims: Union[Size, Tuple[int, int]] = [512, 512],
-    max_dims: Union[Size, Tuple[int, int]] = [512, 512],
-) -> Image.Image:
-    min_x, min_y = min_dims
-    max_x, max_y = max_dims
-
-    if image.width > max_x or image.height > max_y:
-        image = ImageOps.contain(image, (max_x, max_y))
-
-    if image.width < min_x or image.height < min_y:
-        blank = Image.new(image.mode, (min_x, min_y), "black")
-        blank.paste(image)
-        image = blank
-
-    # check for square
-
-    return image
diff --git a/api/onnx_web/main.py b/api/onnx_web/main.py
index 7a43d06f6..9e47cc277 100644
--- a/api/onnx_web/main.py
+++ b/api/onnx_web/main.py
@@ -21,6 +21,7 @@
     load_models,
     load_params,
     load_platforms,
+    load_wildcards,
 )
 from .server.static import register_static_routes
 from .server.utils import check_paths
@@ -46,6 +47,7 @@ def main():
     load_models(server)
     load_params(server)
     load_platforms(server)
+    load_wildcards(server)
 
     if is_debug():
         gc.set_debug(gc.DEBUG_STATS)
diff --git a/api/onnx_web/params.py b/api/onnx_web/params.py
index 584ab8864..9b418c50d 100644
--- a/api/onnx_web/params.py
+++ b/api/onnx_web/params.py
@@ -14,6 +14,7 @@
 
 
 class SizeChart(IntEnum):
+    unlimited = 0
     mini = 128  # small tile for very expensive models
     half = 256  # half tile for outpainting
     auto = 512  # auto tile size
@@ -22,6 +23,7 @@ class SizeChart(IntEnum):
     hd4k = 2**12
     hd8k = 2**13
     hd16k = 2**14
+    hd32k = 2**15
     hd64k = 2**16
 
 
@@ -87,14 +89,14 @@ def round_to_tile(self, tile=512):
 
     def tojson(self) -> Dict[str, int]:
         return {
-            "height": self.height,
             "width": self.width,
+            "height": self.height,
         }
 
     def with_args(self, **kwargs):
         return Size(
-            kwargs.get("height", self.height),
             kwargs.get("width", self.width),
+            kwargs.get("height", self.height),
         )
 
 
@@ -184,7 +186,6 @@ class ImageParams:
     steps: int
     seed: int
     negative_prompt: Optional[str]
-    lpw: bool
     eta: float
     batch: int
     control: Optional[NetworkModel]
@@ -419,12 +420,14 @@ def with_args(self, **kwargs):
 class HighresParams:
     def __init__(
         self,
+        enabled: bool,
         scale: int,
         steps: int,
         strength: float,
         method: Literal["bilinear", "lanczos", "upscale"] = "lanczos",
         iterations: int = 1,
     ):
+        self.enabled = enabled
         self.scale = scale
         self.steps = steps
         self.strength = strength
@@ -439,6 +442,7 @@ def resize(self, size: Size) -> Size:
 
     def tojson(self):
         return {
+            "enabled": self.enabled,
             "iterations": self.iterations,
             "method": self.method,
             "scale": self.scale,
diff --git a/api/onnx_web/prompt/grammar.py b/api/onnx_web/prompt/grammar.py
index ae28d23e2..20127030b 100644
--- a/api/onnx_web/prompt/grammar.py
+++ b/api/onnx_web/prompt/grammar.py
@@ -43,6 +43,8 @@ def __eq__(self, other: object) -> bool:
         if isinstance(other, self.__class__):
             return other.tokens == self.tokens and other.weight == self.weight
 
+        return False
+
 
 class OnnxPromptVisitor(PTNodeVisitor):
     def __init__(self, defaults=True, weight=0.5, **kwargs):
diff --git a/api/onnx_web/prompt/parser.py b/api/onnx_web/prompt/parser.py
index 6a0fde223..8bab4dbfe 100644
--- a/api/onnx_web/prompt/parser.py
+++ b/api/onnx_web/prompt/parser.py
@@ -15,7 +15,7 @@ def parse_prompt_compel(pipeline, prompt: str) -> np.ndarray:
 
 
 def parse_prompt_lpw(pipeline, prompt: str, debug=False) -> np.ndarray:
-    pass
+    raise NotImplementedError()
 
 
 def parse_prompt_onnx(pipeline, prompt: str, debug=False) -> np.ndarray:
diff --git a/api/onnx_web/server/__init__.py b/api/onnx_web/server/__init__.py
index f02fa35a6..b74014f1d 100644
--- a/api/onnx_web/server/__init__.py
+++ b/api/onnx_web/server/__init__.py
@@ -4,5 +4,5 @@
     apply_patch_facexlib,
     apply_patches,
 )
-from .model_cache import ModelCache
+from .model_cache import ModelCache, ModelTypes
 from .context import ServerContext
diff --git a/api/onnx_web/server/admin.py b/api/onnx_web/server/admin.py
index 135233b44..bdbc9adeb 100644
--- a/api/onnx_web/server/admin.py
+++ b/api/onnx_web/server/admin.py
@@ -6,7 +6,7 @@
 from ..utils import load_config, load_config_str
 from ..worker.pool import DevicePoolExecutor
 from .context import ServerContext
-from .load import load_extras, load_models
+from .load import load_extras, load_models, load_wildcards
 from .utils import wrap_route
 
 logger = getLogger(__name__)
@@ -88,6 +88,7 @@ def update_extra_models(server: ServerContext):
 
     logger.info("finished converting models, reloading server")
     load_models(server)
+    load_wildcards(server)
     load_extras(server)
 
     conversion_lock = False
diff --git a/api/onnx_web/server/api.py b/api/onnx_web/server/api.py
index 972c23d9b..aa9d5a22f 100644
--- a/api/onnx_web/server/api.py
+++ b/api/onnx_web/server/api.py
@@ -15,7 +15,7 @@
     run_txt2img_pipeline,
     run_upscale_pipeline,
 )
-from ..image import valid_image  # mask filters; noise sources
+from ..diffusers.utils import replace_wildcards
 from ..output import json_params, make_output_name
 from ..params import Border, StageParams, TileOrder, UpscaleParams
 from ..transformers.run import run_txt2txt_pipeline
@@ -45,6 +45,7 @@
     get_noise_sources,
     get_source_filters,
     get_upscaling_models,
+    get_wildcard_data,
 )
 from .params import (
     border_from_request,
@@ -178,6 +179,12 @@ def img2img(server: ServerContext, pool: DevicePoolExecutor):
         get_config_value("strength", "min"),
     )
 
+    params.prompt = replace_wildcards(params.prompt, params.seed, get_wildcard_data())
+    if params.negative_prompt is not None:
+        params.negative_prompt = replace_wildcards(
+            params.negative_prompt, params.seed, get_wildcard_data()
+        )
+
     output_count = params.batch
     if source_filter is not None and source_filter != "none":
         logger.debug(
@@ -189,12 +196,6 @@ def img2img(server: ServerContext, pool: DevicePoolExecutor):
         server, "img2img", params, size, extras=[strength], count=output_count
     )
 
-    if params.get_valid_pipeline("img2img") != "panorama":
-        logger.info(
-            "resizing input image for limited pipeline, use panorama pipeline for full-size"
-        )
-        source = valid_image(source, min_dims=size, max_dims=size)
-
     job_name = output[0]
     pool.submit(
         job_name,
@@ -220,6 +221,12 @@ def txt2img(server: ServerContext, pool: DevicePoolExecutor):
     upscale = upscale_from_request()
     highres = highres_from_request()
 
+    params.prompt = replace_wildcards(params.prompt, params.seed, get_wildcard_data())
+    if params.negative_prompt is not None:
+        params.negative_prompt = replace_wildcards(
+            params.negative_prompt, params.seed, get_wildcard_data()
+        )
+
     output = make_output_name(server, "txt2img", params, size)
 
     job_name = output[0]
@@ -264,6 +271,12 @@ def inpaint(server: ServerContext, pool: DevicePoolExecutor):
         request.args, "tileOrder", [TileOrder.grid, TileOrder.kernel, TileOrder.spiral]
     )
 
+    params.prompt = replace_wildcards(params.prompt, params.seed, get_wildcard_data())
+    if params.negative_prompt is not None:
+        params.negative_prompt = replace_wildcards(
+            params.negative_prompt, params.seed, get_wildcard_data()
+        )
+
     output = make_output_name(
         server,
         "inpaint",
@@ -321,10 +334,13 @@ def upscale(server: ServerContext, pool: DevicePoolExecutor):
     upscale = upscale_from_request()
     highres = highres_from_request()
 
-    output = make_output_name(server, "upscale", params, size)
+    params.prompt = replace_wildcards(params.prompt, params.seed, get_wildcard_data())
+    if params.negative_prompt is not None:
+        params.negative_prompt = replace_wildcards(
+            params.negative_prompt, params.seed, get_wildcard_data()
+        )
 
-    logger.info("resizing source image for limited pipeline")
-    source = valid_image(source, min_dims=size, max_dims=size)
+    output = make_output_name(server, "upscale", params, size)
 
     job_name = output[0]
     pool.submit(
@@ -364,14 +380,20 @@ def chain(server: ServerContext, pool: DevicePoolExecutor):
     output = make_output_name(server, "chain", params, size)
     job_name = output[0]
 
+    params.prompt = replace_wildcards(params.prompt, params.seed, get_wildcard_data())
+    if params.negative_prompt is not None:
+        params.negative_prompt = replace_wildcards(
+            params.negative_prompt, params.seed, get_wildcard_data()
+        )
+
     pipeline = ChainPipeline()
     for stage_data in data.get("stages", []):
-        callback = CHAIN_STAGES[stage_data.get("type")]
+        stage_class = CHAIN_STAGES[stage_data.get("type")]
         kwargs = stage_data.get("params", {})
-        logger.info("request stage: %s, %s", callback.__name__, kwargs)
+        logger.info("request stage: %s, %s", stage_class.__name__, kwargs)
 
         stage = StageParams(
-            stage_data.get("name", callback.__name__),
+            stage_data.get("name", stage_class.__name__),
             tile_size=get_size(kwargs.get("tile_size")),
             outscale=get_and_clamp_int(kwargs, "outscale", 1, 4),
         )
@@ -396,7 +418,6 @@ def chain(server: ServerContext, pool: DevicePoolExecutor):
             source_file = request.files.get(stage_source_name)
             if source_file is not None:
                 source = Image.open(BytesIO(source_file.read())).convert("RGB")
-                source = valid_image(source, max_dims=(size.width, size.height))
                 kwargs["stage_source"] = source
 
         if stage_mask_name in request.files:
@@ -408,10 +429,9 @@ def chain(server: ServerContext, pool: DevicePoolExecutor):
             mask_file = request.files.get(stage_mask_name)
             if mask_file is not None:
                 mask = Image.open(BytesIO(mask_file.read())).convert("RGB")
-                mask = valid_image(mask, max_dims=(size.width, size.height))
                 kwargs["stage_mask"] = mask
 
-        pipeline.append((callback, stage, kwargs))
+        pipeline.append((stage_class(), stage, kwargs))
 
     logger.info("running chain pipeline with %s stages", len(pipeline.stages))
 
@@ -447,7 +467,6 @@ def blend(server: ServerContext, pool: DevicePoolExecutor):
             logger.warning("missing source %s", i)
         else:
             source = Image.open(BytesIO(source_file.read())).convert("RGBA")
-            source = valid_image(source, mask.size, mask.size)
             sources.append(source)
 
     device, params, size = pipeline_from_request(server)
diff --git a/api/onnx_web/server/load.py b/api/onnx_web/server/load.py
index c1a3babf9..e37957452 100644
--- a/api/onnx_web/server/load.py
+++ b/api/onnx_web/server/load.py
@@ -90,6 +90,7 @@
 diffusion_models: List[str] = []
 network_models: List[NetworkModel] = []
 upscaling_models: List[str] = []
+wildcard_data: Dict[str, List[str]] = {}
 
 # Loaded from extra_models
 extra_hashes: Dict[str, str] = {}
@@ -120,6 +121,10 @@ def get_upscaling_models():
     return upscaling_models
 
 
+def get_wildcard_data():
+    return wildcard_data
+
+
 def get_extra_strings():
     return extra_strings
 
@@ -444,3 +449,22 @@ def any_first_cpu_last(a: DeviceParams, b: DeviceParams):
         "available acceleration platforms: %s",
         ", ".join([str(p) for p in available_platforms]),
     )
+
+
+def load_wildcards(server: ServerContext) -> None:
+    global wildcard_data
+
+    # simple wildcards
+    wildcard_files = list_model_globs(
+        server, ["*.txt"], base_path=path.join(server.model_path, "wildcard")
+    )
+
+    for file in wildcard_files:
+        with open(path.join(server.model_path, "wildcard", f"{file}.txt"), "r") as f:
+            lines = f.read().splitlines()
+            lines = [line.strip() for line in lines if not line.startswith("#")]
+            lines = [line for line in lines if len(line) > 0]
+            logger.debug("loading wildcards from %s: %s", file, lines)
+            wildcard_data[file] = lines
+
+    # TODO: structured wildcards
diff --git a/api/onnx_web/server/model_cache.py b/api/onnx_web/server/model_cache.py
index b7207c418..21da25f4e 100644
--- a/api/onnx_web/server/model_cache.py
+++ b/api/onnx_web/server/model_cache.py
@@ -1,3 +1,4 @@
+from enum import Enum
 from logging import getLogger
 from typing import Any, List, Tuple
 
@@ -6,6 +7,13 @@
 cache: List[Tuple[str, Any, Any]] = []
 
 
+class ModelTypes(str, Enum):
+    correction = "correction"
+    diffusion = "diffusion"
+    scheduler = "scheduler"
+    upscaling = "upscaling"
+
+
 class ModelCache:
     # cache: List[Tuple[str, Any, Any]]
     limit: int
diff --git a/api/onnx_web/server/params.py b/api/onnx_web/server/params.py
index caf0ec822..2598e819a 100644
--- a/api/onnx_web/server/params.py
+++ b/api/onnx_web/server/params.py
@@ -283,6 +283,7 @@ def upscale_from_request() -> UpscaleParams:
 
 
 def highres_from_request() -> HighresParams:
+    enabled = get_boolean(request.args, "highres", get_config_value("highres"))
     iterations = get_and_clamp_int(
         request.args,
         "highresIterations",
@@ -313,6 +314,7 @@ def highres_from_request() -> HighresParams:
         get_config_value("highresStrength", "min"),
     )
     return HighresParams(
+        enabled,
         scale,
         steps,
         strength,
diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py
index caeaab2cf..18c2ebbde 100644
--- a/api/onnx_web/utils.py
+++ b/api/onnx_web/utils.py
@@ -6,7 +6,7 @@
 from logging import getLogger
 from os import environ, path
 from platform import system
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Dict, List, Optional, Sequence, TypeVar, Union
 
 import torch
 from yaml import safe_load
@@ -43,9 +43,12 @@ def get_and_clamp_int(
     return min(max(int(args.get(key, default_value)), min_value), max_value)
 
 
+TElem = TypeVar("TElem")
+
+
 def get_from_list(
-    args: Any, key: str, values: Sequence[Any], default_value: Optional[Any] = None
-) -> Optional[Any]:
+    args: Any, key: str, values: Sequence[TElem], default_value: Optional[TElem] = None
+) -> Optional[TElem]:
     selected = args.get(key, default_value)
     if selected in values:
         return selected
@@ -57,7 +60,9 @@ def get_from_list(
     return None
 
 
-def get_from_map(args: Any, key: str, values: Dict[str, Any], default: Any) -> Any:
+def get_from_map(
+    args: Any, key: str, values: Dict[str, TElem], default: TElem
+) -> TElem:
     selected = args.get(key, default)
     if selected in values:
         return values[selected]
@@ -65,7 +70,7 @@ def get_from_map(args: Any, key: str, values: Dict[str, Any], default: Any) -> A
         return values[default]
 
 
-def get_not_empty(args: Any, key: str, default: Any) -> Any:
+def get_not_empty(args: Any, key: str, default: TElem) -> TElem:
     val = args.get(key, default)
 
     if val is None or len(val) == 0:
diff --git a/api/params.json b/api/params.json
index 7f21d78c1..7f1e9a545 100644
--- a/api/params.json
+++ b/api/params.json
@@ -64,6 +64,9 @@
     "max": 8192,
     "step": 8
   },
+  "highres": {
+    "default": false
+  },
   "highresIterations": {
     "default": 1,
     "min": 1,
diff --git a/api/pyproject.toml b/api/pyproject.toml
index b863bbfb4..efed7be31 100644
--- a/api/pyproject.toml
+++ b/api/pyproject.toml
@@ -14,6 +14,7 @@ exclude = [
 
 [[tool.mypy.overrides]]
 module = [
+"arpeggio",
     "basicsr.archs.rrdbnet_arch",
     "basicsr.utils.download_util",
     "basicsr.utils",
@@ -23,19 +24,45 @@ module = [
     "codeformer.facelib.utils.misc",
     "codeformer.facelib.utils",
     "codeformer.facelib",
+    "compel",
+    "controlnet_aux",
+    "cv2",
     "diffusers",
+    "diffusers.configuration_utils",
+    "diffusers.loaders",
+    "diffusers.models.attention_processor",
+    "diffusers.models.autoencoder_kl",
+    "diffusers.models.cross_attention",
+    "diffusers.models.embeddings",
+    "diffusers.models.modeling_utils",
+    "diffusers.models.unet_2d_blocks",
+    "diffusers.models.vae",
+    "diffusers.utils",
     "diffusers.pipelines.latent_diffusion.pipeline_latent_diffusion",
+    "diffusers.pipelines.onnx_utils",
     "diffusers.pipelines.paint_by_example",
     "diffusers.pipelines.stable_diffusion",
+    "diffusers.pipelines.stable_diffusion.convert_from_ckpt",
     "diffusers.pipeline_utils",
+    "diffusers.schedulers",
     "diffusers.utils.logging",
     "facexlib.utils",
     "facexlib",
     "gfpgan",
+    "gi.repository",
+    "huggingface_hub",
+    "huggingface_hub.file_download",
+    "huggingface_hub.utils.tqdm",
+    "mediapipe",
     "onnxruntime",
+    "onnxruntime.transformers.float16",
+    "piexif",
+    "piexif.helper",
     "realesrgan",
     "realesrgan.archs.srvgg_arch",
     "safetensors",
-    "transformers"
+    "timm.models.layers",
+    "transformers",
+    "win10toast"
 ]
 ignore_missing_imports = true
\ No newline at end of file
diff --git a/api/scripts/parse-prompts.py b/api/scripts/parse-prompts.py
new file mode 100644
index 000000000..954b4e011
--- /dev/null
+++ b/api/scripts/parse-prompts.py
@@ -0,0 +1,67 @@
+from argparse import ArgumentParser
+from collections import Counter
+from json import dumps, loads
+from sys import argv
+from typing import List
+
+
+def parse_args(args: List[str]):
+    parser = ArgumentParser(
+        prog="onnx-web prompt parser",
+        description="count phrase frequency in prompt books",
+    )
+    parser.add_argument("file", nargs="+", help="prompt files to parse")
+    return parser.parse_args(args)
+
+
+def load_duck(file: str):
+    import duckdb
+
+    cursor = duckdb.connect()
+    return [p[0] for p in cursor.sql(f"SELECT * FROM '{file}'").fetchall()]
+
+
+def load_json(file: str):
+    with open(file, "r") as f:
+        data = loads(f.read())
+        params = data.get("params", None)
+        if params:
+            prompt = params.get("input_prompt", None)
+            if prompt:
+                return prompt
+
+            prompt = params.get("prompt", None)
+            if prompt:
+                return prompt
+
+    return ""
+
+
+def load_text(file: str):
+    with open(file, "r") as f:
+        return f.readlines()
+
+
+def main():
+    args = parse_args(argv[1:])
+
+    lines: List[str] = []
+    for file in args.file:
+        if file.endswith(".parquet") or file.endswith(".duckdb"):
+            lines.extend(load_duck(file))
+        elif file.endswith(".json"):
+            # json only contains a single prompt
+            lines.append(load_json(file))
+        else:
+            lines.extend(load_text(file))
+
+    phrases = []
+    for line in lines:
+        phrases.extend([p.lower().strip() for p in line.split(",")])
+
+    count = Counter(phrases)
+    print(dumps(dict(count.most_common())))
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/exe/build.bat b/exe/build.bat
index 04a15e888..ee13b07da 100644
--- a/exe/build.bat
+++ b/exe/build.bat
@@ -28,5 +28,8 @@ git rev-parse HEAD > commit.txt
 set /p GIT_SHA=<commit.txt
 set GIT_HEAD=%GIT_SHA:~0,8%
 
+REM create version file
+echo "%BUNDLE_VERSION%-%BUNDLE_TYPE%-%GIT_HEAD%" > .\dist\version.txt
+
 REM create archive
 "C:\Program Files\7-Zip\7z.exe" a .\dist\onnx-web-v%BUNDLE_VERSION%-%BUNDLE_TYPE%-%GIT_HEAD%.zip ".\dist\*"
diff --git a/gui/src/client/api.ts b/gui/src/client/api.ts
index efcfd7dc3..6bed1e46b 100644
--- a/gui/src/client/api.ts
+++ b/gui/src/client/api.ts
@@ -123,6 +123,7 @@ export function appendUpscaleToURL(url: URL, upscale: UpscaleParams) {
 
 export function appendHighresToURL(url: URL, highres: HighresParams) {
   if (highres.enabled) {
+    url.searchParams.append('highres', String(highres.enabled));
     url.searchParams.append('highresIterations', highres.highresIterations.toFixed(FIXED_INTEGER));
     url.searchParams.append('highresMethod', highres.highresMethod);
     url.searchParams.append('highresScale', highres.highresScale.toFixed(FIXED_INTEGER));
diff --git a/gui/src/components/card/ImageCard.tsx b/gui/src/components/card/ImageCard.tsx
index 1e8893a92..5d6978176 100644
--- a/gui/src/components/card/ImageCard.tsx
+++ b/gui/src/components/card/ImageCard.tsx
@@ -28,7 +28,8 @@ export function ImageCard(props: ImageCardProps) {
   const { params, outputs, size } = image;
 
   const [_hash, setHash] = useHash();
-  const [anchor, setAnchor] = useState<Maybe<HTMLElement>>();
+  const [blendAnchor, setBlendAnchor] = useState<Maybe<HTMLElement>>();
+  const [saveAnchor, setSaveAnchor] = useState<Maybe<HTMLElement>>();
 
   const config = mustExist(useContext(ConfigContext));
   const state = mustExist(useContext(StateContext));
@@ -89,10 +90,18 @@ export function ImageCard(props: ImageCardProps) {
 
   function downloadImage() {
     window.open(outputs[index].url, '_blank');
+    close();
+  }
+
+  function downloadMetadata() {
+    window.open(outputs[index].url + '.json', '_blank');
+    close();
   }
 
   function close() {
-    setAnchor(undefined);
+    // TODO: split these up
+    setBlendAnchor(undefined);
+    setSaveAnchor(undefined);
   }
 
   const [index, setIndex] = useState(0);
@@ -151,10 +160,20 @@ export function ImageCard(props: ImageCardProps) {
           </GridItem>
           <GridItem xs={2}>
             <Tooltip title={t('tooltip.save')}>
-              <IconButton onClick={downloadImage}>
+              <IconButton onClick={(event) => {
+                setSaveAnchor(event.currentTarget);
+              }}>
                 <Download />
               </IconButton>
             </Tooltip>
+            <Menu
+              anchorEl={saveAnchor}
+              open={doesExist(saveAnchor)}
+              onClose={close}
+            >
+              <MenuItem key='save-image' onClick={downloadImage}>{t('save.image')}</MenuItem>
+              <MenuItem key='save-metadata' onClick={downloadMetadata}>{t('save.metadata')}</MenuItem>
+            </Menu>
           </GridItem>
           <GridItem xs={2}>
             <Tooltip title={t('tab.img2img')}>
@@ -180,14 +199,14 @@ export function ImageCard(props: ImageCardProps) {
           <GridItem xs={2}>
             <Tooltip title={t('tab.blend')}>
               <IconButton onClick={(event) => {
-                setAnchor(event.currentTarget);
+                setBlendAnchor(event.currentTarget);
               }}>
                 <Blender />
               </IconButton>
             </Tooltip>
             <Menu
-              anchorEl={anchor}
-              open={doesExist(anchor)}
+              anchorEl={blendAnchor}
+              open={doesExist(blendAnchor)}
               onClose={close}
             >
               {range(BLEND_SOURCES).map((idx) => <MenuItem key={idx} onClick={() => {
diff --git a/gui/src/config.json b/gui/src/config.json
index b1cdd4751..97a53d258 100644
--- a/gui/src/config.json
+++ b/gui/src/config.json
@@ -68,6 +68,9 @@
       "max": 8192,
       "step": 8
     },
+    "highres": {
+      "default": false
+    },
     "highresMethod": {
       "default": "lanczos",
       "keys": [
diff --git a/gui/src/strings/de.ts b/gui/src/strings/de.ts
index f95ed5e4c..8116ab133 100644
--- a/gui/src/strings/de.ts
+++ b/gui/src/strings/de.ts
@@ -191,6 +191,10 @@ export const I18N_STRINGS_DE = {
         txt2txt: '',
         upscale: '',
       },
+      save: {
+        image: '',
+        metadata: '',
+      },
       setting: {
         connectServer: 'verbinden zum Server',
         history: 'Bildgeschichte',
diff --git a/gui/src/strings/en.ts b/gui/src/strings/en.ts
index 99665fdbd..d5d9691ad 100644
--- a/gui/src/strings/en.ts
+++ b/gui/src/strings/en.ts
@@ -256,6 +256,10 @@ export const I18N_STRINGS_EN = {
         rocm: 'ROCm',
         tensorrt: 'TensorRT',
       },
+      save: {
+        image: 'Image',
+        metadata: 'Metadata',
+      },
       setting: {
         connectServer: 'Connect',
         history: 'Image History',
diff --git a/gui/src/strings/es.ts b/gui/src/strings/es.ts
index 5a7e00cfc..3af12260d 100644
--- a/gui/src/strings/es.ts
+++ b/gui/src/strings/es.ts
@@ -191,6 +191,10 @@ export const I18N_STRINGS_ES = {
         txt2txt: '',
         upscale: '',
       },
+      save: {
+        image: '',
+        metadata: '',
+      },
       setting: {
         connectServer: 'Conectar al servidor',
         history: 'Historia de la imagen',
diff --git a/gui/src/strings/fr.ts b/gui/src/strings/fr.ts
index 21ebe66f6..7174de556 100644
--- a/gui/src/strings/fr.ts
+++ b/gui/src/strings/fr.ts
@@ -191,6 +191,10 @@ export const I18N_STRINGS_FR = {
         txt2txt: '',
         upscale: '',
       },
+      save: {
+        image: '',
+        metadata: '',
+      },
       setting: {
         connectServer: '',
         history: '',
diff --git a/api/onnx_web/chain/blend_highres.py b/models/wildcard/.gitkeep
similarity index 100%
rename from api/onnx_web/chain/blend_highres.py
rename to models/wildcard/.gitkeep


Test patch:
diff --git a/api/scripts/test-refs/blend-512-muffin-black-0.png b/api/scripts/test-refs/blend-512-muffin-black-0.png
index 9af7aed24..eb1b096a0 100644
--- a/api/scripts/test-refs/blend-512-muffin-black-0.png
+++ b/api/scripts/test-refs/blend-512-muffin-black-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:d4106859b033cf57d0f2b08c36cff38beb5a67f6e99e2eddefad998223be4182
-size 498897
+oid sha256:dcd91064647c1e1ac7c6f30a34a406e9cb5e9865ec017c5134afc724cceab3bf
+size 556765
diff --git a/api/scripts/test-refs/blend-512-muffin-blend-0.png b/api/scripts/test-refs/blend-512-muffin-blend-0.png
index 93a589cd2..3a04b6928 100644
--- a/api/scripts/test-refs/blend-512-muffin-blend-0.png
+++ b/api/scripts/test-refs/blend-512-muffin-blend-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:2166d951fd884868019f50d3f2aba85a6ddf2ad401e5ff0c074503eb6a8f680c
-size 488620
+oid sha256:4d9935b6b30b36130bb556bf1101666a2f990c0a8e5febfce2aa15c5a5732c5c
+size 548579
diff --git a/api/scripts/test-refs/blend-512-muffin-white-0.png b/api/scripts/test-refs/blend-512-muffin-white-0.png
index 7b1a789a5..103f582c8 100644
--- a/api/scripts/test-refs/blend-512-muffin-white-0.png
+++ b/api/scripts/test-refs/blend-512-muffin-white-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:5dee66e2a5b003c6e892e2f5e7b7ba95c471ee301cdad34e11199083544a46d8
-size 502780
+oid sha256:be25e2a6252de2cd830c6421f75313067bf3ef29904138536615899c7169ac57
+size 573669
diff --git a/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png b/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
index b61dab553..28a2a0461 100644
--- a/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
+++ b/api/scripts/test-refs/img2img-sd-v1-5-256-pumpkin-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:18ab09b41c89f20cae4028770cd169e68a36245822c69ec7ce19ce622921b1bf
-size 215130
+oid sha256:294bc6c6aaa08dc5d5b3a831d5bbc092880fbc82b3c3437f8459f991cee3d03e
+size 133168
diff --git a/api/scripts/test-refs/txt2img-sd-v1-5-512-muffin-heun-0.png b/api/scripts/test-refs/txt2img-sd-v1-5-512-muffin-heun-0.png
index d3c2b16b0..dca24a756 100644
--- a/api/scripts/test-refs/txt2img-sd-v1-5-512-muffin-heun-0.png
+++ b/api/scripts/test-refs/txt2img-sd-v1-5-512-muffin-heun-0.png
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:b4f97d2e321fd091272452832350f4cbfd13bc4ba4a8aeeafcfaa396e6bad3f9
-size 497895
+oid sha256:ce2c5e5510c93d5d4e2e7688cd132e1c6e167e7715173e1250e86f341ac47172
+size 503943
diff --git a/api/scripts/test-refs/txt2img-sd-v1-5-tall-muffin-0.png b/api/scripts/test-refs/txt2img-sd-v1-5-tall-muffin-0.png
new file mode 100644
index 000000000..b7d4f9568
--- /dev/null
+++ b/api/scripts/test-refs/txt2img-sd-v1-5-tall-muffin-0.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:084c489e1e4d374c59c8a4a0ed476805815d3602732fe37f66032473cd4595fb
+size 794421
diff --git a/api/scripts/test-refs/upscale-resrgan-x4-tall-muffin-0.png b/api/scripts/test-refs/upscale-resrgan-x4-tall-muffin-0.png
new file mode 100644
index 000000000..2691ccea0
--- /dev/null
+++ b/api/scripts/test-refs/upscale-resrgan-x4-tall-muffin-0.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d3647eeaa605fb16b6c25c7f56d4335562015f43d1a814446d8e95d6ef21ac43
+size 8587429
diff --git a/api/scripts/test-release.py b/api/scripts/test-release.py
index eec65193e..2e6d7fd27 100644
--- a/api/scripts/test-release.py
+++ b/api/scripts/test-release.py
@@ -1,9 +1,11 @@
 import sys
+from argparse import ArgumentParser
 from io import BytesIO
 from logging import getLogger
 from logging.config import dictConfig
 from os import environ, path
 from time import sleep
+from tqdm import tqdm
 from typing import List, Optional, Union
 
 import cv2
@@ -24,15 +26,9 @@
 
 logger = getLogger(__name__)
 
-FAST_TEST = 20
-SLOW_TEST = 50
-
-
-def test_root() -> str:
-    if len(sys.argv) > 1:
-        return sys.argv[1]
-    else:
-        return "http://127.0.0.1:5000"
+FAST_TEST = 10
+SLOW_TEST = 25
+VERY_SLOW_TEST = 75
 
 
 def test_path(relpath: str) -> str:
@@ -222,6 +218,7 @@ def __init__(
         "upscale-sd-x4-2048-muffin",
         "upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-stable-diffusion-x4&scale=4&outscale=4",
         source="txt2img-sd-v1-5-512-muffin-0",
+        max_attempts=VERY_SLOW_TEST,
     ),
     TestCase(
         "outpaint-panorama-even-256",
@@ -231,7 +228,7 @@ def __init__(
         ),
         source="txt2img-sd-v1-5-512-muffin-0",
         mask="mask-black",
-        max_attempts=SLOW_TEST,
+        max_attempts=VERY_SLOW_TEST,
         mse_threshold=0.025,
     ),
     TestCase(
@@ -242,7 +239,7 @@ def __init__(
         ),
         source="txt2img-sd-v1-5-512-muffin-0",
         mask="mask-black",
-        max_attempts=SLOW_TEST,
+        max_attempts=VERY_SLOW_TEST,
         mse_threshold=0.025,
     ),
     TestCase(
@@ -253,7 +250,7 @@ def __init__(
         ),
         source="txt2img-sd-v1-5-512-muffin-0",
         mask="mask-black",
-        max_attempts=SLOW_TEST,
+        max_attempts=VERY_SLOW_TEST,
         mse_threshold=0.025,
     ),
     TestCase(
@@ -295,6 +292,7 @@ def __init__(
             "&correction=correction-codeformer&faces=true&faceOutscale=1&faceStrength=1.0"
         ),
         source="txt2img-sd-v1-5-512-muffin-0",
+        max_attempts=VERY_SLOW_TEST,
     ),
     TestCase(
         "upscale-sd-x4-gfpgan-2048-muffin",
@@ -303,16 +301,32 @@ def __init__(
             "&correction=correction-gfpgan&faces=true&faceOutscale=1&faceStrength=1.0"
         ),
         source="txt2img-sd-v1-5-512-muffin-0",
+        max_attempts=VERY_SLOW_TEST,
     ),
     TestCase(
         "txt2img-panorama-1024x768-muffin",
         "txt2img?prompt=a+giant+muffin&seed=0&scheduler=ddim&width=1024&height=768&pipeline=panorama&tiledVAE=true",
+        max_attempts=VERY_SLOW_TEST,
     ),
     TestCase(
         "img2img-panorama-1024x768-pumpkin",
         "img2img?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&sourceFilter=none&pipeline=panorama&tiledVAE=true",
         source="txt2img-panorama-1024x768-muffin-0",
+        max_attempts=VERY_SLOW_TEST,
+    ),
+    TestCase(
+        "txt2img-sd-v1-5-tall-muffin",
+        "txt2img?prompt=a+giant+muffin&seed=0&scheduler=ddim&width=512&height=768",
     ),
+    TestCase(
+        "upscale-resrgan-x4-tall-muffin",
+        (
+            "upscale?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&upscaling=upscaling-real-esrgan-x4-plus&scale=4&outscale=4"
+            "&correction=correction-gfpgan&faces=false&faceOutscale=1&faceStrength=1.0"
+        ),
+        source="txt2img-sd-v1-5-tall-muffin-0",
+    ),
+    # TODO: non-square controlnet
 ]
 
 
@@ -321,7 +335,18 @@ def __str__(self) -> str:
         return super().__str__()
 
 
-def generate_images(root: str, test: TestCase) -> Optional[str]:
+def parse_args(args: List[str]):
+    parser = ArgumentParser(
+        prog="onnx-web release tests",
+        description="regression tests for onnx-web",
+    )
+    parser.add_argument("--host", default="http://127.0.0.1:5000")
+    parser.add_argument("-n", "--name")
+    parser.add_argument("-m", "--mse", default=1.0, type=float)
+    return parser.parse_args(args)
+
+
+def generate_images(host: str, test: TestCase) -> Optional[str]:
     files = {}
     if test.source is not None:
         if isinstance(test.source, list):
@@ -354,7 +379,7 @@ def generate_images(root: str, test: TestCase) -> Optional[str]:
         files["mask"] = mask_bytes
 
     logger.debug("generating image: %s", test.query)
-    resp = requests.post(f"{root}/api/{test.query}", files=files)
+    resp = requests.post(f"{host}/api/{test.query}", files=files)
     if resp.status_code == 200:
         json = resp.json()
         return json.get("outputs")
@@ -363,8 +388,8 @@ def generate_images(root: str, test: TestCase) -> Optional[str]:
         raise TestError("error generating image")
 
 
-def check_ready(root: str, key: str) -> bool:
-    resp = requests.get(f"{root}/api/ready?output={key}")
+def check_ready(host: str, key: str) -> bool:
+    resp = requests.get(f"{host}/api/ready?output={key}")
     if resp.status_code == 200:
         json = resp.json()
         ready = json.get("ready", False)
@@ -379,10 +404,10 @@ def check_ready(root: str, key: str) -> bool:
         raise TestError("error getting image status")
 
 
-def download_images(root: str, keys: List[str]) -> List[Image.Image]:
+def download_images(host: str, keys: List[str]) -> List[Image.Image]:
     images = []
     for key in keys:
-        resp = requests.get(f"{root}/output/{key}")
+        resp = requests.get(f"{host}/output/{key}")
         if resp.status_code == 200:
             logger.debug("downloading image: %s", key)
             images.append(Image.open(BytesIO(resp.content)))
@@ -413,31 +438,32 @@ def find_mse(result: Image.Image, ref: Image.Image) -> float:
 
 
 def run_test(
-    root: str,
+    host: str,
     test: TestCase,
+    mse_mult: float = 1.0,
 ) -> bool:
     """
     Generate an image, wait for it to be ready, and calculate the MSE from the reference.
     """
 
-    keys = generate_images(root, test)
+    keys = generate_images(host, test)
     if keys is None:
         raise ValueError("could not generate image")
 
-    attempts = 0
-    while attempts < test.max_attempts:
-        if check_ready(root, keys[0]):
+    ready = False
+    for attempt in tqdm(range(test.max_attempts)):
+        if check_ready(host, keys[0]):
             logger.debug("image is ready: %s", keys)
+            ready = True
             break
         else:
             logger.debug("waiting for image to be ready")
-            attempts += 1
             sleep(6)
 
-    if attempts == test.max_attempts:
+    if not ready:
         raise ValueError("image was not ready in time")
 
-    results = download_images(root, keys)
+    results = download_images(host, keys)
     if results is None:
         raise ValueError("could not download image")
 
@@ -451,7 +477,7 @@ def run_test(
 
         mse = find_mse(result, ref)
 
-        if mse < test.mse_threshold:
+        if mse < (test.mse_threshold * mse_mult):
             logger.info("MSE within threshold: %.5f < %.5f", mse, test.mse_threshold)
         else:
             logger.warning("MSE above threshold: %.5f > %.5f", mse, test.mse_threshold)
@@ -461,24 +487,29 @@ def run_test(
 
 
 def main():
-    root = test_root()
-    logger.info("running release tests against API: %s", root)
+    args = parse_args(sys.argv[1:])
+    logger.info("running release tests against API: %s", args.host)
+
+    if args.name is None:
+        tests = TEST_DATA
+    else:
+        tests = [test for test in TEST_DATA if args.name in test.name]
 
     # make sure tests have unique names
-    test_names = [test.name for test in TEST_DATA]
+    test_names = [test.name for test in tests]
     if len(test_names) != len(set(test_names)):
         logger.error("tests must have unique names: %s", test_names)
         sys.exit(1)
 
     passed = []
     failed = []
-    for test in TEST_DATA:
+    for test in tests:
         test_passed = False
 
         for _i in range(3):
             try:
                 logger.info("starting test: %s", test.name)
-                if run_test(root, test):
+                if run_test(args.host, test, mse_mult=args.mse):
                     logger.info("test passed: %s", test.name)
                     test_passed = True
                     break
@@ -492,7 +523,7 @@ def main():
         else:
             failed.append(test.name)
 
-    logger.info("%s of %s tests passed", len(passed), len(TEST_DATA))
+    logger.info("%s of %s tests passed", len(passed), len(tests))
     failed = list(set(failed))
     if len(failed) > 0:
         logger.error("%s tests had errors: %s", len(failed), failed)