fix loop bug in SlicedAttnProcessor

src/diffusers/models/attention_processor.py

@@ -2190,7 +2190,7 @@ def __call__(
             (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
         )
 
-        for i in range(batch_size_attention // self.slice_size):
+        for i in range((batch_size_attention - 1) // self.slice_size + 1):
             start_idx = i * self.slice_size
             end_idx = (i + 1) * self.slice_size
 
@@ -2287,7 +2287,7 @@ def __call__(
             (batch_size_attention, query_tokens, dim // attn.heads), device=query.device, dtype=query.dtype
         )
 
-        for i in range(batch_size_attention // self.slice_size):
+        for i in range((batch_size_attention - 1) // self.slice_size + 1):
             start_idx = i * self.slice_size
             end_idx = (i + 1) * self.slice_size
 

src/diffusers/models/unets/unet_2d_condition.py

@@ -815,7 +815,10 @@ def fn_recursive_retrieve_sliceable_dims(module: torch.nn.Module):
             size = slice_size[i]
             dim = sliceable_head_dims[i]
             if size is not None and size > dim:
-                raise ValueError(f"size {size} has to be smaller or equal to {dim}.")
+                slice_size[i] = dim
+                logger.warning(
+                    f"size {size} has to be smaller or equal to {dim}, and slice_size {size} has been set to {dim}"
+                )
 
         # Recursively walk through all the children.
         # Any children which exposes the set_attention_slice method

tests/pipelines/test_pipelines_common.py

@@ -1351,14 +1351,30 @@ def _test_attention_slicing_forward_pass(
 
         pipe.enable_attention_slicing(slice_size=1)
         inputs = self.get_dummy_inputs(generator_device)
-        output_with_slicing = pipe(**inputs)[0]
+        output_with_slicing1 = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=2)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing2 = pipe(**inputs)[0]
+
+        pipe.enable_attention_slicing(slice_size=3)
+        inputs = self.get_dummy_inputs(generator_device)
+        output_with_slicing3 = pipe(**inputs)[0]
 
         if test_max_difference:
-            max_diff = np.abs(to_np(output_with_slicing) - to_np(output_without_slicing)).max()
-            self.assertLess(max_diff, expected_max_diff, "Attention slicing should not affect the inference results")
+            max_diff1 = np.abs(to_np(output_with_slicing1) - to_np(output_without_slicing)).max()
+            max_diff2 = np.abs(to_np(output_with_slicing2) - to_np(output_without_slicing)).max()
+            max_diff3 = np.abs(to_np(output_with_slicing3) - to_np(output_without_slicing)).max()
+            self.assertLess(
+                max(max_diff1, max_diff2, max_diff3),
+                expected_max_diff,
+                "Attention slicing should not affect the inference results",
+            )
 
         if test_mean_pixel_difference:
-            assert_mean_pixel_difference(to_np(output_with_slicing[0]), to_np(output_without_slicing[0]))
+            assert_mean_pixel_difference(to_np(output_with_slicing1[0]), to_np(output_without_slicing[0]))
+            assert_mean_pixel_difference(to_np(output_with_slicing2[0]), to_np(output_without_slicing[0]))
+            assert_mean_pixel_difference(to_np(output_with_slicing3[0]), to_np(output_without_slicing[0]))
 
     @unittest.skipIf(
         torch_device != "cuda" or not is_accelerate_available() or is_accelerate_version("<", "0.14.0"),