diff --git a/CHANGELOG.md b/CHANGELOG.md index 5043d8280..051b23fc7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,12 @@ - add support for [deep-cache](https://github.com/horseee/DeepCache) model acceleration it can produce massive speedups (2x-5x) with no overhead, but with some loss of quality *settings -> compute -> model compile -> deep-cache* and *settings -> compute -> model compile -> cache interval* +- **diffusers** + - add *Dynamic Attention Slicing* + dynamically slices attention queries based on query size and slice rate in GB + saves VRAM similar to Sub-Quad on Original backend and it is compatible with HyperTile + slicing will not get triggered if the query size is smaller than the slice rate to gain performance + *settings -> diffusers settings -> dynamic attention slicing* - **other**: - improved `clip-skip` value handling in diffusers, thanks @AI-Casanova & @Disty0 now clip-skip range is 0-12 where previously lowest value was 1 (default is still 1) diff --git a/modules/sd_hijack_optimizations.py b/modules/sd_hijack_optimizations.py index 60a0d3f22..620d4d336 100644 --- a/modules/sd_hijack_optimizations.py +++ b/modules/sd_hijack_optimizations.py @@ -2,6 +2,7 @@ from __future__ import annotations import sys import math import psutil +from functools import cache import torch from torch import einsum @@ -14,6 +15,8 @@ from modules.hypernetworks import hypernetwork from .sub_quadratic_attention import efficient_dot_product_attention # pylint: disable=relative-beyond-top-level +from diffusers.utils import USE_PEFT_BACKEND + if shared.opts.cross_attention_optimization == "xFormers": try: @@ -518,3 +521,172 @@ def sub_quad_attnblock_forward(self, x): out = rearrange(out, 'b (h w) c -> b c h w', h=h) out = self.proj_out(out) return x + out + + +@cache +def find_dynamic_v1_slice_size(slice_size, slice_block_size, slice_rate=4): + while (slice_size * slice_block_size) > slice_rate: + slice_size = slice_size // 2 + if slice_size <= 1: + slice_size = 1 + break + return slice_size + +@cache +def find_dynamic_attention_v1_slice_sizes(query_shape, query_element_size, slice_rate=4): + if len(query_shape) == 3: + batch_size_attention, query_tokens, shape_three = query_shape + shape_four = 1 + else: + batch_size_attention, query_tokens, shape_three, shape_four = query_shape + + slice_block_size = query_tokens * shape_three * shape_four / 1024 / 1024 * query_element_size + block_size = batch_size_attention * slice_block_size + + split_slice_size = batch_size_attention + split_2_slice_size = query_tokens + split_3_slice_size = shape_three + + do_split = False + do_split_2 = False + do_split_3 = False + + if block_size > slice_rate: + do_split = True + split_slice_size = find_dynamic_v1_slice_size(split_slice_size, slice_block_size, slice_rate=slice_rate) + if split_slice_size * slice_block_size > slice_rate: + slice_2_block_size = split_slice_size * shape_three * shape_four / 1024 / 1024 * query_element_size + do_split_2 = True + split_2_slice_size = find_dynamic_v1_slice_size(split_2_slice_size, slice_2_block_size, slice_rate=slice_rate) + if split_2_slice_size * slice_2_block_size > slice_rate: + slice_3_block_size = split_slice_size * split_2_slice_size * shape_four / 1024 / 1024 * query_element_size + do_split_3 = True + split_3_slice_size = find_dynamic_v1_slice_size(split_3_slice_size, slice_3_block_size, slice_rate=slice_rate) + + return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size + +class DynamicAttnProcessorV1: + r""" + dynamically slices attention queries based on query size and slice rate in GB + saves VRAM similar to Sub-Quad on Original backend and it is compatible with HyperTile + slicing will not get triggered if the query size is smaller than the slice rate to gain performance + """ + + def __call__(self, attn, hidden_states: torch.FloatTensor, + encoder_hidden_states=None, attention_mask=None, + temb=None, scale: float = 1.0) -> torch.Tensor: # pylint: disable=too-many-statements, too-many-locals, too-many-branches + + residual = hidden_states + + args = () if USE_PEFT_BACKEND else (scale,) + + if attn.spatial_norm is not None: + hidden_states = attn.spatial_norm(hidden_states, temb) + + input_ndim = hidden_states.ndim + + if input_ndim == 4: + batch_size, channel, height, width = hidden_states.shape + hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2) + + batch_size, sequence_length, _ = ( + hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape + ) + attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size) + + if attn.group_norm is not None: + hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2) + + query = attn.to_q(hidden_states, *args) + + if encoder_hidden_states is None: + encoder_hidden_states = hidden_states + elif attn.norm_cross: + encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states) + + key = attn.to_k(encoder_hidden_states, *args) + value = attn.to_v(encoder_hidden_states, *args) + + query = attn.head_to_batch_dim(query) + key = attn.head_to_batch_dim(key) + value = attn.head_to_batch_dim(value) + + #################################################################### + # Slicing parts: + batch_size_attention, query_tokens, shape_three = query.shape[0], query.shape[1], query.shape[2] + hidden_states = torch.zeros(query.shape, device=query.device, dtype=query.dtype) + do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_dynamic_attention_v1_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate) + + if do_split: + for i in range(batch_size_attention // split_slice_size): + start_idx = i * split_slice_size + end_idx = (i + 1) * split_slice_size + if do_split_2: + for i2 in range(query_tokens // split_2_slice_size): # pylint: disable=invalid-name + start_idx_2 = i2 * split_2_slice_size + end_idx_2 = (i2 + 1) * split_2_slice_size + if do_split_3: + for i3 in range(shape_three // split_3_slice_size): # pylint: disable=invalid-name + start_idx_3 = i3 * split_3_slice_size + end_idx_3 = (i3 + 1) * split_3_slice_size + + query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] + key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] + attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] if attention_mask is not None else None + + attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) + del query_slice + del key_slice + del attn_mask_slice + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3]) + + hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = attn_slice + del attn_slice + else: + query_slice = query[start_idx:end_idx, start_idx_2:end_idx_2] + key_slice = key[start_idx:end_idx, start_idx_2:end_idx_2] + attn_mask_slice = attention_mask[start_idx:end_idx, start_idx_2:end_idx_2] if attention_mask is not None else None + + attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) + del query_slice + del key_slice + del attn_mask_slice + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx, start_idx_2:end_idx_2]) + + hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = attn_slice + del attn_slice + else: + query_slice = query[start_idx:end_idx] + key_slice = key[start_idx:end_idx] + attn_mask_slice = attention_mask[start_idx:end_idx] if attention_mask is not None else None + + attn_slice = attn.get_attention_scores(query_slice, key_slice, attn_mask_slice) + del query_slice + del key_slice + del attn_mask_slice + attn_slice = torch.bmm(attn_slice, value[start_idx:end_idx]) + + hidden_states[start_idx:end_idx] = attn_slice + del attn_slice + if devices.backend != "directml": + getattr(torch, query.device.type).synchronize() + else: + attention_probs = attn.get_attention_scores(query, key, attention_mask) + hidden_states = torch.bmm(attention_probs, value) + #################################################################### + hidden_states = attn.batch_to_head_dim(hidden_states) + + # linear proj + hidden_states = attn.to_out[0](hidden_states, *args) + # dropout + hidden_states = attn.to_out[1](hidden_states) + + if input_ndim == 4: + hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width) + + if attn.residual_connection: + hidden_states = hidden_states + residual + + hidden_states = hidden_states / attn.rescale_output_factor + + return hidden_states diff --git a/modules/sd_models.py b/modules/sd_models.py index f1c691314..c4ae53a24 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -704,6 +704,9 @@ def set_diffuser_options(sd_model, vae = None, op: str = 'model'): sd_model.vqvae.to(torch.float32) # vqvae is producing nans in fp16 if shared.opts.cross_attention_optimization == "xFormers" and hasattr(sd_model, 'enable_xformers_memory_efficient_attention'): sd_model.enable_xformers_memory_efficient_attention() + if shared.opts.diffusers_dynamic_attention_slicing: + from modules.sd_hijack_optimizations import DynamicAttnProcessorV1 + set_diffusers_attention(sd_model, DynamicAttnProcessorV1()) if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'fuse_qkv_projections'): shared.log.debug(f'Setting {op}: enable fused projections') sd_model.fuse_qkv_projections() @@ -1155,6 +1158,15 @@ def set_diffuser_pipe(pipe, new_pipe_type): return pipe +def set_diffusers_attention(pipe, attention): + module_names, _ = pipe._get_signature_keys(pipe) + modules = [getattr(pipe, n, None) for n in module_names] + modules = [m for m in modules if isinstance(m, torch.nn.Module) and hasattr(m, "set_attn_processor")] + + for module in modules: + module.set_attn_processor(attention) + + def get_native(pipe: diffusers.DiffusionPipeline): if hasattr(pipe, "vae") and hasattr(pipe.vae.config, "sample_size"): # Stable Diffusion diff --git a/modules/shared.py b/modules/shared.py index 2c76e4204..95db8f170 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -443,6 +443,8 @@ options_templates.update(options_section(('diffusers', "Diffusers Settings"), { "diffusers_vae_slicing": OptionInfo(True, "VAE slicing"), "diffusers_vae_tiling": OptionInfo(False, "VAE tiling"), "diffusers_attention_slicing": OptionInfo(False, "Attention slicing"), + "diffusers_dynamic_attention_slicing": OptionInfo(False, "Dynamic Attention slicing"), + "dynamic_attention_slice_rate": OptionInfo(4, "Slicing rate for Dynamic Attention Slicing in GB", gr.Slider, {"minimum": 0.1, "maximum": 16, "step": 0.1}), "diffusers_model_load_variant": OptionInfo("default", "Preferred Model variant", gr.Radio, {"choices": ['default', 'fp32', 'fp16']}), "diffusers_vae_load_variant": OptionInfo("default", "Preferred VAE variant", gr.Radio, {"choices": ['default', 'fp32', 'fp16']}), "custom_diffusers_pipeline": OptionInfo('', 'Load custom Diffusers pipeline'),