[wip]

zhhezhhe · May 7, 2024 · fa314ed · fa314ed
1 parent c00a06a
commit fa314ed
Show file tree

Hide file tree

Showing 46 changed files with 90 additions and 88 deletions.
diff --git a/README.md b/README.md
@@ -161,7 +161,7 @@ conda activate opensora
 pip install torch torchvision
 
 # install flash attention (optional)
-# set enable_flashattn=False in config to avoid using flash attention
+# set enable_flash_attn=False in config to avoid using flash attention
 pip install packaging ninja
 pip install flash-attn --no-build-isolation
 

diff --git a/configs/dit/train/16x256x256.py b/configs/dit/train/16x256x256.py
@@ -18,7 +18,7 @@
 model = dict(
     type="DiT-XL/2",
     from_pretrained="DiT-XL-2-256x256.pt",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/dit/train/1x256x256.py b/configs/dit/train/1x256x256.py
@@ -19,7 +19,7 @@
 model = dict(
     type="DiT-XL/2",
     no_temporal_pos_emb=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/latte/train/16x256x256.py b/configs/latte/train/16x256x256.py
@@ -17,7 +17,7 @@
 # Define model
 model = dict(
     type="Latte-XL/2",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/benchmark.py b/configs/opensora-v1-1/train/benchmark.py
@@ -66,7 +66,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/image.py b/configs/opensora-v1-1/train/image.py
@@ -30,7 +30,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/image_rflow.py b/configs/opensora-v1-1/train/image_rflow.py
@@ -35,7 +35,7 @@
 #     type="DiT-XL/2",
 #     from_pretrained="/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/PixArt-XL-2-512x512.pth",
 #     # input_sq_size=512,  # pretrained model is trained on 512x512
-#     enable_flashattn=True,
+#     enable_flash_attn=True,
 #     enable_layernorm_kernel=True,
 # )
 model = dict(
@@ -44,7 +44,7 @@
     time_scale=1.0,
     no_temporal_pos_emb=True,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 # model = dict(
@@ -54,7 +54,7 @@
 #     no_temporal_pos_emb=True,
 #     # from_pretrained="PixArt-XL-2-512x512.pth",
 #     from_pretrained="/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/PixArt-XL-2-512x512.pth",
-#     enable_flashattn=True,
+#     enable_flash_attn=True,
 #     enable_layernorm_kernel=True,
 # )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/stage1.py b/configs/opensora-v1-1/train/stage1.py
@@ -42,7 +42,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/stage2.py b/configs/opensora-v1-1/train/stage2.py
@@ -44,7 +44,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/stage3.py b/configs/opensora-v1-1/train/stage3.py
@@ -44,7 +44,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-1/train/video.py b/configs/opensora-v1-1/train/video.py
@@ -32,7 +32,7 @@
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
     qk_norm_legacy=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-2/inference/sample-ref.py b/configs/opensora-v1-2/inference/sample-ref.py
@@ -22,7 +22,7 @@
     type="STDiT3-XL/2",
     from_pretrained=None,
     qk_norm=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-2/inference/sample.py b/configs/opensora-v1-2/inference/sample.py
@@ -14,7 +14,7 @@
     type="STDiT3-XL/2",
     from_pretrained=None,
     qk_norm=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-2/train/stage1-gc.py b/configs/opensora-v1-2/train/stage1-gc.py
@@ -36,7 +36,7 @@
     type="STDiT3-XL/2",
     from_pretrained=None,
     qk_norm=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora-v1-2/train/stage1.py b/configs/opensora-v1-2/train/stage1.py
@@ -34,12 +34,12 @@
     type="STDiT3-XL/2",
     from_pretrained=None,
     qk_norm=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(
     type="VideoAutoencoderPipeline",
-    from_pretrained="pretrained_models/vae-v3",
+    from_pretrained="pretrained_models/vae-v2",
     micro_frame_size=17,
     vae_2d=dict(
         type="VideoAutoencoderKL",

diff --git a/configs/opensora/inference/16x256x256.py b/configs/opensora/inference/16x256x256.py
@@ -7,7 +7,7 @@
     type="STDiT-XL/2",
     space_scale=0.5,
     time_scale=1.0,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

diff --git a/configs/opensora/inference/16x512x512-rflow.py b/configs/opensora/inference/16x512x512-rflow.py
@@ -7,7 +7,7 @@
     type="STDiT-XL/2",
     space_scale=1.0,
     time_scale=1.0,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

diff --git a/configs/opensora/inference/16x512x512.py b/configs/opensora/inference/16x512x512.py
@@ -7,7 +7,7 @@
     type="STDiT-XL/2",
     space_scale=1.0,
     time_scale=1.0,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

diff --git a/configs/opensora/inference/64x512x512.py b/configs/opensora/inference/64x512x512.py
@@ -7,7 +7,7 @@
     type="STDiT-XL/2",
     space_scale=1.0,
     time_scale=2 / 3,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     from_pretrained="PRETRAINED_MODEL",
 )

diff --git a/configs/opensora/train/16x256x256-mask.py b/configs/opensora/train/16x256x256-mask.py
@@ -20,7 +20,7 @@
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 mask_ratios = {

diff --git a/configs/opensora/train/16x256x256-spee-rflow.py b/configs/opensora/train/16x256x256-spee-rflow.py
@@ -22,8 +22,8 @@
     # from_pretrained="PixArt-XL-2-512x512.pth",
     # from_pretrained = "/home/zhaowangbo/wangbo/PixArt-alpha/pretrained_models/OpenSora-v1-HQ-16x512x512.pth",
     # from_pretrained = "OpenSora-v1-HQ-16x512x512.pth",
-    from_pretrained =  "PRETRAINED_MODEL",
-    enable_flashattn=True,
+    from_pretrained="PRETRAINED_MODEL",
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 # mask_ratios = [0.5, 0.29, 0.07, 0.07, 0.07]

diff --git a/configs/opensora/train/16x256x256-spee.py b/configs/opensora/train/16x256x256-spee.py
@@ -20,7 +20,7 @@
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 mask_ratios = {

diff --git a/configs/opensora/train/16x256x256.py b/configs/opensora/train/16x256x256.py
@@ -20,7 +20,7 @@
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora/train/16x512x512.py b/configs/opensora/train/16x512x512.py
@@ -20,7 +20,7 @@
     space_scale=1.0,
     time_scale=1.0,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/opensora/train/360x512x512.py b/configs/opensora/train/360x512x512.py
@@ -26,7 +26,7 @@
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     enable_sequence_parallelism=True,  # enable sq here
 )

diff --git a/configs/opensora/train/64x512x512-sp.py b/configs/opensora/train/64x512x512-sp.py
@@ -20,7 +20,7 @@
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     enable_sequence_parallelism=True,  # enable sq here
 )

diff --git a/configs/opensora/train/64x512x512.py b/configs/opensora/train/64x512x512.py
@@ -20,7 +20,7 @@
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/pixart/inference/1x20481B.py b/configs/pixart/inference/1x20481B.py
@@ -8,7 +8,7 @@
     from_pretrained="PixArt-1B-2.pth",
     space_scale=4,
     no_temporal_pos_emb=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     base_size=2048 // 8,
 )

diff --git a/configs/pixart/inference/1x2048MS.py b/configs/pixart/inference/1x2048MS.py
@@ -8,7 +8,7 @@
     from_pretrained="PixArt-Sigma-XL-2-2K-MS.pth",
     space_scale=4,
     no_temporal_pos_emb=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
     base_size=2048 // 8,
 )

diff --git a/configs/pixart/train/16x256x256.py b/configs/pixart/train/16x256x256.py
@@ -20,7 +20,7 @@
     space_scale=0.5,
     time_scale=1.0,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/pixart/train/1x2048x2048.py b/configs/pixart/train/1x2048x2048.py
@@ -20,7 +20,7 @@
     space_scale=4.0,
     no_temporal_pos_emb=True,
     from_pretrained="PixArt-1B-2.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 

diff --git a/configs/pixart/train/1x512x512-rflow.py b/configs/pixart/train/1x512x512-rflow.py
@@ -21,8 +21,8 @@
     time_scale=1.0,
     no_temporal_pos_emb=True,
     # from_pretrained="PixArt-XL-2-512x512.pth",
-    from_pretrained =  "PRETRAINED_MODEL",
-    enable_flashattn=True,
+    from_pretrained="PRETRAINED_MODEL",
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/pixart/train/1x512x512.py b/configs/pixart/train/1x512x512.py
@@ -21,7 +21,7 @@
     time_scale=1.0,
     no_temporal_pos_emb=True,
     from_pretrained="PixArt-XL-2-512x512.pth",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/configs/pixart/train/64x512x512.py b/configs/pixart/train/64x512x512.py
@@ -21,7 +21,7 @@
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

diff --git a/docs/config.md b/docs/config.md
@@ -26,8 +26,8 @@ model = dict(
     from_pretrained="PRETRAINED_MODEL",  # (Optional) Load from pretrained model
     input_sq_size=512,        # Base spatial position embedding size
     qk_norm=True,             # Normalize query and key in attention
-    enable_flashattn=True,    # (Optional) Speed up training and inference with flash attention
-    # Turn enable_flashattn to False if you skip flashattn installation
+    enable_flash_attn=True,    # (Optional) Speed up training and inference with flash attention
+    # Turn enable_flash_attn to False if you skip flashattn installation
     enable_layernorm_kernel=True, # (Optional) Speed up training and inference with fused kernel
     # Turn enable_layernorm_kernel to False if you skip apex installation
 )
@@ -177,7 +177,7 @@ model = dict(
     from_pretrained=None,              # Load from pretrained model
     input_sq_size=512,                 # Base spatial position embedding size
     qk_norm=True,                      # Normalize query and key in attention
-    enable_flashattn=True,             # (Optional) Speed up training and inference with flash attention
+    enable_flash_attn=True,             # (Optional) Speed up training and inference with flash attention
     enable_layernorm_kernel=True,      # (Optional) Speed up training and inference with fused kernel
 )
 vae = dict(

diff --git a/docs/zh_CN/structure.md b/docs/zh_CN/structure.md
@@ -89,7 +89,7 @@ model = dict(
     type="STDiT-XL/2",        # Select model type (STDiT-XL/2, DiT-XL/2, etc.)
     space_scale=1.0,          # (Optional) Space positional encoding scale (new height / old height)
     time_scale=2 / 3,         # (Optional) Time positional encoding scale (new frame_interval / old frame_interval)
-    enable_flashattn=True,    # (Optional) Speed up training and inference with flash attention
+    enable_flash_attn=True,    # (Optional) Speed up training and inference with flash attention
     enable_layernorm_kernel=True, # (Optional) Speed up training and inference with fused kernel
     from_pretrained="PRETRAINED_MODEL",  # (Optional) Load from pretrained model
     no_temporal_pos_emb=True,  # (Optional) Disable temporal positional encoding (for image)
@@ -144,7 +144,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained="YOUR_PRETRAINED_MODEL",
-    enable_flashattn=True,        # Enable flash attention
+    enable_flash_attn=True,        # Enable flash attention
     enable_layernorm_kernel=True, # Enable layernorm kernel
 )
 vae = dict(