Upload sd-webui-ultimate-tiled-vae using SD-Hub

sd-webui-ultimate-tiled-vae/LICENSE

@@ -0,0 +1,16 @@
+Creative Commons Attribution-NonCommercial-ShareAlike 4.0 International (CC BY-NC-SA 4.0)
+
+This work is based on "Tiled VAE" by LI YI @ Nanyang Technological University.
+Original source: https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111
+
+You are free to:
+  Share — copy and redistribute the material in any medium or format
+  Adapt — remix, transform, and build upon the material
+
+Under the following terms:
+  Attribution — You must give appropriate credit.
+  NonCommercial — You may not use the material for commercial purposes.
+  ShareAlike — If you remix, transform, or build upon the material,
+               you must distribute your contributions under the same license.
+
+Full license text: https://creativecommons.org/licenses/by-nc-sa/4.0/legalcode

sd-webui-ultimate-tiled-vae/README.md

@@ -0,0 +1,70 @@
+# sd-webui-ultimate-tiled-vae  v1.1
+
+Standalone **Tiled VAE** extension for [AUTOMATIC1111 Stable Diffusion WebUI](https://github.com/AUTOMATIC1111/stable-diffusion-webui).
+
+Splits VAE encode/decode into overlapping tiles so giant images (4K, 8K, panoramas) process without OOM, producing **seamless, artefact-free output**.
+
+## Installation
+
+```bash
+# Option A – git (recommended)
+cd /content/A1111/extensions
+git clone https://github.com/yourname/sd-webui-ultimate-tiled-vae
+
+# Option B – manual (Colab / no git)
+mkdir -p /content/A1111/extensions/sd-webui-ultimate-tiled-vae
+# unzip or copy all files into that folder
+```
+
+Restart the webui after installation.
+
+## Quick smoke test
+
+| Test | Expected |
+|------|----------|
+| txt2img 512×512, extension disabled | Normal generation, no change |
+| txt2img 1024×1024, Fast Approx preset | Tiled execution logged, seamless output |
+| txt2img 2048×2048, Exact Quality preset | Slower, higher quality, no seams |
+| img2img upscale 2× | Works cleanly, hook removed in postprocess |
+| hires.fix on/off | No stuck state after generation |
+| Interrupt mid-decode | Hook still removed on next run |
+
+## Features
+
+| Feature | Detail |
+|---------|--------|
+| **3 GroupNorm modes** | *Exact* (full cross-tile sync), *Approx* (from downsampled estimate), *None* (independent) |
+| **GN/fast-mode validation** | Config automatically enforces coherent combinations |
+| **Configurable pad sizes** | Separate encoder/decoder overlap padding |
+| **Smart ignore** | Skips tiling when image fits in one tile |
+| **Auto-adjust tile** | Last tile never larger than first (prevents late OOM) |
+| **NaN-only fp32 retry** | Retries in fp32 only when output actually contains NaN |
+| **Fallback policy** | Abort / Cheap Approx / Original Forward |
+| **Quick Presets** | Exact Quality / Fast Approx / Safe Minimal VRAM |
+| **Standalone attention** | No dependency on `tile_utils.attn` from multidiffusion |
+| **Dynamic scale factor** | Reads from VAE, not hardcoded to 8 |
+| **Safe hook lifecycle** | `try/finally` everywhere; hooks always restored after crash/interrupt |
+
+## GroupNorm Sync modes
+
+| Mode | How | When to use |
+|------|-----|-------------|
+| **Exact** | Collect stats from all tiles, weighted average, apply globally | Maximum quality; disables fast mode automatically |
+| **Approx** | Estimate stats on downsampled image; apply to all tiles | Best balance; requires (and auto-enables) fast mode |
+| **None** | Each tile normalises itself | Fastest; may produce faint seams on large uniform areas |
+
+## Roadmap
+
+- **Phase 6** – per-block skip matrix (Experimental panel) wired into executor
+- **Phase 7** – SDXL / TAESD explicit guardrails
+- **Phase 8** – Extras-tab postprocess bridge
+
+## Credits
+
+- Original Tiled VAE engine: [LI YI / pkuliyi2015](https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111) — CC BY-NC-SA 4.0
+- Pad size / GN mode ideas: [Kahsolt / stable-diffusion-webui-vae-tile-infer](https://github.com/Kahsolt/stable-diffusion-webui-vae-tile-infer)
+- Tile planner improvements: inspired by [ComfyUI-TiledVaeLite](https://github.com/tinymindlab/comfyui-tiled-vae-lite)
+
+## License
+
+CC BY-NC-SA 4.0 — inherited from original engine.

sd-webui-ultimate-tiled-vae/scripts/ultimate_tiled_vae.py

@@ -0,0 +1,263 @@
+"""
+sd-webui-ultimate-tiled-vae  —  v1.1
+Standalone Tiled VAE extension for AUTOMATIC1111 Stable Diffusion WebUI.
+
+Based on the original Tiled VAE by LI YI (CC BY-NC-SA 4.0)
+  https://github.com/pkuliyi2015/multidiffusion-upscaler-for-automatic1111
+Improvements from:
+  • stable-diffusion-webui-vae-tile-infer (Kahsolt)
+  • ComfyUI-TiledVaeLite (tile planner ideas)
+"""
+
+import sys
+import os
+
+_EXT_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+if _EXT_DIR not in sys.path:
+    sys.path.insert(0, _EXT_DIR)
+
+import gradio as gr
+import modules.scripts as scripts
+
+from tiled_vae.config import (
+    TiledVAEConfig, GNSyncMode, FallbackPolicy,
+    preset_exact, preset_fast_approx, preset_safe_minimal,
+)
+from tiled_vae.hook import attach_hooks, detach_hooks, is_hooked
+from tiled_vae.compat import get_input_accordion, get_optimal_device_name
+from tiled_vae.logging_utils import info, warn
+
+
+# ── VRAM-based recommended defaults ──────────────────────────────────────────
+
+def _rcmd_enc() -> int:
+    try:
+        import torch, modules.devices as devices
+        if torch.cuda.is_available() and devices.device not in ["cpu", devices.cpu]:
+            mb = torch.cuda.get_device_properties(devices.device).total_memory // 2**20
+            if   mb > 16_000: return 3072
+            elif mb > 12_000: return 2048
+            elif mb >  8_000: return 1536
+            else:             return 960
+    except Exception: pass
+    return 512
+
+
+def _rcmd_dec() -> int:
+    try:
+        import torch, modules.devices as devices
+        if torch.cuda.is_available() and devices.device not in ["cpu", devices.cpu]:
+            mb = torch.cuda.get_device_properties(devices.device).total_memory // 2**20
+            if   mb > 30_000: return 256
+            elif mb > 16_000: return 192
+            elif mb > 12_000: return 128
+            elif mb >  8_000: return 96
+            else:             return 64
+    except Exception: pass
+    return 64
+
+
+def _uid(tab: str, name: str) -> str:
+    return f"UTV-{tab}-{name}"
+
+
+# ── Script ────────────────────────────────────────────────────────────────────
+
+class Script(scripts.Script):
+
+    def __init__(self):
+        self._hooked = False
+
+    def title(self):
+        return "Ultimate Tiled VAE"
+
+    def show(self, is_img2img):
+        return scripts.AlwaysVisible
+
+    # ── UI ────────────────────────────────────────────────────────────────────
+
+    def ui(self, is_img2img):
+        tab = "i2i" if is_img2img else "t2i"
+        uid = lambda n: _uid(tab, n)
+
+        InputAccordion = get_input_accordion()
+
+        if InputAccordion:
+            acc = InputAccordion(False, label="Ultimate Tiled VAE", elem_id=uid("acc"))
+        else:
+            acc = gr.Accordion("Ultimate Tiled VAE", open=False, elem_id=uid("acc"))
+
+        with acc as enabled_acc:
+
+            if not InputAccordion:
+                enabled = gr.Checkbox(label="Enable Ultimate Tiled VAE", value=False, elem_id=uid("en"))
+            else:
+                enabled = enabled_acc
+
+            # Row: GPU + preset
+            with gr.Row():
+                move_vae_to_gpu = gr.Checkbox(label="Move VAE to GPU", value=True, elem_id=uid("gpu"))
+                preset = gr.Dropdown(
+                    label="Quick Preset",
+                    choices=["Custom", "Exact Quality", "Fast Approx", "Safe Minimal VRAM"],
+                    value="Custom",
+                    elem_id=uid("preset"),
+                )
+
+            # Tile sizes
+            gr.HTML('<p style="margin:0.4em 0 0.2em"><b>Tile Sizes</b> — larger = better quality, more VRAM.</p>')
+            with gr.Row():
+                encoder_tile_size = gr.Slider(label="Encoder Tile Size", minimum=256, maximum=4096, step=16, value=_rcmd_enc(), elem_id=uid("enc_tile"))
+                decoder_tile_size = gr.Slider(label="Decoder Tile Size", minimum=48, maximum=512, step=16, value=_rcmd_dec(), elem_id=uid("dec_tile"))
+                gr.Button("↻ Reset", variant="tool", elem_id=uid("reset")).click(
+                    fn=lambda: [_rcmd_enc(), _rcmd_dec()],
+                    outputs=[encoder_tile_size, decoder_tile_size],
+                    show_progress=False,
+                )
+
+            # Pad sizes
+            with gr.Row():
+                encoder_pad_size = gr.Slider(label="Encoder Pad", minimum=0, maximum=64, step=1, value=32, elem_id=uid("enc_pad"))
+                decoder_pad_size = gr.Slider(label="Decoder Pad", minimum=0, maximum=32, step=1, value=11, elem_id=uid("dec_pad"))
+
+            # Fast mode + color fix
+            with gr.Row():
+                fast_encoder = gr.Checkbox(label="Fast Encoder", value=True, elem_id=uid("fenc"))
+                fast_decoder = gr.Checkbox(label="Fast Decoder", value=True, elem_id=uid("fdec"))
+                color_fix    = gr.Checkbox(label="Encoder Color Fix", value=False, elem_id=uid("cfix"))
+                fast_encoder.change(fn=lambda v: gr.update(visible=v), inputs=fast_encoder, outputs=color_fix, show_progress=False)
+
+            # GN sync
+            gn_sync_mode = gr.Radio(
+                label="GroupNorm Sync  (Approx auto-enables fast mode; Exact auto-disables it)",
+                choices=[m.value for m in GNSyncMode],
+                value=GNSyncMode.APPROX.value,
+                elem_id=uid("gn"),
+            )
+
+            # Auto / smart / log
+            with gr.Row():
+                auto_adjust  = gr.Checkbox(label="Auto Adjust Tile Size", value=True, elem_id=uid("adj"))
+                smart_ignore = gr.Checkbox(label="Smart Ignore Small Images", value=True, elem_id=uid("si"))
+                log_timings  = gr.Checkbox(label="Log Timings", value=True, elem_id=uid("log"))
+
+            # Reliability
+            with gr.Row():
+                retry_fp32 = gr.Checkbox(label="Retry in FP32 on NaN", value=True, elem_id=uid("fp32"))
+                fallback_policy = gr.Dropdown(
+                    label="Fallback on failure",
+                    choices=[p.value for p in FallbackPolicy],
+                    value=FallbackPolicy.CHEAP_APPROX.value,
+                    elem_id=uid("fb"),
+                )
+
+            # Experimental (Phase 6: per-block skip matrix not yet implemented)
+            with gr.Accordion("Experimental", open=False, elem_id=uid("exp")):
+                gr.HTML(
+                    "<p style='color:orange;margin:0'>"
+                    "⚠ Experimental options may produce artefacts. "
+                    "Per-block skip matrix arrives in Phase 6.</p>"
+                )
+                skip_infer = gr.Checkbox(
+                    label="Enable Skip Infer (no per-block control yet)",
+                    value=False,
+                    elem_id=uid("skip"),
+                )
+
+            # ── Preset wiring ─────────────────────────────────────────────────
+            def _apply_preset(name):
+                enc, dec = _rcmd_enc(), _rcmd_dec()
+                if   name == "Exact Quality":      c = preset_exact(enc, dec)
+                elif name == "Fast Approx":        c = preset_fast_approx(enc, dec)
+                elif name == "Safe Minimal VRAM":  c = preset_safe_minimal(enc, dec)
+                else:
+                    return (gr.update(),) * 7
+                return (
+                    c.fast_encoder, c.fast_decoder, c.encoder_color_fix,
+                    c.gn_sync_mode, c.auto_adjust_tile, c.smart_ignore,
+                    c.fallback_policy,
+                )
+
+            preset.change(
+                fn=_apply_preset,
+                inputs=[preset],
+                outputs=[fast_encoder, fast_decoder, color_fix, gn_sync_mode,
+                         auto_adjust, smart_ignore, fallback_policy],
+                show_progress=False,
+            )
+
+        return [
+            enabled, move_vae_to_gpu,
+            encoder_tile_size, decoder_tile_size,
+            encoder_pad_size,  decoder_pad_size,
+            fast_encoder, fast_decoder, color_fix,
+            gn_sync_mode,
+            auto_adjust, smart_ignore,
+            retry_fp32, fallback_policy,
+            log_timings, skip_infer,
+        ]
+
+    # ── process ───────────────────────────────────────────────────────────────
+
+    def process(self, p, *args):
+        (
+            enabled, move_vae_to_gpu,
+            encoder_tile_size, decoder_tile_size,
+            encoder_pad_size,  decoder_pad_size,
+            fast_encoder, fast_decoder, color_fix,
+            gn_sync_mode,
+            auto_adjust, smart_ignore,
+            retry_fp32, fallback_policy,
+            log_timings, skip_infer,
+        ) = args
+
+        vae = p.sd_model.first_stage_model
+
+        if not enabled:
+            if is_hooked(vae):
+                detach_hooks(vae)
+                self._hooked = False
+            return
+
+        # VRAM hint
+        try:
+            import modules.devices as devices
+            if get_optimal_device_name().startswith("cuda") \
+               and vae.device == devices.cpu and not move_vae_to_gpu:
+                warn("VAE is on CPU. Enable 'Move VAE to GPU' for better performance.")
+        except Exception:
+            pass
+
+        cfg = TiledVAEConfig(
+            enabled=True,
+            move_vae_to_gpu=move_vae_to_gpu,
+            encoder_tile_size=int(encoder_tile_size),
+            decoder_tile_size=int(decoder_tile_size),
+            encoder_pad_size=int(encoder_pad_size),
+            decoder_pad_size=int(decoder_pad_size),
+            fast_encoder=fast_encoder,
+            fast_decoder=fast_decoder,
+            encoder_color_fix=color_fix,
+            gn_sync_mode=gn_sync_mode,
+            auto_adjust_tile=auto_adjust,
+            smart_ignore=smart_ignore,
+            retry_fp32_on_nan=retry_fp32,
+            fallback_policy=fallback_policy,
+            log_timings=log_timings,
+            skip_infer=skip_infer,
+        ).validate()
+
+        attach_hooks(vae, cfg)
+        self._hooked = True
+
+    # ── postprocess: always clean up ─────────────────────────────────────────
+
+    def postprocess(self, p, processed, *args):
+        try:
+            vae = p.sd_model.first_stage_model
+            if is_hooked(vae):
+                detach_hooks(vae)
+        except Exception as e:
+            warn(f"postprocess cleanup failed: {e}")
+        finally:
+            self._hooked = False

sd-webui-ultimate-tiled-vae/tiled_vae/__init__.py

@@ -0,0 +1,3 @@
+"""Ultimate Tiled VAE – internal package."""
+from .config  import TiledVAEConfig, GNSyncMode, FallbackPolicy
+from .hook    import attach_hooks, detach_hooks, is_hooked, VAEHook

sd-webui-ultimate-tiled-vae/tiled_vae/compat.py

@@ -0,0 +1,169 @@
+"""
+Compatibility shims for different A1111 / torch / model variants.
+
+IMPORTANT – attention contract:
+  _attn2task() in taskqueue.py decomposes attention into:
+    store_res  ->  pre_norm  ->  attn  ->  add_res
+
+  Therefore the `attn` callable receives an ALREADY-NORMALIZED tensor
+  and must return ONLY the projected attention output — no residual add,
+  no second norm call.  The residual is re-added by the `add_res` task.
+
+  The old fallback violated that contract (called block.norm again and
+  returned x + out).  Fixed in this version.
+"""
+
+from __future__ import annotations
+from typing import Any, Callable
+
+import torch
+
+
+# ── A1111 module helpers ───────────────────────────────────────────────────────
+
+def get_devices_module():
+    import modules.devices as devices
+    return devices
+
+
+def get_optimal_device() -> torch.device:
+    return get_devices_module().get_optimal_device()
+
+
+def get_optimal_device_name() -> str:
+    devs = get_devices_module()
+    if hasattr(devs, "get_optimal_device_name"):
+        return devs.get_optimal_device_name()
+    return str(devs.get_optimal_device())
+
+
+def torch_gc() -> None:
+    import gc
+    get_devices_module().torch_gc()
+    gc.collect()
+
+
+def test_for_nans(tensor: torch.Tensor, where: str = "vae") -> None:
+    get_devices_module().test_for_nans(tensor, where)
+
+
+# ── A1111 UI helpers ──────────────────────────────────────────────────────────
+
+def get_input_accordion():
+    try:
+        from modules.shared import opts
+        if hasattr(opts, "hypertile_enable_unet"):
+            from modules.ui_components import InputAccordion
+            return InputAccordion
+    except Exception:
+        pass
+    return None
+
+
+# ── Model inspection ──────────────────────────────────────────────────────────
+
+def is_taesd(vae_model: Any) -> bool:
+    return "taesd" in type(vae_model).__name__.lower()
+
+
+def is_sdxl(vae_model: Any) -> bool:
+    try:
+        enc = vae_model.encoder
+        if hasattr(enc, "z_channels"):
+            return enc.z_channels == 16
+        if hasattr(enc, "conv_out"):
+            return enc.conv_out.out_channels == 32
+    except Exception:
+        pass
+    return False
+
+
+def get_vae_scale_factor(vae_model: Any) -> int:
+    """Return the spatial scale factor (encoder /scale, decoder *scale). Usually 8."""
+    try:
+        return int(vae_model.downscale_factor)
+    except Exception:
+        pass
+    try:
+        from modules.processing import opt_f
+        return int(opt_f)
+    except Exception:
+        return 8
+
+
+# ── Attention function factory ────────────────────────────────────────────────
+#
+# CONTRACT (must be honoured by every path):
+#   f(block, x) where x is ALREADY GroupNorm-normalised.
+#   Returns ONLY proj_out(attn(x)).
+#   Must NOT call block.norm(x).
+#   Must NOT add the residual (the add_res task does that).
+#
+# This module is fully self-contained; tile_utils.attn from multidiffusion
+# is no longer a dependency.
+
+def _vanilla_attn_output(block: Any, x: torch.Tensor) -> torch.Tensor:
+    """Scaled dot-product attention on an already-normalised x. No norm, no residual."""
+    B, C, H, W = x.shape
+    q = block.q(x).reshape(B, C, H * W).permute(0, 2, 1)   # (B, HW, C)
+    k = block.k(x).reshape(B, C, H * W)                     # (B, C, HW)
+    v = block.v(x).reshape(B, C, H * W)                     # (B, C, HW)
+
+    w = torch.bmm(q, k) * (int(C) ** -0.5)                  # (B, HW, HW)
+    w = torch.softmax(w, dim=2).permute(0, 2, 1)
+
+    out = torch.bmm(v, w).reshape(B, C, H, W)
+    return block.proj_out(out)
+
+
+def _xformers_attn_output(block: Any, x: torch.Tensor) -> torch.Tensor:
+    """
+    xformers memory-efficient attention on an already-normalised x.
+    No norm, no residual.
+
+    Full axis layout (both directions must be correct and symmetric):
+
+      INPUT:  block.q/k/v return (B, C, H, W)
+              xformers expects (B, seq, C)
+              Convert: (B, C, H, W) -> permute(0,2,3,1) -> (B, H, W, C)
+                                    -> reshape           -> (B, H*W, C)
+              A plain .reshape(B, H*W, C) from (B, C, H, W) silently
+              interleaves channel and spatial bytes — wrong.
+
+      OUTPUT: xformers returns (B, H*W, C)
+              We need (B, C, H, W)
+              Convert: (B, H*W, C) -> reshape  -> (B, H, W, C)
+                                   -> permute(0,3,1,2) -> (B, C, H, W)
+              A plain .reshape(B, C, H, W) from (B, HW, C) is equally
+              wrong for the same reason.
+    """
+    try:
+        import xformers.ops as xops
+        B, C, H, W = x.shape
+        # input: (B, C, H, W) → (B, HW, C)  [permute THEN reshape]
+        q = block.q(x).permute(0, 2, 3, 1).reshape(B, H * W, C)
+        k = block.k(x).permute(0, 2, 3, 1).reshape(B, H * W, C)
+        v = block.v(x).permute(0, 2, 3, 1).reshape(B, H * W, C)
+        # out: (B, HW, C)
+        out = xops.memory_efficient_attention(q, k, v)
+        # output: (B, HW, C) → (B, C, H, W)  [reshape THEN permute]
+        out = out.reshape(B, H, W, C).permute(0, 3, 1, 2).contiguous()
+        return block.proj_out(out)
+    except Exception:
+        # xformers unavailable or call failed – fall back to vanilla
+        return _vanilla_attn_output(block, x)
+
+
+def get_attn_func(block: Any) -> Callable:
+    """
+    Return the correct attention callable for `block`.
+    All paths honour the contract: pre-normed input, proj_out output only.
+    No tile_utils dependency.
+    """
+    try:
+        from ldm.modules.diffusionmodules.model import MemoryEfficientAttnBlock
+        if isinstance(block, MemoryEfficientAttnBlock):
+            return _xformers_attn_output
+    except ImportError:
+        pass
+    return _vanilla_attn_output

sd-webui-ultimate-tiled-vae/tiled_vae/config.py

@@ -0,0 +1,160 @@
+"""
+TiledVAEConfig – dataclass + validation + presets.
+
+GN mode / fast-mode contract (enforced in validate()):
+  Approx  ->  requires fast estimation  ->  forces fast_encoder/fast_decoder True
+  Exact   ->  requires exact path       ->  forces fast_encoder/fast_decoder False
+  None    ->  each tile is independent  ->  fast_* has no effect on GN, allowed either way
+"""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class GNSyncMode(str, Enum):
+    EXACT  = "Exact"    # collect stats from all tiles, weighted average
+    APPROX = "Approx"   # stats from downsampled full-image estimate
+    NONE   = "None"     # each tile normalises itself (fastest, may seam)
+
+
+class FallbackPolicy(str, Enum):
+    ABORT        = "Abort"
+    CHEAP_APPROX = "Cheap Approx"
+    ORIGINAL     = "Original Forward"
+
+
+@dataclass
+class TiledVAEConfig:
+    # ── Core toggles ──────────────────────────────────────────────────────────
+    enabled:            bool = False
+    move_vae_to_gpu:    bool = True
+
+    # ── Tile geometry ─────────────────────────────────────────────────────────
+    encoder_tile_size:  int  = 960
+    decoder_tile_size:  int  = 64
+    encoder_pad_size:   int  = 32
+    decoder_pad_size:   int  = 11
+
+    # ── Quality / speed knobs ─────────────────────────────────────────────────
+    fast_encoder:       bool = True
+    fast_decoder:       bool = True
+    encoder_color_fix:  bool = False
+    gn_sync_mode:       str  = GNSyncMode.APPROX.value
+
+    # ── Auto-sizing ───────────────────────────────────────────────────────────
+    auto_adjust_tile:   bool = True
+    smart_ignore:       bool = True
+
+    # ── Reliability ───────────────────────────────────────────────────────────
+    retry_fp32_on_nan:  bool = True
+    fallback_policy:    str  = FallbackPolicy.CHEAP_APPROX.value
+
+    # ── Logging ───────────────────────────────────────────────────────────────
+    log_timings:        bool = True
+
+    # ── Experimental ─────────────────────────────────────────────────────────
+    skip_infer: bool = False
+    # Per-block skip dicts — populated in Phase 6
+    skip_infer_enc: dict = field(default_factory=dict)
+    skip_infer_dec: dict = field(default_factory=dict)
+
+    # ─────────────────────────────────────────────────────────────────────────
+    def validate(self) -> "TiledVAEConfig":
+        # Clamp numeric fields
+        self.encoder_tile_size = max(256, int(self.encoder_tile_size))
+        self.decoder_tile_size = max(48,  int(self.decoder_tile_size))
+        self.encoder_pad_size  = max(0,   int(self.encoder_pad_size))
+        self.decoder_pad_size  = max(0,   int(self.decoder_pad_size))
+
+        # Validate enum strings
+        valid_gn = {m.value for m in GNSyncMode}
+        if self.gn_sync_mode not in valid_gn:
+            self.gn_sync_mode = GNSyncMode.APPROX.value
+
+        valid_fb = {p.value for p in FallbackPolicy}
+        if self.fallback_policy not in valid_fb:
+            self.fallback_policy = FallbackPolicy.CHEAP_APPROX.value
+
+        # ── GN mode / fast-mode coherence ────────────────────────────────────
+        # Approx mode is implemented by running estimate_group_norm on a
+        # downsampled image BEFORE tile execution.  That path is only activated
+        # when fast_encoder / fast_decoder is True.  If the user picks Approx
+        # but disables fast mode, the executor would silently fall through to
+        # the layer's own forward (effectively None).  We prevent that by
+        # auto-enabling fast mode.
+        if self.gn_sync_mode == GNSyncMode.APPROX.value:
+            if not self.fast_encoder:
+                self.fast_encoder = True
+                _warn_coerce("GN=Approx requires fast encoder — enabling fast_encoder.")
+            if not self.fast_decoder:
+                self.fast_decoder = True
+                _warn_coerce("GN=Approx requires fast decoder — enabling fast_decoder.")
+
+        # Exact mode relies on the EXACT task-queue path (no pre-estimation).
+        # If fast mode is on, estimate_group_norm replaces pre_norm tasks with
+        # apply_norm and the exact cross-tile accumulation never happens.
+        # Force fast_* off so the semantics are correct.
+        if self.gn_sync_mode == GNSyncMode.EXACT.value:
+            if self.fast_encoder:
+                self.fast_encoder = False
+                _warn_coerce("GN=Exact requires fast_encoder=False — disabling.")
+            if self.fast_decoder:
+                self.fast_decoder = False
+                _warn_coerce("GN=Exact requires fast_decoder=False — disabling.")
+            # color_fix is also a fast-encoder sub-feature
+            self.encoder_color_fix = False
+
+        return self
+
+
+def _warn_coerce(msg: str) -> None:
+    print(f"[Ultimate Tiled VAE] config: {msg}")
+
+
+# ── Preset factory ────────────────────────────────────────────────────────────
+
+def preset_exact(enc_tile: int, dec_tile: int) -> TiledVAEConfig:
+    """Maximum quality – full GN sync, no fast paths (auto-validated)."""
+    return TiledVAEConfig(
+        enabled=True,
+        encoder_tile_size=enc_tile,
+        decoder_tile_size=dec_tile,
+        fast_encoder=False,
+        fast_decoder=False,
+        encoder_color_fix=False,
+        gn_sync_mode=GNSyncMode.EXACT.value,
+        retry_fp32_on_nan=True,
+        fallback_policy=FallbackPolicy.CHEAP_APPROX.value,
+    ).validate()
+
+
+def preset_fast_approx(enc_tile: int, dec_tile: int) -> TiledVAEConfig:
+    """Good balance of quality and speed (auto-validated)."""
+    return TiledVAEConfig(
+        enabled=True,
+        encoder_tile_size=enc_tile,
+        decoder_tile_size=dec_tile,
+        fast_encoder=True,
+        fast_decoder=True,
+        encoder_color_fix=True,
+        gn_sync_mode=GNSyncMode.APPROX.value,
+        move_vae_to_gpu=True,
+        retry_fp32_on_nan=True,
+        fallback_policy=FallbackPolicy.CHEAP_APPROX.value,
+    ).validate()
+
+
+def preset_safe_minimal(enc_tile: int, dec_tile: int) -> TiledVAEConfig:
+    """Minimal VRAM – independent tile norm, immediate fp32 retry."""
+    return TiledVAEConfig(
+        enabled=True,
+        encoder_tile_size=enc_tile,
+        decoder_tile_size=max(48, dec_tile // 2),
+        fast_encoder=True,
+        fast_decoder=True,
+        encoder_color_fix=False,
+        gn_sync_mode=GNSyncMode.NONE.value,
+        smart_ignore=True,
+        retry_fp32_on_nan=True,
+        fallback_policy=FallbackPolicy.ORIGINAL.value,
+    ).validate()

sd-webui-ultimate-tiled-vae/tiled_vae/executor.py

@@ -0,0 +1,251 @@
+"""
+Tile execution engine for Ultimate Tiled VAE.
+
+Zigzag order minimises GPU<->CPU transfers.
+Supports all three GN sync modes.
+scale parameter is threaded through so no hardcoded 8.
+"""
+
+from __future__ import annotations
+from typing import List, Optional, TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+
+from .config import GNSyncMode
+from .groupnorm import GroupNormParam
+from .tiling import crop_valid_region
+from .logging_utils import info
+
+if TYPE_CHECKING:
+    from .config import TiledVAEConfig
+
+
+@torch.no_grad()
+def estimate_group_norm(
+    z: torch.Tensor,
+    task_queue: List,
+    color_fix: bool,
+) -> bool:
+    """
+    Fast-mode: run forward on downsampled z to collect GN parameters.
+    Replaces 'pre_norm' tasks in-place with 'apply_norm'.
+    Returns True on success, False on NaN (caller falls back to exact mode).
+    """
+    tile = z
+    last_id = len(task_queue) - 1
+    while last_id >= 0 and task_queue[last_id][0] != "pre_norm":
+        last_id -= 1
+    if last_id <= 0:
+        raise ValueError("No GroupNorm found in task queue")
+
+    for i in range(last_id + 1):
+        task = task_queue[i]
+        name = task[0]
+
+        if name == "pre_norm":
+            gn_func = GroupNormParam.from_tile(tile, task[1])
+            task_queue[i] = ["apply_norm", gn_func]
+            if i == last_id:
+                return True
+            tile = gn_func(tile)
+
+        elif name in ("store_res", "store_res_cpu"):
+            j = i + 1
+            while j < last_id and task_queue[j][0] != "add_res":
+                j += 1
+            if j < last_id:
+                res = task[1](tile)
+                task_queue[j][1] = res.cpu() if name == "store_res_cpu" else res
+
+        elif name == "add_res":
+            tile = tile + task[1].to(tile.device)
+            task[1] = None
+
+        elif color_fix and name == "downsample":
+            # encoder color fix: estimate only up to first downsample
+            for k in range(i, last_id + 1):
+                if task_queue[k][0] == "store_res":
+                    task_queue[k][0] = "store_res_cpu"
+            return True
+
+        else:
+            tile = task[1](tile)
+
+        # NaN guard during estimation
+        if torch.isnan(tile).any():
+            info("NaN during fast-mode estimation — disabling fast mode.")
+            return False
+
+    raise IndexError("estimate_group_norm: should not reach here")
+
+
+def _prepare_fast_mode(
+    z: torch.Tensor,
+    tile_size: int,
+    task_queue_template: List,
+    color_fix: bool,
+    device: torch.device,
+) -> List:
+    """
+    Downsample z, fix distribution shift, run estimate_group_norm.
+    Returns updated queue on success, original template on failure.
+    """
+    H, W = z.shape[2], z.shape[3]
+    scale_factor = tile_size / max(H, W)
+    z_dev  = z.to(device)
+    z_down = F.interpolate(z_dev, scale_factor=scale_factor, mode="nearest-exact")
+
+    # Restore distribution (downsampling shifts mean/std)
+    std_old, mean_old = torch.std_mean(z_dev,  dim=[0, 2, 3], keepdim=True)
+    std_new, mean_new = torch.std_mean(z_down, dim=[0, 2, 3], keepdim=True)
+    z_down = (z_down - mean_new) / std_new * std_old + mean_old
+    del std_old, mean_old, std_new, mean_new
+    z_down = torch.clamp_(z_down, min=z_dev.min(), max=z_dev.max())
+
+    info(
+        f"Fast mode: GN estimation on "
+        f"{z_down.shape[3]}×{z_down.shape[2]} downsampled image"
+    )
+
+    from .taskqueue import clone_task_queue
+    est_queue = clone_task_queue(task_queue_template)
+    if estimate_group_norm(z_down, est_queue, color_fix=color_fix):
+        return est_queue
+    return task_queue_template   # fall back to exact
+
+
+@torch.no_grad()
+def execute_tiles(
+    tiles: List[Optional[torch.Tensor]],
+    task_queues: List[List],
+    in_bboxes,
+    out_bboxes,
+    result_shape,
+    device: torch.device,
+    dtype: torch.dtype,
+    is_decoder: bool,
+    gn_sync_mode: str,
+    cfg: "TiledVAEConfig",
+    interrupted_flag,
+    scale: int = 8,
+) -> Optional[torch.Tensor]:
+    """
+    Execute all task queues in zigzag order.
+    Returns assembled result tensor, or None if interrupted.
+    """
+    num_tiles     = len(tiles)
+    num_completed = 0
+    result: Optional[torch.Tensor] = None
+    forward = True
+
+    total_steps = num_tiles * len(task_queues[0])
+    pbar = tqdm(
+        total=total_steps,
+        desc=f"[Tiled VAE] {'Decoder' if is_decoder else 'Encoder'}",
+    )
+
+    try:
+        while True:
+            if interrupted_flag.interrupted:
+                return None
+
+            gn_param  = GroupNormParam()
+            tile_iter = range(num_tiles) if forward else reversed(range(num_tiles))
+
+            for i in tile_iter:
+                if interrupted_flag.interrupted:
+                    return None
+
+                tile       = tiles[i].to(device)
+                task_queue = task_queues[i]
+                in_bbox    = in_bboxes[i]
+
+                while task_queue:
+                    if interrupted_flag.interrupted:
+                        return None
+
+                    task = task_queue.pop(0)
+                    name = task[0]
+
+                    if name == "pre_norm":
+                        if gn_sync_mode == GNSyncMode.EXACT.value:
+                            # Suspend this tile until all tiles contribute stats
+                            gn_param.add_tile(tile, task[1])
+                            break
+                        elif gn_sync_mode == GNSyncMode.APPROX.value:
+                            # apply_norm was injected by estimate_group_norm
+                            # If somehow pre_norm slipped through, use layer forward
+                            tile = task[1](tile)
+                        else:  # NONE
+                            tile = task[1](tile)
+
+                    elif name == "apply_norm":
+                        tile = task[1](tile)
+
+                    elif name in ("store_res", "store_res_cpu"):
+                        res = task[1](tile)
+                        if name == "store_res_cpu":
+                            res = res.cpu()
+                        j = 0
+                        while j < len(task_queue) and task_queue[j][0] != "add_res":
+                            j += 1
+                        if j < len(task_queue):
+                            task_queue[j][1] = res
+
+                    elif name == "add_res":
+                        tile = tile + task[1].to(device)
+                        task[1] = None
+
+                    else:
+                        tile = task[1](tile)
+
+                    pbar.update(1)
+
+                # Write completed tile
+                if not task_queue:
+                    # NaN check on completed tile
+                    try:
+                        import modules.devices as devmod
+                        devmod.test_for_nans(tile, "vae")
+                    except Exception:
+                        raise
+
+                    N = result_shape[0]
+                    if result is None:
+                        result = torch.zeros(result_shape, device=device, requires_grad=False)
+
+                    valid = crop_valid_region(tile, in_bbox, out_bboxes[i], is_decoder, scale=scale)
+                    ob = out_bboxes[i]
+                    result[:, :, ob[2]:ob[3], ob[0]:ob[1]] = valid
+                    tiles[i] = None
+                    num_completed += 1
+                    del tile
+
+                elif i == num_tiles - 1 and forward:
+                    forward  = False
+                    tiles[i] = tile
+
+                elif i == 0 and not forward:
+                    forward  = True
+                    tiles[i] = tile
+
+                else:
+                    tiles[i] = tile.cpu()
+                    del tile
+
+            if num_completed == num_tiles:
+                break
+
+            # EXACT mode: inject weighted GN for the next round
+            if gn_sync_mode == GNSyncMode.EXACT.value:
+                gn_func = gn_param.summary()
+                if gn_func is not None:
+                    for tq in task_queues:
+                        tq.insert(0, ["apply_norm", gn_func])
+
+    finally:
+        pbar.close()
+
+    return result

sd-webui-ultimate-tiled-vae/tiled_vae/fallback.py

@@ -0,0 +1,127 @@
+"""
+Fallback and NaN-retry logic for Ultimate Tiled VAE.
+
+retry_fp32_on_nan triggers ONLY when a NaN is actually present in the
+output tensor — not on every exception.  Other exceptions propagate normally
+to the caller (hook.py) which then applies the fallback policy.
+
+Fallback policy (FallbackPolicy):
+  Abort          — re-raise; user sees an error message.
+  Cheap Approx   — return cheap_approximation (decoder only, encoder aborts).
+  Original Forward — call the unhooked forward function.
+"""
+
+from __future__ import annotations
+from typing import Optional, TYPE_CHECKING
+
+import torch
+import torch.nn.functional as F
+
+from .config import FallbackPolicy
+from .logging_utils import warn, info
+
+if TYPE_CHECKING:
+    from .config import TiledVAEConfig
+
+
+def has_nans(tensor: torch.Tensor) -> bool:
+    return bool(torch.isnan(tensor).any())
+
+
+def retry_in_fp32(tiled_forward_fn, x: torch.Tensor, net) -> Optional[torch.Tensor]:
+    """
+    Re-run the tiled forward pass with the network temporarily cast to fp32.
+    Returns the result (cast back to original dtype) on success, or None.
+
+    Called ONLY when NaN is detected in the result — not on generic exceptions.
+    """
+    orig_dtype = next(net.parameters()).dtype
+    if orig_dtype == torch.float32:
+        return None   # Already fp32, no point retrying
+
+    info("NaN detected in output — retrying in fp32 …")
+    original_dtype_for_restore = orig_dtype
+    try:
+        net.float()
+        with torch.no_grad():
+            result_fp32 = tiled_forward_fn(x.float())
+        if has_nans(result_fp32):
+            warn("fp32 retry also produced NaN.")
+            return None
+        return result_fp32.to(original_dtype_for_restore)
+    except Exception as e:
+        warn(f"fp32 retry raised an exception: {e}")
+        return None
+    finally:
+        try:
+            net.to(original_dtype_for_restore)
+        except Exception:
+            pass
+
+
+def apply_fallback(
+    cfg: "TiledVAEConfig",
+    x: torch.Tensor,
+    net,
+    original_forward,
+    result_approx: Optional[torch.Tensor],
+    exc: Exception,
+    is_decoder: bool,
+) -> torch.Tensor:
+    """
+    Apply the configured fallback policy after an unrecoverable tiling error.
+    """
+    policy = cfg.fallback_policy
+
+    if policy == FallbackPolicy.ORIGINAL.value:
+        warn("Falling back to original (un-tiled) VAE forward …")
+        try:
+            with torch.no_grad():
+                # Ensure input is on the same device as the network
+                device = next(net.parameters()).device
+                return original_forward(x.to(device))
+        except Exception as e2:
+            warn(f"Original forward also failed: {e2}")
+            raise exc from e2
+
+    elif policy == FallbackPolicy.CHEAP_APPROX.value:
+        if is_decoder and result_approx is not None:
+            warn("Falling back to cheap approximation output …")
+            try:
+                import modules.devices as devices
+                dtype = next(net.parameters()).dtype
+                return result_approx.to(devices.device, dtype=dtype)
+            except Exception as e2:
+                warn(f"Cheap approx fallback failed: {e2}")
+        # Encoder has no approx, or approx failed → escalate
+        raise exc
+
+    else:  # Abort
+        raise exc
+
+
+def get_cheap_approx_result(
+    z: torch.Tensor,
+    is_decoder: bool,
+) -> Optional[torch.Tensor]:
+    """
+    Pre-compute cheap_approximation as a safety net for the decoder.
+    Returns None for encoder (no equivalent approximation).
+    """
+    if not is_decoder:
+        return None
+    try:
+        from modules.sd_vae_approx import cheap_approximation
+        from modules.processing import opt_f
+        with torch.no_grad():
+            return torch.cat(
+                [F.interpolate(
+                    cheap_approximation(xi.unsqueeze(0)),
+                    scale_factor=opt_f,
+                    mode="nearest-exact",
+                )
+                 for xi in z],
+                dim=0,
+            ).cpu()
+    except Exception:
+        return None

sd-webui-ultimate-tiled-vae/tiled_vae/groupnorm.py

@@ -0,0 +1,164 @@
+"""
+GroupNorm parameter management for tiled VAE.
+
+Three sync modes:
+  EXACT  – collect stats from every tile, compute weighted average, apply globally.
+  APPROX – stats come from a downsampled full-image estimate; each tile uses those.
+  NONE   – each tile normalises itself independently (may produce seams).
+"""
+
+from __future__ import annotations
+from typing import Callable, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+from .config import GNSyncMode
+from .compat import get_devices_module
+
+
+# ── Low-level helpers ─────────────────────────────────────────────────────────
+
+def get_var_mean(
+    x: torch.Tensor,
+    num_groups: int = 32,
+    eps: float = 1e-6,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Compute per-group variance and mean.
+    Falls back to fp32 if fp16 overflows.
+    """
+    b, c = x.size(0), x.size(1)
+    channel_in_group = c // num_groups
+    reshaped = x.contiguous().view(
+        1, b * num_groups, channel_in_group, *x.size()[2:]
+    )
+    var, mean = torch.var_mean(reshaped, dim=[0, 2, 3, 4], unbiased=False)
+    if var.dtype == torch.float16 and var.isinf().any():
+        reshaped_fp32 = reshaped.float()
+        var, mean = torch.var_mean(reshaped_fp32, dim=[0, 2, 3, 4], unbiased=False)
+    return var, mean
+
+
+def custom_group_norm(
+    x: torch.Tensor,
+    num_groups: int,
+    mean: torch.Tensor,
+    var: torch.Tensor,
+    weight: Optional[torch.Tensor] = None,
+    bias: Optional[torch.Tensor] = None,
+    eps: float = 1e-6,
+) -> torch.Tensor:
+    """Apply group norm using pre-computed mean/var."""
+    b, c = x.size(0), x.size(1)
+    channel_in_group = c // num_groups
+    reshaped = x.contiguous().view(
+        1, b * num_groups, channel_in_group, *x.size()[2:]
+    )
+    out = F.batch_norm(
+        reshaped, mean.to(x), var.to(x),
+        weight=None, bias=None,
+        training=False, momentum=0, eps=eps,
+    )
+    out = out.view(b, c, *x.size()[2:])
+    if weight is not None:
+        out = out * weight.view(1, -1, 1, 1)
+    if bias is not None:
+        out = out + bias.view(1, -1, 1, 1)
+    return out
+
+
+# ── EXACT mode: GroupNormParam accumulator ────────────────────────────────────
+
+class GroupNormParam:
+    """
+    Accumulates per-tile GroupNorm statistics and computes a pixel-weighted
+    global mean/var for exact cross-tile normalisation.
+    """
+
+    def __init__(self) -> None:
+        self.var_list:   List[torch.Tensor] = []
+        self.mean_list:  List[torch.Tensor] = []
+        self.pixel_list: List[int]          = []
+        self.weight: Optional[torch.Tensor] = None
+        self.bias:   Optional[torch.Tensor] = None
+
+    def add_tile(self, tile: torch.Tensor, layer) -> None:
+        var, mean = get_var_mean(tile, 32)
+        # fp16 overflow guard
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+            if var.device.type == "mps":
+                var  = torch.clamp(var, 0, 60000).half()
+                mean = mean.half()
+        self.var_list.append(var)
+        self.mean_list.append(mean)
+        self.pixel_list.append(tile.shape[2] * tile.shape[3])
+        if hasattr(layer, "weight"):
+            self.weight = layer.weight
+            self.bias   = layer.bias
+        else:
+            self.weight = None
+            self.bias   = None
+
+    def summary(self) -> Optional[Callable]:
+        """Return a GN function using the weighted average stats, or None."""
+        if not self.var_list:
+            return None
+        devices = get_devices_module()
+        var  = torch.vstack(self.var_list)
+        mean = torch.vstack(self.mean_list)
+        max_px = max(self.pixel_list)
+        pix = torch.tensor(
+            self.pixel_list, dtype=torch.float32, device=devices.device
+        ) / max_px
+        pix = pix.unsqueeze(1) / pix.sum()
+        var  = (var  * pix).sum(dim=0)
+        mean = (mean * pix).sum(dim=0)
+        w, b = self.weight, self.bias
+        return lambda x: custom_group_norm(x, 32, mean, var, w, b)
+
+    @staticmethod
+    def from_tile(tile: torch.Tensor, norm_layer) -> Callable:
+        """Create a single-tile GN function without accumulation."""
+        var, mean = get_var_mean(tile, 32)
+        if var.dtype == torch.float16 and var.isinf().any():
+            fp32_tile = tile.float()
+            var, mean = get_var_mean(fp32_tile, 32)
+            if var.device.type == "mps":
+                var  = torch.clamp(var, 0, 60000).half()
+                mean = mean.half()
+        weight = getattr(norm_layer, "weight", None)
+        bias   = getattr(norm_layer, "bias",   None)
+
+        def _fn(x, _mean=mean, _var=var, _w=weight, _b=bias):
+            return custom_group_norm(x, 32, _mean, _var, _w, _b)
+        return _fn
+
+
+# ── APPROX mode: pre-computed plan applied round-robin ───────────────────────
+
+class ApproxGNPlan:
+    """
+    Stores (var, mean) pairs collected on the downsampled image.
+    Each tile pops them in order during execution.
+    """
+
+    def __init__(self) -> None:
+        self._plan:  List[Tuple[torch.Tensor, torch.Tensor]] = []
+        self._index: int = 0
+
+    def record(self, var: torch.Tensor, mean: torch.Tensor) -> None:
+        self._plan.append((var, mean))
+
+    def next_stats(self) -> Tuple[torch.Tensor, torch.Tensor]:
+        var, mean = self._plan[self._index % len(self._plan)]
+        self._index += 1
+        return var, mean
+
+    def reset_counter(self) -> None:
+        self._index = 0
+
+    def __len__(self) -> int:
+        return len(self._plan)

sd-webui-ultimate-tiled-vae/tiled_vae/hook.py

@@ -0,0 +1,256 @@
+"""
+VAEHook – replaces encoder.forward / decoder.forward with tiled processing.
+
+Key invariants:
+  1. net.original_forward is saved exactly once per attach; never lost.
+  2. Reentrancy guard prevents recursive hook invocation.
+  3. VAE device is always restored in a finally block.
+  4. NaN retry fires ONLY when has_nans(result) is True.
+  5. scale factor comes from get_vae_scale_factor(), not hardcoded 8.
+"""
+
+from __future__ import annotations
+from typing import Optional, TYPE_CHECKING
+
+import torch
+
+from .config import GNSyncMode
+from .tiling import split_tiles, crop_valid_region
+from .taskqueue import build_task_queue, clone_task_queue
+from .executor import execute_tiles, _prepare_fast_mode
+from .fallback import (
+    has_nans, retry_in_fp32, apply_fallback, get_cheap_approx_result
+)
+from .logging_utils import info, warn, timed_block
+from .compat import get_devices_module, get_optimal_device, get_vae_scale_factor
+
+if TYPE_CHECKING:
+    from .config import TiledVAEConfig
+
+
+# ── Interrupt flag ────────────────────────────────────────────────────────────
+
+class _InterruptFlag:
+    @property
+    def interrupted(self) -> bool:
+        try:
+            from modules.shared import state
+            return state.interrupted
+        except Exception:
+            return False
+
+
+_INTERRUPT = _InterruptFlag()
+
+
+# ── VAEHook ───────────────────────────────────────────────────────────────────
+
+class VAEHook:
+    """Callable that replaces encoder.forward / decoder.forward."""
+
+    _active: set = set()   # reentrancy guard (uses id(self))
+
+    def __init__(self, net, cfg: "TiledVAEConfig", is_decoder: bool, scale: int = 8):
+        self.net        = net
+        self.cfg        = cfg
+        self.is_decoder = is_decoder
+        self.scale      = scale   # e.g. 8 for SD1.x/SDXL standard VAE
+
+        self.tile_size  = cfg.decoder_tile_size if is_decoder else cfg.encoder_tile_size
+        self.pad        = cfg.decoder_pad_size  if is_decoder else cfg.encoder_pad_size
+        self.fast_mode  = (cfg.fast_decoder and is_decoder) or \
+                          (cfg.fast_encoder and not is_decoder)
+        self.color_fix  = cfg.encoder_color_fix and not is_decoder
+
+    # ── Entry point ───────────────────────────────────────────────────────────
+
+    def __call__(self, x: torch.Tensor) -> torch.Tensor:
+        if id(self) in VAEHook._active:
+            return self.net.original_forward(x)
+        VAEHook._active.add(id(self))
+
+        orig_device = next(self.net.parameters()).device
+        try:
+            if self.cfg.move_vae_to_gpu:
+                target = get_optimal_device()
+                if orig_device != target:
+                    self.net = self.net.to(target)
+            return self._forward_with_nan_retry(x)
+        finally:
+            VAEHook._active.discard(id(self))
+            if self.cfg.move_vae_to_gpu:
+                try:
+                    self.net = self.net.to(orig_device)
+                except Exception:
+                    pass
+
+    # ── Smart-ignore small images ─────────────────────────────────────────────
+
+    def _is_too_small(self, H: int, W: int) -> bool:
+        if not self.cfg.smart_ignore:
+            return False
+        if max(H, W) <= self.pad * 2 + self.tile_size:
+            info("Image fits in one tile – skipping tiled VAE.")
+            return True
+        return False
+
+    # ── NaN-aware retry wrapper ───────────────────────────────────────────────
+
+    def _forward_with_nan_retry(self, x: torch.Tensor) -> torch.Tensor:
+        cfg = self.cfg
+        label = "Decoder" if self.is_decoder else "Encoder"
+
+        result = None
+        try:
+            with timed_block(label, enabled=cfg.log_timings):
+                result = self._tiled_forward(x)
+        except Exception as exc:
+            # Non-NaN exception: go straight to fallback policy
+            approx = get_cheap_approx_result(x, self.is_decoder)
+            return apply_fallback(cfg, x, self.net, self.net.original_forward,
+                                  approx, exc, self.is_decoder)
+
+        # Result obtained – check for NaN
+        if result is not None and cfg.retry_fp32_on_nan and has_nans(result):
+            retried = retry_in_fp32(self._tiled_forward, x, self.net)
+            if retried is not None:
+                return retried
+            # retry also NaN or failed – apply fallback
+            approx = get_cheap_approx_result(x, self.is_decoder)
+            exc = RuntimeError("NaN in tiled VAE output even after fp32 retry.")
+            return apply_fallback(cfg, x, self.net, self.net.original_forward,
+                                  approx, exc, self.is_decoder)
+
+        return result
+
+    # ── Core tiled forward ────────────────────────────────────────────────────
+
+    @torch.no_grad()
+    def _tiled_forward(self, z: torch.Tensor) -> torch.Tensor:
+        cfg        = self.cfg
+        net        = self.net
+        device     = next(net.parameters()).device
+        dtype      = next(net.parameters()).dtype
+        is_decoder = self.is_decoder
+        scale      = self.scale
+
+        z = z.detach()
+        N, C, H, W = z.shape
+
+        if self._is_too_small(H, W):
+            return net.original_forward(z)
+
+        # ── Scale-alignment guard ─────────────────────────────────────────────
+        # For the encoder: H and W must be divisible by `scale` so that H//scale
+        # produces the correct latent dimensions.  A1111 normally ensures this,
+        # but we guard explicitly to surface any mismatch rather than produce a
+        # silent shape error downstream.
+        if not is_decoder and (H % scale != 0 or W % scale != 0):
+            warn(
+                f"Encoder input {W}x{H} not divisible by scale={scale}. "
+                "Falling back to original forward."
+            )
+            return net.original_forward(z)
+
+        info(
+            f"Input {W}x{H}px | tile={self.tile_size} pad={self.pad} "
+            f"scale={scale} mode={'fast' if self.fast_mode else 'exact'} "
+            f"GN={cfg.gn_sync_mode}"
+        )
+
+        # ── Tile plan ─────────────────────────────────────────────────────────
+        in_bboxes, out_bboxes = split_tiles(
+            H, W,
+            self.tile_size,
+            self.pad,
+            is_decoder,
+            auto_adjust=cfg.auto_adjust_tile,
+            scale=scale,
+        )
+        num_tiles = len(in_bboxes)
+
+        # ── Tiles in CPU RAM ──────────────────────────────────────────────────
+        tiles = [
+            z[:, :, bb[2]:bb[3], bb[0]:bb[1]].cpu()
+            for bb in in_bboxes
+        ]
+
+        # ── Safety net for fallback ───────────────────────────────────────────
+        result_approx = get_cheap_approx_result(z, is_decoder)
+
+        # ── Task queues ───────────────────────────────────────────────────────
+        base_queue = build_task_queue(net, is_decoder)
+
+        if self.fast_mode:
+            base_queue = _prepare_fast_mode(
+                z, self.tile_size, base_queue,
+                self.color_fix, device,
+            )
+
+        task_queues = [clone_task_queue(base_queue) for _ in range(num_tiles)]
+        del z
+
+        # ── Output shape ──────────────────────────────────────────────────────
+        if is_decoder:
+            # Decoder output channels are always 3 (RGB)
+            out_shape = (N, 3, H * scale, W * scale)
+        else:
+            out_ch = net.conv_out.out_channels
+            out_shape = (N, out_ch, H // scale, W // scale)
+
+        # ── Execute ───────────────────────────────────────────────────────────
+        result = execute_tiles(
+            tiles=tiles,
+            task_queues=task_queues,
+            in_bboxes=in_bboxes,
+            out_bboxes=out_bboxes,
+            result_shape=out_shape,
+            device=device,
+            dtype=dtype,
+            is_decoder=is_decoder,
+            gn_sync_mode=cfg.gn_sync_mode,
+            cfg=cfg,
+            interrupted_flag=_INTERRUPT,
+            scale=scale,
+        )
+
+        if result is None:
+            info("Generation interrupted during tiled VAE.")
+            if result_approx is not None:
+                return result_approx.to(device, dtype=dtype)
+            raise RuntimeError("Interrupted; no approximation available.")
+
+        return result.to(dtype)
+
+
+# ── Hook lifecycle ────────────────────────────────────────────────────────────
+
+def attach_hooks(vae, cfg: "TiledVAEConfig") -> None:
+    """Attach VAEHook to encoder and decoder.  Safe to call multiple times."""
+    scale   = get_vae_scale_factor(vae)
+    encoder = vae.encoder
+    decoder = vae.decoder
+
+    # Save original forwards exactly once (survives re-attach)
+    if not hasattr(encoder, "original_forward"):
+        encoder.original_forward = encoder.forward
+    if not hasattr(decoder, "original_forward"):
+        decoder.original_forward = decoder.forward
+
+    encoder.forward = VAEHook(encoder, cfg, is_decoder=False, scale=scale)
+    decoder.forward = VAEHook(decoder, cfg, is_decoder=True,  scale=scale)
+    info(f"Hooks attached (scale={scale}).")
+
+
+def detach_hooks(vae) -> None:
+    """Restore original forwards.  Safe to call even if not hooked."""
+    for part in (vae.encoder, vae.decoder):
+        if isinstance(part.forward, VAEHook):
+            part.forward.net = None   # break ref cycle
+        if hasattr(part, "original_forward"):
+            part.forward = part.original_forward
+    info("Hooks removed.")
+
+
+def is_hooked(vae) -> bool:
+    return isinstance(vae.encoder.forward, VAEHook)

sd-webui-ultimate-tiled-vae/tiled_vae/logging_utils.py

@@ -0,0 +1,70 @@
+"""
+Logging utilities for Ultimate Tiled VAE.
+All print/tqdm calls go through here so they can be silenced cleanly.
+"""
+
+import gc
+import time
+from contextlib import contextmanager
+from typing import Optional
+
+
+PREFIX = "[Ultimate Tiled VAE]"
+
+
+def info(msg: str) -> None:
+    print(f"{PREFIX} {msg}")
+
+
+def warn(msg: str) -> None:
+    print(f"{PREFIX} WARNING: {msg}")
+
+
+def error(msg: str) -> None:
+    print(f"{PREFIX} ERROR: {msg}")
+
+
+def debug(msg: str, enabled: bool = False) -> None:
+    if enabled:
+        print(f"{PREFIX} DEBUG: {msg}")
+
+
+@contextmanager
+def timed_block(label: str, enabled: bool = True):
+    """Context manager that prints elapsed time + peak VRAM when done."""
+    if not enabled:
+        yield
+        return
+
+    try:
+        import torch
+        import modules.devices as devices
+        has_cuda = torch.cuda.is_available()
+        if has_cuda:
+            torch.cuda.reset_peak_memory_stats(devices.device)
+        gc.collect()
+    except Exception:
+        has_cuda = False
+
+    t0 = time.time()
+    yield
+
+    elapsed = time.time() - t0
+    try:
+        if has_cuda:
+            import torch
+            import modules.devices as devices
+            vram_mb = torch.cuda.max_memory_allocated(devices.device) / 2 ** 20
+            info(f"{label}: {elapsed:.3f}s, peak VRAM {vram_mb:.1f} MB")
+        else:
+            info(f"{label}: {elapsed:.3f}s")
+    except Exception:
+        info(f"{label}: {elapsed:.3f}s")
+
+
+def tile_summary(num_h: int, num_w: int, real_h: int, real_w: int, tile_size: int) -> None:
+    total = num_h * num_w
+    info(
+        f"Splitting into {num_h}×{num_w} = {total} tiles. "
+        f"Effective tile {real_w}×{real_h}px (requested {tile_size}×{tile_size}px)"
+    )

sd-webui-ultimate-tiled-vae/tiled_vae/taskqueue.py

@@ -0,0 +1,105 @@
+"""
+Task-queue builder for Ultimate Tiled VAE.
+
+Each task is a 2-element list [name, callable_or_data] so that
+executor.py can mutate the second element in-place (e.g. stash residuals).
+"""
+
+from __future__ import annotations
+from typing import Any, Callable, List
+
+import torch
+import torch.nn.functional as F
+
+from .compat import get_attn_func
+
+
+# ── Inplace non-linearity ─────────────────────────────────────────────────────
+
+def _silu(x: torch.Tensor) -> torch.Tensor:
+    return F.silu(x, inplace=True)
+
+
+# ── Block → tasks ─────────────────────────────────────────────────────────────
+
+def _attn2task(queue: List, block: Any) -> None:
+    attn_fn = get_attn_func(block)
+    queue.append(["store_res", lambda x: x])
+    queue.append(["pre_norm",  block.norm])
+    queue.append(["attn",      lambda x, _b=block, _fn=attn_fn: _fn(_b, x)])
+    queue.append(["add_res",   None])
+
+
+def _resblock2task(queue: List, block: Any) -> None:
+    # Shortcut
+    if block.in_channels != block.out_channels:
+        if block.use_conv_shortcut:
+            queue.append(["store_res", block.conv_shortcut])
+        else:
+            queue.append(["store_res", block.nin_shortcut])
+    else:
+        queue.append(["store_res", lambda x: x])
+
+    queue.append(["pre_norm", block.norm1])
+    queue.append(["silu",     _silu])
+    queue.append(["conv1",    block.conv1])
+    queue.append(["pre_norm", block.norm2])
+    queue.append(["silu",     _silu])
+    queue.append(["conv2",    block.conv2])
+    queue.append(["add_res",  None])
+
+
+def _build_sampling(queue: List, net: Any, is_decoder: bool) -> None:
+    """Append the sampling/upsampling ladder tasks to the queue."""
+    if is_decoder:
+        _resblock2task(queue, net.mid.block_1)
+        _attn2task(queue, net.mid.attn_1)
+        _resblock2task(queue, net.mid.block_2)
+        res_iter   = reversed(range(net.num_resolutions))
+        block_ids  = net.num_res_blocks + 1
+        edge_level = 0
+        module     = net.up
+        sample_key = "upsample"
+    else:
+        res_iter   = range(net.num_resolutions)
+        block_ids  = net.num_res_blocks
+        edge_level = net.num_resolutions - 1
+        module     = net.down
+        sample_key = "downsample"
+
+    for i_level in res_iter:
+        for i_block in range(block_ids):
+            _resblock2task(queue, module[i_level].block[i_block])
+        if i_level != edge_level:
+            queue.append([sample_key, getattr(module[i_level], sample_key)])
+
+    if not is_decoder:
+        _resblock2task(queue, net.mid.block_1)
+        _attn2task(queue, net.mid.attn_1)
+        _resblock2task(queue, net.mid.block_2)
+
+
+# ── Public API ────────────────────────────────────────────────────────────────
+
+def build_task_queue(net: Any, is_decoder: bool) -> List:
+    """
+    Build a full task queue for one forward pass of `net`.
+    Each item is [task_name, callable_or_tensor_or_None].
+    """
+    queue: List = []
+    queue.append(["conv_in", net.conv_in])
+    _build_sampling(queue, net, is_decoder)
+
+    if not is_decoder or not net.give_pre_end:
+        queue.append(["pre_norm", net.norm_out])
+        queue.append(["silu",     _silu])
+        queue.append(["conv_out", net.conv_out])
+        if is_decoder and net.tanh_out:
+            queue.append(["tanh", torch.tanh])
+
+    return queue
+
+
+def clone_task_queue(queue: List) -> List:
+    """Deep-copy a task queue (so each tile gets its own mutable copy)."""
+    return [[item for item in task] for task in queue]

sd-webui-ultimate-tiled-vae/tiled_vae/tiling.py

@@ -0,0 +1,148 @@
+"""
+Tile planner for Ultimate Tiled VAE.
+
+Key improvements over original multidiffusion:
+  - scale factor passed explicitly (no hardcoded 8)
+  - _even_partition() ensures last tile <= first tile (prevents late OOM)
+  - Long/thin image edge cases handled
+"""
+
+from __future__ import annotations
+import math
+from typing import List, Tuple
+
+from .logging_utils import tile_summary
+
+# (x1, x2, y1, y2) — x1/y1 inclusive, x2/y2 exclusive
+BBox = Tuple[int, int, int, int]
+
+
+def _best_tile_size(lower: int, upper: int) -> int:
+    """Snap `lower` up to nearest power-of-2 multiple within `upper`."""
+    for div in (32, 16, 8, 4, 2):
+        rem = lower % div
+        if rem == 0:
+            return lower
+        candidate = lower - rem + div
+        if candidate <= upper:
+            return candidate
+    return lower
+
+
+def _even_partition(total: int, n_tiles: int, pad: int) -> List[int]:
+    """
+    Partition (total - 2*pad) pixels into n_tiles segments so that
+    no segment is larger than any earlier one.
+    Extra pixels go to the FIRST tiles, never the last.
+    """
+    usable = total - 2 * pad
+    if usable <= 0:
+        return [total] * n_tiles
+    base  = usable // n_tiles
+    extra = usable  % n_tiles
+    return [base + (1 if i < extra else 0) for i in range(n_tiles)]
+
+
+def split_tiles(
+    h: int,
+    w: int,
+    tile_size: int,
+    pad: int,
+    is_decoder: bool,
+    auto_adjust: bool = True,
+    scale: int = 8,
+) -> Tuple[List[BBox], List[BBox]]:
+    """
+    Compute input and output bounding boxes for all tiles.
+
+    Parameters
+    ----------
+    h, w        : height / width of the input tensor
+    tile_size   : target tile edge length (pixels, in input space)
+    pad         : overlap padding around each tile (input space)
+    is_decoder  : True → output = input × scale; False → output = input ÷ scale
+    auto_adjust : shrink tiles evenly so last tile <= first tile
+    scale       : VAE spatial scale factor (default 8; use get_vae_scale_factor())
+
+    Returns
+    -------
+    in_bboxes  : padded input regions
+    out_bboxes : valid output regions (in output space)
+    """
+    usable_h = max(h - 2 * pad, 1)
+    usable_w = max(w - 2 * pad, 1)
+    num_h = max(math.ceil(usable_h / tile_size), 1)
+    num_w = max(math.ceil(usable_w / tile_size), 1)
+
+    if auto_adjust:
+        raw_th = math.ceil(usable_h / num_h)
+        raw_tw = math.ceil(usable_w / num_w)
+        real_th = _best_tile_size(raw_th, tile_size)
+        real_tw = _best_tile_size(raw_tw, tile_size)
+    else:
+        real_th = tile_size
+        real_tw = tile_size
+
+    tile_summary(num_h, num_w, real_th, real_tw, tile_size)
+
+    # Even partition: last tile never larger than first
+    seg_h = _even_partition(h, num_h, pad)
+    seg_w = _even_partition(w, num_w, pad)
+
+    starts_h = [pad + sum(seg_h[:i]) for i in range(num_h)]
+    starts_w = [pad + sum(seg_w[:i]) for i in range(num_w)]
+
+    in_bboxes:  List[BBox] = []
+    out_bboxes: List[BBox] = []
+
+    for i, (y0, th) in enumerate(zip(starts_h, seg_h)):
+        for j, (x0, tw) in enumerate(zip(starts_w, seg_w)):
+            y1 = min(y0 + th, h)
+            x1 = min(x0 + tw, w)
+
+            # Output region: extend to image boundary at edges
+            oy0 = 0 if y0 <= pad     else y0
+            oy1 = h if y1 >= h - pad else y1
+            ox0 = 0 if x0 <= pad     else x0
+            ox1 = w if x1 >= w - pad else x1
+
+            if is_decoder:
+                out_bboxes.append((ox0 * scale, ox1 * scale, oy0 * scale, oy1 * scale))
+            else:
+                out_bboxes.append((ox0 // scale, ox1 // scale, oy0 // scale, oy1 // scale))
+
+            in_bboxes.append((
+                max(0, x0 - pad), min(w, x1 + pad),
+                max(0, y0 - pad), min(h, y1 + pad),
+            ))
+
+    return in_bboxes, out_bboxes
+
+
+def crop_valid_region(
+    tile: "torch.Tensor",
+    in_bbox: BBox,
+    out_bbox: BBox,
+    is_decoder: bool,
+    scale: int = 8,
+) -> "torch.Tensor":
+    """
+    Trim padded tile output to the valid (output) region.
+    in_bbox / out_bbox are in the coordinate spaces produced by split_tiles().
+    """
+    if is_decoder:
+        padded = (in_bbox[0] * scale, in_bbox[1] * scale,
+                  in_bbox[2] * scale, in_bbox[3] * scale)
+    else:
+        padded = (in_bbox[0] // scale, in_bbox[1] // scale,
+                  in_bbox[2] // scale, in_bbox[3] // scale)
+
+    mx0 = out_bbox[0] - padded[0]   # left trim
+    mx1 = out_bbox[1] - padded[1]   # right  (≤0: from right)
+    my0 = out_bbox[2] - padded[2]   # top trim
+    my1 = out_bbox[3] - padded[3]   # bottom (≤0: from bottom)
+
+    H, W = tile.shape[2], tile.shape[3]
+    y_end = H + my1 if my1 < 0 else H
+    x_end = W + mx1 if mx1 < 0 else W
+    return tile[:, :, my0:y_end, mx0:x_end]