346 lines
14 KiB
Python
346 lines
14 KiB
Python
"""
|
|
Hypertile module for splitting attention layers in SD-1.5 U-Net and SD-1.5 VAE
|
|
Warn: The patch works well only if the input image has a width and height that are multiples of 128
|
|
Original author: @tfernd Github: https://github.com/tfernd/HyperTile
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import functools
|
|
from dataclasses import dataclass
|
|
from typing import Callable
|
|
|
|
from functools import wraps, cache
|
|
|
|
import math
|
|
import torch.nn as nn
|
|
import random
|
|
|
|
from einops import rearrange
|
|
|
|
|
|
@dataclass
|
|
class HypertileParams:
|
|
depth = 0
|
|
layer_name = ""
|
|
tile_size: int = 0
|
|
swap_size: int = 0
|
|
aspect_ratio: float = 1.0
|
|
forward = None
|
|
enabled = False
|
|
|
|
|
|
|
|
# TODO add SD-XL layers
|
|
DEPTH_LAYERS = {
|
|
0: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1",
|
|
"down_blocks.0.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.0.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.2.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"input_blocks.1.1.transformer_blocks.0.attn1",
|
|
"input_blocks.2.1.transformer_blocks.0.attn1",
|
|
"output_blocks.9.1.transformer_blocks.0.attn1",
|
|
"output_blocks.10.1.transformer_blocks.0.attn1",
|
|
"output_blocks.11.1.transformer_blocks.0.attn1",
|
|
# SD 1.5 VAE
|
|
"decoder.mid_block.attentions.0",
|
|
"decoder.mid.attn_1",
|
|
],
|
|
1: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"down_blocks.1.attentions.0.transformer_blocks.0.attn1",
|
|
"down_blocks.1.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.2.attentions.0.transformer_blocks.0.attn1",
|
|
"up_blocks.2.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.2.attentions.2.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"input_blocks.4.1.transformer_blocks.0.attn1",
|
|
"input_blocks.5.1.transformer_blocks.0.attn1",
|
|
"output_blocks.6.1.transformer_blocks.0.attn1",
|
|
"output_blocks.7.1.transformer_blocks.0.attn1",
|
|
"output_blocks.8.1.transformer_blocks.0.attn1",
|
|
],
|
|
2: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"down_blocks.2.attentions.0.transformer_blocks.0.attn1",
|
|
"down_blocks.2.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.1.attentions.0.transformer_blocks.0.attn1",
|
|
"up_blocks.1.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.1.attentions.2.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"input_blocks.7.1.transformer_blocks.0.attn1",
|
|
"input_blocks.8.1.transformer_blocks.0.attn1",
|
|
"output_blocks.3.1.transformer_blocks.0.attn1",
|
|
"output_blocks.4.1.transformer_blocks.0.attn1",
|
|
"output_blocks.5.1.transformer_blocks.0.attn1",
|
|
],
|
|
3: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"mid_block.attentions.0.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"middle_block.1.transformer_blocks.0.attn1",
|
|
],
|
|
}
|
|
# XL layers, thanks for GitHub@gel-crabs for the help
|
|
DEPTH_LAYERS_XL = {
|
|
0: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"down_blocks.0.attentions.0.transformer_blocks.0.attn1",
|
|
"down_blocks.0.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.0.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.1.transformer_blocks.0.attn1",
|
|
"up_blocks.3.attentions.2.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"input_blocks.4.1.transformer_blocks.0.attn1",
|
|
"input_blocks.5.1.transformer_blocks.0.attn1",
|
|
"output_blocks.3.1.transformer_blocks.0.attn1",
|
|
"output_blocks.4.1.transformer_blocks.0.attn1",
|
|
"output_blocks.5.1.transformer_blocks.0.attn1",
|
|
# SD 1.5 VAE
|
|
"decoder.mid_block.attentions.0",
|
|
"decoder.mid.attn_1",
|
|
],
|
|
1: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
#"down_blocks.1.attentions.0.transformer_blocks.0.attn1",
|
|
#"down_blocks.1.attentions.1.transformer_blocks.0.attn1",
|
|
#"up_blocks.2.attentions.0.transformer_blocks.0.attn1",
|
|
#"up_blocks.2.attentions.1.transformer_blocks.0.attn1",
|
|
#"up_blocks.2.attentions.2.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"input_blocks.4.1.transformer_blocks.1.attn1",
|
|
"input_blocks.5.1.transformer_blocks.1.attn1",
|
|
"output_blocks.3.1.transformer_blocks.1.attn1",
|
|
"output_blocks.4.1.transformer_blocks.1.attn1",
|
|
"output_blocks.5.1.transformer_blocks.1.attn1",
|
|
"input_blocks.7.1.transformer_blocks.0.attn1",
|
|
"input_blocks.8.1.transformer_blocks.0.attn1",
|
|
"output_blocks.0.1.transformer_blocks.0.attn1",
|
|
"output_blocks.1.1.transformer_blocks.0.attn1",
|
|
"output_blocks.2.1.transformer_blocks.0.attn1",
|
|
"input_blocks.7.1.transformer_blocks.1.attn1",
|
|
"input_blocks.8.1.transformer_blocks.1.attn1",
|
|
"output_blocks.0.1.transformer_blocks.1.attn1",
|
|
"output_blocks.1.1.transformer_blocks.1.attn1",
|
|
"output_blocks.2.1.transformer_blocks.1.attn1",
|
|
"input_blocks.7.1.transformer_blocks.2.attn1",
|
|
"input_blocks.8.1.transformer_blocks.2.attn1",
|
|
"output_blocks.0.1.transformer_blocks.2.attn1",
|
|
"output_blocks.1.1.transformer_blocks.2.attn1",
|
|
"output_blocks.2.1.transformer_blocks.2.attn1",
|
|
"input_blocks.7.1.transformer_blocks.3.attn1",
|
|
"input_blocks.8.1.transformer_blocks.3.attn1",
|
|
"output_blocks.0.1.transformer_blocks.3.attn1",
|
|
"output_blocks.1.1.transformer_blocks.3.attn1",
|
|
"output_blocks.2.1.transformer_blocks.3.attn1",
|
|
"input_blocks.7.1.transformer_blocks.4.attn1",
|
|
"input_blocks.8.1.transformer_blocks.4.attn1",
|
|
"output_blocks.0.1.transformer_blocks.4.attn1",
|
|
"output_blocks.1.1.transformer_blocks.4.attn1",
|
|
"output_blocks.2.1.transformer_blocks.4.attn1",
|
|
"input_blocks.7.1.transformer_blocks.5.attn1",
|
|
"input_blocks.8.1.transformer_blocks.5.attn1",
|
|
"output_blocks.0.1.transformer_blocks.5.attn1",
|
|
"output_blocks.1.1.transformer_blocks.5.attn1",
|
|
"output_blocks.2.1.transformer_blocks.5.attn1",
|
|
"input_blocks.7.1.transformer_blocks.6.attn1",
|
|
"input_blocks.8.1.transformer_blocks.6.attn1",
|
|
"output_blocks.0.1.transformer_blocks.6.attn1",
|
|
"output_blocks.1.1.transformer_blocks.6.attn1",
|
|
"output_blocks.2.1.transformer_blocks.6.attn1",
|
|
"input_blocks.7.1.transformer_blocks.7.attn1",
|
|
"input_blocks.8.1.transformer_blocks.7.attn1",
|
|
"output_blocks.0.1.transformer_blocks.7.attn1",
|
|
"output_blocks.1.1.transformer_blocks.7.attn1",
|
|
"output_blocks.2.1.transformer_blocks.7.attn1",
|
|
"input_blocks.7.1.transformer_blocks.8.attn1",
|
|
"input_blocks.8.1.transformer_blocks.8.attn1",
|
|
"output_blocks.0.1.transformer_blocks.8.attn1",
|
|
"output_blocks.1.1.transformer_blocks.8.attn1",
|
|
"output_blocks.2.1.transformer_blocks.8.attn1",
|
|
"input_blocks.7.1.transformer_blocks.9.attn1",
|
|
"input_blocks.8.1.transformer_blocks.9.attn1",
|
|
"output_blocks.0.1.transformer_blocks.9.attn1",
|
|
"output_blocks.1.1.transformer_blocks.9.attn1",
|
|
"output_blocks.2.1.transformer_blocks.9.attn1",
|
|
],
|
|
2: [
|
|
# SD 1.5 U-Net (diffusers)
|
|
"mid_block.attentions.0.transformer_blocks.0.attn1",
|
|
# SD 1.5 U-Net (ldm)
|
|
"middle_block.1.transformer_blocks.0.attn1",
|
|
"middle_block.1.transformer_blocks.1.attn1",
|
|
"middle_block.1.transformer_blocks.2.attn1",
|
|
"middle_block.1.transformer_blocks.3.attn1",
|
|
"middle_block.1.transformer_blocks.4.attn1",
|
|
"middle_block.1.transformer_blocks.5.attn1",
|
|
"middle_block.1.transformer_blocks.6.attn1",
|
|
"middle_block.1.transformer_blocks.7.attn1",
|
|
"middle_block.1.transformer_blocks.8.attn1",
|
|
"middle_block.1.transformer_blocks.9.attn1",
|
|
],
|
|
3 : [] # TODO - separate layers for SD-XL
|
|
}
|
|
|
|
|
|
RNG_INSTANCE = random.Random()
|
|
|
|
|
|
def random_divisor(value: int, min_value: int, /, max_options: int = 1) -> int:
|
|
"""
|
|
Returns a random divisor of value that
|
|
x * min_value <= value
|
|
if max_options is 1, the behavior is deterministic
|
|
"""
|
|
min_value = min(min_value, value)
|
|
|
|
# All big divisors of value (inclusive)
|
|
divisors = [i for i in range(min_value, value + 1) if value % i == 0] # divisors in small -> big order
|
|
|
|
ns = [value // i for i in divisors[:max_options]] # has at least 1 element # big -> small order
|
|
|
|
idx = RNG_INSTANCE.randint(0, len(ns) - 1)
|
|
|
|
return ns[idx]
|
|
|
|
|
|
def set_hypertile_seed(seed: int) -> None:
|
|
RNG_INSTANCE.seed(seed)
|
|
|
|
|
|
@functools.cache
|
|
def largest_tile_size_available(width: int, height: int) -> int:
|
|
"""
|
|
Calculates the largest tile size available for a given width and height
|
|
Tile size is always a power of 2
|
|
"""
|
|
gcd = math.gcd(width, height)
|
|
largest_tile_size_available = 1
|
|
while gcd % (largest_tile_size_available * 2) == 0:
|
|
largest_tile_size_available *= 2
|
|
return largest_tile_size_available
|
|
|
|
|
|
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 self_attn_forward(params: HypertileParams, scale_depth=True) -> Callable:
|
|
|
|
@wraps(params.forward)
|
|
def wrapper(*args, **kwargs):
|
|
if not params.enabled:
|
|
return params.forward(*args, **kwargs)
|
|
|
|
latent_tile_size = max(128, params.tile_size) // 8
|
|
x = args[0]
|
|
|
|
# VAE
|
|
if x.ndim == 4:
|
|
b, c, h, w = x.shape
|
|
|
|
nh = random_divisor(h, latent_tile_size, params.swap_size)
|
|
nw = random_divisor(w, latent_tile_size, params.swap_size)
|
|
|
|
if nh * nw > 1:
|
|
x = rearrange(x, "b c (nh h) (nw w) -> (b nh nw) c h w", nh=nh, nw=nw) # split into nh * nw tiles
|
|
|
|
out = params.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)
|
|
|
|
# U-Net
|
|
else:
|
|
hw: int = x.size(1)
|
|
h, w = find_hw_candidates(hw, params.aspect_ratio)
|
|
assert h * w == hw, f"Invalid aspect ratio {params.aspect_ratio} for input of shape {x.shape}, hw={hw}, h={h}, w={w}"
|
|
|
|
factor = 2 ** params.depth if scale_depth else 1
|
|
nh = random_divisor(h, latent_tile_size * factor, params.swap_size)
|
|
nw = random_divisor(w, latent_tile_size * factor, params.swap_size)
|
|
|
|
if nh * nw > 1:
|
|
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 = params.forward(x, *args[1:], **kwargs)
|
|
|
|
if nh * nw > 1:
|
|
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
|
|
|
|
|
|
def hypertile_hook_model(model: nn.Module, width, height, *, enable=False, tile_size_max=128, swap_size=1, max_depth=3, is_sdxl=False):
|
|
hypertile_layers = getattr(model, "__webui_hypertile_layers", None)
|
|
if hypertile_layers is None:
|
|
if not enable:
|
|
return
|
|
|
|
hypertile_layers = {}
|
|
layers = DEPTH_LAYERS_XL if is_sdxl else DEPTH_LAYERS
|
|
|
|
for depth in range(4):
|
|
for layer_name, module in model.named_modules():
|
|
if any(layer_name.endswith(try_name) for try_name in layers[depth]):
|
|
params = HypertileParams()
|
|
module.__webui_hypertile_params = params
|
|
params.forward = module.forward
|
|
params.depth = depth
|
|
params.layer_name = layer_name
|
|
module.forward = self_attn_forward(params)
|
|
|
|
hypertile_layers[layer_name] = 1
|
|
|
|
model.__webui_hypertile_layers = hypertile_layers
|
|
|
|
aspect_ratio = width / height
|
|
tile_size = min(largest_tile_size_available(width, height), tile_size_max)
|
|
|
|
for layer_name, module in model.named_modules():
|
|
if layer_name in hypertile_layers:
|
|
params = module.__webui_hypertile_params
|
|
|
|
params.tile_size = tile_size
|
|
params.swap_size = swap_size
|
|
params.aspect_ratio = aspect_ratio
|
|
params.enabled = enable and params.depth <= max_depth
|