已合并
[Refactor] Dataset Refactor: Modify the configuration file data.json and delete redundant code #1118
zs-derrick创建于 2025年6月16日
[Refactor] Dataset Refactor: Modify the configuration file data.json and delete redundant code #1118
已合并
从refs/pull/1118/head合入到master
共 51 个文件变更+368-508
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "./examples/hunyuanvideo/feature_extract/data.txt", | 5 | "data_path": "./examples/hunyuanvideo/feature_extract/data.txt", |
| 7 | "data_storage_mode": "combine" | 6 | "data_storage_mode": "combine" |
| @@ -67,7 +66,7 @@ | |||
| 67 | "dataloader_param": { | 66 | "dataloader_param": { |
| 68 | "dataloader_mode": "sampler", | 67 | "dataloader_mode": "sampler", |
| 69 | "sampler_type": "LengthGroupedSampler", | 68 | "sampler_type": "LengthGroupedSampler", |
| 70 | - "shuffle": false, | 69 | + "shuffle": true, |
| 71 | "drop_last": true, | 70 | "drop_last": true, |
| 72 | "pin_memory": true, | 71 | "pin_memory": true, |
| 73 | "group_frame": false, | 72 | "group_frame": false, |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "./examples/hunyuanvideo/feature_extract/data.txt", | 5 | "data_path": "./examples/hunyuanvideo/feature_extract/data.txt", |
| 7 | "data_storage_mode": "combine" | 6 | "data_storage_mode": "combine" |
| @@ -70,7 +69,7 @@ | |||
| 70 | "dataloader_param": { | 69 | "dataloader_param": { |
| 71 | "dataloader_mode": "sampler", | 70 | "dataloader_mode": "sampler", |
| 72 | "sampler_type": "LengthGroupedSampler", | 71 | "sampler_type": "LengthGroupedSampler", |
| 73 | - "shuffle": false, | 72 | + "shuffle": true, |
| 74 | "drop_last": true, | 73 | "drop_last": true, |
| 75 | "pin_memory": true, | 74 | "pin_memory": true, |
| 76 | "group_frame": false, | 75 | "group_frame": false, |
| @@ -115,7 +115,7 @@ def extract_feature(): | |||
| 115 | }) | 115 | }) |
| 116 | elif data_storage_mode == "sorafeatured": | 116 | elif data_storage_mode == "sorafeatured": |
| 117 | data_info.update({ | 117 | data_info.update({ |
| 118 | - 'path': f"features/{pt_name}" | 118 | + FILE_INFO: f"features/{pt_name}" |
| 119 | }) | 119 | }) |
| 120 | json_file.write(json.dumps(data_info) + '\n') | 120 | json_file.write(json.dumps(data_info) + '\n') |
| 121 | 121 | ||
| @@ -185,9 +185,9 @@ def extract_feature(): | |||
| 185 | prompt_masks_i = prompt_mask[i] | 185 | prompt_masks_i = prompt_mask[i] |
| 186 | 186 | ||
| 187 | data_to_save = { | 187 | data_to_save = { |
| 188 | - "latents": latent_i, | 188 | + VIDEO: latent_i, |
| 189 | - "prompt": prompts_i, | 189 | + PROMPT_IDS: prompts_i, |
| 190 | - "prompt_mask": prompt_masks_i | 190 | + PROMPT_MASK: prompt_masks_i |
| 191 | } | 191 | } |
| 192 | if latents_dict: | 192 | if latents_dict: |
| 193 | for key in latents_dict.keys(): | 193 | for key in latents_dict.keys(): |
| @@ -1 +0,0 @@ | |||
| 1 | -./hunyuanvideo/sora_featuers/,./hunyuanvideo/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/hunyuanvideo/i2v/data.txt", | 5 | + "data_path": "./hunyuanvideo/sora_featuers/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./hunyuanvideo/sora_featuers/" |
ZC 将data.txt的数据配置移到data.json中，readme是否要做同步修改   | |||
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -71,7 +68,7 @@ | |||
| 71 | "dataloader_param": { | 68 | "dataloader_param": { |
| 72 | "dataloader_mode": "sampler", | 69 | "dataloader_mode": "sampler", |
| 73 | "sampler_type": "LengthGroupedSampler", | 70 | "sampler_type": "LengthGroupedSampler", |
| 74 | - "shuffle": false, | 71 | + "shuffle": true, |
| 75 | "drop_last": true, | 72 | "drop_last": true, |
| 76 | "pin_memory": true, | 73 | "pin_memory": true, |
| 77 | "group_frame": false, | 74 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./hunyuanvideo/sora_featuers/,./hunyuanvideo/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/hunyuanvideo/t2v/data.txt", | 5 | + "data_path": "./hunyuanvideo/sora_featuers/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./hunyuanvideo/sora_featuers/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -69,7 +66,7 @@ | |||
| 69 | "dataloader_param": { | 66 | "dataloader_param": { |
| 70 | "dataloader_mode": "sampler", | 67 | "dataloader_mode": "sampler", |
| 71 | "sampler_type": "LengthGroupedSampler", | 68 | "sampler_type": "LengthGroupedSampler", |
| 72 | - "shuffle": false, | 69 | + "shuffle": true, |
| 73 | "drop_last": true, | 70 | "drop_last": true, |
| 74 | "pin_memory": true, | 71 | "pin_memory": true, |
| 75 | "group_frame": false, | 72 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./hunyuanvideo/sora_featuers/,./hunyuanvideo/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/hunyuanvideo/t2v_A3/data.txt", | 5 | + "data_path": "./hunyuanvideo/sora_featuers/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./hunyuanvideo/sora_featuers/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -69,7 +66,7 @@ | |||
| 69 | "dataloader_param": { | 66 | "dataloader_param": { |
| 70 | "dataloader_mode": "sampler", | 67 | "dataloader_mode": "sampler", |
| 71 | "sampler_type": "LengthGroupedSampler", | 68 | "sampler_type": "LengthGroupedSampler", |
| 72 | - "shuffle": false, | 69 | + "shuffle": true, |
| 73 | "drop_last": true, | 70 | "drop_last": true, |
| 74 | "pin_memory": true, | 71 | "pin_memory": true, |
| 75 | "group_frame": false, | 72 | "group_frame": false, |
| @@ -1,11 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | "basic_parameters": { | 4 | "basic_parameters": { |
| 5 | - "data_path": "examples/opensoraplan1.3/t2v_A3/feature_data.txt", | 5 | + "data_path": "./osp/sora_featuers/data.jsonl", |
| 6 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./osp/sora_featuers/" |
| 7 | - "use_text_feature": true, | ||
| 8 | - "use_video_feature": true | ||
| 9 | }, | 7 | }, |
| 10 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 11 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -1 +0,0 @@ | |||
| 1 | -/home/dataset/osp_feature/sora_features,/home/dataset/osp_feature/sora_features/data.jsonl | ||
| @@ -113,7 +113,7 @@ def extract_feature(): | |||
| 113 | }) | 113 | }) |
| 114 | elif data_storage_mode == "sorafeatured": | 114 | elif data_storage_mode == "sorafeatured": |
| 115 | data_info.update({ | 115 | data_info.update({ |
| 116 | - 'path': f"features/{pt_name}" | 116 | + FILE_INFO: f"features/{pt_name}" |
| 117 | }) | 117 | }) |
| 118 | json_file.write(json.dumps(data_info) + '\n') | 118 | json_file.write(json.dumps(data_info) + '\n') |
| 119 | 119 | ||
| @@ -183,10 +183,10 @@ def extract_feature(): | |||
| 183 | prompt_masks_i = prompt_mask[i] | 183 | prompt_masks_i = prompt_mask[i] |
| 184 | 184 | ||
| 185 | data_to_save = { | 185 | data_to_save = { |
| 186 | - "latents": latent_i, | 186 | + VIDEO: latent_i, |
| 187 | - "prompt": prompts_i, | 187 | + PROMPT_IDS: prompts_i, |
| 188 | - "prompt_mask": prompt_masks_i, | 188 | + PROMPT_MASK: prompt_masks_i, |
| 189 | - "video_mask": video_mask | 189 | + VIDEO_MASK: video_mask |
| 190 | } | 190 | } |
| 191 | pt_name = get_pt_name(file_names[i]) | 191 | pt_name = get_pt_name(file_names[i]) |
| 192 | torch.save(data_to_save, os.path.join(save_path, "features", pt_name)) | 192 | torch.save(data_to_save, os.path.join(save_path, "features", pt_name)) |
| @@ -83,7 +83,7 @@ | |||
| 83 | "dataloader_param": { | 83 | "dataloader_param": { |
| 84 | "dataloader_mode": "sampler", | 84 | "dataloader_mode": "sampler", |
| 85 | "sampler_type": "LengthGroupedSampler", | 85 | "sampler_type": "LengthGroupedSampler", |
| 86 | - "shuffle": false, | 86 | + "shuffle": true, |
| 87 | "drop_last": true, | 87 | "drop_last": true, |
| 88 | "pin_memory": true, | 88 | "pin_memory": true, |
| 89 | "group_frame": false, | 89 | "group_frame": false, |
| @@ -121,7 +121,7 @@ def extract_feature(): | |||
| 121 | }) | 121 | }) |
| 122 | elif data_storage_mode == "sorafeatured": | 122 | elif data_storage_mode == "sorafeatured": |
| 123 | data_info.update({ | 123 | data_info.update({ |
| 124 | - 'path': f"features/{pt_name}" | 124 | + FILE_INFO: f"features/{pt_name}" |
| 125 | }) | 125 | }) |
| 126 | json_file.write(json.dumps(data_info) + '\n') | 126 | json_file.write(json.dumps(data_info) + '\n') |
| 127 | 127 | ||
| @@ -189,9 +189,9 @@ def extract_feature(): | |||
| 189 | prompt_masks_i = prompt_mask[i] | 189 | prompt_masks_i = prompt_mask[i] |
| 190 | 190 | ||
| 191 | data_to_save = { | 191 | data_to_save = { |
| 192 | - "latents": latent_i, | 192 | + VIDEO: latent_i, |
| 193 | - "prompt": prompts_i, | 193 | + PROMPT_IDS: prompts_i, |
| 194 | - "prompt_mask": prompt_masks_i | 194 | + PROMPT_MASK: prompt_masks_i |
| 195 | } | 195 | } |
| 196 | if latents_dict: | 196 | if latents_dict: |
| 197 | for key in latents_dict.keys(): | 197 | for key in latents_dict.keys(): |
| @@ -1 +0,0 @@ | |||
| 1 | -./stepvideo/sora_featuers/,./stepvideo/sora_featuers/data.jsonl | ||
| @@ -1,11 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | "basic_parameters": { | 4 | "basic_parameters": { |
| 5 | - "data_path": "examples/stepvideo/i2v/data.txt", | 5 | + "data_path": "./stepvideo/sora_featuers/data.jsonl", |
| 6 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./stepvideo/sora_featuers/" |
| 7 | - "use_text_feature": true, | ||
| 8 | - "use_video_feature": true | ||
| 9 | }, | 7 | }, |
| 10 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 11 | "video_processor_type": "CogVideoXProcessor", | 9 | "video_processor_type": "CogVideoXProcessor", |
| @@ -1 +0,0 @@ | |||
| 1 | -./stepvideo/sora_featuers/,./stepvideo/sora_featuers/data.jsonl | ||
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/data_path/data.jsonl", | 5 | "data_path": "/data_path/data.jsonl", |
| 7 | "data_folder": "/data_path", | 6 | "data_folder": "/data_path", |
| @@ -1,12 +1,10 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "basic_parameters": { | 4 | + "basic_parameters": { |
| 5 | - "data_path": "examples/stepvideo/t2v/data.txt", | 5 | + "data_path": "./stepvideo/sora_featuers/data.jsonl", |
| 6 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./stepvideo/sora_featuers/" |
| 7 | - "use_text_feature": true, | 7 | + }, |
| 8 | - "use_video_feature": true | ||
| 9 | - }, | ||
| 10 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 11 | "video_processor_type": "CogVideoXProcessor", | 9 | "video_processor_type": "CogVideoXProcessor", |
| 12 | "video_reader_type": "DecordVideo", | 10 | "video_reader_type": "DecordVideo", |
| @@ -1 +0,0 @@ | |||
| 1 | -./wan2.1/sora_featuers/,./wan2.1/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/wan2.1/1.3b/i2v/data.txt", | 5 | + "data_path": "./wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./wan2.1/sora_featuers/,./wan2.1/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/wan2.1/1.3b/t2v/data.txt", | 5 | + "data_path": "./wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./wan2.1/sora_featuers/,./wan2.1/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/wan2.1/14b/i2v/data.txt", | 5 | + "data_path": "./wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./wan2.1/sora_featuers/,./wan2.1/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/wan2.1/14b/t2v/data.txt", | 5 | + "data_path": "./wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
| @@ -1 +0,0 @@ | |||
| 1 | -./wan2.1/sora_featuers/,./wan2.1/sora_featuers/data.jsonl | ||
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "./examples/wan2.1/14b/t2v_A3/data.txt", | 5 | + "data_path": "./wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "./wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "./examples/wan2.1/feature_extract/data.txt", | 5 | "data_path": "./examples/wan2.1/feature_extract/data.txt", |
| 7 | "data_storage_mode": "combine" | 6 | "data_storage_mode": "combine" |
| @@ -117,7 +116,7 @@ | |||
| 117 | "dataloader_param": { | 116 | "dataloader_param": { |
| 118 | "dataloader_mode": "sampler", | 117 | "dataloader_mode": "sampler", |
| 119 | "sampler_type": "LengthGroupedSampler", | 118 | "sampler_type": "LengthGroupedSampler", |
| 120 | - "shuffle": false, | 119 | + "shuffle": true, |
| 121 | "drop_last": true, | 120 | "drop_last": true, |
| 122 | "pin_memory": true, | 121 | "pin_memory": true, |
| 123 | "group_frame": false, | 122 | "group_frame": false, |
| @@ -82,13 +82,7 @@ class WanTextVideoDataset(T2VDataset): | |||
| 82 | raise AssertionError(f"file {file_path} do not exist!") | 82 | raise AssertionError(f"file {file_path} do not exist!") |
| 83 | 83 | ||
| 84 | vframes = self.video_reader(file_path) | 84 | vframes = self.video_reader(file_path) |
| 85 | - start_frame_idx = sample.get("start_frame_idx", 0) | 85 | + video = self.video_processer(vframes=vframes, **sample) |
| 86 | - frame_indice = sample["sample_frame_index"] | ||
| 87 | - video = self.video_processer( | ||
| 88 | - vframes, | ||
| 89 | - predefine_num_frames=len(frame_indice), | ||
| 90 | - start_frame_idx=start_frame_idx, | ||
| 91 | - ) | ||
| 92 | 86 | ||
| 93 | if self.task == "i2v": | 87 | if self.task == "i2v": |
| 94 | first_frame = video[:, 0, :, :] # c t h w | 88 | first_frame = video[:, 0, :, :] # c t h w |
| @@ -203,7 +197,7 @@ def extract_feature(): | |||
| 203 | }) | 197 | }) |
| 204 | elif data_storage_mode == "sorafeatured": | 198 | elif data_storage_mode == "sorafeatured": |
| 205 | data_info.update({ | 199 | data_info.update({ |
| 206 | - 'path': f"features/{pt_name}" | 200 | + FILE_INFO: f"features/{pt_name}" |
| 207 | }) | 201 | }) |
| 208 | json_file.write(json.dumps(data_info) + '\n') | 202 | json_file.write(json.dumps(data_info) + '\n') |
| 209 | 203 | ||
| @@ -270,9 +264,9 @@ def extract_feature(): | |||
| 270 | prompt_masks_i = prompt_mask[i] | 264 | prompt_masks_i = prompt_mask[i] |
| 271 | 265 | ||
| 272 | data_to_save = { | 266 | data_to_save = { |
| 273 | - "latents": latent_i, | 267 | + VIDEO: latent_i, |
| 274 | - "prompt": prompts_i, | 268 | + PROMPT_IDS: prompts_i, |
| 275 | - "prompt_mask": prompt_masks_i | 269 | + PROMPT_MASK: prompt_masks_i |
| 276 | } | 270 | } |
| 277 | 271 | ||
| 278 | if latents_dict: | 272 | if latents_dict: |
| @@ -2,7 +2,7 @@ | |||
| 2 | "sorafeature": { | 2 | "sorafeature": { |
| 3 | "extract_video_feature": true, | 3 | "extract_video_feature": true, |
| 4 | "extract_text_feature": true, | 4 | "extract_text_feature": true, |
| 5 | - "save_path": "./wan2.1/sora_featuers/", | 5 | + "save_path": "./wan2.1/sora_features/", |
| 6 | "data_storage_mode": "sorafeatured" | 6 | "data_storage_mode": "sorafeatured" |
| 7 | }, | 7 | }, |
| 8 | "task": "t2v" | 8 | "task": "t2v" |
| @@ -278,14 +278,21 @@ class TextProcesser: | |||
| 278 | 278 | ||
| 279 | def __call__(self, texts): | 279 | def __call__(self, texts): |
| 280 | if self.enable_text_preprocessing: | 280 | if self.enable_text_preprocessing: |
| 281 | - texts_info = [ | 281 | + if isinstance(texts, tuple) or isinstance(texts, list): |
| 282 | - TextProcesser.text_preprocessing( | 282 | + texts_info = [ |
| 283 | - text, | 283 | + TextProcesser.text_preprocessing( |
| 284 | + text, | ||
| 285 | + self.use_clean_caption, | ||
| 286 | + text_preprocess_methods=self.text_preprocess_methods | ||
| 287 | + ) | ||
| 288 | + for text in texts | ||
| 289 | + ] | ||
| 290 | + else: | ||
| 291 | + texts_info = TextProcesser.text_preprocessing( | ||
| 292 | + texts, | ||
| 284 | self.use_clean_caption, | 293 | self.use_clean_caption, |
| 285 | text_preprocess_methods=self.text_preprocess_methods | 294 | text_preprocess_methods=self.text_preprocess_methods |
| 286 | ) | 295 | ) |
| 287 | - for text in texts | ||
| 288 | - ] | ||
| 289 | texts_info = texts_info if random.random() > self.cfg else [""] | 296 | texts_info = texts_info if random.random() > self.cfg else [""] |
| 290 | else: | 297 | else: |
| 291 | texts_info = texts | 298 | texts_info = texts |
| @@ -354,29 +354,46 @@ class OpensoraplanVideoProcessor(AbstractVideoProcessor): | |||
| 354 | def __call__( | 354 | def __call__( |
| 355 | self, | 355 | self, |
| 356 | vframes, | 356 | vframes, |
| 357 | - predefine_num_frames=13, | 357 | + sample_num_frames=13, |
| 358 | start_frame_idx=0, | 358 | start_frame_idx=0, |
| 359 | - clip_total_frames=-1, | 359 | + num_frames=-1, |
| 360 | - resolution_crop=(None, None, None, None), | 360 | + crop=(None, None, None, None), |
| 361 | **kwargs | 361 | **kwargs |
| 362 | ): | 362 | ): |
| 363 | - """Process video with temporal speed adjustment and spatial validation""" | 363 | + """Process video frames with temporal speed adjustment and spatial validation. |
| 364 | - total_frames = vframes.get_len() if clip_total_frames == -1 else clip_total_frames | 364 | + |
| 365 | - fps = vframes.get_video_fps() if vframes.get_video_fps() > 0 else 30.0 | 365 | + Args: |
| 366 | - s_x, e_x, s_y, e_y = resolution_crop | 366 | + vframes: Video frames container object with frame access methods |
| 367 | + sample_num_frames: Expected number of output frames for validation | ||
| 368 | + start_frame_idx: Starting index for frame sampling | ||
| 369 | + num_frames: Total available frames (-1 = auto-detect from vframes) | ||
| 370 | + crop: Spatial crop coordinates (start_x, end_x, start_y, end_y) | ||
| 367 | 371 | ||
| 368 | - # Frame interval calculation | 372 | + Returns: |
| 373 | + torch.Tensor: Processed video tensor in CTHW format | ||
| 374 | + | ||
| 375 | + Raises: | ||
| 376 | + IndexError: When video is too short for required processing | ||
| 377 | + ValueError: When sampled frames mismatch predefined count | ||
| 378 | + AssertionError: When aspect ratio validation fails | ||
| 379 | + """ | ||
| 380 | + # Frame count and FPS initialization | ||
| 381 | + total_frames = vframes.get_len() if num_frames == -1 else num_frames | ||
| 382 | + fps = vframes.get_video_fps() if vframes.get_video_fps() > 0 else 30.0 | ||
| 383 | + s_x, e_x, s_y, e_y = crop | ||
| 384 | + | ||
| 385 | + # Temporal sampling interval calculation | ||
| 369 | if self.auto_interval: | 386 | if self.auto_interval: |
| 370 | # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24) | 387 | # resample in case high fps, such as 50/60/90/144 -> train_fps(e.g, 24) |
| 371 | frame_interval = 1.0 if abs(fps - self.train_fps) < 0.1 else fps / self.train_fps | 388 | frame_interval = 1.0 if abs(fps - self.train_fps) < 0.1 else fps / self.train_fps |
| 372 | else: | 389 | else: |
| 373 | frame_interval = self.frame_interval | 390 | frame_interval = self.frame_interval |
| 374 | 391 | ||
| 375 | - # Temporal sampling | 392 | + # Generate initial frame indices |
| 376 | frame_indices = np.arange(start_frame_idx, start_frame_idx + total_frames, frame_interval).astype(int) | 393 | frame_indices = np.arange(start_frame_idx, start_frame_idx + total_frames, frame_interval).astype(int) |
| 377 | frame_indices = frame_indices[frame_indices < start_frame_idx + total_frames] | 394 | frame_indices = frame_indices[frame_indices < start_frame_idx + total_frames] |
| 378 | 395 | ||
| 379 | - # speed up | 396 | + # speed up through temporal subsampling |
| 380 | max_speed_factor = len(frame_indices) / self.num_frames | 397 | max_speed_factor = len(frame_indices) / self.num_frames |
| 381 | if self.speed_factor > 1 and max_speed_factor > 1: | 398 | if self.speed_factor > 1 and max_speed_factor > 1: |
| 382 | speed_factor = min(self.speed_factor, max_speed_factor) | 399 | speed_factor = min(self.speed_factor, max_speed_factor) |
| @@ -386,7 +403,7 @@ class OpensoraplanVideoProcessor(AbstractVideoProcessor): | |||
| 386 | ) | 403 | ) |
| 387 | frame_indices = frame_indices[speed_frame_idx] | 404 | frame_indices = frame_indices[speed_frame_idx] |
| 388 | 405 | ||
| 389 | - # too long video will be temporal-crop randomly | 406 | + # Random temporal cropping for long sequences |
| 390 | if len(frame_indices) > self.num_frames: | 407 | if len(frame_indices) > self.num_frames: |
| 391 | begin_index, end_index = self.temporal_sample(len(frame_indices)) | 408 | begin_index, end_index = self.temporal_sample(len(frame_indices)) |
| 392 | frame_indices = frame_indices[begin_index:end_index] | 409 | frame_indices = frame_indices[begin_index:end_index] |
| @@ -402,9 +419,9 @@ class OpensoraplanVideoProcessor(AbstractVideoProcessor): | |||
| 402 | frame_indices = frame_indices[:end_frame_idx] | 419 | frame_indices = frame_indices[:end_frame_idx] |
| 403 | 420 | ||
| 404 | # Frame validation | 421 | # Frame validation |
| 405 | - if predefine_num_frames != len(frame_indices): | 422 | + if sample_num_frames != len(frame_indices): |
| 406 | raise ValueError( | 423 | raise ValueError( |
| 407 | - f"predefine_num_frames ({predefine_num_frames}) is not equal with frame_indices ({len(frame_indices)})" | 424 | + f"sample_num_frames ({sample_num_frames}) is not equal with frame_indices ({len(frame_indices)})" |
| 408 | ) | 425 | ) |
| 409 | if len(frame_indices) < self.num_frames and self.drop_short_ratio >= 1: | 426 | if len(frame_indices) < self.num_frames and self.drop_short_ratio >= 1: |
| 410 | raise IndexError( | 427 | raise IndexError( |
| @@ -175,13 +175,18 @@ def prepare_sampler_dataloader( | |||
| 175 | sample_sizes = [] | 175 | sample_sizes = [] |
| 176 | for data_sample in dataset.data_samples: | 176 | for data_sample in dataset.data_samples: |
| 177 | sample_num_frames.append(data_sample["sample_num_frames"]) | 177 | sample_num_frames.append(data_sample["sample_num_frames"]) |
| 178 | - sample_sizes.append(data_sample["sample_size"]) | 178 | + if data_sample.get("sample_size", None): |
| 179 | + sample_size = data_sample["sample_size"] | ||
| 180 | + else: | ||
| 181 | + sample_size = f'{len(data_sample["sample_frame_index"])}x{data_sample["resolution"]["sample_height"]}x{data_sample["resolution"]["sample_width"]}' | ||
| 182 | + sample_sizes.append(sample_size) | ||
| 179 | sampler = ( | 183 | sampler = ( |
| 180 | LengthGroupedSampler( | 184 | LengthGroupedSampler( |
| 181 | batch_size, | 185 | batch_size, |
| 182 | world_size=process_group.size(), | 186 | world_size=process_group.size(), |
| 183 | num_replicas=process_group.size(), | 187 | num_replicas=process_group.size(), |
| 184 | rank=process_group.rank(), | 188 | rank=process_group.rank(), |
| 189 | + shuffle=shuffle, | ||
| 185 | gradient_accumulation_size=gradient_accumulation_size, | 190 | gradient_accumulation_size=gradient_accumulation_size, |
| 186 | initial_global_step=initial_global_step_for_sampler, | 191 | initial_global_step=initial_global_step_for_sampler, |
| 187 | lengths=sample_num_frames if not group_data else sample_sizes, | 192 | lengths=sample_num_frames if not group_data else sample_sizes, |
| @@ -58,7 +58,7 @@ def split_to_even_chunks(indices, lengths, num_chunks, batch_size): | |||
| 58 | return pad_chunks | 58 | return pad_chunks |
| 59 | 59 | ||
| 60 | 60 | ||
| 61 | -def split_data_to_even_chunks(megabatch, lengths, world_size, batch_size): | 61 | +def split_data_to_even_chunks(megabatch, lengths, world_size, batch_size, shuffle=True): |
| 62 | """ | 62 | """ |
| 63 | Split a list of indices into `chunks` chunks of roughly equal lengths. | 63 | Split a list of indices into `chunks` chunks of roughly equal lengths. |
| 64 | """ | 64 | """ |
| @@ -74,9 +74,16 @@ def split_data_to_even_chunks(megabatch, lengths, world_size, batch_size): | |||
| 74 | if batch_size <= len(chunk): | 74 | if batch_size <= len(chunk): |
| 75 | raise AssertionError("batch_size must greater than len_chunk !") | 75 | raise AssertionError("batch_size must greater than len_chunk !") |
| 76 | if len(chunk) != 0: # [[1, 2], [3]] -> [[1, 2], [3, 3]] | 76 | if len(chunk) != 0: # [[1, 2], [3]] -> [[1, 2], [3, 3]] |
| 77 | - chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))] | 77 | + if shuffle: |
| 78 | - else: | 78 | + chunk = chunk + [random.choice(chunk) for _ in range(batch_size - len(chunk))] |
| 79 | - chunk = random.choice(pad_chunks) # [[1], []] -> [[1], [1]] | 79 | + else: |
| 80 | + chunk = chunk + [chunk[0] for _ in range(batch_size - len(chunk))] | ||
| 81 | + else: # [[1], []] -> [[1], [1]] | ||
| 82 | + if shuffle: | ||
| 83 | + chunk = random.choice(pad_chunks) | ||
| 84 | + else: | ||
| 85 | + chunk = pad_chunks[0] | ||
| 86 | + | ||
| 80 | pad_chunks.append(chunk) | 87 | pad_chunks.append(chunk) |
| 81 | return pad_chunks | 88 | return pad_chunks |
| 82 | 89 | ||
| @@ -95,9 +102,9 @@ def group_frame_and_resolution_fun(indices): | |||
| 95 | raise NotImplementedError | 102 | raise NotImplementedError |
| 96 | 103 | ||
| 97 | 104 | ||
| 98 | -def last_group_frame_fun(shuffled_megabatches, lengths): | 105 | +def last_group_frame_fun(megabatches, lengths, shuffle=True): |
| 99 | - re_shuffled_megabatches = [] | 106 | + re_megabatches = [] |
| 100 | - for megabatch in shuffled_megabatches: | 107 | + for megabatch in megabatches: |
| 101 | re_megabatch = [] | 108 | re_megabatch = [] |
| 102 | for batch in megabatch: | 109 | for batch in megabatch: |
| 103 | if len(batch) == 0: | 110 | if len(batch) == 0: |
| @@ -113,14 +120,22 @@ def last_group_frame_fun(shuffled_megabatches, lengths): | |||
| 113 | for idx, length in idx_length_dict.items() | 120 | for idx, length in idx_length_dict.items() |
| 114 | if length == pick_length | 121 | if length == pick_length |
| 115 | ] | 122 | ] |
| 116 | - random_select_batch = [ | 123 | + |
| 117 | - random.choice(candidate_batch) | 124 | + if shuffle: |
| 118 | - for i in range(len(len_each_batch) - len(candidate_batch)) | 125 | + random_select_batch = [ |
| 119 | - ] | 126 | + random.choice(candidate_batch) |
| 127 | + for _ in range(len(len_each_batch) - len(candidate_batch)) | ||
| 128 | + ] | ||
| 129 | + else: | ||
| 130 | + random_select_batch = [ | ||
| 131 | + candidate_batch[0] | ||
| 132 | + for _ in range(len(len_each_batch) - len(candidate_batch)) | ||
| 133 | + ] | ||
| 134 | + | ||
| 120 | batch = candidate_batch + random_select_batch | 135 | batch = candidate_batch + random_select_batch |
| 121 | re_megabatch.append(batch) | 136 | re_megabatch.append(batch) |
| 122 | - re_shuffled_megabatches.append(re_megabatch) | 137 | + re_megabatches.append(re_megabatch) |
| 123 | - return re_shuffled_megabatches | 138 | + return re_megabatches |
| 124 | 139 | ||
| 125 | 140 | ||
| 126 | def last_group_resolution_fun(indices): | 141 | def last_group_resolution_fun(indices): |
| @@ -188,7 +203,7 @@ def get_length_grouped_indices( | |||
| 188 | return out_list | 203 | return out_list |
| 189 | 204 | ||
| 190 | 205 | ||
| 191 | -def group_data_fun(lengths, generator=None): | 206 | +def group_data_fun(lengths, generator=None, shuffle=True): |
| 192 | # counter is decrease order | 207 | # counter is decrease order |
| 193 | counter = Counter(lengths) # counter {'1x256x256': 3, ''} lengths ['1x256x256', '1x256x256', '1x256x256', ...] | 208 | counter = Counter(lengths) # counter {'1x256x256': 3, ''} lengths ['1x256x256', '1x256x256', '1x256x256', ...] |
| 194 | grouped_indices = defaultdict(list) | 209 | grouped_indices = defaultdict(list) |
| @@ -199,11 +214,19 @@ def group_data_fun(lengths, generator=None): | |||
| 199 | sorted_indices = [grouped_indices[item] for (item, _) in sorted(counter.items(), key=lambda x: x[1], reverse=True)] | 214 | sorted_indices = [grouped_indices[item] for (item, _) in sorted(counter.items(), key=lambda x: x[1], reverse=True)] |
| 200 | 215 | ||
| 201 | # shuffle in each group | 216 | # shuffle in each group |
| 202 | - shuffle_sorted_indices = [] | 217 | + if shuffle: |
| 203 | - for indice in sorted_indices: | 218 | + shuffle_sorted_indices = [] |
| 204 | - shuffle_idx = torch.randperm(len(indice), generator=generator).tolist() | 219 | + for indice in sorted_indices: |
| 205 | - shuffle_sorted_indices.extend([indice[idx] for idx in shuffle_idx]) | 220 | + shuffle_idx = torch.randperm(len(indice), generator=generator).tolist() |
| 206 | - return shuffle_sorted_indices | 221 | + shuffle_sorted_indices.extend([indice[idx] for idx in shuffle_idx]) |
| 222 | + | ||
| 223 | + return shuffle_sorted_indices | ||
| 224 | + else: | ||
| 225 | + unshuffle_sorted_indices = [] | ||
| 226 | + for indice in sorted_indices: | ||
| 227 | + unshuffle_sorted_indices.extend(indice) | ||
| 228 | + | ||
| 229 | + return unshuffle_sorted_indices | ||
| 207 | 230 | ||
| 208 | 231 | ||
| 209 | def get_length_grouped_data_indices( | 232 | def get_length_grouped_data_indices( |
| @@ -214,7 +237,8 @@ def get_length_grouped_data_indices( | |||
| 214 | initial_global_step, | 237 | initial_global_step, |
| 215 | generator=None, | 238 | generator=None, |
| 216 | group_data=False, | 239 | group_data=False, |
| 217 | - seed=42): | 240 | + seed=42, |
| 241 | + shuffle=True): | ||
| 218 | # We need to use torch for the random part as a distributed sampler will set the random seed for torch. | 242 | # We need to use torch for the random part as a distributed sampler will set the random seed for torch. |
| 219 | if generator is None: | 243 | if generator is None: |
| 220 | if world_size == 1: | 244 | if world_size == 1: |
| @@ -223,27 +247,33 @@ def get_length_grouped_data_indices( | |||
| 223 | generator = torch.Generator() # every rank will generate a fixed order but random index | 247 | generator = torch.Generator() # every rank will generate a fixed order but random index |
| 224 | 248 | ||
| 225 | if group_data: | 249 | if group_data: |
| 226 | - indices = group_data_fun(lengths, generator) | 250 | + indices = group_data_fun(lengths, generator, shuffle) |
| 227 | else: | 251 | else: |
| 228 | - indices = torch.randperm(len(lengths), generator=generator).tolist() | 252 | + if shuffle: |
| 253 | + indices = torch.randperm(len(lengths), generator=generator).tolist() | ||
| 254 | + else: | ||
| 255 | + indices = list(range(len(lengths))) | ||
| 229 | 256 | ||
| 230 | megabatch_size = world_size * batch_size | 257 | megabatch_size = world_size * batch_size |
| 231 | megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)] | 258 | megabatches = [indices[i: i + megabatch_size] for i in range(0, len(lengths), megabatch_size)] |
| 232 | 259 | ||
| 233 | - megabatches = [split_data_to_even_chunks(megabatch, lengths, world_size, batch_size) for megabatch in megabatches] | 260 | + megabatches = [split_data_to_even_chunks(megabatch, lengths, world_size, batch_size, shuffle) for megabatch in megabatches] |
| 234 | 261 | ||
| 235 | - indices_mega = torch.randperm(len(megabatches), generator=generator).tolist() | 262 | + if shuffle: |
| 263 | + indices_mega = torch.randperm(len(megabatches), generator=generator).tolist() | ||
| 264 | + else: | ||
| 265 | + indices_mega = list(range(len(megabatches))) | ||
| 236 | 266 | ||
| 237 | - shuffled_megabatches = [megabatches[i] for i in indices_mega] | 267 | + megabatches = [megabatches[i] for i in indices_mega] |
| 238 | 268 | ||
| 239 | if group_data: | 269 | if group_data: |
| 240 | - shuffled_megabatches = last_group_frame_fun(shuffled_megabatches, lengths) | 270 | + megabatches = last_group_frame_fun(megabatches, lengths, shuffle) |
| 241 | 271 | ||
| 242 | initial_global_step = initial_global_step * gradient_accumulation_size | 272 | initial_global_step = initial_global_step * gradient_accumulation_size |
| 243 | - shuffled_megabatches = shuffled_megabatches[initial_global_step:] | 273 | + megabatches = megabatches[initial_global_step:] |
| 244 | 274 | ||
| 245 | out_list = [] | 275 | out_list = [] |
| 246 | - for megabatch in shuffled_megabatches: | 276 | + for megabatch in megabatches: |
| 247 | for batch in megabatch: | 277 | for batch in megabatch: |
| 248 | for i in batch: | 278 | for i in batch: |
| 249 | out_list.append(i) | 279 | out_list.append(i) |
| @@ -262,6 +292,7 @@ class LengthGroupedSampler(DistributedSampler): | |||
| 262 | world_size: int, | 292 | world_size: int, |
| 263 | num_replicas: Optional[int] = None, | 293 | num_replicas: Optional[int] = None, |
| 264 | rank: Optional[int] = None, | 294 | rank: Optional[int] = None, |
| 295 | + shuffle: bool = True, | ||
| 265 | gradient_accumulation_size: int = 1, | 296 | gradient_accumulation_size: int = 1, |
| 266 | initial_global_step: int = 0, | 297 | initial_global_step: int = 0, |
| 267 | lengths: Optional[List[int]] = None, | 298 | lengths: Optional[List[int]] = None, |
| @@ -279,6 +310,7 @@ class LengthGroupedSampler(DistributedSampler): | |||
| 279 | raise ValueError("world_size must be provided.") | 310 | raise ValueError("world_size must be provided.") |
| 280 | self.batch_size = batch_size | 311 | self.batch_size = batch_size |
| 281 | self.world_size = world_size | 312 | self.world_size = world_size |
| 313 | + self.shuffle = shuffle | ||
| 282 | self.initial_global_step = initial_global_step | 314 | self.initial_global_step = initial_global_step |
| 283 | self.gradient_accumulation_size = gradient_accumulation_size | 315 | self.gradient_accumulation_size = gradient_accumulation_size |
| 284 | self.lengths = lengths | 316 | self.lengths = lengths |
| @@ -312,6 +344,7 @@ class LengthGroupedSampler(DistributedSampler): | |||
| 312 | self.initial_global_step, | 344 | self.initial_global_step, |
| 313 | group_data=self.group_data, | 345 | group_data=self.group_data, |
| 314 | generator=self.generator, | 346 | generator=self.generator, |
| 347 | + shuffle=self.shuffle | ||
| 315 | ) | 348 | ) |
| 316 | 349 | ||
| 317 | # start sampling from the consumed samples point to continue training from where it left off | 350 | # start sampling from the consumed samples point to continue training from where it left off |
| @@ -58,7 +58,7 @@ class FeatureDataset(MMBaseDataset): | |||
| 58 | 58 | ||
| 59 | # Extract video features | 59 | # Extract video features |
| 60 | examples[VIDEO] = feature_data.pop(VIDEO, None) | 60 | examples[VIDEO] = feature_data.pop(VIDEO, None) |
| 61 | - | 61 | + |
| 62 | # Extract text prompt IDs and masks | 62 | # Extract text prompt IDs and masks |
| 63 | examples[PROMPT_IDS] = feature_data.pop(PROMPT_IDS, None) | 63 | examples[PROMPT_IDS] = feature_data.pop(PROMPT_IDS, None) |
| 64 | examples[PROMPT_MASK] = feature_data.pop(PROMPT_MASK, None) | 64 | examples[PROMPT_MASK] = feature_data.pop(PROMPT_MASK, None) |
| @@ -52,7 +52,6 @@ class I2VDataset(T2VDataset): | |||
| 52 | use_clean_caption: bool = True, | 52 | use_clean_caption: bool = True, |
| 53 | support_chinese: bool = False, | 53 | support_chinese: bool = False, |
| 54 | tokenizer_config: Union[dict, None] = None, | 54 | tokenizer_config: Union[dict, None] = None, |
| 55 | - use_feature_data: bool = False, | ||
| 56 | vid_img_fusion_by_splicing: bool = False, | 55 | vid_img_fusion_by_splicing: bool = False, |
| 57 | use_img_num: int = 0, | 56 | use_img_num: int = 0, |
| 58 | use_img_from_vid: bool = True, | 57 | use_img_from_vid: bool = True, |
| @@ -98,7 +97,6 @@ class I2VDataset(T2VDataset): | |||
| 98 | use_clean_caption=use_clean_caption, | 97 | use_clean_caption=use_clean_caption, |
| 99 | support_chinese=support_chinese, | 98 | support_chinese=support_chinese, |
| 100 | tokenizer_config=tokenizer_config, | 99 | tokenizer_config=tokenizer_config, |
| 101 | - use_feature_data=use_feature_data, | ||
| 102 | vid_img_fusion_by_splicing=vid_img_fusion_by_splicing, | 100 | vid_img_fusion_by_splicing=vid_img_fusion_by_splicing, |
| 103 | use_img_num=use_img_num, | 101 | use_img_num=use_img_num, |
| 104 | use_img_from_vid=use_img_from_vid, | 102 | use_img_from_vid=use_img_from_vid, |
| @@ -110,107 +108,65 @@ class I2VDataset(T2VDataset): | |||
| 110 | examples = copy.deepcopy(I2VOutputData) | 108 | examples = copy.deepcopy(I2VOutputData) |
| 111 | 109 | ||
| 112 | if self.data_storage_mode == "combine": | 110 | if self.data_storage_mode == "combine": |
| 113 | - examples = self.get_merge_data(examples, index) | 111 | + sample = self.data_samples[index] |
| 112 | + file_path = sample["path"] | ||
| 113 | + texts = sample["cap"] | ||
| 114 | elif self.data_storage_mode == "standard": | 114 | elif self.data_storage_mode == "standard": |
| 115 | sample = self.data_samples[index] | 115 | sample = self.data_samples[index] |
| 116 | - if self.use_feature_data: | 116 | + file_path, texts = sample[FILE_INFO], sample[CAPTIONS] |
| 117 | - video_path, masked_video_path, text_path = sample[FILE_INFO], sample[MASKED_VIDEO], sample[CAPTIONS] | 117 | + if self.data_folder: |
| 118 | - if self.data_folder: | 118 | + file_path = os.path.join(self.data_folder, file_path) |
| 119 | - video_path = os.path.join(self.data_folder, video_path) | ||
| 120 | - masked_video_path = os.path.join(self.data_folder, masked_video_path) | ||
| 121 | - text_path = os.path.join(self.data_folder, text_path) | ||
| 122 | - video_value = self.get_data_from_feature_data(video_path) | ||
| 123 | - masked_video_value = self.get_data_from_feature_data(masked_video_path) | ||
| 124 | - texts = self.get_data_from_feature_data(text_path) | ||
| 125 | - examples[VIDEO] = video_value | ||
| 126 | - examples[MASKED_VIDEO] = masked_video_value | ||
| 127 | - examples[TEXT] = texts | ||
| 128 | - examples[PROMPT_IDS] = texts | ||
| 129 | - examples[PROMPT_MASK] = texts | ||
| 130 | - else: | ||
| 131 | - raise NotImplementedError( | ||
| 132 | - f"Not support now: data_storage_mode={self.data_storage_mode} and use_feature_data=false" | ||
| 133 | - ) | ||
| 134 | else: | 119 | else: |
| 135 | raise NotImplementedError( | 120 | raise NotImplementedError( |
| 136 | f"Not support now: data_storage_mode={self.data_storage_mode}." | 121 | f"Not support now: data_storage_mode={self.data_storage_mode}." |
| 137 | ) | 122 | ) |
| 138 | - return examples | 123 | + |
| 139 | - | 124 | + # get video or img |
| 140 | - def drop(self, text, is_video=True): | ||
| 141 | - rand_num = random.random() | ||
| 142 | - rand_num_text = random.random() | ||
| 143 | - | ||
| 144 | - if rand_num < self.cfg: | ||
| 145 | - if rand_num_text < self.default_text_ratio: | ||
| 146 | - if not is_video: | ||
| 147 | - text = "The image showcases a scene with coherent and clear visuals." | ||
| 148 | - else: | ||
| 149 | - text = "The video showcases a scene with coherent and clear visuals." | ||
| 150 | - else: | ||
| 151 | - text = '' | ||
| 152 | - | ||
| 153 | - return dict(text=text) | ||
| 154 | - | ||
| 155 | - def get_merge_data(self, examples, index): | ||
| 156 | - sample = self.data_samples[index] | ||
| 157 | - file_path = sample["path"] | ||
| 158 | - if not os.path.exists(file_path): | ||
| 159 | - raise AssertionError(f"file {file_path} do not exist!") | ||
| 160 | file_type = self.get_type(file_path) | 125 | file_type = self.get_type(file_path) |
| 161 | - if file_type == "video": | 126 | + if file_type == "image": |
| 162 | - frame_indice = sample["sample_frame_index"] | 127 | + video_value = self.image_processer(file_path) |
| 128 | + video_value = video_value.transpose(0, 1) | ||
| 129 | + transforms_after_resize = self.image_transforms_after_resize | ||
| 130 | + elif file_type == "video": | ||
| 163 | vframes = self.video_reader(file_path) | 131 | vframes = self.video_reader(file_path) |
| 164 | - start_frame_idx = sample.get("start_frame_idx", 0) | 132 | + video_value = self.video_processer(vframes=vframes, **sample) |
| 165 | - clip_total_frames = sample.get("num_frames", -1) | 133 | + if self.vid_img_fusion_by_splicing: |
| 166 | - resolution_crop = tuple(sample.get("crop", (None, None, None, None))) | 134 | + video_value = self.get_vid_img_fusion(video_value) |
| 167 | - video = self.video_processer( | 135 | + video_value = video_value.permute(1, 0, 2, 3) |
| 168 | - vframes, | 136 | + transforms_after_resize = self.video_transforms_after_resize |
| 169 | - predefine_num_frames=len(frame_indice), | 137 | + |
| 170 | - start_frame_idx=start_frame_idx, | 138 | + inpaint_cond_data = self.mask_processor(video_value, mask_type_ratio_dict=self.mask_type_ratio_dict_video) |
| 171 | - clip_total_frames=clip_total_frames, | 139 | + mask, masked_video = inpaint_cond_data['mask'], inpaint_cond_data['masked_pixel_values'] |
| 172 | - resolution_crop=resolution_crop | ||
| 173 | - ) | ||
| 174 | - video = video.permute(1, 0, 2, 3) | ||
| 175 | - inpaint_cond_data = self.mask_processor(video, mask_type_ratio_dict=self.mask_type_ratio_dict_video) | ||
| 176 | - mask, masked_video = inpaint_cond_data['mask'], inpaint_cond_data['masked_pixel_values'] | ||
| 177 | 140 | ||
| 178 | - video = self.video_transforms_after_resize(video) # T C H W -> T C H W | 141 | + video_value = transforms_after_resize(video_value) # T C H W -> T C H W |
| 179 | - masked_video = self.video_transforms_after_resize(masked_video) # T C H W -> T C H W | 142 | + masked_video = transforms_after_resize(masked_video) # T C H W -> T C H W |
| 180 | 143 | ||
| 181 | - video = torch.cat([video, masked_video, mask], dim=1) # T 2C+1 H W | 144 | + video_value = torch.cat([video_value, masked_video, mask], dim=1) # T 2C+1 H W |
| 182 | 145 | ||
| 183 | - video = video.transpose(0, 1) # T C H W -> C T H W | 146 | + video_value = video_value.transpose(0, 1) # T C H W -> C T H W |
| 184 | 147 | ||
| 185 | - examples[VIDEO] = video | 148 | + examples[VIDEO] = video_value |
| 186 | - elif file_type == "image": | ||
| 187 | - image = self.image_processer(file_path) | ||
| 188 | - image = image.transpose(0, 1) | ||
| 189 | - inpaint_cond_data = self.mask_processor(image, mask_type_ratio_dict=self.mask_type_ratio_dict_image) | ||
| 190 | - mask, masked_image = inpaint_cond_data['mask'], inpaint_cond_data['masked_pixel_values'] | ||
| 191 | 149 | ||
| 192 | - image = self.image_transforms_after_resize(image) | 150 | + # get text tokens |
| 193 | - masked_image = self.image_transforms_after_resize(masked_image) | 151 | + if (isinstance(texts, list) or isinstance(texts, tuple)) and len(texts) > 1: |
| 152 | + texts = random.choice(texts) | ||
| 194 | 153 | ||
| 195 | - image = torch.cat([image, masked_image, mask], dim=1) # [1 2C+1 H W] | ||
| 196 | - | ||
| 197 | - image = image.transpose(0, 1) # [1 C H W] -> [C 1 H W] | ||
| 198 | - examples[VIDEO] = image | ||
| 199 | - | ||
| 200 | - text = sample["cap"] | ||
| 201 | - if not isinstance(text, list): | ||
| 202 | - text = [text] | ||
| 203 | - text = [random.choice(text)] | ||
| 204 | if self.use_aesthetic: | 154 | if self.use_aesthetic: |
| 205 | if sample.get('aesthetic', None) is not None or sample.get('aes', None) is not None: | 155 | if sample.get('aesthetic', None) is not None or sample.get('aes', None) is not None: |
| 206 | aes = sample.get('aesthetic', None) or sample.get('aes', None) | 156 | aes = sample.get('aesthetic', None) or sample.get('aes', None) |
| 207 | if file_type == "video": | 157 | if file_type == "video": |
| 208 | - text = [add_aesthetic_notice_video(text[0], aes)] | 158 | + texts = add_aesthetic_notice_video(texts, aes) |
| 209 | elif file_type == "image": | 159 | elif file_type == "image": |
| 210 | - text = [add_aesthetic_notice_image(text[0], aes)] | 160 | + texts = add_aesthetic_notice_image(texts, aes) |
| 211 | - prompt_ids, prompt_mask = self.get_text_processer(text) | ||
| 212 | - examples[PROMPT_IDS], examples[PROMPT_MASK] = prompt_ids, prompt_mask | ||
| 213 | 161 | ||
| 162 | + if self.use_text_processer: | ||
| 163 | + prompt_ids, prompt_mask = self.get_text_processer(texts) | ||
| 164 | + examples[PROMPT_IDS], examples[PROMPT_MASK] = ( | ||
| 165 | + prompt_ids, | ||
| 166 | + prompt_mask, | ||
| 167 | + ) | ||
| 168 | + | ||
| 214 | # for feature extract, trace source file name | 169 | # for feature extract, trace source file name |
| 215 | examples[FILE_INFO] = file_path | 170 | examples[FILE_INFO] = file_path |
| 171 | + | ||
| 216 | return examples | 172 | return examples |
| @@ -12,9 +12,8 @@ import sys | |||
| 12 | import torch | 12 | import torch |
| 13 | import numpy as np | 13 | import numpy as np |
| 14 | from megatron.core import mpu | 14 | from megatron.core import mpu |
| 15 | + | ||
| 15 | from mindspeed_mm.data.data_utils.utils import map_target_fps | 16 | from mindspeed_mm.data.data_utils.utils import map_target_fps |
| 16 | - | ||
| 17 | - | ||
| 18 | from mindspeed_mm.data.data_utils.constants import ( | 17 | from mindspeed_mm.data.data_utils.constants import ( |
| 19 | CAPTIONS, | 18 | CAPTIONS, |
| 20 | FILE_INFO, | 19 | FILE_INFO, |
| @@ -62,7 +61,6 @@ class T2VDataset(MMBaseDataset): | |||
| 62 | vid_img_process(dict): some data preprocessing parameters | 61 | vid_img_process(dict): some data preprocessing parameters |
| 63 | use_text_processer(bool): whether text preprocessing | 62 | use_text_processer(bool): whether text preprocessing |
| 64 | tokenizer_config(dict or list(dict)): the config of tokenizer or a list of configs for multi tokenizers | 63 | tokenizer_config(dict or list(dict)): the config of tokenizer or a list of configs for multi tokenizers |
| 65 | - use_feature_data(bool): use vae feature instead of raw video data or use text feature instead of raw text. | ||
| 66 | vid_img_fusion_by_splicing(bool): videos and images are fused by splicing | 64 | vid_img_fusion_by_splicing(bool): videos and images are fused by splicing |
| 67 | use_img_num(int): the number of fused images | 65 | use_img_num(int): the number of fused images |
| 68 | use_img_from_vid(bool): sampling some images from video | 66 | use_img_from_vid(bool): sampling some images from video |
| @@ -78,7 +76,6 @@ class T2VDataset(MMBaseDataset): | |||
| 78 | use_clean_caption: bool = True, | 76 | use_clean_caption: bool = True, |
| 79 | support_chinese: bool = False, | 77 | support_chinese: bool = False, |
| 80 | tokenizer_config: Optional[Union[dict, List[dict]]] = None, | 78 | tokenizer_config: Optional[Union[dict, List[dict]]] = None, |
| 81 | - use_feature_data: bool = False, | ||
| 82 | vid_img_fusion_by_splicing: bool = False, | 79 | vid_img_fusion_by_splicing: bool = False, |
| 83 | use_img_num: int = 0, | 80 | use_img_num: int = 0, |
| 84 | use_img_from_vid: bool = True, | 81 | use_img_from_vid: bool = True, |
| @@ -87,7 +84,6 @@ class T2VDataset(MMBaseDataset): | |||
| 87 | super().__init__(**basic_param) | 84 | super().__init__(**basic_param) |
| 88 | self.use_text_processer = use_text_processer | 85 | self.use_text_processer = use_text_processer |
| 89 | self.enable_text_preprocessing = enable_text_preprocessing | 86 | self.enable_text_preprocessing = enable_text_preprocessing |
| 90 | - self.use_feature_data = use_feature_data | ||
| 91 | self.vid_img_fusion_by_splicing = vid_img_fusion_by_splicing | 87 | self.vid_img_fusion_by_splicing = vid_img_fusion_by_splicing |
| 92 | self.use_img_num = use_img_num | 88 | self.use_img_num = use_img_num |
| 93 | self.use_img_from_vid = use_img_from_vid | 89 | self.use_img_from_vid = use_img_from_vid |
| @@ -137,22 +133,14 @@ class T2VDataset(MMBaseDataset): | |||
| 137 | support_chinese=support_chinese, | 133 | support_chinese=support_chinese, |
| 138 | cfg=self.cfg, | 134 | cfg=self.cfg, |
| 139 | ) | 135 | ) |
| 140 | - | ||
| 141 | - # Feature usage flags | ||
| 142 | - self.use_text_feature = basic_param.get("use_text_feature", False) | ||
| 143 | - self.use_video_feature = basic_param.get("use_video_feature", False) | ||
| 144 | 136 | ||
| 145 | # Thread pool configuration for data loading | 137 | # Thread pool configuration for data loading |
| 146 | max_workers = kwargs.get("max_workers", 1) | 138 | max_workers = kwargs.get("max_workers", 1) |
| 147 | self.executor = ThreadPoolExecutor(max_workers=max_workers) | 139 | self.executor = ThreadPoolExecutor(max_workers=max_workers) |
| 148 | self.timeout = kwargs.get("timeout", 60) | 140 | self.timeout = kwargs.get("timeout", 60) |
| 149 | 141 | ||
| 150 | - # Data validation for specific storage modes | 142 | + # Data validation |
| 151 | - if self.data_storage_mode == "combine" or self.data_storage_mode == "sorafeatured": | 143 | + self.data_samples = self.video_processer.select_valid_data(self.data_samples) |
| 152 | - if basic_param.get("skip_define_frame_index", False): | ||
| 153 | - return | ||
| 154 | - | ||
| 155 | - self.data_samples = self.video_processer.select_valid_data(self.data_samples) | ||
| 156 | 144 | ||
| 157 | def __getitem__(self, index): | 145 | def __getitem__(self, index): |
| 158 | try: | 146 | try: |
| @@ -176,89 +164,71 @@ class T2VDataset(MMBaseDataset): | |||
| 176 | 164 | ||
| 177 | if self.data_storage_mode == "standard": | 165 | if self.data_storage_mode == "standard": |
| 178 | sample = self.data_samples[index] | 166 | sample = self.data_samples[index] |
| 179 | - if self.use_feature_data: | 167 | + file_path, texts = sample[FILE_INFO], sample[CAPTIONS] |
| 180 | - path, texts_path = sample[FILE_INFO], sample[CAPTIONS] | 168 | + if self.data_folder: |
| 181 | - if self.data_folder: | 169 | + file_path = os.path.join(self.data_folder, file_path) |
| 182 | - path = os.path.join(self.data_folder, path) | ||
| 183 | - texts_path = os.path.join(self.data_folder, texts_path) | ||
| 184 | - video_value = self.get_data_from_feature_data(path) | ||
| 185 | - texts = self.get_data_from_feature_data(texts_path) | ||
| 186 | - if "image_latent" in sample: | ||
| 187 | - image_latent_path = sample["image_latent"] | ||
| 188 | - if self.data_folder: | ||
| 189 | - image_latent_path = os.path.join(self.data_folder, image_latent_path) | ||
| 190 | - examples.update(self.get_data_from_feature_data(image_latent_path)) | ||
| 191 | - examples[VIDEO] = video_value | ||
| 192 | - examples[TEXT] = texts | ||
| 193 | - examples[PROMPT_IDS] = texts | ||
| 194 | - examples[PROMPT_MASK] = texts | ||
| 195 | - else: | ||
| 196 | - path, texts = sample[FILE_INFO], sample[CAPTIONS] | ||
| 197 | - if self.data_folder: | ||
| 198 | - path = os.path.join(self.data_folder, path) | ||
| 199 | - examples[TEXT] = texts | ||
| 200 | - video_value = ( | ||
| 201 | - self.get_vid_img_fusion(path) | ||
| 202 | - if self.vid_img_fusion_by_splicing | ||
| 203 | - else self.get_value_from_vid_or_img(path) | ||
| 204 | - ) | ||
| 205 | - examples[VIDEO] = video_value | ||
| 206 | - if FILE_REJECTED_INFO in sample.keys(): | ||
| 207 | - video_rejected_path = os.path.join(self.data_folder, sample[FILE_REJECTED_INFO]) | ||
| 208 | - video_rejected_value = ( | ||
| 209 | - self.get_vid_img_fusion(video_rejected_path) | ||
| 210 | - if self.vid_img_fusion_by_splicing | ||
| 211 | - else self.get_value_from_vid_or_img(video_rejected_path) | ||
| 212 | - ) | ||
| 213 | - examples[VIDEO_REJECTED] = video_rejected_value | ||
| 214 | - examples[SCORE] = sample[SCORE] | ||
| 215 | - examples[SCORE_REJECTED] = sample[SCORE_REJECTED] | ||
| 216 | - if self.use_text_processer: | ||
| 217 | - prompt_ids, prompt_mask = self.get_text_processer(texts) | ||
| 218 | - examples[PROMPT_IDS], examples[PROMPT_MASK] = ( | ||
| 219 | - prompt_ids, | ||
| 220 | - prompt_mask, | ||
| 221 | - ) | ||
| 222 | - # for feature extract, trace source file name | ||
| 223 | - examples[FILE_INFO] = sample[FILE_INFO] | ||
| 224 | 170 | ||
| 225 | - elif self.data_storage_mode == "sorafeatured": | ||
| 226 | - sample = self.data_samples[index] | ||
| 227 | - | ||
| 228 | - if self.use_video_feature: | ||
| 229 | - video_path = sample['path'] | ||
| 230 | - if self.data_folder: | ||
| 231 | - video_path = os.path.join(self.data_folder, video_path) | ||
| 232 | - video_value = self.get_data_from_feature_data(video_path) | ||
| 233 | - examples[VIDEO] = video_value['latents'] | ||
| 234 | - if 'video_mask' in video_value.keys(): | ||
| 235 | - examples[VIDEO_MASK] = video_value['video_mask'] | ||
| 236 | - | ||
| 237 | - # update extra info, and avoid critical key values from being overwritten | ||
| 238 | - for key in SORA_MODEL_PROTECTED_KEYS: | ||
| 239 | - video_value.pop(key, None) | ||
| 240 | - examples.update(video_value) | ||
| 241 | - | ||
| 242 | - file_type = 'video' | ||
| 243 | - else: | ||
| 244 | - examples, file_type = self.get_video_data(examples, index) | ||
| 245 | - | ||
| 246 | - if self.use_text_feature: | ||
| 247 | - texts_path = sample['path'] | ||
| 248 | - if self.data_folder: | ||
| 249 | - texts_path = os.path.join(self.data_folder, texts_path) | ||
| 250 | - texts_value = self.get_data_from_feature_data(texts_path) | ||
| 251 | - examples[PROMPT_IDS] = texts_value['prompt'] | ||
| 252 | - examples[PROMPT_MASK] = texts_value['prompt_mask'] | ||
| 253 | - else: | ||
| 254 | - examples = self.get_text_data(examples, sample, file_type) | ||
| 255 | - | ||
| 256 | elif self.data_storage_mode == "combine": | 171 | elif self.data_storage_mode == "combine": |
| 257 | - examples = self.get_merge_data(examples, index) | 172 | + sample = self.data_samples[index] |
| 173 | + file_path = sample["path"] | ||
| 174 | + texts = sample["cap"] | ||
| 175 | + | ||
| 258 | else: | 176 | else: |
| 259 | raise NotImplementedError( | 177 | raise NotImplementedError( |
| 260 | f"Not support now: data_storage_mode={self.data_storage_mode}." | 178 | f"Not support now: data_storage_mode={self.data_storage_mode}." |
| 261 | ) | 179 | ) |
| 180 | + | ||
| 181 | + # Generic media processing pipeline | ||
| 182 | + file_type = self.get_type(file_path) | ||
| 183 | + | ||
| 184 | + # Image/video processing | ||
| 185 | + if file_type == "image": | ||
| 186 | + video_value = self.image_processer(file_path) | ||
| 187 | + elif file_type == "video": | ||
| 188 | + vframes = self.video_reader(file_path) | ||
| 189 | + video_value = self.video_processer(vframes=vframes, **sample) | ||
| 190 | + if self.vid_img_fusion_by_splicing: | ||
| 191 | + video_value = self.get_vid_img_fusion(video_value) | ||
| 192 | + examples[VIDEO] = video_value | ||
| 193 | + | ||
| 194 | + # Text processing | ||
| 195 | + if isinstance(texts, (list, tuple)) and len(texts) > 1: | ||
| 196 | + texts = random.choice(texts) # Random selection from multiple options | ||
| 197 | + | ||
| 198 | + # Handle aesthetic scoring | ||
| 199 | + if self.use_aesthetic: | ||
| 200 | + aes = sample.get('aesthetic') or sample.get('aes') | ||
| 201 | + if aes is not None: | ||
| 202 | + if file_type == "video": | ||
| 203 | + texts = add_aesthetic_notice_video(texts, aes) | ||
| 204 | + elif file_type == "image": | ||
| 205 | + texts = add_aesthetic_notice_image(texts, aes) | ||
| 206 | + | ||
| 207 | + # Text tokenization | ||
| 208 | + if self.use_text_processer: | ||
| 209 | + prompt_ids, prompt_mask = self.get_text_processer(texts) | ||
| 210 | + examples[PROMPT_IDS], examples[PROMPT_MASK] = prompt_ids, prompt_mask | ||
| 211 | + | ||
| 212 | + # DPO (Direct Preference Optimization) handling | ||
| 213 | + if FILE_REJECTED_INFO in sample.keys(): | ||
| 214 | + rejected_video_path = os.path.join(self.data_folder, sample[FILE_REJECTED_INFO]) | ||
| 215 | + | ||
| 216 | + rejected_file_type = self.get_type(rejected_video_path) | ||
| 217 | + if rejected_file_type == "image": | ||
| 218 | + rejected_video_value = self.image_processer(rejected_video_path) | ||
| 219 | + elif rejected_file_type == "video": | ||
| 220 | + rejected_vframes = self.video_reader(rejected_video_path) | ||
| 221 | + rejected_video_value = self.video_processer(vframes=rejected_vframes, **sample) | ||
| 222 | + if self.vid_img_fusion_by_splicing: | ||
| 223 | + rejected_video_value = self.get_vid_img_fusion(rejected_video_value) | ||
| 224 | + | ||
| 225 | + examples[VIDEO_REJECTED] = rejected_video_value | ||
| 226 | + examples[SCORE] = sample[SCORE] | ||
| 227 | + examples[SCORE_REJECTED] = sample[SCORE_REJECTED] | ||
| 228 | + | ||
| 229 | + # for feature extract, trace source file name | ||
| 230 | + examples[FILE_INFO] = file_path | ||
| 231 | + | ||
| 262 | return examples | 232 | return examples |
| 263 | 233 | ||
| 264 | def get_data_from_feature_data(self, feature_path): | 234 | def get_data_from_feature_data(self, feature_path): |
| @@ -266,83 +236,6 @@ class T2VDataset(MMBaseDataset): | |||
| 266 | return torch.load(feature_path, map_location=torch.device('cpu')) | 236 | return torch.load(feature_path, map_location=torch.device('cpu')) |
| 267 | raise NotImplementedError("Not implemented.") | 237 | raise NotImplementedError("Not implemented.") |
| 268 | 238 | ||
| 269 | - def get_video_data(self, examples, index): | ||
| 270 | - sample = self.data_samples[index] | ||
| 271 | - file_path = sample["path"] | ||
| 272 | - if not os.path.exists(file_path): | ||
| 273 | - raise AssertionError(f"file {file_path} do not exist!") | ||
| 274 | - file_type = self.get_type(file_path) | ||
| 275 | - if file_type == "video": | ||
| 276 | - frame_indice = sample["sample_frame_index"] | ||
| 277 | - vframes = self.video_reader(file_path) | ||
| 278 | - video = self.video_processer( | ||
| 279 | - vframes, | ||
| 280 | - predefine_num_frames=len(frame_indice), | ||
| 281 | - ) | ||
| 282 | - examples[VIDEO] = video | ||
| 283 | - elif file_type == "image": | ||
| 284 | - image = self.image_processer(file_path) | ||
| 285 | - examples[VIDEO] = image | ||
| 286 | - return examples, file_type | ||
| 287 | - | ||
| 288 | - def get_text_data(self, examples, sample, file_type='video'): | ||
| 289 | - text = sample["cap"] | ||
| 290 | - if not isinstance(text, list): | ||
| 291 | - text = [text] | ||
| 292 | - text = [random.choice(text)] | ||
| 293 | - if self.use_aesthetic: | ||
| 294 | - if sample.get('aesthetic', None) is not None or sample.get('aes', None) is not None: | ||
| 295 | - aes = sample.get('aesthetic', None) or sample.get('aes', None) | ||
| 296 | - if file_type == "video": | ||
| 297 | - text = [add_aesthetic_notice_video(text[0], aes)] | ||
| 298 | - elif file_type == "image": | ||
| 299 | - text = [add_aesthetic_notice_image(text[0], aes)] | ||
| 300 | - prompt_ids, prompt_mask = self.get_text_processer(text) | ||
| 301 | - examples[PROMPT_IDS], examples[PROMPT_MASK] = prompt_ids, prompt_mask | ||
| 302 | - return examples | ||
| 303 | - | ||
| 304 | - def get_merge_data(self, examples, index): | ||
| 305 | - sample = self.data_samples[index] | ||
| 306 | - file_path = sample["path"] | ||
| 307 | - if not os.path.exists(file_path): | ||
| 308 | - raise AssertionError(f"file {file_path} do not exist!") | ||
| 309 | - file_type = self.get_type(file_path) | ||
| 310 | - if file_type == "video": | ||
| 311 | - frame_indice = sample["sample_frame_index"] | ||
| 312 | - vframes = self.video_reader(file_path) | ||
| 313 | - start_frame_idx = sample.get("start_frame_idx", 0) | ||
| 314 | - clip_total_frames = sample.get("num_frames", -1) | ||
| 315 | - resolution_crop = tuple(sample.get("crop", (None, None, None, None))) | ||
| 316 | - video = self.video_processer( | ||
| 317 | - vframes, | ||
| 318 | - predefine_num_frames=len(frame_indice), | ||
| 319 | - start_frame_idx=start_frame_idx, | ||
| 320 | - clip_total_frames=clip_total_frames, | ||
| 321 | - resolution_crop=resolution_crop | ||
| 322 | - ) | ||
| 323 | - examples[VIDEO] = video | ||
| 324 | - elif file_type == "image": | ||
| 325 | - image = self.image_processer(file_path) | ||
| 326 | - examples[VIDEO] = image | ||
| 327 | - | ||
| 328 | - text = sample["cap"] | ||
| 329 | - if not isinstance(text, list): | ||
| 330 | - text = [text] | ||
| 331 | - text = [random.choice(text)] | ||
| 332 | - if self.use_aesthetic: | ||
| 333 | - if sample.get('aesthetic', None) is not None or sample.get('aes', None) is not None: | ||
| 334 | - aes = sample.get('aesthetic', None) or sample.get('aes', None) | ||
| 335 | - if file_type == "video": | ||
| 336 | - text = [add_aesthetic_notice_video(text[0], aes)] | ||
| 337 | - elif file_type == "image": | ||
| 338 | - text = [add_aesthetic_notice_image(text[0], aes)] | ||
| 339 | - prompt_ids, prompt_mask = self.get_text_processer(text) | ||
| 340 | - examples[PROMPT_IDS], examples[PROMPT_MASK] = prompt_ids, prompt_mask | ||
| 341 | - | ||
| 342 | - # for feature extract, trace source file name | ||
| 343 | - examples[FILE_INFO] = file_path | ||
| 344 | - return examples | ||
| 345 | - | ||
| 346 | def get_value_from_vid_or_img(self, path): | 239 | def get_value_from_vid_or_img(self, path): |
| 347 | file_type = self.get_type(path) | 240 | file_type = self.get_type(path) |
| 348 | if file_type == "video": | 241 | if file_type == "video": |
| @@ -352,9 +245,7 @@ class T2VDataset(MMBaseDataset): | |||
| 352 | video_value = self.image_processer(path) | 245 | video_value = self.image_processer(path) |
| 353 | return video_value | 246 | return video_value |
| 354 | 247 | ||
| 355 | - def get_vid_img_fusion(self, path): | 248 | + def get_vid_img_fusion(self, video_value): |
| 356 | - vframes = self.video_reader(path) | ||
| 357 | - video_value = self.video_processer(vframes=vframes) | ||
| 358 | if self.use_img_num != 0 and self.use_img_from_vid: | 249 | if self.use_img_num != 0 and self.use_img_from_vid: |
| 359 | select_image_idx = np.linspace( | 250 | select_image_idx = np.linspace( |
| 360 | 0, self.num_frames - 1, self.use_img_num, dtype=int | 251 | 0, self.num_frames - 1, self.use_img_num, dtype=int |
| @@ -407,7 +298,6 @@ class DynamicVideoTextDataset(MMBaseDataset): | |||
| 407 | vid_img_process(dict): some data preprocessing parameters | 298 | vid_img_process(dict): some data preprocessing parameters |
| 408 | use_text_processer(bool): whether text preprocessing | 299 | use_text_processer(bool): whether text preprocessing |
| 409 | tokenizer_config(dict): the config of tokenizer | 300 | tokenizer_config(dict): the config of tokenizer |
| 410 | - use_feature_data(bool): use vae feature instead of raw video data or use text feature instead of raw text. | ||
| 411 | vid_img_fusion_by_splicing(bool): videos and images are fused by splicing | 301 | vid_img_fusion_by_splicing(bool): videos and images are fused by splicing |
| 412 | use_img_num(int): the number of fused images | 302 | use_img_num(int): the number of fused images |
| 413 | use_img_from_vid(bool): sampling some images from video | 303 | use_img_from_vid(bool): sampling some images from video |
| @@ -421,7 +311,6 @@ class DynamicVideoTextDataset(MMBaseDataset): | |||
| 421 | enable_text_preprocessing: bool = True, | 311 | enable_text_preprocessing: bool = True, |
| 422 | use_clean_caption: bool = True, | 312 | use_clean_caption: bool = True, |
| 423 | tokenizer_config: Union[dict, None] = None, | 313 | tokenizer_config: Union[dict, None] = None, |
| 424 | - use_feature_data: bool = False, | ||
| 425 | vid_img_fusion_by_splicing: bool = False, | 314 | vid_img_fusion_by_splicing: bool = False, |
| 426 | use_img_num: int = 0, | 315 | use_img_num: int = 0, |
| 427 | use_img_from_vid: bool = True, | 316 | use_img_from_vid: bool = True, |
| @@ -432,7 +321,6 @@ class DynamicVideoTextDataset(MMBaseDataset): | |||
| 432 | ): | 321 | ): |
| 433 | super().__init__(**basic_param) | 322 | super().__init__(**basic_param) |
| 434 | self.use_text_processer = use_text_processer | 323 | self.use_text_processer = use_text_processer |
| 435 | - self.use_feature_data = use_feature_data | ||
| 436 | self.vid_img_fusion_by_splicing = vid_img_fusion_by_splicing | 324 | self.vid_img_fusion_by_splicing = vid_img_fusion_by_splicing |
| 437 | self.use_img_num = use_img_num | 325 | self.use_img_num = use_img_num |
| 438 | self.use_img_from_vid = use_img_from_vid | 326 | self.use_img_from_vid = use_img_from_vid |
| @@ -560,4 +448,4 @@ class DynamicVideoTextDataset(MMBaseDataset): | |||
| 560 | sample["text"] = sample["text"] + postfix | 448 | sample["text"] = sample["text"] + postfix |
| 561 | else: | 449 | else: |
| 562 | frame_interval = self.frame_interval | 450 | frame_interval = self.frame_interval |
| 563 | - return frame_interval | 451 | + return frame_interval |
| @@ -2,6 +2,7 @@ import json | |||
| 2 | import os | 2 | import os |
| 3 | import time | 3 | import time |
| 4 | import uuid | 4 | import uuid |
| 5 | +import copy | ||
| 5 | 6 | ||
| 6 | import mindspeed.megatron_adaptor | 7 | import mindspeed.megatron_adaptor |
| 7 | import torch | 8 | import torch |
| @@ -14,10 +15,9 @@ from numpy import save | |||
| 14 | from mindspeed_mm.configs.config import merge_mm_args, mm_extra_args_provider | 15 | from mindspeed_mm.configs.config import merge_mm_args, mm_extra_args_provider |
| 15 | from mindspeed_mm.data import build_mm_dataloader, build_mm_dataset | 16 | from mindspeed_mm.data import build_mm_dataloader, build_mm_dataset |
| 16 | from mindspeed_mm.data.data_utils.constants import ( | 17 | from mindspeed_mm.data.data_utils.constants import ( |
| 18 | + FILE_INFO, | ||
| 17 | PROMPT_IDS, | 19 | PROMPT_IDS, |
| 18 | - PROMPT_IDS_2, | ||
| 19 | PROMPT_MASK, | 20 | PROMPT_MASK, |
| 20 | - PROMPT_MASK_2, | ||
| 21 | VIDEO, | 21 | VIDEO, |
| 22 | VIDEO_MASK, | 22 | VIDEO_MASK, |
| 23 | ) | 23 | ) |
| @@ -38,28 +38,22 @@ def prepare_model(args, device): | |||
| 38 | return vae, text_encoder | 38 | return vae, text_encoder |
| 39 | 39 | ||
| 40 | 40 | ||
| 41 | +def get_pt_name(file_name): | ||
| 42 | + pt_name = os.path.basename(file_name).replace(".", "_") + ".pt" | ||
| 43 | + return pt_name | ||
| 44 | + | ||
| 45 | + | ||
| 41 | def extract_feature(): | 46 | def extract_feature(): |
| 42 | 47 | ||
| 43 | initialize_megatron(extra_args_provider=mm_extra_args_provider, args_defaults={}) | 48 | initialize_megatron(extra_args_provider=mm_extra_args_provider, args_defaults={}) |
| 44 | args = get_args() | 49 | args = get_args() |
| 45 | merge_mm_args(args) | 50 | merge_mm_args(args) |
| 46 | - | ||
| 47 | - extract_video_feature = args.mm.tool.sorafeature.extract_video_feature | ||
| 48 | - extract_text_feature = args.mm.tool.sorafeature.extract_text_feature | ||
| 49 | - data_storage_mode = args.mm.tool.sorafeature.data_storage_mode | ||
| 50 | 51 | ||
| 51 | save_path = args.mm.tool.sorafeature.save_path | 52 | save_path = args.mm.tool.sorafeature.save_path |
| 52 | 53 | ||
| 53 | if torch.distributed.get_rank() == 0: | 54 | if torch.distributed.get_rank() == 0: |
| 54 | - if not os.path.exists(save_path): | 55 | + if not os.path.exists(os.path.join(save_path, 'features')): |
| 55 | - os.makedirs(save_path) | 56 | + os.makedirs(os.path.join(save_path, 'features')) |
| 56 | - if data_storage_mode == 'standard': | ||
| 57 | - if not os.path.exists(os.path.join(save_path, 'videos')): | ||
| 58 | - os.makedirs(os.path.join(save_path, 'videos')) | ||
| 59 | - if not os.path.exists(os.path.join(save_path, 'labels')): | ||
| 60 | - os.makedirs(os.path.join(save_path, 'labels')) | ||
| 61 | - if not os.path.exists(os.path.join(save_path, 'images')): | ||
| 62 | - os.makedirs(os.path.join(save_path, 'images')) | ||
| 63 | 57 | ||
| 64 | 58 | ||
| 65 | set_jit_fusion_options() | 59 | set_jit_fusion_options() |
| @@ -72,63 +66,87 @@ def extract_feature(): | |||
| 72 | train_dataset, | 66 | train_dataset, |
| 73 | args.mm.data.dataloader_param, | 67 | args.mm.data.dataloader_param, |
| 74 | process_group=mpu.get_data_parallel_group(), | 68 | process_group=mpu.get_data_parallel_group(), |
| 69 | + dataset_param=args.mm.data.dataset_param, | ||
| 75 | ) | 70 | ) |
| 71 | + | ||
| 72 | + # master rank, write data info jsonl | ||
| 73 | + if torch.distributed.get_rank() == 0: | ||
| 74 | + with open(os.path.join(save_path, 'data.jsonl'), 'w', encoding="utf-8") as json_file: | ||
| 75 | + for data_sample in train_dataset.data_samples: | ||
| 76 | + source_data_storage_mode = args.mm.data.dataset_param.basic_parameters.data_storage_mode | ||
| 77 | + if source_data_storage_mode == "combine": | ||
| 78 | + source_file_key = "path" | ||
| 79 | + elif source_data_storage_mode == "standard": | ||
| 80 | + source_file_key = FILE_INFO | ||
| 81 | + else: | ||
| 82 | + raise NotImplementedError(f"Extract features from data storage mode {source_data_storage_mode} is not implemented") | ||
| 83 | + | ||
| 84 | + file_name = data_sample[source_file_key] | ||
| 85 | + pt_name = get_pt_name(file_name) | ||
| 86 | + data_info = copy.deepcopy(data_sample) | ||
| 87 | + | ||
| 88 | + data_info.update({ | ||
| 89 | + FILE_INFO: f"features/{pt_name}" | ||
| 90 | + }) | ||
| 91 | + json_file.write(json.dumps(data_info) + '\n') | ||
| 76 | 92 | ||
| 77 | vae, text_encoder = prepare_model(args, device) | 93 | vae, text_encoder = prepare_model(args, device) |
| 78 | 94 | ||
| 79 | start_time = time.time() | 95 | start_time = time.time() |
| 80 | print_rank_0(f"Features extraction begins. {len(train_dataloader)} data in total.") | 96 | print_rank_0(f"Features extraction begins. {len(train_dataloader)} data in total.") |
| 81 | counter = 0 | 97 | counter = 0 |
| 82 | - prof = Profiler(args.mm.tool.profile) | 98 | + if hasattr(args.mm.tool, "profile"): |
| 83 | - prof.start() | 99 | + prof = Profiler(args.mm.tool.profile) |
| 100 | + prof.start() | ||
| 101 | + | ||
| 84 | for batch in train_dataloader: | 102 | for batch in train_dataloader: |
| 85 | - counter += 1 | 103 | + if batch: |
| 86 | - | 104 | + video = batch.pop(VIDEO).to(device, dtype) |
| 87 | - video = batch.pop(VIDEO).to(device, dtype) | 105 | + prompt_ids = batch.pop(PROMPT_IDS) |
| 88 | - prompt_ids = batch.pop(PROMPT_IDS).to(device) | 106 | + prompt_mask = batch.pop(PROMPT_MASK) |
| 89 | - prompt_mask = batch.pop(PROMPT_MASK).to(device) | 107 | + file_names = batch.pop(FILE_INFO) |
| 90 | - | ||
| 91 | - | ||
| 92 | - if extract_video_feature: | ||
| 93 | - latents, latents_dict = vae.encode(video) | ||
| 94 | else: | 108 | else: |
| 95 | - latents = video | 109 | + raise ValueError("Batch is None!") |
| 96 | 110 | ||
| 97 | - if extract_text_feature: | 111 | + # extract feature |
| 98 | - prompt, prompt_mask = text_encoder.encode(prompt_ids, prompt_mask) | 112 | + latents, latents_dict = vae.encode(video, **batch) |
| 99 | - else: | 113 | + prompt, prompt_mask = text_encoder.encode(prompt_ids, prompt_mask) |
| 100 | - prompt = prompt_ids | 114 | + |
| 101 | - | 115 | + bs = video.shape[0] |
| 102 | - if torch.distributed.get_rank() == 0: | 116 | + counter += bs |
| 103 | - if data_storage_mode == 'standard': | 117 | + |
| 104 | - loca = time.strftime("%Y-%m-%d-%H-%M-%S") | 118 | + for i in range(bs): |
| 105 | - pt_name = "feature" + loca + "-" + uuid.uuid4().hex + ".pt" | 119 | + pt_name = get_pt_name(file_names[i]) |
| 106 | - torch.save(latents.squeeze(0), os.path.join(save_path, 'videos', pt_name)) | 120 | + latent_i = latents[i].cpu() |
| 107 | - torch.save(prompt.squeeze(0), os.path.join(save_path, 'labels', pt_name)) | 121 | + if isinstance(prompt_ids, (list, tuple)): |
| 108 | - data_to_save = { | 122 | + prompts_i = [_prompt[i].cpu() for _prompt in prompt] |
| 109 | - "file": os.path.join('videos', pt_name), | 123 | + prompt_masks_i = [_prompt_mask[i].cpu() for _prompt_mask in prompt_mask] |
| 110 | - "captions": os.path.join('labels', pt_name) | ||
| 111 | - } | ||
| 112 | - if latents_dict is not None: | ||
| 113 | - for k in latents_dict: | ||
| 114 | - latents_dict[k] = latents_dict[k].squeeze(0) | ||
| 115 | - torch.save(latents_dict, os.path.join(save_path, 'images', pt_name)) | ||
| 116 | - data_to_save["image_latent"] = os.path.join('images', pt_name) | ||
| 117 | - with open(os.path.join(save_path, 'data.jsonl'), 'a', encoding="utf-8") as json_file: | ||
| 118 | - json_file.write(json.dumps(data_to_save) + '\n') | ||
| 119 | - print_rank_0(f"consumed sample {counter} | elapsed time {(time.time() - start_time):.2f} | file {pt_name}") | ||
| 120 | else: | 124 | else: |
| 121 | - data_to_save = { | 125 | + prompts_i = prompt[i].cpu() |
| 122 | - "latents": latents.squeeze(0), | 126 | + prompt_masks_i = prompt_mask[i].cpu() |
| 123 | - "prompt": prompt.squeeze(0), | 127 | + |
| 124 | - "prompt_mask": prompt_mask.squeeze(0) | 128 | + data_to_save = { |
| 125 | - } | 129 | + VIDEO: latent_i, |
| 126 | - | 130 | + PROMPT_IDS: prompts_i, |
| 127 | - loca = time.strftime('%Y-%m-%d-%H-%M-%S') | 131 | + PROMPT_MASK: prompt_masks_i |
| 128 | - pt_name = "feature" + loca + "-" + uuid.uuid4().hex + ".pt" | 132 | + } |
| 129 | - torch.save(data_to_save, os.path.join(save_path, pt_name)) | 133 | + |
| 130 | - prof.step() | 134 | + # other i2v variables |
| 131 | - prof.stop() | 135 | + if latents_dict: |
| 136 | + for key in latents_dict.keys(): | ||
| 137 | + if isinstance(latents_dict[key][i], torch.Tensor): | ||
| 138 | + data_to_save[key] = latents_dict[key][i].cpu() | ||
| 139 | + else: | ||
| 140 | + data_to_save[key] = latents_dict[key][i] | ||
| 141 | + | ||
| 142 | + torch.save(data_to_save, os.path.join(save_path, 'features', pt_name)) | ||
| 143 | + | ||
| 144 | + print_rank_0(f"consumed sample {counter} | elapsed time {(time.time() - start_time):.2f} | file {file_names}") | ||
| 145 | + | ||
| 146 | + if hasattr(args.mm.tool, "profile"): | ||
| 147 | + prof.step() | ||
| 148 | + if hasattr(args.mm.tool, "profile"): | ||
| 149 | + prof.stop() | ||
| 132 | 150 | ||
| 133 | duration = time.time() - start_time | 151 | duration = time.time() - start_time |
| 134 | print_rank_0(f"{counter} feature vectors extracted in {duration:.2f} seconds.") | 152 | print_rank_0(f"{counter} feature vectors extracted in {duration:.2f} seconds.") |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/home/ci_resource/data/cogvideox1_0/data.jsonl", | 5 | "data_path": "/home/ci_resource/data/cogvideox1_0/data.jsonl", |
| 7 | "data_folder": "/home/ci_resource/data/cogvideox1_0/", | 6 | "data_folder": "/home/ci_resource/data/cogvideox1_0/", |
| @@ -1,7 +1,6 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | "dataset_type": "t2v", | 3 | "dataset_type": "t2v", |
| 4 | - "use_feature_data": false, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | "data_path": "/home/ci_resource/data/cogvideox1_0/data.jsonl", | 5 | "data_path": "/home/ci_resource/data/cogvideox1_0/data.jsonl", |
| 7 | "data_folder": "/home/ci_resource/data/cogvideox1_0/", | 6 | "data_folder": "/home/ci_resource/data/cogvideox1_0/", |
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "/home/ci_resource/models/hunyuanvideo_t2v/data.txt", | 5 | + "data_path": "/home/ci_resource/data/hunyuanvideo_t2v/feature_data/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "/home/ci_resource/data/hunyuanvideo_t2v/feature_data/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -69,7 +66,7 @@ | |||
| 69 | "dataloader_param": { | 66 | "dataloader_param": { |
| 70 | "dataloader_mode": "sampler", | 67 | "dataloader_mode": "sampler", |
| 71 | "sampler_type": "LengthGroupedSampler", | 68 | "sampler_type": "LengthGroupedSampler", |
| 72 | - "shuffle": false, | 69 | + "shuffle": true, |
| 73 | "drop_last": true, | 70 | "drop_last": true, |
| 74 | "pin_memory": true, | 71 | "pin_memory": true, |
| 75 | "group_frame": false, | 72 | "group_frame": false, |
| @@ -67,7 +67,7 @@ | |||
| 67 | "dataloader_param": { | 67 | "dataloader_param": { |
| 68 | "dataloader_mode": "sampler", | 68 | "dataloader_mode": "sampler", |
| 69 | "sampler_type": "LengthGroupedSampler", | 69 | "sampler_type": "LengthGroupedSampler", |
| 70 | - "shuffle": false, | 70 | + "shuffle": true, |
| 71 | "drop_last": true, | 71 | "drop_last": true, |
| 72 | "pin_memory": false, | 72 | "pin_memory": false, |
| 73 | "group_frame": false, | 73 | "group_frame": false, |
| @@ -61,7 +61,7 @@ | |||
| 61 | "dataloader_param": { | 61 | "dataloader_param": { |
| 62 | "dataloader_mode": "sampler", | 62 | "dataloader_mode": "sampler", |
| 63 | "sampler_type": "LengthGroupedSampler", | 63 | "sampler_type": "LengthGroupedSampler", |
| 64 | - "shuffle": false, | 64 | + "shuffle": true, |
| 65 | "drop_last": true, | 65 | "drop_last": true, |
| 66 | "pin_memory": true, | 66 | "pin_memory": true, |
| 67 | "group_frame": false, | 67 | "group_frame": false, |
| @@ -1,12 +1,9 @@ | |||
| 1 | { | 1 | { |
| 2 | "dataset_param": { | 2 | "dataset_param": { |
| 3 | - "dataset_type": "t2v", | 3 | + "dataset_type": "feature", |
| 4 | - "use_feature_data": true, | ||
| 5 | "basic_parameters": { | 4 | "basic_parameters": { |
| 6 | - "data_path": "/home/ci_resource/data/wan2.1/data.txt", | 5 | + "data_path": "/home/ci_resource/models/wan2.1/sora_features/data.jsonl", |
| 7 | - "data_storage_mode": "sorafeatured", | 6 | + "data_folder": "/home/ci_resource/models/wan2.1/sora_features/" |
| 8 | - "use_text_feature": true, | ||
| 9 | - "use_video_feature": true | ||
| 10 | }, | 7 | }, |
| 11 | "preprocess_parameters": { | 8 | "preprocess_parameters": { |
| 12 | "video_processor_type": "OpensoraplanVideoProcessor", | 9 | "video_processor_type": "OpensoraplanVideoProcessor", |
| @@ -39,7 +36,7 @@ | |||
| 39 | "dataloader_param": { | 36 | "dataloader_param": { |
| 40 | "dataloader_mode": "sampler", | 37 | "dataloader_mode": "sampler", |
| 41 | "sampler_type": "LengthGroupedSampler", | 38 | "sampler_type": "LengthGroupedSampler", |
| 42 | - "shuffle": false, | 39 | + "shuffle": true, |
| 43 | "drop_last": true, | 40 | "drop_last": true, |
| 44 | "pin_memory": true, | 41 | "pin_memory": true, |
| 45 | "group_frame": false, | 42 | "group_frame": false, |
会提另一个pr，最后两个pr会同步合入