sd3 xformers quality

Author: CanvaChen

src/diffusers/models/attention_processor.py

@@ -1530,13 +1530,13 @@ def __init__(self, attention_op: Optional[Callable] = None):
         self.attention_op = attention_op
 
     def __call__(
-            self,
-            attn: Attention,
-            hidden_states: torch.FloatTensor,
-            encoder_hidden_states: torch.FloatTensor = None,
-            attention_mask: Optional[torch.FloatTensor] = None,
-            *args,
-            **kwargs,
+        self,
+        attn: Attention,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        *args,
+        **kwargs,
     ) -> torch.FloatTensor:
         residual = hidden_states
 
@@ -1579,7 +1579,7 @@ def __call__(
         # Split the attention outputs.
         hidden_states, encoder_hidden_states = (
             hidden_states[:, : residual.shape[1]],
-            hidden_states[:, residual.shape[1]:],
+            hidden_states[:, residual.shape[1] :],
         )
 
         # linear proj

tests/models/transformers/test_models_transformer_sd3.py

@@ -89,6 +89,5 @@ def test_xformers_enable_works(self):
         model.enable_xformers_memory_efficient_attention()
 
         assert (
-                model.transformer_blocks[0].attn.processor.__class__.__name__
-                == "XFormersJointAttnProcessor"
+            model.transformer_blocks[0].attn.processor.__class__.__name__ == "XFormersJointAttnProcessor"
         ), "xformers is not enabled"