File size: 1,937 Bytes
ec0a9aa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 | import json
import os
def load_prompt_or_image(prompt_source, prompt_idx, prompt, image_path):
"""
Load the prompt or image path based on the prompt source.
"""
if prompt_source == "prompt":
assert prompt_idx == 0, "You have already provided a prompt"
return prompt, image_path
elif prompt_source == "I2V_VBench":
# assert prompt is a json file
assert prompt.endswith(".json"), "Prompt must be a json file"
with open(prompt, "r") as f:
prompts = json.load(f)
prompt_idx = str(prompt_idx)
original_prompt = prompts[prompt_idx]["original"]
improved_prompt = prompts[prompt_idx]["improved"]
image_path = os.path.join(image_path, f"{original_prompt}.jpg")
assert os.path.exists(image_path), "Image path does not exist"
return improved_prompt, image_path
elif prompt_source == "I2V_Wan_Web":
assert prompt == image_path, "Prompt and image path must be the same"
prompt_idx = str(prompt_idx).zfill(3)
prompt_path = os.path.join(prompt, f"{prompt_idx}/prompt.txt")
image_path = os.path.join(image_path, f"{prompt_idx}/image.jpg")
with open(prompt_path, "r") as f:
prompt = f.read()
return prompt, image_path
elif prompt_source in ["T2V_Wan_VBench", "T2V_Hyv_VBench", "T2V_Hyv_Web"]:
assert prompt.endswith(".txt"), "Prompt must be a txt file"
with open(prompt, "r") as f:
prompts = f.readlines()
prompt = prompts[prompt_idx]
return prompt, None
elif prompt_source in ["T2V_Xingyang_Motion", "T2V_Xingyang_VBench"]:
assert prompt.endswith(".txt"), "Prompt must be a txt file"
with open(prompt, "r") as f:
prompts = f.readlines()
prompt = prompts[prompt_idx]
return prompt, None
else:
raise ValueError(f"Invalid prompt source: {prompt_source}")
|