data_utils/extract_training_video.py

import argparse
import os
import torch
import torch.distributed as dist
from torch.utils.data import DataLoader
from tqdm import tqdm  # 导入 tqdm
import logging  # 导入 logging
import sys
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
from data_utils.v2a_utils.vggsound_224 import VGGSound
from data_utils.v2a_utils.feature_utils_224 import FeaturesUtils
import torchaudio
from einops import rearrange
from torch.utils.data.dataloader import default_collate


# 设置日志配置
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def setup(rank, world_size):
    dist.init_process_group("nccl", rank=rank, world_size=world_size)
    torch.cuda.set_device(rank)

def cleanup():
    dist.destroy_process_group()

def error_avoidance_collate(batch):
    batch = list(filter(lambda x: x is not None, batch))
    return default_collate(batch)

def main(args):
    rank = int(os.environ["RANK"])  # 从环境变量中获取 rank
    world_size = int(os.environ["WORLD_SIZE"])  # 从环境变量中获取 world size
    setup(rank, world_size)

    dataset = VGGSound(
        root=args.root,
        tsv_path=args.tsv_path,
        sample_rate=args.sample_rate,
        duration_sec=args.duration_sec,
        audio_samples=args.audio_samples,
        start_row=args.start_row,
        end_row=args.end_row,
        save_dir=args.save_dir
    )
    save_dir = args.save_dir
    os.makedirs(save_dir, exist_ok=True)
    # 使用 DataLoader 加载数据集，增加 batch_size 和 num_workers
    train_sampler = torch.utils.data.distributed.DistributedSampler(dataset, num_replicas=world_size, rank=rank)
    dataloader = DataLoader(dataset, batch_size=1, sampler=train_sampler, num_workers=8, drop_last=False,collate_fn=error_avoidance_collate)

    feature_extractor = FeaturesUtils(
        vae_ckpt=args.vae_ckpt,
        vae_config=args.vae_config,
        enable_conditions=True,
        synchformer_ckpt=args.synchformer_ckpt
    ).eval().cuda(rank)

    # 使用 DistributedDataParallel 支持多显卡
    feature_extractor = torch.nn.parallel.DistributedDataParallel(feature_extractor, device_ids=[rank])

    # 使用 tqdm 显示进度条
    for i, data in enumerate(tqdm(dataloader, desc="Processing", unit="batch")):
        # 使用 torch.no_grad() 来加快推理速度
        ids = data['id']  # 获取当前批次的所有 ID
        with torch.no_grad():
            audio = data['audio'].cuda(rank, non_blocking=True)

            output = {
                'caption': data['caption'],
                'caption_cot': data['caption_cot']
            }
            # logging.info(f'Processing batch {i} with IDs: {ids}')  # 添加日志记录

            latent = feature_extractor.module.encode_audio(audio)
            output['latent'] = latent.detach().cpu()

            clip_video = data['clip_video'].cuda(rank, non_blocking=True)
            # logging.info(f'Processing batch {i} with shape: {clip_video.shape}')  # 添加日志记录
            clip_features = feature_extractor.module.encode_video_with_clip(clip_video)
            output['metaclip_features'] = clip_features.detach().cpu()

            sync_video = data['sync_video'].cuda(rank, non_blocking=True)
            sync_features = feature_extractor.module.encode_video_with_sync(sync_video)
            output['sync_features'] = sync_features.detach().cpu()

            caption = data['caption']
            metaclip_global_text_features, metaclip_text_features = feature_extractor.module.encode_text(caption)
            output['metaclip_global_text_features'] = metaclip_global_text_features.detach().cpu()
            output['metaclip_text_features'] = metaclip_text_features.detach().cpu()

            caption_cot = data['caption_cot']
            t5_features = feature_extractor.module.encode_t5_text(caption_cot)
            output['t5_features'] = t5_features.detach().cpu()


            # 保存每个样本的输出
            for j in range(audio.size(0)):  # 遍历当前批次的每个样本
                sample_output = {
                    'id': ids[j],
                    'caption': output['caption'][j],
                    'caption_cot': output['caption_cot'][j],
                    'latent': output['latent'][j],
                    'metaclip_features': output['metaclip_features'][j],
                    'sync_features': output['sync_features'][j],
                    'metaclip_global_text_features': output['metaclip_global_text_features'][j],
                    'metaclip_text_features': output['metaclip_text_features'][j],
                    't5_features': output['t5_features'][j],
                }
                torch.save(sample_output, f'{save_dir}/{ids[j]}.pth')

        ## test the sync between videos and audios
        # torchaudio.save(f'input_{i}.wav',data['audio'],sample_rate=44100)
        # recon_audio = feature_extractor.decode_audio(latent)
        # recon_audio = rearrange(recon_audio, "b d n -> d (b n)")
        # id = data['id']
        # torchaudio.save(f'recon_{i}.wav',recon_audio.cpu(),sample_rate=44100)
        # os.system(f'ffmpeg -y -i dataset/vggsound/video/train/{id}.mp4 -i recon_{i}.wav -t 9 -map 0:v -map 1:a -c:v copy -c:a aac -strict experimental -shortest out_{i}.mp4')
        # os.system(f'ffmpeg -y -i dataset/vggsound/video/train/{id}.mp4 -i input_{i}.wav -t 9 -map 0:v -map 1:a -c:v copy -c:a aac -strict experimental -shortest input_{i}.mp4')
    cleanup()
    

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Extract Video Training Latents')
    parser.add_argument('--root', type=str, default='dataset/vggsound/video/test', help='Root directory of the video dataset')
    parser.add_argument('--tsv_path', type=str, default='dataset/vggsound/split_txt/test_caption.csv', help='Path to the TSV file')
    parser.add_argument('--save-dir', type=str, default='dataset/vggsound/video_text_latents/test', help='Save Directory')
    parser.add_argument('--sample_rate', type=int, default=44100, help='Sample rate of the audio')
    parser.add_argument('--duration_sec', type=float, default=9.0, help='Duration of the audio in seconds')
    parser.add_argument('--audio_samples', type=int, default=397312, help='Number of audio samples')
    parser.add_argument('--vae_ckpt', type=str, default='ckpts/vae.ckpt', help='Path to the VAE checkpoint')
    parser.add_argument('--vae_config', type=str, default='PrismAudio/configs/model_configs/stable_audio_2_0_vae.json', help='Path to the VAE configuration file')
    parser.add_argument('--synchformer_ckpt', type=str, default='ckpts/synchformer_state_dict.pth', help='Path to the Synchformer checkpoint')
    parser.add_argument('--start-row', type=int, default=0, help='start row')
    parser.add_argument('--end-row', type=int, default=None, help='end row')

    args = parser.parse_args()

    # 直接使用 torch.distributed.launch 启动
    main(args=args)  # 这里的 rank 需要在命令行中指定