Upload STDiT

config.json

@@ -11,6 +11,7 @@
   "depth": 28,
   "drop_path": 0.0,
   "enable_flash_attn": false,
+  "enable_flashattn": false,
   "enable_layernorm_kernel": false,
   "enable_sequence_parallelism": false,
   "freeze": null,
@@ -32,7 +33,7 @@
     2
   ],
   "pred_sigma": true,
-  "space_scale": 0.5,
+  "space_scale": 1.0,
   "time_scale": 1.0,
   "torch_dtype": "float32",
   "transformers_version": "4.38.2"

modeling_stdit.py

@@ -112,7 +112,7 @@ class STDiT(PreTrainedModel):
         x = x.to(self.final_layer.linear.weight.dtype)
         timestep = timestep.to(self.final_layer.linear.weight.dtype)
         y = y.to(self.final_layer.linear.weight.dtype)
-        
+
         # embedding
         x = self.x_embedder(x)  # [B, N, C]
         x = rearrange(x, "B (T S) C -> B T S C", T=self.num_temporal, S=self.num_spatial)
@@ -148,7 +148,8 @@ class STDiT(PreTrainedModel):
                     tpe = self.pos_embed_temporal
             else:
                 tpe = None
-            x = auto_grad_checkpoint(block, x, y, t0, y_lens, tpe)
+            x = block(x, y, t0, y_lens, tpe)
+            # x = auto_grad_checkpoint(block, x, y, t0, y_lens, tpe)
 
         if self.enable_sequence_parallelism:
             x = gather_forward_split_backward(x, get_sequence_parallel_group(), dim=1, grad_scale="up")