Merge pull request #288 from testIgnor/main

Add Custom Tiled VAE Sizes to NeverOOM

Author: Panchovix

extensions-builtin/sd_forge_neveroom/scripts/forge_never_oom.py

@@ -1,4 +1,6 @@
 import gradio as gr
+import torch
+import modules.devices as devices
 
 from modules import scripts
 from ldm_patched.modules import model_management
@@ -17,23 +19,54 @@ def title(self):
     def show(self, is_img2img):
         return scripts.AlwaysVisible
 
+    """
+    The following two functions are pulled directly from
+    pkuliyi2015/multidiffusion-upscaler-for-automatic1111
+    """
+    def get_rcmd_enc_tsize(self):
+        if torch.cuda.is_available() and devices.device not in ['cpu', devices.cpu]:
+            total_memory = torch.cuda.get_device_properties(devices.device).total_memory // 2**20
+            if   total_memory > 16*1000: ENCODER_TILE_SIZE = 3072
+            elif total_memory > 12*1000: ENCODER_TILE_SIZE = 2048
+            elif total_memory >  8*1000: ENCODER_TILE_SIZE = 1536
+            else:                        ENCODER_TILE_SIZE = 960
+        else:                            ENCODER_TILE_SIZE = 512
+        return ENCODER_TILE_SIZE
+
+    def get_rcmd_dec_tsize(self):
+        if torch.cuda.is_available() and devices.device not in ['cpu', devices.cpu]:
+            total_memory = torch.cuda.get_device_properties(devices.device).total_memory // 2**20
+            if   total_memory > 30*1000: DECODER_TILE_SIZE = 256
+            elif total_memory > 16*1000: DECODER_TILE_SIZE = 192
+            elif total_memory > 12*1000: DECODER_TILE_SIZE = 128
+            elif total_memory >  8*1000: DECODER_TILE_SIZE = 96
+            else:                        DECODER_TILE_SIZE = 64
+        else:                            DECODER_TILE_SIZE = 64
+        return DECODER_TILE_SIZE
+
     def ui(self, *args, **kwargs):
         with gr.Accordion(open=False, label=self.title()):
             unet_enabled = gr.Checkbox(label='Enabled for UNet (always maximize offload)', value=False)
             vae_enabled = gr.Checkbox(label='Enabled for VAE (always tiled)', value=False)
-        return unet_enabled, vae_enabled
+            encoder_tile_size = gr.Slider(label='Encoder Tile Size', minimum=256, maximum=4096, step=16, value=self.get_rcmd_enc_tsize())
+            decoder_tile_size = gr.Slider(label='Decoder Tile Size', minimum=48,  maximum=512,  step=16, value=self.get_rcmd_dec_tsize())
+        return unet_enabled, vae_enabled, encoder_tile_size, decoder_tile_size
 
     def process(self, p, *script_args, **kwargs):
-        unet_enabled, vae_enabled = script_args
+        unet_enabled, vae_enabled, encoder_tile_size, decoder_tile_size = script_args
 
         if unet_enabled:
             print('NeverOOM Enabled for UNet (always maximize offload)')
 
         if vae_enabled:
             print('NeverOOM Enabled for VAE (always tiled)')
+            print('With tile sizes')
+            print(f'Encode:\t x:{encoder_tile_size}\t y:{encoder_tile_size}')
+            print(f'Decode:\t x:{decoder_tile_size}\t y:{decoder_tile_size}')
 
         model_management.VAE_ALWAYS_TILED = vae_enabled
-
+        model_management.VAE_ENCODE_TILE_SIZE_X = model_management.VAE_ENCODE_TILE_SIZE_Y = encoder_tile_size
+        model_management.VAE_DECODE_TILE_SIZE_X = model_management.VAE_DECODE_TILE_SIZE_Y = decoder_tile_size
         if self.previous_unet_enabled != unet_enabled:
             model_management.unload_all_models()
             if unet_enabled:

ldm_patched/modules/model_management.py

@@ -204,6 +204,11 @@ def is_nvidia():
     VAE_DTYPE = torch.float32
 
 VAE_ALWAYS_TILED = False
+# include VAE tile size parameters
+VAE_ENCODE_TILE_SIZE_X = 512
+VAE_ENCODE_TILE_SIZE_Y = 512
+VAE_DECODE_TILE_SIZE_X = 64
+VAE_DECODE_TILE_SIZE_Y = 64
 
 def set_fp16_accumulation_if_available():
     if args.allow_fp16_accumulation:
@@ -506,7 +511,7 @@ def load_models_gpu(models, memory_required=0):
 
         if vram_set_state == VRAMState.NO_VRAM:
             async_kept_memory = 0
-        
+
         loaded_model.model_load(async_kept_memory)
         current_loaded_models.insert(0, loaded_model)
 
@@ -692,7 +697,7 @@ def cast_to_device(tensor, device, dtype, copy=False):
             return tensor.to(device, non_blocking=non_blocking).to(dtype, non_blocking=non_blocking)
     else:
         return tensor.to(device, dtype, copy=copy, non_blocking=non_blocking)
-    
+
 def sage_attention_enabled():
     return args.use_sage_attention
 
@@ -758,7 +763,7 @@ def get_free_memory(dev=None, torch_free_too=False):
         return (mem_free_total, mem_free_torch)
     else:
         return mem_free_total
-    
+
 def mac_version():
     try:
         return tuple(int(n) for n in platform.mac_ver()[0].split("."))

ldm_patched/modules/sd.py

@@ -58,7 +58,7 @@ def load_clip_weights(model, sd):
 
 def load_lora_for_models(model, clip, lora, strength_model, strength_clip, filename='default'):
     model_flag = type(model.model).__name__ if model is not None else 'default'
-    
+
     # Only build key maps for components we'll actually use
     key_map = {}
     if model is not None and strength_model != 0:
@@ -72,7 +72,7 @@ def load_lora_for_models(model, clip, lora, strength_model, strength_clip, filen
 
     # Load LoRA weights
     loaded = ldm_patched.modules.lora.load_lora(lora, key_map)
-    
+
     # Only clone and patch if we have relevant weights
     if model is not None and strength_model != 0:
         new_modelpatcher = model.clone()
@@ -261,7 +261,11 @@ def encode_tiled_(self, pixel_samples, tile_x=512, tile_y=512, overlap = 64):
 
     def decode_inner(self, samples_in):
         if model_management.VAE_ALWAYS_TILED:
-            return self.decode_tiled(samples_in).to(self.output_device)
+            return self.decode_tiled(
+                                    samples_in,
+                                    tile_x = model_management.VAE_DECODE_TILE_SIZE_X,
+                                    tile_y = model_management.VAE_DECODE_TILE_SIZE_Y
+                                    ).to(self.output_device)
 
         try:
             memory_used = self.memory_used_decode(samples_in.shape, self.vae_dtype)
@@ -295,8 +299,12 @@ def decode_tiled(self, samples, tile_x=64, tile_y=64, overlap = 16):
 
     def encode_inner(self, pixel_samples):
         if model_management.VAE_ALWAYS_TILED:
-            return self.encode_tiled(pixel_samples)
-        
+            return self.encode_tiled(
+                                    pixel_samples, 
+                                    tile_x = model_management.VAE_ENCODE_TILE_SIZE_X,
+                                    tile_y = model_management.VAE_ENCODE_TILE_SIZE_Y
+                                    )
+
         regulation = self.patcher.model_options.get("model_vae_regulation", None)
 
         pixel_samples = pixel_samples.movedim(-1,1)