sd3 xformers test

Author: CanvaChen

tests/models/transformers/test_models_transformer_sd3.py

@@ -18,6 +18,7 @@
 import torch
 
 from diffusers import SD3Transformer2DModel
+from diffusers.utils.import_utils import is_xformers_available
 from diffusers.utils.testing_utils import (
     enable_full_determinism,
     torch_device,
@@ -76,3 +77,18 @@ def prepare_init_args_and_inputs_for_common(self):
         }
         inputs_dict = self.dummy_input
         return init_dict, inputs_dict
+
+    @unittest.skipIf(
+        torch_device != "cuda" or not is_xformers_available(),
+        reason="XFormers attention is only available with CUDA and `xformers` installed",
+    )
+    def test_xformers_enable_works(self):
+        init_dict, inputs_dict = self.prepare_init_args_and_inputs_for_common()
+        model = self.model_class(**init_dict)
+
+        model.enable_xformers_memory_efficient_attention()
+
+        assert (
+                model.transformer_blocks[0].attn.processor.__class__.__name__
+                == "XFormersJointAttnProcessor"
+        ), "xformers is not enabled"