[wip] debug extract feat

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

@@ -0,0 +1,92 @@
+# 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
+
+# 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,
+    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

opensora/datasets/datasets.py

@@ -119,10 +119,12 @@ class VariableVideoTextDataset(VideoTextDataset):
         frame_interval=1,
         image_size=(None, None),
         transform_name=None,
+        dummy_text_feature=False,
     ):
         super().__init__(data_path, num_frames, frame_interval, image_size, transform_name=None)
         self.transform_name = transform_name
         self.data["id"] = np.arange(len(self.data))
+        self.dummy_text_feature = dummy_text_feature
 
     def get_data_info(self, index):
         T = self.data.iloc[index]["num_frames"]
@@ -176,6 +178,9 @@ class VariableVideoTextDataset(VideoTextDataset):
         }
         if self.get_text:
             ret["text"] = sample["text"]
+        if self.dummy_text_feature:
+            ret["text"] = "dummy text"
+            ret["mask"] = None
         return ret
 
     def __getitem__(self, index):

opensora/models/stdit/stdit3.py

@@ -373,6 +373,7 @@ class STDiT3(PreTrainedModel):
             y_lens = mask.tolist()
         else:
             y, y_lens = self.encode_text(y, mask)
+        breakpoint()
 
         # === get x embed ===
         x = self.x_embedder(x)  # [B, N, C]

scripts/misc/search_bs.py

@@ -13,8 +13,8 @@ from tqdm import tqdm
 
 from opensora.acceleration.checkpoint import set_grad_checkpoint
 from opensora.acceleration.parallel_states import get_data_parallel_group
-from opensora.datasets import prepare_variable_dataloader
 from opensora.datasets.aspect import get_num_frames
+from opensora.datasets.dataloader import prepare_dataloader
 from opensora.registry import DATASETS, MODELS, SCHEDULERS, build_module
 from opensora.utils.ckpt_utils import model_sharding
 from opensora.utils.config_utils import parse_configs
@@ -155,11 +155,11 @@ def main():
             pin_memory=True,
             process_group=get_data_parallel_group(),
         )
-        dataloader = prepare_variable_dataloader(
+        dataloader, sampler = prepare_dataloader(
             bucket_config=bucket_config,
             **dataloader_args,
         )
-        num_batch = dataloader.batch_sampler.get_num_batch()
+        num_batch = sampler.get_num_batch()
         num_steps_per_epoch = num_batch // dist.get_world_size()
 
         dataloader_iter = iter(dataloader)

scripts/train.py

@@ -226,9 +226,15 @@ def main():
                 # == visual and text encoding ==
                 with torch.no_grad():
                     # Prepare visual inputs
-                    x = vae.encode(x)  # [B, C, T, H/P, W/P]
+                    if cfg.get("load_video_features", False):
+                        x = x.to(device, dtype)
+                    else:
+                        x = vae.encode(x)  # [B, C, T, H/P, W/P]
                     # Prepare text inputs
-                    model_args = text_encoder.encode(y)
+                    if cfg.get("load_text_features", False):
+                        model_args = {"y": y.to(device, dtype), "mask": batch.pop("mask").to(device, dtype)}
+                    else:
+                        model_args = text_encoder.encode(y)
 
                 # == mask ==
                 mask = None