Upload STDiT2

config.json

@@ -10,7 +10,7 @@
   "class_dropout_prob": 0.1,
   "depth": 28,
   "drop_path": 0.0,
-  "enable_flashattn": false,
+  "enable_flash_attn": false,
   "enable_layernorm_kernel": false,
   "enable_sequence_parallelism": false,
   "freeze": null,

configuration_stdit2.py

@@ -24,7 +24,7 @@ class STDiT2Config(PretrainedConfig):
         model_max_length=120,
         freeze=None,
         qk_norm=False,
-        enable_flashattn=False,
+        enable_flash_attn=False,
         enable_layernorm_kernel=False,
         enable_sequence_parallelism=False,
         **kwargs,
@@ -45,7 +45,7 @@ class STDiT2Config(PretrainedConfig):
         self.model_max_length = model_max_length
         self.freeze = freeze
         self.qk_norm = qk_norm
-        self.enable_flashattn = enable_flashattn
+        self.enable_flash_attn = enable_flash_attn
         self.enable_layernorm_kernel = enable_layernorm_kernel
         self.enable_sequence_parallelism = enable_sequence_parallelism
         super().__init__(**kwargs)

layers.py

@@ -30,7 +30,7 @@ class STDiT2Block(nn.Module):
         num_heads,
         mlp_ratio=4.0,
         drop_path=0.0,
-        enable_flashattn=False,
+        enable_flash_attn=False,
         enable_layernorm_kernel=False,
         enable_sequence_parallelism=False,
         rope=None,
@@ -38,7 +38,7 @@ class STDiT2Block(nn.Module):
     ):
         super().__init__()
         self.hidden_size = hidden_size
-        self.enable_flashattn = enable_flashattn
+        self.enable_flash_attn = enable_flash_attn
         self._enable_sequence_parallelism = enable_sequence_parallelism
 
         assert not self._enable_sequence_parallelism, "Sequence parallelism is not supported."
@@ -55,7 +55,7 @@ class STDiT2Block(nn.Module):
             hidden_size,
             num_heads=num_heads,
             qkv_bias=True,
-            enable_flashattn=enable_flashattn,
+            enable_flash_attn=enable_flash_attn,
             qk_norm=qk_norm,
         )
         self.scale_shift_table = nn.Parameter(torch.randn(6, hidden_size) / hidden_size**0.5)
@@ -76,7 +76,7 @@ class STDiT2Block(nn.Module):
             hidden_size,
             num_heads=num_heads,
             qkv_bias=True,
-            enable_flashattn=self.enable_flashattn,
+            enable_flash_attn=self.enable_flash_attn,
             rope=rope,
             qk_norm=qk_norm,
         )
@@ -196,7 +196,7 @@ class Attention(nn.Module):
         attn_drop: float = 0.0,
         proj_drop: float = 0.0,
         norm_layer: nn.Module = LlamaRMSNorm,
-        enable_flashattn: bool = False,
+        enable_flash_attn: bool = False,
         rope=None,
     ) -> None:
         super().__init__()
@@ -205,7 +205,7 @@ class Attention(nn.Module):
         self.num_heads = num_heads
         self.head_dim = dim // num_heads
         self.scale = self.head_dim**-0.5
-        self.enable_flashattn = enable_flashattn
+        self.enable_flash_attn = enable_flash_attn
 
         self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
         self.q_norm = norm_layer(self.head_dim) if qk_norm else nn.Identity()
@@ -222,7 +222,7 @@ class Attention(nn.Module):
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         B, N, C = x.shape
         # flash attn is not memory efficient for small sequences, this is empirical
-        enable_flashattn = self.enable_flashattn and (N > B)
+        enable_flash_attn = self.enable_flash_attn and (N > B)
         qkv = self.qkv(x)
         qkv_shape = (B, N, 3, self.num_heads, self.head_dim)
 
@@ -233,7 +233,7 @@ class Attention(nn.Module):
             k = self.rotary_emb(k)
         q, k = self.q_norm(q), self.k_norm(k)
 
-        if enable_flashattn:
+        if enable_flash_attn:
             from flash_attn import flash_attn_func
 
             # (B, #heads, N, #dim) -> (B, N, #heads, #dim)
@@ -258,7 +258,7 @@ class Attention(nn.Module):
             x = attn @ v
 
         x_output_shape = (B, N, C)
-        if not enable_flashattn:
+        if not enable_flash_attn:
             x = x.transpose(1, 2)
         x = x.reshape(x_output_shape)
         x = self.proj(x)

model.safetensors

@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:e777ae49713478957c48f97eda4405e392e3fd12580e01be944465b741c6521c
+oid sha256:033319fd49c2ff9bc57836e2fcaacea5ac64f6efa357f603f91f06a9164d0e1c
 size 3071846872

modeling_stdit2.py

@@ -38,7 +38,7 @@ class STDiT2(PreTrainedModel):
         self.no_temporal_pos_emb = config.no_temporal_pos_emb
         self.depth = config.depth
         self.mlp_ratio = config.mlp_ratio
-        self.enable_flashattn = config.enable_flashattn
+        self.enable_flash_attn = config.enable_flash_attn
         self.enable_layernorm_kernel = config.enable_layernorm_kernel
         self.enable_sequence_parallelism = config.enable_sequence_parallelism
 
@@ -69,7 +69,7 @@ class STDiT2(PreTrainedModel):
                     self.num_heads,
                     mlp_ratio=self.mlp_ratio,
                     drop_path=drop_path[i],
-                    enable_flashattn=self.enable_flashattn,
+                    enable_flash_attn=self.enable_flash_attn,
                     enable_layernorm_kernel=self.enable_layernorm_kernel,
                     enable_sequence_parallelism=self.enable_sequence_parallelism,
                     rope=self.rope.rotate_queries_or_keys,