update memory efficient attention

animatediff/models/motion_module.py

@@ -458,6 +458,10 @@ class CrossAttention(nn.Module):
         hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
         return hidden_states
 
+    def set_use_memory_efficient_attention_xformers(self, *args, **kwargs):
+        print('Set Xformers for MotionModule\'s Attention.')
+        self._use_memory_efficient_attention_xformers = True
+
     def _memory_efficient_attention_xformers(self, query, key, value, attention_mask):
         # TODO attention_mask
         query = query.contiguous()