add examples for openvid1M dataset, update README.md, remove .yaml files

Author: Huy Vu2

examples/diffusion/recipes/wan/README.md

@@ -74,6 +74,8 @@ energon prepare "${DATASET_PATH}"
 
 In the training config, point `dataset.path` to the processed data output directory: `dataset.path=${DATASET_PATH}`.
 
+**Note**: We provide an example instruction to process and train a small portion of OpenVid-1M dataset, at `examples/diffusion/recipes/wan/prepare_dataset/openvid1M_dataset`.
+
 ---
 
 ## 2. Checkpoint Conversion
@@ -136,41 +138,33 @@ This recipe leverages sequence packing to maximize throughput. When a batch cont
 
 ### Training mode
 
-WAN uses different flow-matching hyperparameters for pretraining vs fine-tuning. Select the appropriate step function via `--step_func`:
-- `wan_pretrain_step`: logit-normal sampling, higher `logit_std`, lower `flow_shift` — for stability and broad learning
-- `wan_finetune_step`: uniform sampling, lower `logit_std`, higher `flow_shift` — to refine details and improve quality
+WAN uses different flow-matching hyperparameters for pretraining vs fine-tuning. Always pass `--step_func wan_step`; the mode is automatically inferred from the recipe name:
+- **pretrain** (recipe contains `pretrain`): logit-normal sampling, higher `logit_std` (1.5), lower `flow_shift` (2.5) — for stability and broad learning
+- **finetune** (recipe contains `finetune`/`sft`/`peft`): uniform sampling, lower `logit_std` (1.0), higher `flow_shift` (3.0) — to refine details and improve quality
 
 ### WAN 1.3B — Mock data (no dataset path):
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
   --recipe wan_1_3B_pretrain_config \
-  --step_func wan_pretrain_step
+  --step_func wan_step
 ```
 
 ### WAN 1.3B — Real data (WebDataset path):
 
 ```bash
 uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
   --recipe wan_1_3B_pretrain_config \
-  --step_func wan_pretrain_step \
+  --step_func wan_step \
   dataset.path=${WORKSPACE}/datasets/wan
 ```
 
-### WAN 14B — Mock data:
-
-```bash
-uv run torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
-  --recipe wan_14B_pretrain_config \
-  --step_func wan_pretrain_step
-```
-
 ### With CLI overrides (iters, LR, batch size, parallelism, etc.):
 
 ```bash
 uv run torchrun --nproc_per_node=$NUM_GPUS scripts/training/run_recipe.py \
   --recipe wan_1_3B_pretrain_config \
-  --step_func wan_pretrain_step \
+  --step_func wan_step \
   dataset.path=${WORKSPACE}/datasets/wan \
   train.global_batch_size=8 \
   train.micro_batch_size=1 \

examples/diffusion/recipes/wan/conf/wan_14B.yaml

@@ -0,0 +1,110 @@
+# Preparing OpenVid-1M Dataset for Wan Fine-tuning
+
+This guide walks through downloading a subset of the [OpenVid-1M](https://huggingface.co/datasets/nkp37/OpenVid-1M) dataset, preprocessing it, and converting it into the Energon WebDataset format for Wan training.
+
+## Prerequisites
+
+- `wget`, `unzip` available on your system
+- A Hugging Face account with access token (`HF_TOKEN`)
+- Python environment with Megatron-Bridge installed
+
+## Step 1: Download a Small Portion of OpenVid-1M
+
+```bash
+PROCESSED_DATA_PATH=<path/to/dataset_dir>
+cd ${PROCESSED_DATA_PATH}
+
+export HF_TOKEN=<HF_API_KEY>
+wget https://huggingface.co/datasets/nkp37/OpenVid-1M/resolve/main/OpenVidHD/OpenVidHD_part_1.zip
+wget https://huggingface.co/datasets/nkp37/OpenVid-1M/resolve/main/data/train/OpenVidHD.csv
+```
+
+## Step 2: Extract a Subset of Videos
+
+Extract 200 videos as an example (adjust the number as needed):
+
+```bash
+unzip -Z1 OpenVidHD_part_1.zip | head -n 200 | xargs -I {} unzip OpenVidHD_part_1.zip "{}" -d OpenVidHD_part_1
+```
+
+## Step 3: Create Per-Video Caption JSON Files
+
+The preprocessing script reads the CSV caption file and creates a sidecar `.json` file next to each `.mp4`:
+
+```bash
+cd ${MBRIDGE_PATH}
+python examples/diffusion/recipes/wan/prepare_dataset/openvid1M_preprocess.py \
+  --csv_path ${PROCESSED_DATA_PATH}/OpenVidHD.csv \
+  --video_dir ${PROCESSED_DATA_PATH}/OpenVidHD_part_1
+```
+
+Each video will have a corresponding `.json` file with `video` and `caption` fields.
+
+## Step 4: Convert to Energon WebDataset Format
+
+Run the dataset preparation script using `torchrun` with 8 GPUs:
+
+```bash
+DATASET_SRC=${PROCESSED_DATA_PATH}/OpenVidHD_part_1
+DATASET_PATH=${PROCESSED_DATA_PATH}/prepared_dataset_wds
+
+torchrun --nproc_per_node=8 \
+  examples/diffusion/recipes/wan/prepare_dataset/prepare_dataset_wan.py \
+  --video_folder "${DATASET_SRC}" \
+  --output_dir "${DATASET_PATH}" \
+  --output_format energon \
+  --model "Wan-AI/Wan2.1-T2V-1.3B-Diffusers" \
+  --mode video \
+  --height 480 \
+  --width 832 \
+  --resize_mode bilinear \
+  --center-crop
+```
+
+The output directory `${DATASET_PATH}` can then be passed directly as `dataset.path` when launching Wan training.
+
+## Step 5: Launch Wan Fine-tuning
+
+Set the required environment variables and run the training script:
+
+```bash
+DATASET_PATH=${PROCESSED_DATA_PATH}/prepared_dataset_wds
+CHECKPOINT_DIR=<path/to/save/checkpoints>
+INIT_CHECKPOINT_DIR=<path/to/pretrained/wan/checkpoint>
+EXP_NAME=<experiment_name>
+
+NVTE_FUSED_ATTN=1 torchrun --nproc_per_node=8 scripts/training/run_recipe.py \
+    --recipe wan_1_3B_pretrain_config \
+    --step_func wan_step \
+    model.tensor_model_parallel_size=1 \
+    model.pipeline_model_parallel_size=1 \
+    model.context_parallel_size=4 \
+    model.sequence_parallel=false \
+    model.qkv_format=thd \
+    model.bias_activation_fusion=false \
+    dataset.path=${DATASET_PATH} \
+    dataset.packing_buffer_size=200 \
+    dataset.num_workers=10 \
+    checkpoint.save=${CHECKPOINT_DIR} \
+    checkpoint.load=${INIT_CHECKPOINT_DIR} \
+    checkpoint.load_optim=false \
+    checkpoint.save_interval=200 \
+    optimizer.lr=5e-5 \
+    optimizer.min_lr=5e-5 \
+    train.eval_iters=0 \
+    scheduler.lr_decay_style=constant \
+    scheduler.lr_warmup_iters=1000 \
+    model.seq_length=40000 \
+    dataset.seq_length=40000 \
+    train.global_batch_size=8 \
+    train.micro_batch_size=1 \
+    dataset.global_batch_size=8 \
+    dataset.micro_batch_size=1 \
+    logger.log_interval=1 \
+    logger.wandb_project="wan" \
+    logger.wandb_exp_name=${EXP_NAME} \
+    logger.wandb_save_dir=${CHECKPOINT_DIR}
+```
+
+> **Note:** With a small dataset (e.g. 200 videos), set `dataset.num_workers=1` to avoid an
+> `IndexError` in the Energon sharder caused by too few samples in the val split shard.

examples/diffusion/recipes/wan/conf/wan_1_3B.yaml

@@ -0,0 +1,52 @@
+import argparse
+import pandas as pd
+import json
+import os
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--csv_path", default="OpenVidHD.csv")
+    parser.add_argument("--video_dir", default="OpenVidHD_part_1")
+    args = parser.parse_args()
+
+    csv_path = args.csv_path
+    video_dir = args.video_dir
+
+    # 1. Load the CSV into a dictionary for O(1) lookup speed
+    print(f"Loading {csv_path}...")
+    df = pd.read_csv(csv_path)
+    
+    # Create a mapping of filename -> caption
+    # We strip any whitespace just in case
+    caption_map = dict(zip(df['video'].str.strip(), df['caption']))
+
+    # 2. List all mp4 files in the directory
+    video_files = [f for f in os.listdir(video_dir) if f.endswith('.mp4')]
+    print(f"Found {len(video_files)} videos in {video_dir}. Starting conversion...")
+
+    count = 0
+    for video_name in video_files:
+        if video_name in caption_map:
+            # Construct the JSON content
+            content = {
+                "video": video_name,
+                "caption": caption_map[video_name]
+            }
+            
+            # Create the output filename (replace .mp4 with .json)
+            json_name = os.path.splitext(video_name)[0] + '.json'
+            json_path = os.path.join(video_dir, json_name)
+            
+            # Write the individual JSON file
+            with open(json_path, 'w', encoding='utf-8') as f:
+                json.dump(content, f, indent=4, ensure_ascii=False)
+            
+            count += 1
+        else:
+            print(f"Warning: No caption found for {video_name}")
+
+    print(f"Finished! Created {count} JSON files in {video_dir}.")
+
+if __name__ == "__main__":
+    main()