diff --git a/CHANGELOG.md b/CHANGELOG.md index 2325c721e..7094553fe 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/launch.py b/launch.py index 52452136b..d5fe9a7d1 100644 --- a/launch.py +++ b/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() diff --git a/modules/devices.py b/modules/devices.py index 22ca4e667..5fe08b73d 100644 --- a/modules/devices.py +++ b/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(): diff --git a/modules/processing.py b/modules/processing.py index d40f1d33e..73f78b3bc 100644 --- a/modules/processing.py +++ b/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}") diff --git a/modules/sd_hijack_freeu.py b/modules/sd_hijack_freeu.py index 51448561a..2fa2dc10e 100644 --- a/modules/sd_hijack_freeu.py +++ b/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() diff --git a/modules/sd_hijack_hypertile.py b/modules/sd_hijack_hypertile.py new file mode 100644 index 000000000..7b6a2c760 --- /dev/null +++ b/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) diff --git a/modules/shared.py b/modules/shared.py index 15960445b..d8136c11b 100644 --- a/modules/shared.py +++ b/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("

HyperTile

", "", 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("

Inference mode

", "", 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)"),