diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000000_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000000_action.npy new file mode 100644 index 000000000..8a6ab2ab1 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000000_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000001_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000001_action.npy new file mode 100644 index 000000000..e649b6c72 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000001_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000002_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000002_action.npy new file mode 100644 index 000000000..3ee9a9c38 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000002_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000003_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000003_action.npy new file mode 100644 index 000000000..882381a54 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000003_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000004_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000004_action.npy new file mode 100644 index 000000000..9ce523037 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000004_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000005_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000005_action.npy new file mode 100644 index 000000000..02cd7ab2e Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000005_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000006_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000006_action.npy new file mode 100644 index 000000000..f06f6883b Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000006_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000007_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000007_action.npy new file mode 100644 index 000000000..b0f8a4d3e Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000007_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000008_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000008_action.npy new file mode 100644 index 000000000..8a754f020 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000008_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000009_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000009_action.npy new file mode 100644 index 000000000..b1b94d93d Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000009_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000010_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000010_action.npy new file mode 100644 index 000000000..e97638cfc Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000010_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000011_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000011_action.npy new file mode 100644 index 000000000..c4565bf73 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000011_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000012_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000012_action.npy new file mode 100644 index 000000000..0392aeeb8 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000012_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000013_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000013_action.npy new file mode 100644 index 000000000..95208043c Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000013_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000014_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000014_action.npy new file mode 100644 index 000000000..5a251d101 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000014_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000015_action.npy b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000015_action.npy new file mode 100644 index 000000000..a8c7fba01 Binary files /dev/null and b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/000015_action.npy differ diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/action/README.md b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/README.md new file mode 100644 index 000000000..cdddcfab6 --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/action/README.md @@ -0,0 +1,18 @@ +Total Files: 16 + +00. Hold [W] + Static +01. Hold [S] + Static +02. Hold [A] + Static +03. Hold [D] + Static +04. Hold [WA] + Static +05. Hold [WD] + Static +06. Hold [SA] + Static +07. Hold [SD] + Static +08. No Key + Hold [up] +09. No Key + Hold [down] +10. No Key + Hold [left] +11. No Key + Hold [right] +12. No Key + Hold [up_right] +13. No Key + Hold [up_left] +14. No Key + Hold [down_right] +15. No Key + Hold [down_left] diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/finetune_i2v.slurm b/examples/training/finetune/WanGame2.1_1.3b_i2v/finetune_i2v.slurm index 18a77d2da..fc4eb79d2 100644 --- a/examples/training/finetune/WanGame2.1_1.3b_i2v/finetune_i2v.slurm +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/finetune_i2v.slurm @@ -35,7 +35,7 @@ export HOME="/mnt/weka/home/hao.zhang/wei" MODEL_PATH="weizhou03/Wan2.1-Game-Fun-1.3B-InP-Diffusers" DATA_DIR="mc_wasd_10/preprocessed/combined_parquet_dataset" -VALIDATION_DATASET_FILE="mc_wasd_10/validation.json" +VALIDATION_DATASET_FILE="examples/training/finetune/WanGame2.1_1.3b_i2v/validation.json" # Configs NUM_GPUS=8 @@ -43,7 +43,7 @@ NUM_GPUS=8 training_args=( --tracker_project_name "wangame_1.3b_overfit" --output_dir "wangame_1.3b_overfit" - --max_train_steps 1500 + --max_train_steps 15000 --train_batch_size 1 --train_sp_batch_size 1 --gradient_accumulation_steps 1 @@ -88,8 +88,8 @@ validation_args=( optimizer_args=( --learning_rate 2e-5 --mixed_precision "bf16" - --weight_only_checkpointing_steps 1000 - --training_state_checkpointing_steps 1000 + --weight_only_checkpointing_steps 1000000 + --training_state_checkpointing_steps 10000000 --weight_decay 1e-4 --max_grad_norm 1.0 ) diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/generate_actions.py b/examples/training/finetune/WanGame2.1_1.3b_i2v/generate_actions.py new file mode 100644 index 000000000..edc3ad6dd --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/generate_actions.py @@ -0,0 +1,193 @@ +import os +import numpy as np + +# Configuration +SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +BASE_OUTPUT_DIR = os.path.join(SCRIPT_DIR, 'action') +VIDEO_OUTPUT_DIR = BASE_OUTPUT_DIR +os.makedirs(VIDEO_OUTPUT_DIR, exist_ok=True) + +FRAME_COUNT = 81 +CAM_VALUE = 0.1 + +# Action Mapping +KEY_TO_INDEX = { + 'W': 0, 'S': 1, 'A': 2, 'D': 3, +} + +VIEW_ACTION_TO_MOUSE = { + "stop": [0.0, 0.0], + "up": [CAM_VALUE, 0.0], + "down": [-CAM_VALUE, 0.0], + "left": [0.0, -CAM_VALUE], + "right": [0.0, CAM_VALUE], + "up_right": [CAM_VALUE, CAM_VALUE], + "up_left": [CAM_VALUE, -CAM_VALUE], + "down_right": [-CAM_VALUE, CAM_VALUE], + "down_left": [-CAM_VALUE, -CAM_VALUE], +} + +def get_multihot_vector(keys_str): + """Convert string like 'WA' to [1, 0, 1, 0, 0, 0]""" + vector = [0.0] * 6 + if not keys_str: + return vector + for char in keys_str.upper(): + if char in KEY_TO_INDEX: + vector[KEY_TO_INDEX[char]] = 1.0 + return vector + +def get_mouse_vector(view_str): + """Convert view string to [x, y]""" + return VIEW_ACTION_TO_MOUSE.get(view_str.lower(), [0.0, 0.0]) + +def generate_sequence(key_seq, mouse_seq): + """ + Generates action arrays based on sequences. + """ + keyboard_arr = np.zeros((FRAME_COUNT, 6), dtype=np.float32) + mouse_arr = np.zeros((FRAME_COUNT, 2), dtype=np.float32) + + mid_point = FRAME_COUNT // 2 + + # First Half + k_vec1 = get_multihot_vector(key_seq[0]) + m_vec1 = get_mouse_vector(mouse_seq[0]) + keyboard_arr[:mid_point] = k_vec1 + mouse_arr[:mid_point] = m_vec1 + + # Second Half + k_vec2 = get_multihot_vector(key_seq[1]) + m_vec2 = get_mouse_vector(mouse_seq[1]) + keyboard_arr[mid_point:] = k_vec2 + mouse_arr[mid_point:] = m_vec2 + + return keyboard_arr, mouse_arr + +def save_action(index, keyboard_arr, mouse_arr): + filename = f"{index:06d}_action.npy" + filepath = os.path.join(VIDEO_OUTPUT_DIR, filename) + + action_dict = { + 'keyboard': keyboard_arr, + 'mouse': mouse_arr + } + np.save(filepath, action_dict) + return filename + +def generate_description(key_seq, mouse_seq): + """Generates a human-readable string for the combination.""" + k1, k2 = key_seq + m1, m2 = mouse_seq + + # Format Keyboard Description + if not k1 and not k2: + k_desc = "No Key" + elif k1 == k2: + k_desc = f"Hold [{k1}]" + else: + k_desc = f"Switch [{k1}]->[{k2}]" + + # Format Mouse Description + if m1 == "stop" and m2 == "stop": + m_desc = "Static" + elif m1 == m2: + m_desc = f"Hold [{m1}]" + else: + m_desc = f"Switch [{m1}]->[{m2}]" + + return f"{k_desc} + {m_desc}" + +# ========================================== +# Main Generation Logic +# ========================================== + +configs = [] +readme_content = [] + +# Group 1: Constant Keyboard, No Mouse (0-7) +keys_basic = ['W', 'S', 'A', 'D', 'WA', 'WD', 'SA', 'SD'] +for k in keys_basic: + configs.append(((k, k), ("stop", "stop"))) + +# Group 2: No Keyboard, Constant Mouse (8-15) +mouse_basic = ['up', 'down', 'left', 'right', 'up_right', 'up_left', 'down_right', 'down_left'] +for m in mouse_basic: + configs.append((("", ""), (m, m))) + +# Group 3: Split Keyboard, No Mouse (16-23) +split_keys = [ + ('W', 'S'), ('S', 'W'), + ('A', 'D'), ('D', 'A'), + ('W', 'A'), ('W', 'D'), + ('S', 'A'), ('S', 'D') +] +for k1, k2 in split_keys: + configs.append(((k1, k2), ("stop", "stop"))) + +# Group 4: No Keyboard, Split Mouse (24-31) +split_mouse = [ + ('left', 'right'), ('right', 'left'), + ('up', 'down'), ('down', 'up'), + ('up_left', 'up_right'), ('up_right', 'up_left'), + ('left', 'up'), ('right', 'down') +] +for m1, m2 in split_mouse: + configs.append((("", ""), (m1, m2))) + +# Group 5: Constant Keyboard + Constant Mouse (32-47) +combo_keys = ['W', 'S', 'W', 'S', 'A', 'D', 'WA', 'WD', 'W', 'S', 'W', 'S', 'A', 'D', 'WA', 'WD'] +combo_mice = ['left', 'left', 'right', 'right', 'up', 'up', 'down', 'down', 'up_left', 'up_left', 'up_right', 'up_right', 'down_left', 'down_right', 'right', 'left'] +for i in range(16): + configs.append(((combo_keys[i], combo_keys[i]), (combo_mice[i], combo_mice[i]))) + +# Group 6: Constant Keyboard, Split Mouse (48-55) +complex_1_keys = ['W'] * 8 +complex_1_mice = [ + ('left', 'right'), ('right', 'left'), + ('up', 'down'), ('down', 'up'), + ('left', 'up'), ('right', 'up'), + ('left', 'down'), ('right', 'down') +] +for i in range(8): + configs.append(((complex_1_keys[i], complex_1_keys[i]), complex_1_mice[i])) + +# Group 7: Split Keyboard, Constant Mouse (56-63) +complex_2_keys = [ + ('W', 'S'), ('S', 'W'), + ('A', 'D'), ('D', 'A'), + ('W', 'A'), ('W', 'D'), + ('S', 'A'), ('S', 'D') +] +complex_2_mouse = 'up' +for k1, k2 in complex_2_keys: + configs.append(((k1, k2), (complex_2_mouse, complex_2_mouse))) + + +# Execution +print(f"Preparing to generate {len(configs)} action files...") + +for i, (key_seq, mouse_seq) in enumerate(configs): + if i >= 16: break + + # Generate Data + kb_arr, ms_arr = generate_sequence(key_seq, mouse_seq) + filename = save_action(i, kb_arr, ms_arr) + + # Generate Description for README + description = generate_description(key_seq, mouse_seq) + readme_entry = f"{i:02d}. {description}" + readme_content.append(readme_entry) + + print(f"Generated {filename} -> {description}") + +# Write README +readme_path = os.path.join(BASE_OUTPUT_DIR, 'README.md') +with open(readme_path, 'w', encoding='utf-8') as f: + f.write(f"Total Files: {len(readme_content)}\n\n") + for line in readme_content: + f.write(line + '\n') + +print(f"\nProcessing complete.") +print(f"64 .npy files generated in {VIDEO_OUTPUT_DIR}") +print(f"Manifest saved to {readme_path}") \ No newline at end of file diff --git a/examples/training/finetune/WanGame2.1_1.3b_i2v/validation.json b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation.json new file mode 100644 index 000000000..5af4cca74 --- /dev/null +++ b/examples/training/finetune/WanGame2.1_1.3b_i2v/validation.json @@ -0,0 +1,404 @@ +{ + "data": [ + { + "caption": "0", + "image_path": "../../../../mc_wasd_10/validate/000000.jpg", + "action_path": "../../../../mc_wasd_10/videos/000000_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "1", + "image_path": "../../../../mc_wasd_10/validate/000001.jpg", + "action_path": "../../../../mc_wasd_10/videos/000001_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "2", + "image_path": "../../../../mc_wasd_10/validate/000002.jpg", + "action_path": "../../../../mc_wasd_10/videos/000002_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "3", + "image_path": "../../../../mc_wasd_10/validate/000003.jpg", + "action_path": "../../../../mc_wasd_10/videos/000003_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "4", + "image_path": "../../../../mc_wasd_10/validate/000004.jpg", + "action_path": "../../../../mc_wasd_10/videos/000004_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "5", + "image_path": "../../../../mc_wasd_10/validate/000005.jpg", + "action_path": "../../../../mc_wasd_10/videos/000005_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "6", + "image_path": "../../../../mc_wasd_10/validate/000006.jpg", + "action_path": "../../../../mc_wasd_10/videos/000006_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "7", + "image_path": "../../../../mc_wasd_10/validate/000007.jpg", + "action_path": "../../../../mc_wasd_10/videos/000007_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "00. Hold [W] + Static", + "image_path": "../../../../mc_wasd_10/validate/000000.jpg", + "action_path": "action/000000_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01. Hold [S] + Static", + "image_path": "../../../../mc_wasd_10/validate/000001.jpg", + "action_path": "action/000001_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02. Hold [A] + Static", + "image_path": "../../../../mc_wasd_10/validate/000002.jpg", + "action_path": "action/000002_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03. Hold [D] + Static", + "image_path": "../../../../mc_wasd_10/validate/000003.jpg", + "action_path": "action/000003_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "04. Hold [WA] + Static", + "image_path": "../../../../mc_wasd_10/validate/000004.jpg", + "action_path": "action/000004_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "05. Hold [WD] + Static", + "image_path": "../../../../mc_wasd_10/validate/000005.jpg", + "action_path": "action/000005_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "06. Hold [SA] + Static", + "image_path": "../../../../mc_wasd_10/validate/000006.jpg", + "action_path": "action/000006_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "07. Hold [SD] + Static", + "image_path": "../../../../mc_wasd_10/validate/000007.jpg", + "action_path": "action/000007_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "08. No Key + Hold [up]", + "image_path": "../../../../mc_wasd_10/validate/000000.jpg", + "action_path": "action/000008_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "09. No Key + Hold [down]", + "image_path": "../../../../mc_wasd_10/validate/000001.jpg", + "action_path": "action/000009_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "10. No Key + Hold [left]", + "image_path": "../../../../mc_wasd_10/validate/000002.jpg", + "action_path": "action/000010_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "11. No Key + Hold [right]", + "image_path": "../../../../mc_wasd_10/validate/000003.jpg", + "action_path": "action/000011_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "12. No Key + Hold [up_right]", + "image_path": "../../../../mc_wasd_10/validate/000004.jpg", + "action_path": "action/000012_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "13. No Key + Hold [up_left]", + "image_path": "../../../../mc_wasd_10/validate/000005.jpg", + "action_path": "action/000013_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "14. No Key + Hold [down_right]", + "image_path": "../../../../mc_wasd_10/validate/000006.jpg", + "action_path": "action/000014_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "15. No Key + Hold [down_left]", + "image_path": "../../../../mc_wasd_10/validate/000007.jpg", + "action_path": "action/000015_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "00. Hold [W] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000000.jpg", + "action_path": "action/000000_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "01. Hold [S] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000001.jpg", + "action_path": "action/000001_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "02. Hold [A] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000002.jpg", + "action_path": "action/000002_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "03. Hold [D] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000003.jpg", + "action_path": "action/000003_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "04. Hold [WA] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000004.jpg", + "action_path": "action/000004_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "05. Hold [WD] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000005.jpg", + "action_path": "action/000005_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "06. Hold [SA] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000006.jpg", + "action_path": "action/000006_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "07. Hold [SD] + Static", + "image_path": "../../../../mc_wasd_10/validate/gen_000007.jpg", + "action_path": "action/000007_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "08. No Key + Hold [up]", + "image_path": "../../../../mc_wasd_10/validate/gen_000000.jpg", + "action_path": "action/000008_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "09. No Key + Hold [down]", + "image_path": "../../../../mc_wasd_10/validate/gen_000001.jpg", + "action_path": "action/000009_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "10. No Key + Hold [left]", + "image_path": "../../../../mc_wasd_10/validate/gen_000002.jpg", + "action_path": "action/000010_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "11. No Key + Hold [right]", + "image_path": "../../../../mc_wasd_10/validate/gen_000003.jpg", + "action_path": "action/000011_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "12. No Key + Hold [up_right]", + "image_path": "../../../../mc_wasd_10/validate/gen_000004.jpg", + "action_path": "action/000012_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "13. No Key + Hold [up_left]", + "image_path": "../../../../mc_wasd_10/validate/gen_000005.jpg", + "action_path": "action/000013_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "14. No Key + Hold [down_right]", + "image_path": "../../../../mc_wasd_10/validate/gen_000006.jpg", + "action_path": "action/000014_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + }, + { + "caption": "15. No Key + Hold [down_left]", + "image_path": "../../../../mc_wasd_10/validate/gen_000007.jpg", + "action_path": "action/000015_action.npy", + "video_path": null, + "num_inference_steps": 40, + "height": 352, + "width": 640, + "num_frames": 77 + } + ] +} \ No newline at end of file diff --git a/fastvideo/dataset/validation_dataset.py b/fastvideo/dataset/validation_dataset.py index 5ab467d75..8d831d5df 100644 --- a/fastvideo/dataset/validation_dataset.py +++ b/fastvideo/dataset/validation_dataset.py @@ -4,6 +4,7 @@ import os import pathlib import datasets +import numpy as np from torch.utils.data import IterableDataset from fastvideo.distributed import (get_sp_world_size, get_world_rank, @@ -160,5 +161,25 @@ class ValidationDataset(IterableDataset): else: sample["control_video"] = load_video(control_video_path) + if sample.get("action_path", None) is not None: + action_path = sample["action_path"] + action_path = os.path.join(self.dir, action_path) + sample["action_path"] = action_path + if not pathlib.Path(action_path).is_file(): + logger.warning("Action file %s does not exist.", action_path) + else: + try: + action_data = np.load(action_path, allow_pickle=True) + num_frames = sample["num_frames"] + if action_data.dtype == object: action_data = action_data.item() + if isinstance(action_data, dict): + sample["keyboard_cond"] = action_data["keyboard"][:num_frames] + sample["mouse_cond"] = action_data["mouse"][:num_frames] + else: + sample["keyboard_cond"] = action_data[:num_frames] + except Exception as e: + logger.error("Error loading action file %s: %s", + action_path, e) + sample = {k: v for k, v in sample.items() if v is not None} yield sample