hypertile autodetect optimal value

CHANGELOG.md

@@ -6,6 +6,9 @@
   - allow deployment without git clone  
     for example, zip of the sdnext folder can be used  
   - control: configurable output folder in settings
+  - hypertile: enable vae tiling
+  - hypertile: add autodetect optimial value  
+    set tile size to 0 to use autodetected value  
   - cli: sdapi.py allow manual api invoke  
     example: `python cli/sdapi.py /sdapi/v1/sd-models`  
 - **Fixes**:

modules/control/units/reference.py

@@ -1,3 +1,4 @@
+from typing import Union
 import time
 from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline
 from modules.control.proc.reference_sd15 import StableDiffusionReferencePipeline
@@ -13,7 +14,7 @@ def list_models():
 
 
 class ReferencePipeline():
-    def __init__(self, pipeline: StableDiffusionXLPipeline | StableDiffusionPipeline, dtype = None):
+    def __init__(self, pipeline: Union[StableDiffusionXLPipeline, StableDiffusionPipeline], dtype = None):
         t0 = time.time()
         self.orig_pipeline = pipeline
         self.pipeline = None

modules/sd_hijack_hypertile.py

@@ -11,6 +11,7 @@ import torch
 import torch.nn as nn
 from einops import rearrange
 from modules.shared import log
+from functools import cache
 
 
 # global variables to keep track of changing image size in multiple passes
@@ -20,6 +21,45 @@ max_h = 0
 max_w = 0
 error_reported = False
 reset_needed = False
+RNG_INSTANCE = random.Random()
+
+
+def set_seed(seed: int) -> None:
+    RNG_INSTANCE.seed(seed)
+
+
+def iterative_closest_divisors(hw:int, aspect_ratio:float) -> tuple[int, int]:
+    """
+    Finds h and w such that h*w = hw and h/w = aspect_ratio
+    We check all possible divisors of hw and return the closest to the aspect ratio
+    """
+    divisors = [i for i in range(2, hw + 1) if hw % i == 0] # all divisors of hw
+    pairs = [(i, hw // i) for i in divisors] # all pairs of divisors of hw
+    ratios = [w/h for h, w in pairs] # all ratios of pairs of divisors of hw
+    closest_ratio = min(ratios, key=lambda x: abs(x - aspect_ratio)) # closest ratio to aspect_ratio
+    closest_pair = pairs[ratios.index(closest_ratio)] # closest pair of divisors to aspect_ratio
+    return closest_pair
+
+@cache
+def find_hw_candidates(hw:int, aspect_ratio:float) -> tuple[int, int]:
+    """
+    Finds h and w such that h*w = hw and h/w = aspect_ratio
+    """
+    h, w = round(math.sqrt(hw * aspect_ratio)), round(math.sqrt(hw / aspect_ratio))
+    # find h and w such that h*w = hw and h/w = aspect_ratio
+    if h * w != hw:
+        w_candidate = hw / h
+        # check if w is an integer
+        if not w_candidate.is_integer():
+            h_candidate = hw / w
+            # check if h is an integer
+            if not h_candidate.is_integer():
+                return iterative_closest_divisors(hw, aspect_ratio)
+            else:
+                h = int(h_candidate)
+        else:
+            w = int(w_candidate)
+    return h, w
 
 
 def possible_tile_sizes(dimension: int, tile_size: int, min_tile_size: int, tile_options: int) -> list[int]:
@@ -41,7 +81,7 @@ def parse_list(x: list[int], /) -> str:
 
 
 @contextmanager
-def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256, swap_size: int=1, depth: int=0):
+def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=128, swap_size: int=1, depth: int=0):
     # hijacks AttnBlock from ldm and attention from diffusers
     global reset_needed # pylint: disable=global-statement
     ar = height / width # Aspect ratio
@@ -65,8 +105,10 @@ def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256
             global height, width, max_h, max_w, reset_needed, error_reported # pylint: disable=global-statement
             x = args[0]
             try:
-                nh = nhs[random.randint(0, len(nhs) - 1)]
-                nw = nws[random.randint(0, len(nws) - 1)]
+                nh = RNG_INSTANCE.randint(0, len(nhs) - 1)
+                nw = RNG_INSTANCE.randint(0, len(nws) - 1)
+                # nh = nhs[random.randint(0, len(nhs) - 1)]
+                # nw = nws[random.randint(0, len(nws) - 1)]
             except Exception as e:
                 if not error_reported:
                     error_reported = True
@@ -83,6 +125,7 @@ def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256
             else: # Unet
                 hw = x.size(1)
                 h, w = round(math.sqrt(ar * hw)), round(math.sqrt(hw / ar))
+                h, w = find_hw_candidates(hw, ar)
                 # dynamic height/width based on fact that first two forward calls contain actual height/width
                 # and reset if latest hw is larger since we're never downscaling in 2nd pass
                 if reset_needed:
@@ -155,9 +198,11 @@ def context_hypertile_vae(p):
         # shared.log.warning('Hypertile VAE is enabled but no VAE model was found')
         return nullcontext()
     else:
-        shared.log.info(f'Applying hypertile: vae={shared.opts.hypertile_vae_tile}')
-        p.extra_generation_params['Hypertile VAE'] = shared.opts.hypertile_vae_tile
-        return split_attention(vae, tile_size=shared.opts.hypertile_vae_tile, min_tile_size=128, swap_size=1)
+        tile_size = shared.opts.hypertile_vae_tile if shared.opts.hypertile_vae_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128))
+        shared.log.info(f'Applying hypertile: vae={tile_size}')
+        p.extra_generation_params['Hypertile VAE'] = tile_size
+        return split_attention(vae, tile_size=tile_size, min_tile_size=128, swap_size=1)
+
 
 
 def context_hypertile_unet(p):
@@ -179,9 +224,10 @@ def context_hypertile_unet(p):
         # shared.log.warning('Hypertile UNet is enabled but no Unet model was found')
         return nullcontext()
     else:
-        shared.log.info(f'Applying hypertile: unet={shared.opts.hypertile_unet_tile}')
-        p.extra_generation_params['Hypertile UNet'] = shared.opts.hypertile_unet_tile
-        return split_attention(unet, tile_size=shared.opts.hypertile_unet_tile, min_tile_size=128, swap_size=1)
+        tile_size = shared.opts.hypertile_unet_tile if shared.opts.hypertile_unet_tile > 0 else max(256, 64 * min(p.width // 128, p.height // 128))
+        shared.log.info(f'Applying hypertile: unet={tile_size}')
+        p.extra_generation_params['Hypertile UNet'] = tile_size
+        return split_attention(unet, tile_size=tile_size, min_tile_size=128, swap_size=1)
 
 
 def hypertile_set(p, hr=False):
@@ -198,4 +244,5 @@ def hypertile_set(p, hr=False):
     else:
         width=p.width
         height=p.height
+    set_seed(p.all_seeds[0])
     reset_needed = True

modules/shared.py

@@ -357,10 +357,10 @@ options_templates.update(options_section(('advanced', "Inference Settings"), {
     "freeu_s2": OptionInfo(0.2, "2nd stage skip factor", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
 
     "hypertile_sep": OptionInfo("<h2>HyperTile</h2>", "", gr.HTML),
-    "hypertile_vae_enabled": OptionInfo(False, "HyperTile for VAE enabled", gr.Checkbox, {"visible": False}),
-    "hypertile_vae_tile": OptionInfo(128, "HyperTile for VAE tile size", gr.Slider, {"minimum": 128, "maximum": 512, "step": 8, "visible": False}),
     "hypertile_unet_enabled": OptionInfo(False, "HyperTile for UNet enabled"),
-    "hypertile_unet_tile": OptionInfo(256, "HyperTile for UNet tile size", gr.Slider, {"minimum": 256, "maximum": 1024, "step": 8}),
+    "hypertile_unet_tile": OptionInfo(256, "HyperTile for UNet tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}),
+    "hypertile_vae_enabled": OptionInfo(False, "HyperTile for VAE enabled", gr.Checkbox),
+    "hypertile_vae_tile": OptionInfo(128, "HyperTile for VAE tile size", gr.Slider, {"minimum": 0, "maximum": 1024, "step": 8}),
 
     "inference_other_sep": OptionInfo("<h2>Other</h2>", "", gr.HTML),
     "batch_frame_mode": OptionInfo(False, "Process multiple images in batch in parallel"),