add hypertile

CHANGELOG.md

@@ -145,8 +145,17 @@ Upgrades are still possible and supported, but above is recommended for best exp
 - **Inference**  
   - new section in **settings**  
     - [Token Merging](https://github.com/dbolya/tomesd): not new, but updated  
+      available for *diffusers* and *original* backends  
+      speed-up your generations by merging redundant tokens  
+      speed up will depend on how aggressive you want to be with token merging  
     - [Free-U](https://github.com/ChenyangSi/FreeU): new!  
+      available for *diffusers* and *original* backends  
+      improve generations quality at no cost (other than finding params that work for you)  
+      thanks @ljleb
     - [HyperTile](https://github.com/tfernd/HyperTile): new!  
+      available for *diffusers* and *original* backends  
+      2x speed-up your generations for free :)
+      thanks @tfernd
 - **General**  
   - **Startup**  
     - all main CLI parameters can now be set as environment variable as well  

launch.py

@@ -189,7 +189,7 @@ if __name__ == "__main__":
         installer.log.info('Skipping all checks')
         installer.quick_allowed = True
     elif installer.check_timestamp():
-        installer.log.info('No changes detected: Quick launch active')
+        installer.log.info('No changes detected: quick launch active')
         installer.install_requirements()
         installer.install_packages()
         installer.check_extensions()

modules/devices.py

@@ -15,7 +15,7 @@ def has_mps() -> bool:
     if sys.platform != "darwin":
         return False
     else:
-        return mac_specific.has_mps
+        return mac_specific.has_mps # pylint: disable=used-before-assignment
 
 
 def get_gpu_info():

modules/processing.py

@@ -35,6 +35,7 @@ import modules.sd_models
 import modules.sd_vae
 import modules.sd_vae_approx
 import modules.generation_parameters_copypaste
+from modules.sd_hijack_hypertile import context_hypertile_vae, context_hypertile_unet
 
 
 opt_C = 4
@@ -79,7 +80,7 @@ def txt2img_image_conditioning(sd_model, x, width, height):
         image_conditioning = torch.zeros(x.shape[0], 3, height, width, device=x.device)
         image_conditioning = sd_model.get_first_stage_encoding(sd_model.encode_first_stage(image_conditioning))
         # Add the fake full 1s mask to the first dimension.
-        image_conditioning = torch.nn.functional.pad(image_conditioning, (0, 0, 0, 0, 1, 0), value=1.0)
+        image_conditioning = torch.nn.functional.pad(image_conditioning, (0, 0, 0, 0, 1, 0), value=1.0) # pylint: disable=not-callable
         image_conditioning = image_conditioning.to(x.dtype)
         return image_conditioning
     elif sd_model.model.conditioning_key == "crossattn-adm": # UnCLIP models
@@ -656,7 +657,8 @@ def process_images(p: StableDiffusionProcessing) -> Processed:
             res = process_images_inner(p)
             print_profile(pr, 'Torch')
         else:
-            res = process_images_inner(p)
+            with context_hypertile_vae(p), context_hypertile_unet(p):
+                res = process_images_inner(p)
     finally:
         if not shared.opts.cuda_compile:
             modules.sd_models.apply_token_merging(p.sd_model, 0)
@@ -815,7 +817,6 @@ def process_images_inner(p: StableDiffusionProcessing) -> Processed:
             elif shared.backend == shared.Backend.DIFFUSERS:
                 from modules.processing_diffusers import process_diffusers
                 x_samples_ddim = process_diffusers(p, p.seeds, p.prompts, p.negative_prompts)
-
             else:
                 raise ValueError(f"Unknown backend {shared.backend}")
 

modules/sd_hijack_freeu.py

@@ -47,7 +47,6 @@ def lerp(a, b, r):
 
 
 def free_u_cat_hijack(hs, *args, original_function, **kwargs):
-    print('HERE', state.sampling_step, state.sampling_steps)
     if not opts.freeu_enabled:
         return original_function(hs, *args, **kwargs)
     schedule_ratio = get_schedule_ratio()

modules/sd_hijack_hypertile.py

@@ -0,0 +1,145 @@
+# credits: @tfernd https://github.com/tfernd/HyperTile
+# based on: https://github.com/tfernd/HyperTile/tree/main/hyper_tile/utils.py + https://github.com/tfernd/HyperTile/tree/main/hyper_tile/hyper_tile.py
+
+from __future__ import annotations
+from typing import Callable
+from functools import wraps
+from contextlib import contextmanager, nullcontext
+import random
+import math
+import torch
+import torch.nn as nn
+from einops import rearrange
+
+
+# global variables to keep track of changing image size in multiple passes
+height = None
+width = None
+max_h = 0
+max_w = 0
+
+
+def possible_tile_sizes(dimension: int, tile_size: int, min_tile_size: int, tile_options: int) -> list[int]:
+    assert tile_options >= 1
+    min_tile_size = min(min_tile_size, tile_size, dimension)
+    # all divisors that are themselves divisible by 8 and give tile-size above min
+    n = torch.arange(1, dimension + 1)
+    n = n[dimension // n // 8 * 8 * n == dimension]
+    n = n[dimension // n >= min_tile_size]
+    pos = (dimension // n).sub(tile_size).abs().argsort()
+    pos = pos[:tile_options]
+    return n[pos].tolist()
+
+
+def parse_list(x: list[int], /) -> str:
+    if len(x) == 0:
+        return str(x[0])
+    return str(x)
+
+
+@contextmanager
+def split_attention(layer: nn.Module, tile_size: int=256, min_tile_size: int=256, swap_size: int=1, depth: int=0):
+    # hijacks AttnBlock from ldm and attention from diffusers
+    ar = height / width # Aspect ratio
+    nhs = possible_tile_sizes(height, tile_size, min_tile_size, swap_size) # possible sub-grids that fit into the image
+    nws = possible_tile_sizes(width, tile_size, min_tile_size, swap_size)
+    # random sub-grid indices # TODO remove randomness. seed?
+    make_ns = lambda: (nhs[random.randint(0, len(nhs) - 1)], nws[random.randint(0, len(nws) - 1)]) # pylint: disable=unnecessary-lambda-assignment
+
+    def reset_nhs():
+        nonlocal nhs
+        nhs = possible_tile_sizes(height, tile_size, min_tile_size, swap_size)
+
+    def reset_nws():
+        nonlocal nws
+        nws = possible_tile_sizes(width, tile_size, min_tile_size, swap_size)
+
+    def self_attn_forward(forward: Callable) -> Callable:
+        @wraps(forward)
+        def wrapper(*args, **kwargs):
+            global height, width, max_h, max_w # pylint: disable=global-statement
+            nh, nw = make_ns()
+            x = args[0]
+            if x.ndim == 4: # VAE
+                # TODO: VAE breaks for diffusers when using non-standard sizes
+                if nh * nw > 1:
+                    x = rearrange(x, "b c (nh h) (nw w) -> (b nh nw) c h w", nh=nh, nw=nw)
+                out = forward(x, *args[1:], **kwargs)
+                if nh * nw > 1:
+                    out = rearrange(out, "(b nh nw) c h w -> b c (nh h) (nw w)", nh=nh, nw=nw)
+            else: # Unet
+                hw = x.size(1)
+                h, w = round(math.sqrt(ar * hw)), round(math.sqrt(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 h > max_h:
+                    height = 8 * h
+                    max_h = max(max_h, h)
+                    reset_nhs()
+                if w > max_w:
+                    width = 8 * w
+                    max_w = max(max_w, w)
+                    reset_nws()
+                down_ratio = height // 8 // h
+                curr_depth = round(math.log(down_ratio, 2))
+                # scale-up the tile-size the deeper we go
+                nh = max(1, nh // down_ratio)
+                nw = max(1, nw // down_ratio)
+                do_split = curr_depth <= depth and h % nh == 0 and w % nw == 0 and nh * nw > 1
+                if do_split:
+                    x = rearrange(x, "b (nh h nw w) c -> (b nh nw) (h w) c", h=h // nh, w=w // nw, nh=nh, nw=nw)
+                out = forward(x, *args[1:], **kwargs)
+                if do_split:
+                    out = rearrange(out, "(b nh nw) hw c -> b nh nw hw c", nh=nh, nw=nw)
+                    out = rearrange(out, "b nh nw (h w) c -> b (nh h nw w) c", h=h // nh, w=w // nw)
+            return out
+        return wrapper
+    try: # hijack forward method and restore
+        for name, module in layer.named_modules():
+            if module.__class__.__qualname__ in ("Attention", "CrossAttention", "AttnBlock"):
+                if name.endswith("attn2") or name.endswith("attn_2"): # skip cross-attention layers
+                    continue
+                setattr(module, "_original_forward", module.forward) # save original forward for recovery later # noqa: B010
+                setattr(module, "forward", self_attn_forward(module.forward)) # noqa: B010
+        yield
+    finally:
+        for _name, module in layer.named_modules():
+            if hasattr(module, "_original_forward"): # remove hijack
+                setattr(module, "forward", module._original_forward) # pylint: disable=protected-access # noqa: B010
+                del module._original_forward
+
+
+def context_hypertile_vae(p):
+    global height, width, max_h, max_w # pylint: disable=global-statement
+    height=p.height
+    width=p.width
+    max_h = 0
+    max_w = 0
+    from modules import shared
+    if p.sd_model is None or not shared.opts.hypertile_vae_enabled:
+        return nullcontext()
+    vae = getattr(p.sd_model, "vae", None) if shared.backend == shared.Backend.DIFFUSERS else getattr(p.sd_model, "first_stage_model", None)
+    if vae is None:
+        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}')
+        return split_attention(vae, tile_size=shared.opts.hypertile_vae_tile, min_tile_size=128, swap_size=1)
+
+
+def context_hypertile_unet(p):
+    global height, width, max_h, max_w # pylint: disable=global-statement
+    height=p.height
+    width=p.width
+    max_h = 0
+    max_w = 0
+    from modules import shared
+    if p.sd_model is None or not shared.opts.hypertile_unet_enabled:
+        return nullcontext()
+    unet = getattr(p.sd_model, "unet", None) if shared.backend == shared.Backend.DIFFUSERS else getattr(p.sd_model.model, "diffusion_model", None)
+    if unet is None:
+        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}')
+        return split_attention(unet, tile_size=shared.opts.hypertile_unet_tile, min_tile_size=128, swap_size=1)

modules/shared.py

@@ -457,6 +457,12 @@ options_templates.update(options_section(('advanced', "Inference Settings"), {
     "freeu_s1": OptionInfo(0.9, "1st stage skip factor", gr.Slider, {"minimum": 0.0, "maximum": 1.0, "step": 0.01}),
     "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, "2nd stage skip factor", gr.Slider, {"minimum": 128, "maximum": 512, "step": 8, "visible": False}),
+    "hypertile_unet_enabled": OptionInfo(False, "HyperTile for UNet enabled"),
+    "hypertile_unet_tile": OptionInfo(256, "2nd stage skip factor", gr.Slider, {"minimum": 256, "maximum": 1024, "step": 8}),
+
     "inference_mode_sep": OptionInfo("<h2>Inference mode</h2>", "", gr.HTML),
     "inference_mode": OptionInfo("no-grad", "Torch inference mode", gr.Radio, lambda: {"choices": ["no-grad", "inference-mode", "none"]}),
     "sd_vae_sliced_encode": OptionInfo(False, "VAE Slicing (original)"),