feat: split up UNet and VAE tile size and overlap/stride params

api/onnx_web/chain/highres.py

@@ -43,7 +43,7 @@ def stage_highres(
                     outscale=highres.scale,
                 ),
                 chain=chain,
-                overlap=params.overlap,
+                overlap=params.vae_overlap,
             )
         else:
             logger.debug("using simple upscaling for highres")
@@ -51,14 +51,14 @@ def stage_highres(
                 UpscaleSimpleStage(),
                 stage,
                 method=highres.method,
-                overlap=params.overlap,
+                overlap=params.vae_overlap,
                 upscale=upscale.with_args(scale=highres.scale, outscale=highres.scale),
             )
 
         chain.stage(
             BlendImg2ImgStage(),
             stage,
-            overlap=params.overlap,
+            overlap=params.vae_overlap,
             prompt_index=prompt_index + i,
             strength=highres.strength,
         )

api/onnx_web/chain/source_txt2img.py

@@ -60,9 +60,9 @@ class SourceTxt2ImgStage(BaseStage):
         )
 
         if params.is_xl():
-            tile_size = max(stage.tile_size, params.tiles)
+            tile_size = max(stage.tile_size, params.unet_tile)
         else:
-            tile_size = params.tiles
+            tile_size = params.unet_tile
 
         # this works for panorama as well, because tile_size is already max(tile_size, *size)
         latent_size = size.min(tile_size, tile_size)

api/onnx_web/chain/upscale_bsrgan.py

@@ -106,5 +106,5 @@ class UpscaleBSRGANStage(BaseStage):
         params: ImageParams,
         size: Size,
     ) -> int:
-        tile = min(params.tiles, self.max_tile)
+        tile = min(params.unet_tile, self.max_tile)
         return size.width // tile * size.height // tile

api/onnx_web/chain/upscale_outpaint.py

@@ -71,7 +71,7 @@ class UpscaleOutpaintStage(BaseStage):
                 outputs.append(source)
                 continue
 
-            tile_size = params.tiles
+            tile_size = params.unet_tile
             size = Size(*source.size)
             latent_size = size.min(tile_size, tile_size)
 

api/onnx_web/diffusers/load.py

@@ -266,14 +266,13 @@ def load_pipeline(
 
     # update panorama params
     if params.is_panorama():
-        latent_window = params.tiles // 8
+        unet_stride = (params.unet_tile * (1 - params.unet_overlap)) // 8
-        latent_stride = params.stride // 8
+        logger.debug("setting panorama window parameters: %s/%s for UNet, %s/%s for VAE", params.unet_tile, unet_stride, params.vae_tile, params.vae_overlap)
-
+        pipe.set_window_size(params.unet_tile // 8, unet_stride)
-        pipe.set_window_size(latent_window, latent_stride)
 
         for vae in VAE_COMPONENTS:
             if hasattr(pipe, vae):
-                getattr(pipe, vae).set_window_size(latent_window, params.overlap)
+                getattr(pipe, vae).set_window_size(params.vae_tile // 8, params.vae_overlap)
 
     run_gc([device])
 
@@ -626,8 +625,8 @@ def patch_pipeline(
             server,
             original_decoder,
             decoder=True,
-            window=params.tiles,
+            window=params.unet_tile,
-            overlap=params.overlap,
+            overlap=params.vae_overlap,
         )
         logger.debug("patched VAE decoder with wrapper")
 
@@ -637,8 +636,8 @@ def patch_pipeline(
             server,
             original_encoder,
             decoder=False,
-            window=params.tiles,
+            window=params.unet_tile,
-            overlap=params.overlap,
+            overlap=params.vae_overlap,
         )
         logger.debug("patched VAE encoder with wrapper")
 

api/onnx_web/diffusers/run.py

@@ -44,9 +44,9 @@ def run_txt2img_pipeline(
 ) -> None:
     # if using panorama, the pipeline will tile itself (views)
     if params.is_panorama() or params.is_xl():
-        tile_size = max(params.tiles, size.width, size.height)
+        tile_size = max(params.unet_tile, size.width, size.height)
     else:
-        tile_size = params.tiles
+        tile_size = params.unet_tile
 
     # prepare the chain pipeline and first stage
     chain = ChainPipeline()
@@ -57,11 +57,11 @@ def run_txt2img_pipeline(
         ),
         size=size,
         prompt_index=0,
-        overlap=params.overlap,
+        overlap=params.vae_overlap,
     )
 
     # apply upscaling and correction, before highres
-    stage = StageParams(tile_size=params.tiles)
+    stage = StageParams(tile_size=params.unet_tile)
     first_upscale, after_upscale = split_upscale(upscale)
     if first_upscale:
         stage_upscale_correction(
@@ -139,14 +139,14 @@ def run_img2img_pipeline(
     # prepare the chain pipeline and first stage
     chain = ChainPipeline()
     stage = StageParams(
-        tile_size=params.tiles,
+        tile_size=params.unet_tile,
     )
     chain.stage(
         BlendImg2ImgStage(),
         stage,
         prompt_index=0,
         strength=strength,
-        overlap=params.overlap,
+        overlap=params.vae_overlap,
     )
 
     # apply upscaling and correction, before highres
@@ -236,7 +236,7 @@ def run_inpaint_pipeline(
     full_res_inpaint_padding: float,
 ) -> None:
     logger.debug("building inpaint pipeline")
-    tile_size = params.tiles
+    tile_size = params.unet_tile
 
     if mask is None:
         # if no mask was provided, keep the full source image
@@ -332,7 +332,7 @@ def run_inpaint_pipeline(
         fill_color=fill_color,
         mask_filter=mask_filter,
         noise_source=noise_source,
-        overlap=params.overlap,
+        overlap=params.vae_overlap,
         prompt_index=0,
     )
 
@@ -410,7 +410,7 @@ def run_upscale_pipeline(
 ) -> None:
     # set up the chain pipeline, no base stage for upscaling
     chain = ChainPipeline()
-    stage = StageParams(tile_size=params.tiles)
+    stage = StageParams(tile_size=params.unet_tile)
 
     # apply upscaling and correction, before highres
     first_upscale, after_upscale = split_upscale(upscale)

api/onnx_web/params.py

@@ -204,8 +204,10 @@ class ImageParams:
     input_negative_prompt: str
     loopback: int
     tiled_vae: bool
-    tiles: int
+    unet_tile: int
-    overlap: float
+    unet_overlap: float
+    vae_tile: int
+    vae_overlap: float
 
     def __init__(
         self,
@@ -224,9 +226,10 @@ class ImageParams:
         input_negative_prompt: Optional[str] = None,
         loopback: int = 0,
         tiled_vae: bool = False,
-        tiles: int = 512,
+        unet_overlap: float = 0.25,
-        overlap: float = 0.25,
+        unet_tile: int = 512,
-        stride: int = 64,
+        vae_overlap: float = 0.25,
+        vae_tile: int = 512,
     ) -> None:
         self.model = model
         self.pipeline = pipeline
@@ -243,9 +246,10 @@ class ImageParams:
         self.input_negative_prompt = input_negative_prompt or negative_prompt
         self.loopback = loopback
         self.tiled_vae = tiled_vae
-        self.tiles = tiles
+        self.unet_overlap = unet_overlap
-        self.overlap = overlap
+        self.unet_tile = unet_tile
-        self.stride = stride
+        self.vae_overlap = vae_overlap
+        self.vae_tile = vae_tile
 
     def do_cfg(self):
         return self.cfg > 1.0
@@ -312,9 +316,10 @@ class ImageParams:
             "input_negative_prompt": self.input_negative_prompt,
             "loopback": self.loopback,
             "tiled_vae": self.tiled_vae,
-            "tiles": self.tiles,
+            "unet_overlap": self.unet_overlap,
-            "overlap": self.overlap,
+            "unet_tile": self.unet_tile,
-            "stride": self.stride,
+            "vae_overlap": self.vae_overlap,
+            "vae_tile": self.vae_tile,
         }
 
     def with_args(self, **kwargs):
@@ -334,9 +339,10 @@ class ImageParams:
             kwargs.get("input_negative_prompt", self.input_negative_prompt),
             kwargs.get("loopback", self.loopback),
             kwargs.get("tiled_vae", self.tiled_vae),
-            kwargs.get("tiles", self.tiles),
+            kwargs.get("unet_overlap", self.unet_overlap),
-            kwargs.get("overlap", self.overlap),
+            kwargs.get("unet_tile", self.unet_tile),
-            kwargs.get("stride", self.stride),
+            kwargs.get("vae_overlap", self.vae_overlap),
+            kwargs.get("vae_tile", self.vae_tile),
         )
 
 

api/onnx_web/server/params.py

@@ -117,32 +117,35 @@ def build_params(
         get_config_value("steps", "max"),
         get_config_value("steps", "min"),
     )
-    tiled_vae = get_boolean(data, "tiledVAE", get_config_value("tiledVAE"))
+    tiled_vae = get_boolean(data, "tiled_vae", get_config_value("tiled_vae"))
-    tiles = get_and_clamp_int(
+    unet_overlap = get_and_clamp_float(
         data,
-        "tiles",
+        "unet_overlap",
-        get_config_value("tiles"),
+        get_config_value("unet_overlap"),
-        get_config_value("tiles", "max"),
+        get_config_value("unet_overlap", "max"),
-        get_config_value("tiles", "min"),
+        get_config_value("unet_overlap", "min"),
     )
-    overlap = get_and_clamp_float(
+    unet_tile = get_and_clamp_int(
         data,
-        "overlap",
+        "unet_tile",
-        get_config_value("overlap"),
+        get_config_value("unet_tile"),
-        get_config_value("overlap", "max"),
+        get_config_value("unet_tile", "max"),
-        get_config_value("overlap", "min"),
+        get_config_value("unet_tile", "min"),
     )
-    stride = get_and_clamp_int(
+    vae_overlap = get_and_clamp_float(
         data,
-        "stride",
+        "vae_overlap",
-        get_config_value("stride"),
+        get_config_value("vae_overlap"),
-        get_config_value("stride", "max"),
+        get_config_value("vae_overlap", "max"),
-        get_config_value("stride", "min"),
+        get_config_value("vae_overlap", "min"),
+    )
+    vae_tile = get_and_clamp_int(
+        data,
+        "vae_tile",
+        get_config_value("vae_tile"),
+        get_config_value("vae_tile", "max"),
+        get_config_value("vae_tile", "min"),
     )
-
-    if stride > tiles:
-        logger.info("limiting stride to tile size, %s > %s", stride, tiles)
-        stride = tiles
 
     seed = int(data.get("seed", -1))
     if seed == -1:
@@ -163,9 +166,10 @@ def build_params(
         control=control,
         loopback=loopback,
         tiled_vae=tiled_vae,
-        tiles=tiles,
+        unet_overlap=unet_overlap,
-        overlap=overlap,
+        unet_tile=unet_tile,
-        stride=stride,
+        vae_overlap=vae_overlap,
+        vae_tile=vae_tile,
     )
 
     return params

api/params.json

@@ -141,12 +141,6 @@
     "max": 4,
     "step": 1
   },
-  "overlap": {
-    "default": 0.25,
-    "min": 0.0,
-    "max": 0.9,
-    "step": 0.01
-  },
   "pipeline": {
     "default": "",
     "keys": [
@@ -197,21 +191,9 @@
     "max": 1,
     "step": 0.01
   },
-  "stride": {
+  "tiled_vae": {
-    "default": 128,
-    "min": 64,
-    "max": 512,
-    "step": 64
-  },
-  "tiledVAE": {
     "default": false
   },
-  "tiles": {
-    "default": 512,
-    "min": 128,
-    "max": 2048,
-    "step": 128
-  },
   "tileOrder": {
     "default": "spiral",
     "keys": [
@@ -225,6 +207,18 @@
     "max": 1024,
     "step": 8
   },
+  "unet_overlap": {
+    "default": 0.25,
+    "min": 0.0,
+    "max": 0.9,
+    "step": 0.01
+  },
+  "unet_tile": {
+    "default": 512,
+    "min": 128,
+    "max": 2048,
+    "step": 128
+  },
   "upscaleOrder": {
     "default": "correction-first",
     "keys": [
@@ -237,6 +231,18 @@
     "default": "",
     "keys": []
   },
+  "vae_overlap": {
+    "default": 0.25,
+    "min": 0.0,
+    "max": 0.9,
+    "step": 0.01
+  },
+  "vae_tile": {
+    "default": 512,
+    "min": 256,
+    "max": 1024,
+    "step": 128
+  },
   "width": {
     "default": 512,
     "min": 128,

api/scripts/test-release.py

@@ -305,12 +305,12 @@ TEST_DATA = [
     ),
     TestCase(
         "txt2img-panorama-1024x768-muffin",
-        "txt2img?prompt=a+giant+muffin&seed=0&scheduler=ddim&width=1024&height=768&pipeline=panorama&tiledVAE=true",
+        "txt2img?prompt=a+giant+muffin&seed=0&scheduler=ddim&width=1024&height=768&pipeline=panorama&tiled_vae=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",
+        "img2img?prompt=a+giant+pumpkin&seed=0&scheduler=ddim&sourceFilter=none&pipeline=panorama&tiled_vae=true",
         source="txt2img-panorama-1024x768-muffin-0",
         max_attempts=VERY_SLOW_TEST,
     ),

gui/src/client/api.ts

@@ -70,10 +70,11 @@ export function makeImageURL(root: string, type: string, params: BaseImgParams):
   url.searchParams.append('cfg', params.cfg.toFixed(FIXED_FLOAT));
   url.searchParams.append('eta', params.eta.toFixed(FIXED_FLOAT));
   url.searchParams.append('steps', params.steps.toFixed(FIXED_INTEGER));
-  url.searchParams.append('tiledVAE', String(params.tiledVAE));
+  url.searchParams.append('tiled_vae', String(params.tiled_vae));
-  url.searchParams.append('tiles', params.tiles.toFixed(FIXED_INTEGER));
+  url.searchParams.append('unet_overlap', params.unet_overlap.toFixed(FIXED_FLOAT));
-  url.searchParams.append('overlap', params.overlap.toFixed(FIXED_FLOAT));
+  url.searchParams.append('unet_tile', params.unet_tile.toFixed(FIXED_INTEGER));
-  url.searchParams.append('stride', params.stride.toFixed(FIXED_INTEGER));
+  url.searchParams.append('vae_overlap', params.vae_overlap.toFixed(FIXED_FLOAT));
+  url.searchParams.append('vae_tile', params.vae_tile.toFixed(FIXED_INTEGER));
 
   if (doesExist(params.scheduler)) {
     url.searchParams.append('scheduler', params.scheduler);

gui/src/components/control/ImageControl.tsx

@@ -1,3 +1,4 @@
+/* eslint-disable camelcase */
 import { doesExist, mustDefault, mustExist } from '@apextoaster/js-utils';
 import { Casino } from '@mui/icons-material';
 import { Button, Checkbox, FormControlLabel, Stack } from '@mui/material';
@@ -47,9 +48,6 @@ export function ImageControl(props: ImageControlProps) {
     staleTime: STALE_TIME,
   });
 
-  // max stride is the lesser of tile size and server's max stride
-  const maxStride = Math.min(state.tiles, params.stride.max);
-
   return <Stack spacing={2}>
     <Stack direction='row' spacing={4}>
       <QueryList
@@ -156,58 +154,74 @@ export function ImageControl(props: ImageControlProps) {
         }}
       />
       <NumericField
-        label={t('parameter.tiles')}
+        label={t('parameter.unet_tile')}
-        min={params.tiles.min}
+        min={params.unet_tile.min}
-        max={params.tiles.max}
+        max={params.unet_tile.max}
-        step={params.tiles.step}
+        step={params.unet_tile.step}
-        value={state.tiles}
+        value={state.unet_tile}
-        onChange={(tiles) => {
+        onChange={(unet_tile) => {
           props.onChange({
             ...state,
-            tiles,
+            unet_tile,
+          });
+        }}
+      />
+      <NumericField
+        label={t('parameter.unet_overlap')}
+        min={params.unet_overlap.min}
+        max={params.unet_overlap.max}
+        step={params.unet_overlap.step}
+        value={state.unet_overlap}
+        onChange={(unet_overlap) => {
+          props.onChange({
+            ...state,
+            unet_overlap,
+          });
+        }}
+      />
+      <FormControlLabel
+        label={t('parameter.tiled_vae')}
+        control={<Checkbox
+          checked={state.tiled_vae}
+          value='check'
+          onChange={(event) => {
+            props.onChange({
+              ...state,
+              tiled_vae: state.tiled_vae === false,
+            });
+          }}
+        />}
+      />
+      <NumericField
+        decimal
+        disabled={state.tiled_vae === false}
+        label={t('parameter.vae_overlap')}
+        min={params.vae_overlap.min}
+        max={params.vae_overlap.max}
+        step={params.vae_overlap.step}
+        value={state.vae_overlap}
+        onChange={(vae_tile) => {
+          props.onChange({
+            ...state,
+            vae_tile,
           });
         }}
       />
       <NumericField
         decimal
-        label={t('parameter.overlap')}
+        disabled={state.tiled_vae === false}
-        min={params.overlap.min}
+        label={t('parameter.vae_tile')}
-        max={params.overlap.max}
+        min={params.vae_tile.min}
-        step={params.overlap.step}
+        max={params.vae_tile.max}
-        value={state.overlap}
+        step={params.vae_tile.step}
-        onChange={(overlap) => {
+        value={state.vae_tile}
+        onChange={(vae_tile) => {
           props.onChange({
             ...state,
-            overlap,
+            vae_tile,
           });
         }}
       />
-      <NumericField
-        label={t('parameter.stride')}
-        min={params.stride.min}
-        max={maxStride}
-        step={params.stride.step}
-        value={state.stride}
-        onChange={(stride) => {
-          props.onChange({
-            ...state,
-            stride,
-          });
-        }}
-      />
-      <FormControlLabel
-        label={t('parameter.tiledVAE')}
-        control={<Checkbox
-          checked={state.tiledVAE}
-          value='check'
-          onChange={(event) => {
-            props.onChange({
-              ...state,
-              tiledVAE: state.tiledVAE === false,
-            });
-          }}
-        />}
-      />
     </Stack>
     <PromptInput
       selector={selector}

gui/src/config.json

@@ -131,12 +131,6 @@
       "max": 4,
       "step": 1
     },
-    "overlap": {
-      "default": 0.25,
-      "min": 0.0,
-      "max": 0.9,
-      "step": 0.01
-    },
     "pipeline": {
       "default": "",
       "keys": [
@@ -193,15 +187,9 @@
       "max": 512,
       "step": 64
     },
-    "tiledVAE": {
+    "tiled_vae": {
       "default": false
     },
-    "tiles": {
-      "default": 512,
-      "min": 128,
-      "max": 1024,
-      "step": 128
-    },
     "tileOrder": {
       "default": "spiral",
       "keys": [
@@ -215,6 +203,18 @@
       "max": 1024,
       "step": 8
     },
+    "unet_overlap": {
+      "default": 0.25,
+      "min": 0.0,
+      "max": 0.9,
+      "step": 0.01
+    },
+    "unet_tile": {
+      "default": 512,
+      "min": 128,
+      "max": 1024,
+      "step": 128
+    },
     "upscaleOrder": {
       "default": "correction-first",
       "keys": [
@@ -227,6 +227,12 @@
       "default": "",
       "keys": []
     },
+    "vae_overlap": {
+      "default": 0.25,
+      "min": 0.0,
+      "max": 0.9,
+      "step": 0.01
+    },
     "width": {
       "default": 512,
       "min": 256,

gui/src/state.ts

@@ -1,3 +1,4 @@
+/* eslint-disable camelcase */
 /* eslint-disable max-lines */
 /* eslint-disable no-null/no-null */
 import { Maybe } from '@apextoaster/js-utils';
@@ -267,10 +268,11 @@ export function baseParamsFromServer(defaults: ServerParams): Required<BaseImgPa
     scheduler: defaults.scheduler.default,
     steps: defaults.steps.default,
     seed: defaults.seed.default,
-    tiledVAE: defaults.tiledVAE.default,
+    tiled_vae: defaults.tiled_vae.default,
-    tiles: defaults.tiles.default,
+    unet_overlap: defaults.unet_overlap.default,
-    overlap: defaults.overlap.default,
+    unet_tile: defaults.unet_tile.default,
-    stride: defaults.stride.default,
+    vae_overlap: defaults.vae_overlap.default,
+    vae_tile: defaults.vae_tile.default,
   };
 }
 

gui/src/strings/de.ts

@@ -162,9 +162,11 @@ export const I18N_STRINGS_DE = {
         steps: 'Schritte',
         strength: 'Stärke',
         stride: '',
-        tiledVAE: '',
+        tiled_vae: '',
         tiles: '',
         tileOrder: '',
+        unet_overlap: '',
+        unet_tile: '',
         upscale: {
           label: '',
           denoise: 'Entrauschen',
@@ -172,6 +174,8 @@ export const I18N_STRINGS_DE = {
           order: '',
           outscale: 'Ausgangsskala',
         },
+        vae_overlap: '',
+        vae_tile: '',
         width: 'Breite',
         correction: {
           label: 'Gesichtskorrektur',

gui/src/strings/en.ts

@@ -215,10 +215,10 @@ export const I18N_STRINGS_EN = {
         sourceFilter: 'Source Filter',
         steps: 'Steps',
         strength: 'Strength',
-        stride: 'UNet Stride',
+        tiled_vae: 'Tiled VAE',
-        tiledVAE: 'Tiled VAE',
-        tiles: 'Tile Size',
         tileOrder: 'Tile Order',
+        unet_overlap: 'UNet Overlap',
+        unet_tile: 'UNet Tile Size',
         upscale: {
           label: 'Upscale',
           denoise: 'Denoise',
@@ -226,6 +226,8 @@ export const I18N_STRINGS_EN = {
           order: 'Upscale Order',
           outscale: 'Outscale',
         },
+        vae_overlap: 'UNet Overlap',
+        vae_tile: 'VAE Tile Size',
         width: 'Width',
         correction: {
           label: 'Face Correction',

gui/src/strings/es.ts

@@ -162,9 +162,11 @@ export const I18N_STRINGS_ES = {
         steps: 'Pasos',
         strength: 'Fuerza',
         stride: '',
-        tiledVAE: '',
+        tiled_vae: '',
         tiles: '',
         tileOrder: 'Orden de secciones',
+        unet_overlap: '',
+        unet_tile: '',
         upscale: {
           label: 'Aumento',
           denoise: '',
@@ -172,6 +174,8 @@ export const I18N_STRINGS_ES = {
           order: '',
           outscale: 'Escala de producción',
         },
+        vae_overlap: '',
+        vae_tile: '',
         width: 'Anchura',
         correction: {
           label: 'Corrección facial',

gui/src/strings/fr.ts

@@ -162,9 +162,11 @@ export const I18N_STRINGS_FR = {
         steps: '',
         strength: '',
         stride: '',
-        tiledVAE: '',
+        tiled_vae: '',
         tiles: '',
         tileOrder: '',
+        unet_overlap: '',
+        unet_tile: '',
         upscale: {
           label: '',
           denoise: '',
@@ -172,6 +174,8 @@ export const I18N_STRINGS_FR = {
           order: '',
           outscale: '',
         },
+        vae_overlap: '',
+        vae_tile: '',
         width: '',
         correction: {
           label: '',

gui/src/types/params.ts

@@ -50,10 +50,11 @@ export interface BaseImgParams {
   negativePrompt?: string;
 
   batch: number;
-  tiledVAE: boolean;
+  tiled_vae: boolean;
-  tiles: number;
+  vae_overlap: number;
-  overlap: number;
+  vae_tile: number;
-  stride: number;
+  unet_overlap: number;
+  unet_tile: number;
 
   cfg: number;
   steps: number;