From 779bbdf9428a25d463fa749c5c83f3b7c0e1e05e Mon Sep 17 00:00:00 2001
From: Oscar Anghammar <oscar.anghammar@gmail.com>
Date: Sat, 3 May 2025 17:31:36 +0200
Subject: [PATCH] Add ForcedTiledVAE subclass of VAE

---
 comfy/sd.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/comfy/sd.py b/comfy/sd.py
index 5b95cf75a..54530e11f 100644
--- a/comfy/sd.py
+++ b/comfy/sd.py
@@ -718,6 +718,26 @@ class VAE:
         except:
             return None
 
+class ForcedTiledVAE(VAE):
+    def __init__(self, tile_size, overlap, temporal_size=64, temporal_overlap=8, sd=None, device=None, config=None, dtype=None, metadata=None):
+        super().__init__(sd, device, config, dtype, metadata)
+        self.tile_size = tile_size
+        self.overlap = overlap
+        self.temporal_size = temporal_size
+        self.temporal_overlap = temporal_overlap
+
+    @classmethod
+    def cast(cls, vae: VAE, tile_size, overlap, temporal_size=64, temporal_overlap=8):
+        vae.__class__ = cls
+        vae.tile_size = tile_size
+        vae.overlap = overlap
+        vae.temporal_size = temporal_size
+        vae.temporal_overlap = temporal_overlap
+        return vae
+
+    def encode(self, pixel_samples):
+        return super().encode_tiled(pixel_samples, self.tile_size, self.tile_size, self.overlap, self.temporal_size, self.temporal_overlap)
+
 class StyleModel:
     def __init__(self, model, device="cpu"):
         self.model = model