diff --git a/modules/devices.py b/modules/devices.py
index 638e9df0c..ad008add4 100644
--- a/modules/devices.py
+++ b/modules/devices.py
@@ -4,6 +4,7 @@ import sys
import time
import contextlib
import torch
+from functools import wraps
from modules.errors import log
from modules import cmd_args, shared, memstats, errors
@@ -12,7 +13,6 @@ if sys.platform == "darwin":
previous_oom = 0
-backup_sdpa = None
debug = os.environ.get('SD_DEVICE_DEBUG', None) is not None
@@ -250,18 +250,17 @@ def set_cuda_params():
except Exception:
pass
try:
- if shared.opts.cross_attention_optimization == "Scaled-Dot-Product" or shared.opts.cross_attention_optimization == "Dynamic Attention SDP":
+ if shared.opts.cross_attention_optimization == "Scaled-Dot-Product":
torch.backends.cuda.enable_flash_sdp('Flash attention' in shared.opts.sdp_options)
torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in shared.opts.sdp_options)
torch.backends.cuda.enable_math_sdp('Math attention' in shared.opts.sdp_options)
if backend == "rocm":
- global backup_sdpa # pylint: disable=global-statement
if 'Flash attention' in shared.opts.sdp_options:
try:
# https://github.com/huggingface/diffusers/discussions/7172
from flash_attn import flash_attn_func
- if backup_sdpa is None:
- backup_sdpa = torch.nn.functional.scaled_dot_product_attention
+ backup_sdpa = torch.nn.functional.scaled_dot_product_attention
+ @wraps(torch.nn.functional.scaled_dot_product_attention)
def sdpa_hijack(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None):
if query.shape[3] <= 128 and attn_mask is None and query.dtype != torch.float32:
return flash_attn_func(q=query.transpose(1, 2), k=key.transpose(1, 2), v=value.transpose(1, 2), dropout_p=dropout_p, causal=is_causal, softmax_scale=scale).transpose(1, 2)
@@ -271,8 +270,9 @@ def set_cuda_params():
shared.log.debug('ROCm Flash Attention Hijacked')
except Exception as err:
log.error(f'ROCm Flash Attention failed: {err}')
- elif backup_sdpa is not None: # Restore original SDPA
- torch.nn.functional.scaled_dot_product_attention = backup_sdpa
+ if 'Dynamic attention' in shared.opts.sdp_options:
+ from modules.sd_hijack_dynamic_atten import sliced_scaled_dot_product_attention
+ torch.nn.functional.scaled_dot_product_attention = sliced_scaled_dot_product_attention
except Exception:
pass
if shared.cmd_opts.profile:
diff --git a/modules/sd_hijack_dynamic_atten.py b/modules/sd_hijack_dynamic_atten.py
index b2d6fdc42..fb2befc18 100644
--- a/modules/sd_hijack_dynamic_atten.py
+++ b/modules/sd_hijack_dynamic_atten.py
@@ -1,7 +1,6 @@
-from functools import cache
+from functools import cache, wraps
import torch
-import torch.nn.functional as F
from diffusers.utils import USE_PEFT_BACKEND
from modules import shared, devices
@@ -49,6 +48,8 @@ def find_slice_sizes(query_shape, query_element_size, slice_rate=4):
return do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size
+backup_sdpa = torch.nn.functional.scaled_dot_product_attention
+@wraps(torch.nn.functional.scaled_dot_product_attention)
def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, **kwargs):
do_split, do_split_2, do_split_3, split_slice_size, split_2_slice_size, split_3_slice_size = find_slice_sizes(query.shape, query.element_size(), slice_rate=shared.opts.dynamic_attention_slice_rate)
@@ -67,7 +68,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
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
- hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = F.scaled_dot_product_attention(
+ hidden_states[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3] = backup_sdpa(
query[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
key[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
value[start_idx:end_idx, start_idx_2:end_idx_2, start_idx_3:end_idx_3],
@@ -75,7 +76,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
else:
- hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = F.scaled_dot_product_attention(
+ hidden_states[start_idx:end_idx, start_idx_2:end_idx_2] = backup_sdpa(
query[start_idx:end_idx, start_idx_2:end_idx_2],
key[start_idx:end_idx, start_idx_2:end_idx_2],
value[start_idx:end_idx, start_idx_2:end_idx_2],
@@ -83,7 +84,7 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
dropout_p=dropout_p, is_causal=is_causal, **kwargs
)
else:
- hidden_states[start_idx:end_idx] = F.scaled_dot_product_attention(
+ hidden_states[start_idx:end_idx] = backup_sdpa(
query[start_idx:end_idx],
key[start_idx:end_idx],
value[start_idx:end_idx],
@@ -93,93 +94,10 @@ def sliced_scaled_dot_product_attention(query, key, value, attn_mask=None, dropo
if devices.backend != "directml":
getattr(torch, query.device.type).synchronize()
else:
- return F.scaled_dot_product_attention(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
+ return backup_sdpa(query, key, value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, **kwargs)
return hidden_states
-class DynamicAttnProcessorSDP:
- r"""
- dynamically slices attention queries in order to keep them under the slice rate
- slicing will not get triggered if the query size is smaller than the slice rate to gain performance
-
- slice rate is in GB
- based on AttnProcessor V2
- """
-
- def __init__(self):
- if not hasattr(F, "scaled_dot_product_attention"):
- raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")
-
- def __call__(self, attn, hidden_states: torch.Tensor, encoder_hidden_states=None, attention_mask=None, temb=None, *args, **kwargs) -> torch.Tensor:
-
- residual = hidden_states
- 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
- )
-
- if attention_mask is not None:
- attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
- # scaled_dot_product_attention expects attention_mask shape to be
- # (batch, heads, source_length, target_length)
- attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
-
- 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)
-
- 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)
- value = attn.to_v(encoder_hidden_states)
-
- inner_dim = key.shape[-1]
- head_dim = inner_dim // attn.heads
-
- query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
- key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
- value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
-
- # the output of sdp = (batch, num_heads, seq_len, head_dim)
- # -: add support for attn.scale when we move to Torch 2.1
- ####################################################################
- # Slicing part:
- hidden_states = sliced_scaled_dot_product_attention(
- query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
- )
- ####################################################################
-
- hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
- hidden_states = hidden_states.to(query.dtype)
-
- # linear proj
- hidden_states = attn.to_out[0](hidden_states)
- # 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
-
class DynamicAttnProcessorBMM:
r"""
dynamically slices attention queries in order to keep them under the slice rate
diff --git a/modules/sd_models.py b/modules/sd_models.py
index 272587a8e..12533bdbe 100644
--- a/modules/sd_models.py
+++ b/modules/sd_models.py
@@ -1525,9 +1525,6 @@ def set_diffusers_attention(pipe):
elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM":
from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM
set_attn(pipe, DynamicAttnProcessorBMM())
- elif shared.opts.cross_attention_optimization == "Dynamic Attention SDP":
- from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorSDP
- set_attn(pipe, DynamicAttnProcessorSDP())
pipe.current_attn_name = shared.opts.cross_attention_optimization
diff --git a/modules/shared.py b/modules/shared.py
index bdc6a7916..cd136689c 100644
--- a/modules/shared.py
+++ b/modules/shared.py
@@ -359,7 +359,7 @@ cpu_memory = psutil.virtual_memory().total / 1024 / 1024 / 1024
mem_stat = memory_stats()
if "gpu" in mem_stat:
- gpu_memory = mem_stat['gpu']['total']
+ gpu_memory = mem_stat['gpu']['total'] + 0.1
if not (cmd_opts.lowvram or cmd_opts.medvram):
if "gpu" in mem_stat:
@@ -382,13 +382,11 @@ elif cmd_opts.lowvram:
if devices.backend == "directml": # Force BMM for DirectML instead of SDP
cross_attention_optimization_default = "Dynamic Attention BMM" if native else "Sub-quadratic"
-elif native and (cmd_opts.lowvram or cmd_opts.medvram):
- cross_attention_optimization_default = "Dynamic Attention SDP"
elif devices.backend == "cpu":
cross_attention_optimization_default = "Scaled-Dot-Product" if native else "Doggettx's"
elif devices.backend == "mps":
cross_attention_optimization_default = "Scaled-Dot-Product" if native else "Doggettx's"
-else: # cuda, rocm, ipex
+else: # cuda, rocm, ipex, openvino
cross_attention_optimization_default ="Scaled-Dot-Product"
@@ -399,6 +397,9 @@ if devices.backend == "rocm":
else:
sdp_options_default = ['Flash attention', 'Memory attention', 'Math attention']
+if (cmd_opts.lowvram or cmd_opts.medvram) and 'Flash attention' not in sdp_options_default:
+ sdp_options_default.append('Dynamic attention')
+
options_templates.update(options_section(('sd', "Execution & Models"), {
"sd_backend": OptionInfo(default_backend, "Execution backend", gr.Radio, {"choices": ["diffusers", "original"] }),
"sd_model_checkpoint": OptionInfo(default_checkpoint, "Base model", gr.Dropdown, lambda: {"choices": list_checkpoint_tiles()}, refresh=refresh_checkpoints),
@@ -437,9 +438,9 @@ options_templates.update(options_section(('cuda', "Compute Settings"), {
"cross_attention_sep": OptionInfo("Cross Attention
", "", gr.HTML),
"cross_attention_optimization": OptionInfo(cross_attention_optimization_default, "Attention optimization method", gr.Radio, lambda: {"choices": shared_items.list_crossattention(native) }),
- "sdp_options": OptionInfo(sdp_options_default, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention'] }),
+ "sdp_options": OptionInfo(sdp_options_default, "SDP options", gr.CheckboxGroup, {"choices": ['Flash attention', 'Memory attention', 'Math attention', 'Dynamic attention'] }),
"xformers_options": OptionInfo(['Flash attention'], "xFormers options", gr.CheckboxGroup, {"choices": ['Flash attention'] }),
- "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": 24, "step": 0.1, "visible": native}),
+ "dynamic_attention_slice_rate": OptionInfo(4, "Dynamic Attention slicing rate in GB", gr.Slider, {"minimum": 0.1, "maximum": gpu_memory, "step": 0.1, "visible": native}),
"sub_quad_sep": OptionInfo("Sub-quadratic options
", "", gr.HTML, {"visible": not native}),
"sub_quad_q_chunk_size": OptionInfo(512, "Attention query chunk size", gr.Slider, {"minimum": 16, "maximum": 8192, "step": 8, "visible": not native}),
"sub_quad_kv_chunk_size": OptionInfo(512, "Attention kv chunk size", gr.Slider, {"minimum": 0, "maximum": 8192, "step": 8, "visible": not native}),
diff --git a/modules/shared_items.py b/modules/shared_items.py
index 0c76d59bd..1b0077ef8 100644
--- a/modules/shared_items.py
+++ b/modules/shared_items.py
@@ -31,8 +31,7 @@ def list_crossattention(diffusers=False):
"xFormers",
"Batch matrix-matrix",
"Split attention",
- "Dynamic Attention BMM",
- "Dynamic Attention SDP"
+ "Dynamic Attention BMM"
]
else:
return [