import yaml  
import os  
# 假设对应关系存储在名为 "file_pairs.txt" 的文本文件中  
file_pairs_file = "./dataset/ViViD/upper_body/test_pairs.txt"  
output_yaml_path = "./configs/prompts/upper_body2.yaml"  # 输出的 YAML 文件路径  
videos_dir = "./dataset/ViViD/upper_body/videos"  
images_dir = "./dataset/ViViD/upper_body/images"  
# 准备要写入 YAML 的数据结构  
yaml_data = {  
    "pretrained_base_model_path": "ckpts/sd-image-variations-diffusers",  
    "pretrained_vae_path": "ckpts/sd-vae-ft-mse",  
    "image_encoder_path": "ckpts/sd-image-variations-diffusers/image_encoder",  
    "denoising_unet_path": "ckpts/ViViD/denoising_unet.pth",  
    "reference_unet_path": "ckpts/ViViD/reference_unet.pth",  
    "pose_guider_path": "ckpts/ViViD/pose_guider.pth",  
    "motion_module_path": "ckpts/MotionModule/mm_sd_v15_v2.ckpt",  
    "inference_config": "./configs/inference/inference.yaml",  
    "weight_dtype": "fp16",  
    "model_video_paths": [],  
    "cloth_image_paths": []  
}  


# 读取文本文件并填充 YAML 数据结构  
with open(file_pairs_file, 'r') as file:  
    for line in file:  
        # 每行可能是 "视频文件路径 对应图像文件路径"  
        video_file_name, image_file_name = line.strip().split()  # 假设用空格分隔  
        # 构建完整的路径  
        video_path = os.path.join(videos_dir, video_file_name)  # 完整视频文件路径  
        image_path = os.path.join(images_dir, image_file_name)  # 完整图像文件路径  
        yaml_data["model_video_paths"].append(video_path)  # 添加视频文件路径  
        yaml_data["cloth_image_paths"].append(image_path)  # 添加图像文件路径  

# 将数据写入 YAML 文件  
with open(output_yaml_path, 'w') as yaml_file:  
    yaml.dump(yaml_data, yaml_file, default_flow_style=False)  

print(f"YAML 文件已生成: {output_yaml_path}")