Make AttentionModuleMixin.set_attention_slice callable

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]`.

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