update

configs/dit/train/16x256x256.py

@@ -18,7 +18,7 @@ sp_size = 1
 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(

configs/dit/train/1x256x256.py

@@ -19,7 +19,7 @@ sp_size = 1
 model = dict(
     type="DiT-XL/2",
     no_temporal_pos_emb=True,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/latte/train/16x256x256.py

@@ -17,7 +17,7 @@ sp_size = 1
 # Define model
 model = dict(
     type="Latte-XL/2",
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/inference/sample-ref.py

@@ -7,33 +7,35 @@ multi_resolution = "STDiT2"
 # Condition
 prompt_path = None
 prompt = [
-    "A car driving on the ocean.",
-    'Drone view of waves crashing against the rugged cliffs along Big Sur\'s garay point beach. The crashing blue waters create white-tipped waves, while the golden light of the setting sun illuminates the rocky shore. A small island with a lighthouse sits in the distance, and green shrubbery covers the cliff\'s edge. The steep drop from the road down to the beach is a dramatic feat, with the cliff\'s edges jutting out over the sea. This is a view that captures the raw beauty of the coast and the rugged landscape of the Pacific Coast Highway.{"reference_path": "assets/images/condition/cliff.png", "mask_strategy": "0"}',
-    "In an ornate, historical hall, a massive tidal wave peaks and begins to crash. Two surfers, seizing the moment, skillfully navigate the face of the wave.",
+    'Drone view of waves crashing against the rugged cliffs along Big Sur\'s garay point beach. {"reference_path": "assets/images/condition/cliff.png", "mask_strategy": "0"}',
+    'A breathtaking sunrise scene.{"reference_path": "assets/images/condition/sunset1.png","mask_strategy": "0"}',
+    'A car driving on the ocean.{"reference_path": "https://cdn.openai.com/tmp/s/interp/d0.mp4","mask_strategy": "0,0,-8,0,8"}',
+    'A snowy forest.{"reference_path": "https://cdn.pixabay.com/video/2021/04/25/72171-542991404_large.mp4","mask_strategy": "0,0,0,0,15,0.8"}',
+    'A breathtaking sunrise scene.{"reference_path": "assets/images/condition/sunset1.png;assets/images/condition/sunset2.png","mask_strategy": "0;0,1,0,-1,1"}',
+    '|0|a white jeep equipped with a roof rack driving on a dirt road in a coniferous forest.|2|a white jeep equipped with a roof rack driving on a dirt road in the desert.|4|a white jeep equipped with a roof rack driving on a dirt road in a mountain.|6|A white jeep equipped with a roof rack driving on a dirt road in a city.|8|a white jeep equipped with a roof rack driving on a dirt road on the surface of a river.|10|a white jeep equipped with a roof rack driving on a dirt road under the lake.|12|a white jeep equipped with a roof rack flying into the sky.|14|a white jeep equipped with a roof rack driving in the universe. Earth is the background.{"reference_path": "https://cdn.openai.com/tmp/s/interp/d0.mp4", "mask_strategy": "0,0,0,0,15"}',
 ]
 
 loop = 2
 condition_frame_length = 4
-reference_path = [
-    "https://cdn.openai.com/tmp/s/interp/d0.mp4",
-    None,
-    "assets/images/condition/wave.png",
-]
-# valid when reference_path is not None
-# (loop id, ref id, ref start, length, target start)
-mask_strategy = [
-    "0,0,0,0,8,0.3",
-    None,
-    "0",
-]
+# (
+#   loop id, [the loop index of the condition image or video]
+#   reference id, [the index of the condition image or video in the reference_path]
+#   reference start, [the start frame of the condition image or video]
+#   target start, [the location to insert]
+#   length, [the number of frames to insert]
+#   edit_ratio [the edit rate of the condition image or video]
+# )
+# See https://github.com/hpcaitech/Open-Sora/blob/main/docs/config.md#advanced-inference-config for more details
+# See https://github.com/hpcaitech/Open-Sora/blob/main/docs/commands.md#inference-with-open-sora-11 for more examples
 
 # Define model
 model = dict(
     type="STDiT2-XL/2",
-    from_pretrained=None,
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v2-stage3",
     input_sq_size=512,
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/inference/sample.py

@@ -7,10 +7,11 @@ multi_resolution = "STDiT2"
 # Define model
 model = dict(
     type="STDiT2-XL/2",
-    from_pretrained=None,
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v2-stage3",
     input_sq_size=512,
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/benchmark.py

@@ -65,7 +65,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/image.py

@@ -29,7 +29,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/image_rflow.py

@@ -0,0 +1,88 @@
+# Define dataset
+# dataset = dict(
+#     type="VariableVideoTextDataset",
+#     data_path=None,
+#     num_frames=None,
+#     frame_interval=3,
+#     image_size=(None, None),
+#     transform_name="resize_crop",
+# )
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=1,
+    frame_interval=1,
+    image_size=(256, 256),
+    transform_name="center",
+)
+bucket_config = {  # 6s/it
+    "256": {1: (1.0, 256)},
+    "512": {1: (1.0, 80)},
+    "480p": {1: (1.0, 52)},
+    "1024": {1: (1.0, 20)},
+    "1080p": {1: (1.0, 8)},
+}
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+# model = dict(
+#     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_flash_attn=True,
+#     enable_layernorm_kernel=True,
+# )
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PixArt-XL-2-512x512.pth",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+# model = dict(
+#     type="DiT-XL/2",
+#     # space_scale=1.0,
+#     # time_scale=1.0,
+#     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_flash_attn=True,
+#     enable_layernorm_kernel=True,
+# )
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=200,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 10
+log_every = 10
+ckpt_every = 500
+load = None
+
+batch_size = 100  # only for logging
+lr = 2e-5
+grad_clip = 1.0

configs/opensora-v1-1/train/stage1.py

@@ -16,15 +16,15 @@ bucket_config = {  # 1s/it
     "1024": {1: (0.3, 3)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -41,7 +41,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/stage2.py

@@ -18,15 +18,15 @@ bucket_config = {  # 7s/it
     "1080p": {1: (0.4, 8)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -43,7 +43,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/stage3.py

@@ -18,15 +18,15 @@ bucket_config = {  # 13s/it
     "1024": {1: (0.3, 40)},
 }
 mask_ratios = {
-    "mask_no": 0.75,
-    "mask_quarter_random": 0.025,
-    "mask_quarter_head": 0.025,
-    "mask_quarter_tail": 0.025,
-    "mask_quarter_head_tail": 0.05,
-    "mask_image_random": 0.025,
-    "mask_image_head": 0.025,
-    "mask_image_tail": 0.025,
-    "mask_image_head_tail": 0.05,
+    "identity": 0.75,
+    "quarter_random": 0.025,
+    "quarter_head": 0.025,
+    "quarter_tail": 0.025,
+    "quarter_head_tail": 0.05,
+    "image_random": 0.025,
+    "image_head": 0.025,
+    "image_tail": 0.025,
+    "image_head_tail": 0.05,
 }
 
 # Define acceleration
@@ -43,7 +43,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-1/train/video.py

@@ -31,7 +31,8 @@ model = dict(
     from_pretrained=None,
     input_sq_size=512,  # pretrained model is trained on 512x512
     qk_norm=True,
-    enable_flashattn=True,
+    qk_norm_legacy=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora-v1-2/inference/sample.py

@@ -0,0 +1,42 @@
+resolution = "240p"
+aspect_ratio = "9:16"
+num_frames = 51
+fps = 24
+frame_interval = 1
+save_fps = 24
+
+save_dir = "./samples/samples/"
+seed = 42
+batch_size = 1
+multi_resolution = "STDiT2"
+dtype = "bf16"
+condition_frame_length = 5
+align = 5
+
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained="hpcai-tech/OpenSora-STDiT-v3",
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="hpcai-tech/OpenSora-VAE-v1.2",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    num_sampling_steps=30,
+    cfg_scale=7.0,
+)
+
+aes = 6.5
+flow = None

configs/opensora-v1-2/misc/bs.py

@@ -0,0 +1,117 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# == Config 1: Webvid ==
+# base: (512, 408), 12s/it
+grad_checkpoint = True
+base = ("512", "408")
+base_step_time = 12
+bucket_config = {
+    "144p": {
+        1: (475, 0),
+        51: (51, 0),
+        102: (27, 0),
+        204: (13, 0),
+        408: (6, 0),
+    },
+    # ---
+    "240p": {
+        1: (297, 200),  # 8.25
+        51: (20, 0),
+        102: (10, 0),
+        204: (5, 0),
+        408: (2, 0),
+    },
+    # ---
+    "512": {
+        1: (141, 0),
+        51: (8, 0),
+        102: (4, 0),
+        204: (2, 0),
+        408: (1, 0),
+    },
+    # ---
+    "480p": {
+        1: (89, 0),
+        51: (5, 0),
+        102: (2, 0),
+        204: (1, 0),
+    },
+    # ---
+    "1024": {
+        1: (36, 0),
+        51: (1, 0),
+    },
+    # ---
+    "1080p": {1: (5, 0)},
+    # ---
+    "2048": {1: (5, 0)},
+}
+
+# == Config 1 ==
+# base: (512, 408), 16s/it
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/misc/eval_loss.py

@@ -0,0 +1,49 @@
+num_workers = 8
+dtype = "bf16"
+seed = 42
+num_eval_timesteps = 10
+
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+bucket_config = {
+    "144p": {1: (None, 100), 51: (None, 30), 102: (None, 20), 204: (None, 8), 408: (None, 4)},
+    # ---
+    "240p": {1: (None, 100), 51: (None, 24), 102: (None, 12), 204: (None, 4), 408: (None, 2)},
+    # ---
+    "360p": {1: (None, 60), 51: (None, 12), 102: (None, 6), 204: (None, 2), 408: (None, 1)},
+    # ---
+    "480p": {1: (None, 40), 51: (None, 6), 102: (None, 3), 204: (None, 1)},
+    # ---
+    "720p": {1: (None, 20), 51: (None, 2), 102: (None, 1)},
+    # ---
+    "1080p": {1: (None, 10)},
+    # ---
+    "2048": {1: (None, 5)},
+}
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+    local_files_only=True,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    local_files_only=True,
+)
+scheduler = dict(type="rflow")

configs/opensora-v1-2/misc/extract.py

@@ -0,0 +1,62 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+seed = 42
+outputs = "outputs"
+wandb = False
+
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained="/mnt/nfs-206/zangwei/opensora/outputs/1091-STDiT3-XL-2/epoch0-global_step8500",
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=32,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+
+# feature extraction settings
+save_text_features = True
+save_compressed_text_features = True
+bin_size = 250  # 1GB, 4195 bins
+log_time = False

configs/opensora-v1-2/misc/feat.py

@@ -0,0 +1,94 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+    dummy_text_feature=True,
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+grad_checkpoint = True
+
+load_text_features = True
+
+# Acceleration settings
+num_workers = 0
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+    skip_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 1
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/adapt.py

@@ -0,0 +1,84 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+bucket_config = {  # 2s/it
+    "144p": {1: (0.5, 48), 34: (1.0, 2), 51: (1.0, 4), 102: (1.0, 2), 204: (1.0, 1)},
+    # ---
+    "256": {1: (0.6, 20), 34: (0.5, 2), 51: (0.5, 1), 68: (0.5, 1), 136: (0.0, None)},
+    "240p": {1: (0.6, 20), 34: (0.5, 2), 51: (0.5, 1), 68: (0.5, 1), 136: (0.0, None)},
+    # ---
+    "360p": {1: (0.5, 8), 34: (0.2, 1), 102: (0.0, None)},
+    "512": {1: (0.5, 8), 34: (0.2, 1), 102: (0.0, None)},
+    # ---
+    "480p": {1: (0.2, 4), 17: (0.3, 1), 68: (0.0, None)},
+    # ---
+    "720p": {1: (0.1, 2)},
+    "1024": {1: (0.1, 2)},
+    # ---
+    "1080p": {1: (0.1, 1)},
+}
+grad_checkpoint = False
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="pretrained_models/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/stage1.py

@@ -0,0 +1,111 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# backup
+# bucket_config = {  # 20s/it
+#     "144p": {1: (1.0, 100), 51: (1.0, 30), 102: (1.0, 20), 204: (1.0, 8), 408: (1.0, 4)},
+#     # ---
+#     "256": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
+#     "240p": {1: (0.5, 100), 51: (0.3, 24), 102: (0.3, 12), 204: (0.3, 4), 408: (0.3, 2)},
+#     # ---
+#     "360p": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
+#     "512": {1: (0.5, 60), 51: (0.3, 12), 102: (0.3, 6), 204: (0.3, 2), 408: (0.3, 1)},
+#     # ---
+#     "480p": {1: (0.5, 40), 51: (0.3, 6), 102: (0.3, 3), 204: (0.3, 1), 408: (0.0, None)},
+#     # ---
+#     "720p": {1: (0.2, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
+#     "1024": {1: (0.1, 20), 51: (0.3, 2), 102: (0.3, 1), 204: (0.0, None)},
+#     # ---
+#     "1080p": {1: (0.1, 10)},
+#     # ---
+#     "2048": {1: (0.1, 5)},
+# }
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 0.1), 5), 408: ((0.5, 0.1), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 0.1), 5), 408: ((0.4, 0.1), 2)},
+    # ---
+    "360p": {1: (0.2, 141), 51: (0.15, 8), 102: ((0.15, 0.33), 4), 204: ((0.15, 0.1), 2), 408: ((0.15, 0.1), 1)},
+    "512": {1: (0.1, 141)},
+    # ---
+    "480p": {1: (0.1, 89)},
+    # ---
+    "720p": {1: (0.05, 36)},
+    "1024": {1: (0.05, 36)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.1, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.05,
+    "intepolate": 0.005,
+    "quarter_random": 0.005,
+    "quarter_head": 0.005,
+    "quarter_tail": 0.005,
+    "quarter_head_tail": 0.005,
+    "image_random": 0.025,
+    "image_head": 0.05,
+    "image_tail": 0.025,
+    "image_head_tail": 0.025,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15
+warmup_steps = 1000

configs/opensora-v1-2/train/stage1_feat.py

@@ -0,0 +1,59 @@
+# Dataset settings
+dataset = dict(type="BatchFeatureDataset")
+grad_checkpoint = True
+num_workers = 4
+
+# Acceleration settings
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+    skip_y_embedder=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+vae_out_channels = 4
+model_max_length = 300
+text_encoder_output_dim = 4096
+load_video_features = True
+load_text_features = True
+
+# Mask settings
+mask_ratios = {
+    "random": 0.2,
+    "intepolate": 0.01,
+    "quarter_random": 0.01,
+    "quarter_head": 0.01,
+    "quarter_tail": 0.01,
+    "quarter_head_tail": 0.01,
+    "image_random": 0.05,
+    "image_head": 0.1,
+    "image_tail": 0.05,
+    "image_head_tail": 0.05,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 500
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 2e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/stage2.py

@@ -0,0 +1,90 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 12s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: ((1.0, 0.33), 27), 204: ((1.0, 0.1), 13), 408: ((1.0, 0.1), 6)},
+    # ---
+    "256": {1: (0.4, 297), 51: (0.5, 20), 102: ((0.5, 0.33), 10), 204: ((0.5, 1.0), 5), 408: ((0.5, 1.0), 2)},
+    "240p": {1: (0.3, 297), 51: (0.4, 20), 102: ((0.4, 0.33), 10), 204: ((0.4, 1.0), 5), 408: ((0.4, 1.0), 2)},
+    # ---
+    "360p": {1: (0.5, 141), 51: (0.15, 8), 102: ((0.3, 0.5), 4), 204: ((0.3, 1.0), 2), 408: ((0.5, 0.5), 1)},
+    "512": {1: (0.4, 141), 51: (0.15, 8), 102: ((0.2, 0.4), 4), 204: ((0.2, 1.0), 2), 408: ((0.4, 0.5), 1)},
+    # ---
+    "480p": {1: (0.5, 89), 51: (0.2, 5), 102: (0.2, 2), 204: (0.1, 1)},
+    # ---
+    "720p": {1: (0.1, 36), 51: (0.03, 1)},
+    "1024": {1: (0.1, 36), 51: (0.02, 1)},
+    # ---
+    "1080p": {1: (0.01, 5)},
+    # ---
+    "2048": {1: (0.01, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+mask_ratios = {
+    "random": 0.05,
+    "intepolate": 0.005,
+    "quarter_random": 0.005,
+    "quarter_head": 0.005,
+    "quarter_tail": 0.005,
+    "quarter_head_tail": 0.005,
+    "image_random": 0.025,
+    "image_head": 0.05,
+    "image_tail": 0.025,
+    "image_head_tail": 0.025,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15

configs/opensora-v1-2/train/stage3.py

@@ -0,0 +1,92 @@
+# Dataset settings
+dataset = dict(
+    type="VariableVideoTextDataset",
+    transform_name="resize_crop",
+)
+
+# webvid
+bucket_config = {  # 20s/it
+    "144p": {1: (1.0, 475), 51: (1.0, 51), 102: (1.0, 27), 204: (1.0, 13), 408: (1.0, 6)},
+    # ---
+    "256": {1: (1.0, 297), 51: (0.5, 20), 102: (0.5, 10), 204: (0.5, 5), 408: ((0.5, 0.5), 2)},
+    "240p": {1: (1.0, 297), 51: (0.5, 20), 102: (0.5, 10), 204: (0.5, 5), 408: ((0.5, 0.4), 2)},
+    # ---
+    "360p": {1: (1.0, 141), 51: (0.5, 8), 102: (0.5, 4), 204: (0.5, 2), 408: ((0.5, 0.3), 1)},
+    "512": {1: (1.0, 141), 51: (0.5, 8), 102: (0.5, 4), 204: (0.5, 2), 408: ((0.5, 0.2), 1)},
+    # ---
+    "480p": {1: (1.0, 89), 51: (0.5, 5), 102: (0.5, 3), 204: ((0.5, 0.5), 1), 408: (0.0, None)},
+    # ---
+    "720p": {1: (0.3, 36), 51: (0.2, 2), 102: (0.1, 1), 204: (0.0, None)},
+    "1024": {1: (0.3, 36), 51: (0.1, 2), 102: (0.1, 1), 204: (0.0, None)},
+    # ---
+    "1080p": {1: (0.1, 5)},
+    # ---
+    "2048": {1: (0.05, 5)},
+}
+
+grad_checkpoint = True
+
+# Acceleration settings
+num_workers = 8
+num_bucket_build_workers = 16
+dtype = "bf16"
+plugin = "zero2"
+
+# Model settings
+model = dict(
+    type="STDiT3-XL/2",
+    from_pretrained=None,
+    qk_norm=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    freeze_y_embedder=True,
+)
+vae = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="/mnt/jfs/sora_checkpoints/vae-pipeline",
+    micro_frame_size=17,
+    micro_batch_size=4,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+    shardformer=True,
+    local_files_only=True,
+)
+scheduler = dict(
+    type="rflow",
+    use_timestep_transform=True,
+    sample_method="logit-normal",
+)
+
+# Mask settings
+# 25%
+mask_ratios = {
+    "random": 0.01,
+    "intepolate": 0.002,
+    "quarter_random": 0.002,
+    "quarter_head": 0.002,
+    "quarter_tail": 0.002,
+    "quarter_head_tail": 0.002,
+    "image_random": 0.0,
+    "image_head": 0.22,
+    "image_tail": 0.005,
+    "image_head_tail": 0.005,
+}
+
+# Log settings
+seed = 42
+outputs = "outputs"
+wandb = False
+epochs = 1000
+log_every = 10
+ckpt_every = 200
+
+# optimization settings
+load = None
+grad_clip = 1.0
+lr = 1e-4
+ema_decay = 0.99
+adam_eps = 1e-15
+warmup_steps = 1000

configs/opensora/inference/16x256x256.py

@@ -7,7 +7,7 @@ model = dict(
     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",
 )

configs/opensora/inference/16x512x512-rflow.py

@@ -0,0 +1,35 @@
+num_frames = 16
+fps = 24 // 3
+image_size = (512, 512)
+
+# Define model
+model = dict(
+    type="STDiT-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    from_pretrained="PRETRAINED_MODEL",
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+    micro_batch_size=2,
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+)
+scheduler = dict(
+    type="rflow",
+    num_sampling_steps=10,
+    cfg_scale=7.0,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 2
+seed = 42
+prompt_path = "./assets/texts/t2v_samples.txt"
+save_dir = "./outputs/samples/"

configs/opensora/inference/16x512x512.py

@@ -7,7 +7,7 @@ model = dict(
     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",
 )

configs/opensora/inference/64x512x512.py

@@ -7,7 +7,7 @@ model = dict(
     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",
 )

configs/opensora/train/16x256x256-mask.py

@@ -20,12 +20,12 @@ model = dict(
     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 = {
-    "mask_no": 0.7,
-    "mask_random": 0.15,
+    "identity": 0.7,
+    "random": 0.15,
     "mask_head": 0.05,
     "mask_tail": 0.05,
     "mask_head_tail": 0.05,

configs/opensora/train/16x256x256-spee-rflow.py

@@ -0,0 +1,64 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=16,
+    frame_interval=3,
+    image_size=(256, 256),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="STDiT-XL/2",
+    space_scale=0.5,
+    time_scale=1.0,
+    # 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_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+# mask_ratios = [0.5, 0.29, 0.07, 0.07, 0.07]
+# mask_ratios = {
+#     "identity": 0.9,
+#     "random": 0.06,
+#     "mask_head": 0.01,
+#     "mask_tail": 0.01,
+#     "mask_head_tail": 0.02,
+# }
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = True
+
+epochs = 1
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 16
+lr = 2e-5
+grad_clip = 1.0

configs/opensora/train/16x256x256-spee.py

@@ -20,12 +20,12 @@ model = dict(
     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 = {
-    "mask_no": 0.5,
-    "mask_random": 0.29,
+    "identity": 0.5,
+    "random": 0.29,
     "mask_head": 0.07,
     "mask_tail": 0.07,
     "mask_head_tail": 0.07,

configs/opensora/train/16x256x256.py

@@ -8,7 +8,7 @@ dataset = dict(
 )
 
 # Define acceleration
-num_workers = 4
+num_workers = 0
 dtype = "bf16"
 grad_checkpoint = True
 plugin = "zero2"
@@ -20,7 +20,7 @@ model = dict(
     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(

configs/opensora/train/16x512x512.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=1.0,
     time_scale=1.0,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/opensora/train/360x512x512.py

@@ -26,7 +26,7 @@ model = dict(
     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
 )

configs/opensora/train/64x512x512-sp.py

@@ -20,7 +20,7 @@ model = dict(
     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
 )

configs/opensora/train/64x512x512.py

@@ -20,7 +20,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/pixart/inference/1x20481B.py

@@ -0,0 +1,36 @@
+num_frames = 1
+fps = 1
+image_size = (2560, 1536)
+# image_size = (2048, 2048)
+
+model = dict(
+    type="PixArt-1B/2",
+    from_pretrained="PixArt-1B-2.pth",
+    space_scale=4,
+    no_temporal_pos_emb=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    base_size=2048 // 8,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="dpm-solver",
+    num_sampling_steps=14,
+    cfg_scale=4.5,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 1
+seed = 42
+prompt_path = "./assets/texts/t2i_sigma.txt"
+save_dir = "./samples/samples/"

configs/pixart/inference/1x2048MS.py

@@ -0,0 +1,36 @@
+num_frames = 1
+fps = 1
+image_size = (2560, 1536)
+# image_size = (2048, 2048)
+
+model = dict(
+    type="PixArt-XL/2",
+    from_pretrained="PixArt-Sigma-XL-2-2K-MS.pth",
+    space_scale=4,
+    no_temporal_pos_emb=True,
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+    base_size=2048 // 8,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="dpm-solver",
+    num_sampling_steps=14,
+    cfg_scale=4.5,
+)
+dtype = "bf16"
+
+# Others
+batch_size = 1
+seed = 42
+prompt_path = "./assets/texts/t2i_sigma.txt"
+save_dir = "./samples/samples/"

configs/pixart/inference/1x512x512-rflow.py

@@ -0,0 +1,39 @@
+num_frames = 1
+fps = 1
+image_size = (512, 512)
+
+# Define model
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PRETRAINED_MODEL",
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+)
+scheduler = dict(
+    type="rflow",
+    num_sampling_steps=20,
+    cfg_scale=7.0,
+)
+dtype = "bf16"
+
+# prompt_path = "./assets/texts/t2i_samples.txt"
+prompt = [
+    "Pirate ship trapped in a cosmic maelstrom nebula.",
+    "A small cactus with a happy face in the Sahara desert.",
+    "A small cactus with a sad face in the Sahara desert.",
+]
+
+# Others
+batch_size = 2
+seed = 42
+save_dir = "./outputs/samples2/"

configs/pixart/train/16x256x256.py

@@ -20,7 +20,7 @@ model = dict(
     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(

configs/pixart/train/1x2048x2048.py

@@ -0,0 +1,54 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path="/home/zhaowangbo/data/csv/image-v1_1_ext_noempty_rcp_clean_info.csv",
+    num_frames=1,
+    frame_interval=3,
+    image_size=(2048, 2048),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="PixArt-1B/2",
+    space_scale=4.0,
+    no_temporal_pos_emb=True,
+    from_pretrained="PixArt-1B-2.pth",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+    subfolder="vae",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=300,
+)
+scheduler = dict(
+    type="iddpm",
+    timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = False
+
+epochs = 1000
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 4
+lr = 2e-5
+grad_clip = 1.0

configs/pixart/train/1x512x512-rflow.py

@@ -0,0 +1,55 @@
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=1,
+    frame_interval=3,
+    image_size=(512, 512),
+)
+
+# Define acceleration
+num_workers = 4
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+sp_size = 1
+
+# Define model
+model = dict(
+    type="PixArt-XL/2",
+    space_scale=1.0,
+    time_scale=1.0,
+    no_temporal_pos_emb=True,
+    # from_pretrained="PixArt-XL-2-512x512.pth",
+    from_pretrained="PRETRAINED_MODEL",
+    enable_flash_attn=True,
+    enable_layernorm_kernel=True,
+)
+vae = dict(
+    type="VideoAutoencoderKL",
+    from_pretrained="stabilityai/sd-vae-ft-ema",
+)
+text_encoder = dict(
+    type="t5",
+    from_pretrained="DeepFloyd/t5-v1_1-xxl",
+    model_max_length=120,
+    shardformer=True,
+)
+scheduler = dict(
+    type="rflow",
+    # timestep_respacing="",
+)
+
+# Others
+seed = 42
+outputs = "outputs"
+wandb = True
+
+epochs = 2
+log_every = 10
+ckpt_every = 1000
+load = None
+
+batch_size = 64
+lr = 2e-5
+grad_clip = 1.0

configs/pixart/train/1x512x512.py

@@ -21,7 +21,7 @@ model = dict(
     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(

configs/pixart/train/64x512x512.py

@@ -21,7 +21,7 @@ model = dict(
     space_scale=1.0,
     time_scale=2 / 3,
     from_pretrained=None,
-    enable_flashattn=True,
+    enable_flash_attn=True,
     enable_layernorm_kernel=True,
 )
 vae = dict(

configs/vae/inference/image.py

@@ -0,0 +1,32 @@
+image_size = (256, 256)
+num_frames = 1
+
+dtype = "bf16"
+batch_size = 1
+seed = 42
+save_dir = "samples/vae_video"
+cal_stats = True
+log_stats_every = 100
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    image_size=image_size,
+)
+num_samples = 100
+num_workers = 4
+
+# Define model
+model = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="hpcai-tech/OpenSora-VAE-v1.2",
+    micro_frame_size=None,
+    micro_batch_size=4,
+    cal_loss=True,
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6

configs/vae/inference/video.py

@@ -0,0 +1,32 @@
+image_size = (256, 256)
+num_frames = 17
+
+dtype = "bf16"
+batch_size = 1
+seed = 42
+save_dir = "samples/vae_video"
+cal_stats = True
+log_stats_every = 100
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    image_size=image_size,
+)
+num_samples = 100
+num_workers = 4
+
+# Define model
+model = dict(
+    type="OpenSoraVAE_V1_2",
+    from_pretrained="hpcai-tech/OpenSora-VAE-v1.2",
+    micro_frame_size=None,
+    micro_batch_size=4,
+    cal_loss=True,
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6

configs/vae/train/stage1.py

@@ -0,0 +1,59 @@
+num_frames = 17
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=True,
+    from_pretrained=None,
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_strategy = "mixed_video_image"
+mixed_image_ratio = 0.2
+use_real_rec_loss = False
+use_z_rec_loss = True
+use_image_identity_loss = True
+
+# Others
+seed = 42
+outputs = "outputs/vae_stage1"
+wandb = False
+
+epochs = 100  # NOTE: adjust accordingly w.r.t dataset size
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0

configs/vae/train/stage2.py

@@ -0,0 +1,59 @@
+num_frames = 17
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=False,
+    from_pretrained="outputs/vae_stage1",
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_strategy = "mixed_video_image"
+mixed_image_ratio = 0.2
+use_real_rec_loss = False
+use_z_rec_loss = True
+use_image_identity_loss = False
+
+# Others
+seed = 42
+outputs = "outputs/vae_stage2"
+wandb = False
+
+epochs = 100  # NOTE: adjust accordingly w.r.t dataset size
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0

configs/vae/train/stage3.py

@@ -0,0 +1,58 @@
+num_frames = 33
+image_size = (256, 256)
+
+# Define dataset
+dataset = dict(
+    type="VideoTextDataset",
+    data_path=None,
+    num_frames=num_frames,
+    frame_interval=1,
+    image_size=image_size,
+)
+
+# Define acceleration
+num_workers = 16
+dtype = "bf16"
+grad_checkpoint = True
+plugin = "zero2"
+
+# Define model
+model = dict(
+    type="VideoAutoencoderPipeline",
+    freeze_vae_2d=False,
+    from_pretrained="outputs/vae_stage2",
+    cal_loss=True,
+    vae_2d=dict(
+        type="VideoAutoencoderKL",
+        from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers",
+        subfolder="vae",
+        local_files_only=True,
+    ),
+    vae_temporal=dict(
+        type="VAE_Temporal_SD",
+        from_pretrained=None,
+    ),
+)
+
+# loss weights
+perceptual_loss_weight = 0.1  # use vgg is not None and more than 0
+kl_loss_weight = 1e-6
+
+mixed_strategy = "mixed_video_random"
+use_real_rec_loss = True
+use_z_rec_loss = False
+use_image_identity_loss = False
+
+# Others
+seed = 42
+outputs = "outputs/vae_stage3"
+wandb = False
+
+epochs = 100  # NOTE: adjust accordingly w.r.t dataset size
+log_every = 1
+ckpt_every = 1000
+load = None
+
+batch_size = 1
+lr = 1e-5
+grad_clip = 1.0