file_path = '/home/lizhikai/webvid_prompt100.txt'
str_list = []
with open(file_path, 'r', encoding='utf-8') as file:
    for line in file:
        str_list.append(line.strip())
        if len(str_list) == 100:
            break

def generate_image_ig_api(prompt, model_name):
    model_source, model_name, model_type = model_name.split("_")
    pipe = load_replicate_model(model_name, model_type)
    result = pipe(prompt=prompt)
    return result
model_names = ['replicate_Zeroscope-v2-xl_text2video',
                    #    'replicate_Damo-Text-to-Video_text2video',
                        'replicate_Animate-Diff_text2video',
                        'replicate_OpenSora_text2video',
                        'replicate_LaVie_text2video',
                        'replicate_VideoCrafter2_text2video',
                        'replicate_Stable-Video-Diffusion_text2video',
                        ]
save_names = []
for name in model_names:
    model_source, model_name, model_type = name.split("_")
    save_names.append(model_name)

for i, prompt in enumerate(str_list):
    print("save the {} prompt".format(i+1))
    # if i+1 < 97:
    #     continue
    with concurrent.futures.ThreadPoolExecutor() as executor:
        futures = [executor.submit(generate_image_ig_api, prompt, model) for model in model_names]
        results = [future.result() for future in futures]

    root_dir = '/mnt/data/lizhikai/ksort_video_cache/'
    save_dir = os.path.join(root_dir, f'cache_{i+1}')
    if not os.path.exists(save_dir):
        os.makedirs(save_dir, exist_ok=True)
    with open(os.path.join(save_dir, "prompt.txt"), 'w', encoding='utf-8') as file:
        file.write(prompt)

    # 下载视频并保存
    repeat_num = 5
    for j, url in enumerate(results):
        while 1:
            time.sleep(1)
            response = requests.get(url, stream=True)
            if response.status_code == 200:
                file_path = os.path.join(save_dir, f'{save_names[j]}.mp4')
                with open(file_path, 'wb') as file:
                    for chunk in response.iter_content(chunk_size=8192):
                        file.write(chunk)
                print(f"视频 {j} 已保存到 {file_path}")
                break
            else:
                repeat_num = repeat_num - 1
                if repeat_num == 0:
                    print(f"视频 {j} 保存失败")
                    # raise ValueError("Video request failed.")
                continue