remove attn_mode and xformers

Author: brianfitzgerald

mmditx.py

@@ -217,15 +217,13 @@ def optimized_attention(qkv, num_heads):
 
 
 class SelfAttention(nn.Module):
-    ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
 
     def __init__(
         self,
         dim: int,
         num_heads: int = 8,
         qkv_bias: bool = False,
         qk_scale: Optional[float] = None,
-        attn_mode: str = "xformers",
         pre_only: bool = False,
         qk_norm: Optional[str] = None,
         rmsnorm: bool = False,
@@ -239,8 +237,6 @@ def __init__(
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias, dtype=dtype, device=device)
         if not pre_only:
             self.proj = nn.Linear(dim, dim, dtype=dtype, device=device)
-        assert attn_mode in self.ATTENTION_MODES
-        self.attn_mode = attn_mode
         self.pre_only = pre_only
 
         if qk_norm == "rms":
@@ -294,7 +290,7 @@ def post_attention(self, x: torch.Tensor) -> torch.Tensor:
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         (q, k, v) = self.pre_attention(x)
-        x = attention(q, k, v, self.num_heads, self.attn_mode)
+        x = attention(q, k, v, self.num_heads)
         x = self.post_attention(x)
         return x
 
@@ -391,14 +387,11 @@ def forward(self, x):
 class DismantledBlock(nn.Module):
     """A DiT block with gated adaptive layer norm (adaLN) conditioning."""
 
-    ATTENTION_MODES = ("xformers", "torch", "torch-hb", "math", "debug")
-
     def __init__(
         self,
         hidden_size: int,
         num_heads: int,
         mlp_ratio: float = 4.0,
-        attn_mode: str = "xformers",
         qkv_bias: bool = False,
         pre_only: bool = False,
         rmsnorm: bool = False,
@@ -411,7 +404,6 @@ def __init__(
         **block_kwargs,
     ):
         super().__init__()
-        assert attn_mode in self.ATTENTION_MODES
         if not rmsnorm:
             self.norm1 = nn.LayerNorm(
                 hidden_size,
@@ -426,7 +418,6 @@ def __init__(
             dim=hidden_size,
             num_heads=num_heads,
             qkv_bias=qkv_bias,
-            attn_mode=attn_mode,
             pre_only=pre_only,
             qk_norm=qk_norm,
             rmsnorm=rmsnorm,
@@ -441,7 +432,6 @@ def __init__(
                 dim=hidden_size,
                 num_heads=num_heads,
                 qkv_bias=qkv_bias,
-                attn_mode=attn_mode,
                 pre_only=False,
                 qk_norm=qk_norm,
                 rmsnorm=rmsnorm,
@@ -716,7 +706,6 @@ def __init__(
         adm_in_channels: Optional[int] = None,
         context_embedder_config: Optional[Dict] = None,
         register_length: int = 0,
-        attn_mode: str = "torch",
         rmsnorm: bool = False,
         scale_mod_only: bool = False,
         swiglu: bool = False,
@@ -735,7 +724,7 @@ def __init__(
         super().__init__()
         if verbose:
             print(
-                f"mmdit initializing with: {input_size=}, {patch_size=}, {in_channels=}, {depth=}, {mlp_ratio=}, {learn_sigma=}, {adm_in_channels=}, {context_embedder_config=}, {register_length=}, {attn_mode=}, {rmsnorm=}, {scale_mod_only=}, {swiglu=}, {out_channels=}, {pos_embed_scaling_factor=}, {pos_embed_offset=}, {pos_embed_max_size=}, {num_patches=}, {qk_norm=}, {qkv_bias=}, {dtype=}, {device=}"
+                f"mmdit initializing with: {input_size=}, {patch_size=}, {in_channels=}, {depth=}, {mlp_ratio=}, {learn_sigma=}, {adm_in_channels=}, {context_embedder_config=}, {register_length=}, {rmsnorm=}, {scale_mod_only=}, {swiglu=}, {out_channels=}, {pos_embed_scaling_factor=}, {pos_embed_offset=}, {pos_embed_max_size=}, {num_patches=}, {qk_norm=}, {qkv_bias=}, {dtype=}, {device=}"
             )
         self.dtype = dtype
         self.learn_sigma = learn_sigma
@@ -805,7 +794,6 @@ def __init__(
                     num_heads,
                     mlp_ratio=mlp_ratio,
                     qkv_bias=qkv_bias,
-                    attn_mode=attn_mode,
                     pre_only=i == depth - 1,
                     rmsnorm=rmsnorm,
                     scale_mod_only=scale_mod_only,

other_impls.py

@@ -9,32 +9,20 @@
 from transformers import CLIPTokenizer, T5TokenizerFast
 from einops import rearrange
 
-try:
-    import xformers.ops
-except ImportError:
-    xformers.ops = None
-    print("xformers not found, attn_mode='xformers' will not work")
-
 #################################################################################################
 ### Core/Utility
 #################################################################################################
 
 
-def attention(q, k, v, heads, mask=None, attn_mode: str = "torch"):
+def attention(q, k, v, heads, mask=None):
     """Convenience wrapper around a basic attention operation"""
     b, _, dim_head = q.shape
     dim_head //= heads
     q, k, v = map(lambda t: t.view(b, -1, heads, dim_head).transpose(1, 2), (q, k, v))
-    if attn_mode == "torch":
-        out = torch.nn.functional.scaled_dot_product_attention(
-            q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
-        )
-        return out.transpose(1, 2).reshape(b, -1, heads * dim_head)
-    elif attn_mode == "xformers":
-        x = xformers.ops.memory_efficient_attention(q, k, v)
-        x = rearrange(x, "b h n d -> b n (h d)")
-        return x
-
+    out = torch.nn.functional.scaled_dot_product_attention(
+        q, k, v, attn_mask=mask, dropout_p=0.0, is_causal=False
+    )
+    return out.transpose(1, 2).reshape(b, -1, heads * dim_head)
 
 class Mlp(nn.Module):
     """MLP as used in Vision Transformer, MLP-Mixer and related networks"""