K-Sort-Arena / model /models /generate_video_cache.py
YangZhoumill's picture
pre-release
b257e01
raw
history blame
2.48 kB
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