diff --git a/CHANGELOG.md b/CHANGELOG.md index 76624b18e..d1a3aecf2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -5,12 +5,12 @@ - **Features** - allow recursive inline wildcards using curly braces syntax - simplify SDNQ pre-quantization saved config + - refactor settings and improve handling of attention mechanisms - **Fixes** - hires strength save/load in metadata, thanks @awsr - fix imgi2img initial scale tab, thanks @awsr - fix pony-v7 text-encoder - detailer with face-restorers - - fix sage-attention checks on sm86 ## Update for 2025-11-06 diff --git a/modules/attention.py b/modules/attention.py index 6ed2bfb24..1ad1b1d4b 100644 --- a/modules/attention.py +++ b/modules/attention.py @@ -16,41 +16,40 @@ def set_dynamic_attention(): log.error(f'Torch attention: type="dynamic attention" {err}') return None -def set_triton_flash_attention(backend: str): +def set_triton_flash_attention(): try: - if backend in {"zluda", "rocm"}: - from modules.flash_attn_triton_amd import interface_fa - sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention - @wraps(sdpa_pre_triton_flash_atten) - def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: Optional[torch.FloatTensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor: - if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32: - if scale is None: - scale = query.shape[-1] ** (-0.5) - head_size_og = query.size(3) - if head_size_og % 8 != 0: - query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8]) - key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8]) - value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8]) - query = query.transpose(1, 2) - key = key.transpose(1, 2) - value = value.transpose(1, 2) - out_padded = torch.zeros_like(query) - interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal) - return out_padded[..., :head_size_og].transpose(1, 2) - else: - if enable_gqa: - kwargs["enable_gqa"] = enable_gqa - return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs) - torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten - log.debug('Torch attention: type="triton flash attention"') + from modules.flash_attn_triton_amd import interface_fa + sdpa_pre_triton_flash_atten = torch.nn.functional.scaled_dot_product_attention + @wraps(sdpa_pre_triton_flash_atten) + def sdpa_triton_flash_atten(query: torch.FloatTensor, key: torch.FloatTensor, value: torch.FloatTensor, attn_mask: Optional[torch.FloatTensor] = None, dropout_p: float = 0.0, is_causal: bool = False, scale: Optional[float] = None, enable_gqa: bool = False, **kwargs) -> torch.FloatTensor: + if query.shape[-1] <= 128 and attn_mask is None and query.dtype != torch.float32: + if scale is None: + scale = query.shape[-1] ** (-0.5) + head_size_og = query.size(3) + if head_size_og % 8 != 0: + query = torch.nn.functional.pad(query, [0, 8 - head_size_og % 8]) + key = torch.nn.functional.pad(key, [0, 8 - head_size_og % 8]) + value = torch.nn.functional.pad(value, [0, 8 - head_size_og % 8]) + query = query.transpose(1, 2) + key = key.transpose(1, 2) + value = value.transpose(1, 2) + out_padded = torch.zeros_like(query) + interface_fa.fwd(query, key, value, out_padded, dropout_p, scale, is_causal) + return out_padded[..., :head_size_og].transpose(1, 2) + else: + if enable_gqa: + kwargs["enable_gqa"] = enable_gqa + return sdpa_pre_triton_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs) + torch.nn.functional.scaled_dot_product_attention = sdpa_triton_flash_atten + log.debug('Torch attention: type="Triton Flash attention"') except Exception as err: - log.error(f'Torch attention: type="triton flash attention" {err}') + log.error(f'Torch attention: type="Triton Flash attention" {err}') def set_ck_flash_attention(backend: str, device: torch.device): try: if backend == "rocm": if not installed('flash-attn'): - log.info('Building CK Flash attention...') + log.info('Torch attention: type="CK Flash" building...') agent = rocm.Agent(getattr(torch.cuda.get_device_properties(device), "gcnArchName", "gfx0000")) install(rocm.get_flash_attention_command(agent), reinstall=True) else: @@ -83,9 +82,9 @@ def set_ck_flash_attention(backend: str, device: torch.device): kwargs["enable_gqa"] = enable_gqa return sdpa_pre_flash_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs) torch.nn.functional.scaled_dot_product_attention = sdpa_flash_atten - log.debug('Torch attention: type="ck flash attention"') + log.debug('Torch attention: type="CK Flash attention"') except Exception as err: - log.error(f'Torch attention: type="ck flash attention" {err}') + log.error(f'Torch attention: type="CK Flash attention" {err}') def set_sage_attention(backend: str, device: torch.device): try: @@ -100,7 +99,6 @@ def set_sage_attention(backend: str, device: torch.device): use_cuda_backend = False if use_cuda_backend: - log.debug('Torch attention: type=SageAttention backend=cuda') from sageattention import sageattn_qk_int8_pv_fp16_cuda def sage_attn_impl(query, key, value, is_causal, scale): return sageattn_qk_int8_pv_fp16_cuda( @@ -112,7 +110,6 @@ def set_sage_attention(backend: str, device: torch.device): pv_accum_dtype="fp32", ) else: - log.debug('Torch attention: type=SageAttention backend=auto') from sageattention import sageattn def sage_attn_impl(query, key, value, is_causal, scale): return sageattn( @@ -138,6 +135,60 @@ def set_sage_attention(backend: str, device: torch.device): kwargs["enable_gqa"] = enable_gqa return sdpa_pre_sage_atten(query=query, key=key, value=value, attn_mask=attn_mask, dropout_p=dropout_p, is_causal=is_causal, scale=scale, **kwargs) torch.nn.functional.scaled_dot_product_attention = sdpa_sage_atten - log.debug('Torch attention: type="sage attention"') + log.debug(f'Torch attention: type="Sage attention" backend={"cuda" if use_cuda_backend else "auto"}') except Exception as err: - log.error(f'Torch attention: type="sage attention" {err}') + log.error(f'Torch attention: type="Sage attention" {err}') + + +def set_diffusers_attention(pipe, quiet:bool=False): + from modules import shared + import diffusers.models.attention_processor as p + + def set_attn(pipe, attention, name:str=None): + if attention is None: + return + # other models uses their own attention processor + if getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"): + try: + pipe.unet.set_attn_processor(attention) + except Exception as e: + if 'Nunchaku' in pipe.unet.__class__.__name__: + pass + else: + shared.log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}') + """ # each transformer typically has its own attention processor + if getattr(pipe, "transformer", None) is not None and hasattr(pipe.transformer, "set_attn_processor"): + try: + pipe.transformer.set_attn_processor(attention) + except Exception as e: + if 'Nunchaku' in pipe.transformer.__class__.__name__: + pass + else: + shared.log.error(f'Torch attention: type="{name}" cls={attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}') + """ + + shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"') + if shared.opts.cross_attention_optimization == "Disabled": + pass # do nothing + elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers + # set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product") + pass + elif shared.opts.cross_attention_optimization == "xFormers": + if hasattr(pipe, 'enable_xformers_memory_efficient_attention'): + pipe.enable_xformers_memory_efficient_attention() + else: + shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}") + elif shared.opts.cross_attention_optimization == "Batch matrix-matrix": + set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix") + elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM": + from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM + set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM") + + if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"): + if shared.opts.attention_slicing: + pipe.enable_attention_slicing() + else: + pipe.disable_attention_slicing() + shared.log.debug(f"Torch attention: slicing={shared.opts.attention_slicing}") + + pipe.current_attn_name = shared.opts.cross_attention_optimization diff --git a/modules/devices.py b/modules/devices.py index ba988837d..f8ef7d3d1 100644 --- a/modules/devices.py +++ b/modules/devices.py @@ -458,29 +458,29 @@ def set_sdpa_params(): log.warning(f'Torch attention: type="sdpa" {err}') try: - torch.backends.cuda.enable_flash_sdp('Flash attention' in opts.sdp_options) - torch.backends.cuda.enable_mem_efficient_sdp('Memory attention' in opts.sdp_options) - torch.backends.cuda.enable_math_sdp('Math attention' in opts.sdp_options) + torch.backends.cuda.enable_flash_sdp('Flash' in opts.sdp_options) + torch.backends.cuda.enable_mem_efficient_sdp('Memory' in opts.sdp_options) + torch.backends.cuda.enable_math_sdp('Math' in opts.sdp_options) if hasattr(torch.backends.cuda, "allow_fp16_bf16_reduction_math_sdp"): # only valid for torch >= 2.5 torch.backends.cuda.allow_fp16_bf16_reduction_math_sdp(True) - log.debug(f'Torch attention: type="sdpa" opts={opts.sdp_options}') + log.debug(f'Torch attention: type="sdpa" kernels={opts.sdp_options} overrides={opts.sdp_overrides}') except Exception as err: log.warning(f'Torch attention: type="sdpa" {err}') # Stack hijcaks in reverse order. This gives priority to the last added hijack. # If the last hijack is not compatible, it will use the one before it and so on. - if 'Dynamic attention' in opts.sdp_options: + if 'Dynamic attention' in opts.sdp_overrides: global sdpa_pre_dyanmic_atten # pylint: disable=global-statement sdpa_pre_dyanmic_atten = attention.set_dynamic_attention() - if 'Triton Flash attention' in opts.sdp_options: - attention.set_triton_flash_attention(backend) + if 'Triton Flash attention' in opts.sdp_overrides: + attention.set_triton_flash_attention() - if 'CK Flash attention' in opts.sdp_options: + if 'CK Flash attention' in opts.sdp_overrides: attention.set_ck_flash_attention(backend, device) - if 'Sage attention' in opts.sdp_options: + if 'Sage attention' in opts.sdp_overrides: attention.set_sage_attention(backend, device) from importlib.metadata import version diff --git a/modules/processing_diffusers.py b/modules/processing_diffusers.py index c2d852f25..6bca30a3f 100644 --- a/modules/processing_diffusers.py +++ b/modules/processing_diffusers.py @@ -5,7 +5,7 @@ import numpy as np import torch import torchvision.transforms.functional as TF from PIL import Image -from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks +from modules import shared, devices, processing, sd_models, errors, sd_hijack_hypertile, processing_vae, sd_models_compile, timer, modelstats, extra_networks, attention from modules.processing_helpers import resize_hires, calculate_base_steps, calculate_hires_steps, calculate_refiner_steps, save_intermediate, update_sampler, is_txt2img, is_refiner_enabled, get_job_name from modules.processing_args import set_pipeline_args from modules.onnx_impl import preprocess_pipeline as preprocess_onnx_pipeline, check_parameters_changed as olive_check_parameters_changed @@ -497,7 +497,7 @@ def update_pipeline(sd_model, p: processing.StableDiffusionProcessing): orig_pipeline = sd_model # processed ONNX pipeline should not be replaced with original pipeline. if getattr(sd_model, "current_attn_name", None) != shared.opts.cross_attention_optimization: shared.log.info(f"Setting attention optimization: {shared.opts.cross_attention_optimization}") - sd_models.set_diffusers_attention(sd_model) + attention.set_diffusers_attention(sd_model) return sd_model diff --git a/modules/sd_models.py b/modules/sd_models.py index b40cf4aea..2c3e9f0f0 100644 --- a/modules/sd_models.py +++ b/modules/sd_models.py @@ -10,7 +10,7 @@ import diffusers.loaders.single_file_utils import torch import huggingface_hub as hf from installer import log -from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors +from modules import timer, paths, shared, shared_items, modelloader, devices, script_callbacks, sd_vae, sd_unet, errors, sd_models_compile, sd_detect, model_quant, sd_hijack_te, sd_hijack_accelerate, sd_hijack_safetensors, attention from modules.memstats import memory_stats from modules.modeldata import model_data from modules.sd_checkpoint import CheckpointInfo, select_checkpoint, list_models, checkpoints_list, checkpoint_titles, get_closest_checkpoint_match, model_hash, update_model_hashes, setup_model, write_metadata, read_metadata_from_safetensors # pylint: disable=unused-import @@ -130,7 +130,7 @@ def set_diffuser_options(sd_model, vae=None, op:str='model', offload:bool=True, clear_caches() set_vae_options(sd_model, vae, op, quiet) - set_diffusers_attention(sd_model, quiet) + attention.set_diffusers_attention(sd_model, quiet) if shared.opts.diffusers_fuse_projections and hasattr(sd_model, 'fuse_qkv_projections'): try: @@ -1157,60 +1157,6 @@ def set_diffuser_pipe(pipe, new_pipe_type): return pipe -def set_diffusers_attention(pipe, quiet:bool=False): - import diffusers.models.attention_processor as p - - def set_attn(pipe, attention, name:str=None, quiet:bool=False): - if attention is None: - return - # other models uses their own attention processor - if pipe.__class__.__name__.startswith("StableDiffusion") and getattr(pipe, "unet", None) is not None and hasattr(pipe.unet, "set_attn_processor"): - try: - pipe.unet.set_attn_processor(attention) - except Exception as e: - if 'Nunchaku' in pipe.unet.__class__.__name__: - pass - else: - shared.log.error(f"Attention: {name if name is not None else attention.__class__.__name__} pipe={pipe.__class__.__name__} {e}") - elif not quiet: - shared.log.warning(f"Attention: {name if name is not None else attention.__class__.__name__} is not compatible with {pipe.__class__.__name__}") - - # if hasattr(pipe, 'pipe'): - # set_diffusers_attention(pipe.pipe) - - if 'Control' in pipe.__class__.__name__ or 'Adapter' in pipe.__class__.__name__ or not (pipe.__class__.__name__.startswith("StableDiffusion") and hasattr(pipe, "unet")): - if shared.opts.cross_attention_optimization not in {"Scaled-Dot-Product", "Disabled"}: - shared.log.warning(f"Attention: {shared.opts.cross_attention_optimization} is not compatible with {pipe.__class__.__name__}") - else: - pipe.current_attn_name = shared.opts.cross_attention_optimization - return - - shared.log.quiet(quiet, f'Setting model: attention="{shared.opts.cross_attention_optimization}"') - if shared.opts.cross_attention_optimization == "Disabled": - pass # do nothing - elif shared.opts.cross_attention_optimization == "Scaled-Dot-Product": # The default set by Diffusers - set_attn(pipe, p.AttnProcessor2_0(), name="Scaled-Dot-Product", quiet=True) - elif shared.opts.cross_attention_optimization == "xFormers": - if hasattr(pipe, 'enable_xformers_memory_efficient_attention'): - pipe.enable_xformers_memory_efficient_attention() - else: - shared.log.warning(f"Attention: xFormers is not compatible with {pipe.__class__.__name__}") - elif shared.opts.cross_attention_optimization == "Batch matrix-matrix": - set_attn(pipe, p.AttnProcessor(), name="Batch matrix-matrix") - elif shared.opts.cross_attention_optimization == "Dynamic Attention BMM": - from modules.sd_hijack_dynamic_atten import DynamicAttnProcessorBMM - set_attn(pipe, DynamicAttnProcessorBMM(), name="Dynamic Attention BMM") - - if shared.opts.attention_slicing != "Default" and hasattr(pipe, "enable_attention_slicing") and hasattr(pipe, "disable_attention_slicing"): - if shared.opts.attention_slicing: - pipe.enable_attention_slicing() - else: - pipe.disable_attention_slicing() - shared.log.debug(f"Attention: slicing={shared.opts.attention_slicing}") - - pipe.current_attn_name = shared.opts.cross_attention_optimization - - def add_noise_pred_to_diffusers_callback(pipe): if not hasattr(pipe, "_callback_tensor_inputs"): return pipe diff --git a/modules/shared.py b/modules/shared.py index 24aad5271..66165e385 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -136,7 +136,7 @@ def list_samplers(): return modules.sd_samplers.all_samplers -startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_sdp_choices, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat) +startup_offload_mode, startup_offload_min_gpu, startup_offload_max_gpu, startup_cross_attention, startup_sdp_options, startup_sdp_choices, startup_sdp_override_options, startup_sdp_override_choices, startup_offload_always, startup_offload_never = get_default_modes(cmd_opts=cmd_opts, mem_stat=mem_stat) options_templates.update(options_section(('sd', "Model Loading"), { "sd_backend": OptionInfo('diffusers', "Execution backend", gr.Radio, {"choices": ['diffusers', 'original'], "visible": False }), @@ -296,13 +296,13 @@ options_templates.update(options_section(('cuda', "Compute Settings"), { "diffusers_generator_device": OptionInfo("GPU", "Generator device", gr.Radio, {"choices": ["GPU", "CPU", "Unset"]}), "cross_attention_sep": OptionInfo("