From 862c6619fadb6b2f89b7dfdcdb62069e71c3b3e0 Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Thu, 23 Apr 2026 18:08:23 +0800
Subject: [PATCH] vae

---
 backend/args.py       |  2 ++
 backend/loader.py     |  6 ++---
 backend/operations.py | 53 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/backend/args.py b/backend/args.py
index 94cfaa02..781c33d7 100644
--- a/backend/args.py
+++ b/backend/args.py
@@ -102,6 +102,8 @@ parser.add_argument("--autotune", action="store_true", help="torch.backends.cudn
 parser.add_argument("--mmap-torch-files", action="store_true", help="Use mmap when loading ckpt/pt files")
 parser.add_argument("--disable-mmap", action="store_true", help="Don't use mmap when loading safetensors")
 
+parser.add_argument("--tiled-conv2d", type=int, default=0, metavar="TILE_SIZE", choices=[0, 64, 128, 256, 512], help="reduce VAE memory usage ; increase processing time")
+
 
 class SageAttentionFuncs(enum.Enum):
     auto = "auto"
diff --git a/backend/loader.py b/backend/loader.py
index b616ff3b..cdd4cfc0 100644
--- a/backend/loader.py
+++ b/backend/loader.py
@@ -67,7 +67,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
             config = IntegratedAutoencoderKL.load_config(config_path)
 
             with no_init_weights():
-                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype()):
+                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype(), bnb_dtype="vae"):
                     model = IntegratedAutoencoderKL.from_config(config)
 
             load_state_dict(model, state_dict, ignore_start="loss.")
@@ -79,7 +79,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
             config = AutoencoderKLFlux2.load_config(config_path)
 
             with no_init_weights():
-                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype()):
+                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype(), bnb_dtype="vae"):
                     model = AutoencoderKLFlux2.from_config(config)
 
             load_state_dict(model, state_dict, ignore_start="loss.")
@@ -91,7 +91,7 @@ def load_huggingface_component(guess, component_name, lib_name, cls_name, repo_p
             config = WanVAE.load_config(config_path)
 
             with no_init_weights():
-                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype()):
+                with using_forge_operations(device=memory_management.cpu, dtype=memory_management.vae_dtype(), bnb_dtype="vae"):
                     model = WanVAE.from_config(config)
 
             load_state_dict(model, state_dict)
diff --git a/backend/operations.py b/backend/operations.py
index 3d79f407..77ae44e6 100644
--- a/backend/operations.py
+++ b/backend/operations.py
@@ -773,6 +773,56 @@ class ForgeOperationsFP8(ForgeOperations):
             return super().forward(x)
 
 
+# region Tiled
+
+
+class TiledOperations(ForgeOperations):
+    class Conv2d(ForgeOperations.Conv2d):
+        tile_size: int
+
+        def __init__(self, *arg, **kwargs):
+            super().__init__(*arg, **kwargs)
+            self._3x1x1: bool = self.kernel_size == (3, 3) and self.stride == (1, 1) and self.padding == (1, 1)
+            self.tile_size = args.tiled_conv2d
+
+        @torch.inference_mode()
+        def forward(self, x: torch.Tensor):
+            if not self._3x1x1:
+                return super().forward(x)
+
+            B, C, H, W = x.shape
+
+            if H <= self.tile_size and W <= self.tile_size:
+                return super().forward(x)
+
+            orig_forward = super().forward
+            out_channels = self.out_channels if self.out_channels is not None else C
+
+            out = torch.empty((B, out_channels, H, W), device=x.device, dtype=x.dtype, memory_format=torch.contiguous_format)
+            non_blocking = memory_management.device_supports_non_blocking(x.device)
+
+            for i in range(0, H, self.tile_size):
+                i0 = max(i - 1, 0)
+                i1 = min(i + self.tile_size + 1, H)
+                pi = i - i0
+                ph = min(self.tile_size, H - i)
+
+                for j in range(0, W, self.tile_size):
+                    j0 = max(j - 1, 0)
+                    j1 = min(j + self.tile_size + 1, W)
+
+                    tile = x[:, :, i0:i1, j0:j1]
+                    tile_conv = orig_forward(tile)
+
+                    pj = j - j0
+                    pw = min(self.tile_size, W - j)
+
+                    out[:, :, i : i + ph, j : j + pw].copy_(tile_conv[:, :, pi : pi + ph, pj : pj + pw], non_blocking=non_blocking)
+                    del tile_conv
+
+            return out
+
+
 # region Pick OPs
 
 
@@ -829,6 +879,9 @@ def using_forge_operations(operations=None, device=None, dtype=None, manual_cast
         elif bnb_dtype in ["nf4", "fp4"]:
             assert memory_management.bnb_enabled(), 'Install the "bitsandbytes" package with --bnb'
             operations = ForgeOperationsBNB4bits
+        elif bnb_dtype in ["vae"] and args.tiled_conv2d:
+            memory_management.logger.info(f"Using TiledOperations ({args.tiled_conv2d}) for VAE")
+            operations = TiledOperations
         elif dtype is torch.float8_e4m3fn and args.fast_fp8 and memory_management.supports_fp8_compute(memory_management.get_torch_device()):
             operations = ForgeOperationsFP8
         else: