From a5ae3e31dc97367786be323644415f06084f7cc7 Mon Sep 17 00:00:00 2001 From: ayo0la <168016617+ayo0la@users.noreply.github.com> Date: Thu, 10 Sep 2026 09:11:08 -0400 Subject: [PATCH] fix(attention): make AttentionModuleMixin.set_attention_slice callable set_attention_slice on the mixin referenced two names that did not exist, so it raised AttributeError for every argument: default_processor_cls (the attribute is _default_processor_cls) and _get_compatible_processor (never defined). Nothing routes into the method today, so it was latent, but any mixin-based model wired into enable_attention_slicing would have hit it. Fix the attribute name, add _get_compatible_processor which instantiates the first entry in _available_processors whose class name matches the requested type, and warn instead of silently no-oping when slicing is requested on a module that lists no sliced processor. Fixes #14729 --- src/diffusers/models/attention.py | 22 ++++++- tests/models/test_attention_processor.py | 77 +++++++++++++++++++++++- 2 files changed, 96 insertions(+), 3 deletions(-) diff --git a/src/diffusers/models/attention.py b/src/diffusers/models/attention.py index 65289e4b5f16..cdc0763d5ae4 100644 --- a/src/diffusers/models/attention.py +++ b/src/diffusers/models/attention.py @@ -352,14 +352,32 @@ def set_attention_slice(self, slice_size: int) -> None: # Try to get a compatible processor for sliced attention if slice_size is not None: - processor = self._get_compatible_processor("sliced") + processor = self._get_compatible_processor("sliced", slice_size=slice_size) + if processor is None: + logger.warning( + f"Attention slicing was requested but `{type(self).__name__}` has no sliced attention " + "processor in `_available_processors`. Falling back to the default processor, so attention " + "will not be sliced for this module." + ) # If no processor was found or slice_size is None, use default processor if processor is None: - processor = self.default_processor_cls() + processor = self._default_processor_cls() self.set_processor(processor) + def _get_compatible_processor(self, processor_type: str, **init_kwargs) -> "AttentionProcessor | None": + """ + Instantiate the first processor in `_available_processors` whose class name contains `processor_type` + (case-insensitive), e.g. `"sliced"` for a `SlicedAttnProcessor`-style class. Returns `None` if the module lists + no such processor. + """ + processor_type = processor_type.lower() + for processor_cls in self._available_processors: + if processor_type in processor_cls.__name__.lower(): + return processor_cls(**init_kwargs) + return None + def batch_to_head_dim(self, tensor: torch.Tensor) -> torch.Tensor: """ Reshape the tensor from `[batch_size, seq_len, dim]` to `[batch_size // heads, seq_len, dim * heads]`. diff --git a/tests/models/test_attention_processor.py b/tests/models/test_attention_processor.py index a2b02b56692c..f2d55a9a3c2b 100644 --- a/tests/models/test_attention_processor.py +++ b/tests/models/test_attention_processor.py @@ -5,11 +5,14 @@ import pytest import torch from packaging import version +from torch import nn from diffusers import DiffusionPipeline +from diffusers.models.attention import AttentionModuleMixin from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor +from diffusers.utils import logging -from ..testing_utils import torch_device +from ..testing_utils import CaptureLogger, torch_device class TestAttnAddedKVProcessor: @@ -132,3 +135,75 @@ def test_conversion_when_using_device_map(self): assert np.allclose(pre_conversion, conversion, atol=1e-3) assert np.allclose(conversion, after_conversion, atol=1e-3) + + +class _MixinTestProcessor: + def __call__(self, attn, hidden_states, *args, **kwargs): + return hidden_states + + +class _MixinTestSlicedProcessor: + def __init__(self, slice_size: int): + self.slice_size = slice_size + + def __call__(self, attn, hidden_states, *args, **kwargs): + return hidden_states + + +class _MixinAttention(nn.Module, AttentionModuleMixin): + _default_processor_cls = _MixinTestProcessor + _available_processors = [_MixinTestProcessor] + _supports_qkv_fusion = False + + def __init__(self, sliceable_head_dim: int | None = None): + super().__init__() + if sliceable_head_dim is not None: + self.sliceable_head_dim = sliceable_head_dim + self.set_processor(self._default_processor_cls()) + + +class _MixinAttentionWithSliced(_MixinAttention): + _available_processors = [_MixinTestProcessor, _MixinTestSlicedProcessor] + + +class TestAttentionModuleMixinSetAttentionSlice: + def test_none_restores_default_processor(self): + attn = _MixinAttention() + attn.set_processor(_MixinTestSlicedProcessor(2)) + + attn.set_attention_slice(None) + + assert isinstance(attn.processor, _MixinTestProcessor) + + def test_uses_sliced_processor_when_available(self): + attn = _MixinAttentionWithSliced() + + attn.set_attention_slice(2) + + assert isinstance(attn.processor, _MixinTestSlicedProcessor) + assert attn.processor.slice_size == 2 + + def test_falls_back_to_default_and_warns_when_no_sliced_processor(self): + attn = _MixinAttention() + attn_logger = logging.get_logger("diffusers.models.attention") + attn_logger.setLevel(logging.WARNING) + + with CaptureLogger(attn_logger) as cap_logger: + attn.set_attention_slice(2) + + assert isinstance(attn.processor, _MixinTestProcessor) + assert "sliced" in cap_logger.out + assert "_MixinAttention" in cap_logger.out + + def test_slice_size_larger_than_sliceable_head_dim_raises(self): + attn = _MixinAttentionWithSliced(sliceable_head_dim=4) + + with pytest.raises(ValueError, match="has to be smaller or equal to 4"): + attn.set_attention_slice(8) + + def test_get_compatible_processor(self): + attn = _MixinAttentionWithSliced() + assert isinstance(attn._get_compatible_processor("sliced", slice_size=3), _MixinTestSlicedProcessor) + + attn = _MixinAttention() + assert attn._get_compatible_processor("sliced", slice_size=3) is None