返回 VideoClaw
test_cot.py
根目录 / FilmAgent / FilmAgent / test_cot.py
1 from FilmAgent_root.FilmAgent.util import *
2 from FilmAgent_root.FilmAgent.LLMCaller import *
3 from typing import Dict, List, Union
4
5 # TO DO
6 ROOT_PATH = "/path/to/FilmAgent"
7 model = "gpt-4o"
8 # TO DO
9
10 topics=["Reconcilation in a friend reunion", "A quarrel and breakup scene", "Casual meet-up with an old friend", "Emergency meeting after a security breach", "Late night brainstorming for a startup", "Family argument during dinner", "Emotional farewell at the roadside", "Heated debate over investments in the office", "Heated family discussion ending in a heartfelt apology", "Office gossip turning into a major understanding", "Celebratory end of project cheers with team members", "Planning a secret escape from a mundane routine", "Unexpected guest crashes a small house party", "An employee's emotional breakdown after being terminated", "Confession of a long-held secret among close friends"]
11
12 class FilmCrafter:
13
14 def __init__(self, topic: str, ID) -> None:
15 self.topic = topic
16 self.store_path = os.path.join(ROOT_PATH, f"store\cot\{ID}")
17 self.log_path = os.path.join(self.store_path, "prompt.txt")
18 self.profile_path = os.path.join(self.store_path, "actors_profile.json")
19 self.action_description_path = os.path.join(ROOT_PATH, "Locations\\actions.txt")
20 self.shot_description_path = os.path.join(ROOT_PATH, "Locations\\shots.txt")
21 self.script_path = os.path.join(self.store_path, "script.json")
22
23 # The maximum number of characters in a film
24 self.character_limit = 4
25
26 if not os.path.exists(self.store_path):
27 os.makedirs(self.store_path)
28
29
30 def call(self, identity: str, params: Dict, trans2json: bool = True) -> Union[str, dict, list]:
31 prompt = read_prompt(os.path.join(ROOT_PATH, f"Prompt\COT_Prompt\{identity}.txt") )
32 prompt = prompt_format(prompt, params)
33 log_prompt(self.log_path, prompt)
34 result = LLMCall(prompt, model)
35 if trans2json:
36 result = clean_text(result)
37 result = GPTResponse2JSON(result)
38 log_prompt(self.log_path, result)
39 return result
40
41
42 def casting(self):
43 params = {"{topic}": self.topic, "{character_limit}": self.character_limit}
44 result = self.call("director_1", params)
45 write_json(self.profile_path, result)
46
47
48 def script_(self):
49 profile = read_json(self.profile_path)
50 male_characters = ", ".join(list(map(lambda x: x['name'],
51 filter(lambda x: x['gender'].lower() == 'male', profile))))
52 female_characters = ", ".join(list(map(lambda x: x['name'],
53 filter(lambda x: x['gender'].lower() == 'female', profile))))
54 optional_positions = ""
55 base_path = os.path.join(ROOT_PATH, "Locations")
56 for entry in os.listdir(base_path):
57 full_path = os.path.join(base_path, entry)
58 if os.path.isdir(full_path):
59 ps = json.dumps(read_json(os.path.join(full_path, "position.json")))
60 optional_positions = optional_positions + f"**{entry}**: {ps}\n"
61
62 all_actions = read_prompt(self.action_description_path)
63 all_shots = read_prompt(self.shot_description_path)
64
65 params = {"{topic}": self.topic,
66 "{male_characters}": male_characters,
67 "{female_characters}": female_characters,
68 "{profiles}": profile,
69 "{optional_positions}": optional_positions,
70 "{all_actions}": all_actions,
71 "{all_shots}": all_shots
72 }
73 result = self.call("script", params)
74 write_json(self.script_path, result)
75
76
77 if __name__ == '__main__':
78 id=14
79 f = FilmCrafter(topic = topics[id-1], ID=id)
80 f.casting()
81 f.script_()
82
83
83 lines PYTHON