返回 ViMax
idea2video_pipeline.py
根目录 / pipelines / idea2video_pipeline.py
1 import os
2 import shutil
3 import logging
4 from agents import Screenwriter, CharacterExtractor, CharacterPortraitsGenerator
5 from pipelines.script2video_pipeline import Script2VideoPipeline
6 from interfaces import CharacterInScene
7 from typing import List, Dict, Optional
8 import asyncio
9 import json
10 import yaml
11 from langchain.chat_models import init_chat_model
12 from tools.render_backend import RenderBackend
13 from utils.provider_presets import resolve_chat_model_config
14 from utils.text import safe_path_component
15 from utils.video import concatenate_video_files
16
17
18 def _pipeline_print(quiet: bool, message: str) -> None:
19 if not quiet:
20 print(message)
21
22
23 class Idea2VideoPipeline:
24 def __init__(
25 self,
26 chat_model: str,
27 image_generator: str,
28 video_generator: str,
29 working_dir: str,
30 ):
31 self.chat_model = chat_model
32 self.image_generator = image_generator
33 self.video_generator = video_generator
34 self.working_dir = working_dir
35 os.makedirs(self.working_dir, exist_ok=True)
36
37 self.screenwriter = Screenwriter(chat_model=self.chat_model)
38 self.character_extractor = CharacterExtractor(
39 chat_model=self.chat_model)
40 self.character_portraits_generator = CharacterPortraitsGenerator(
41 image_generator=self.image_generator)
42
43 @classmethod
44 def init_from_config(cls, config_path: str):
45 with open(config_path, "r") as f:
46 config = yaml.safe_load(f)
47
48 chat_model_args = resolve_chat_model_config(config["chat_model"]["init_args"])
49 chat_model = init_chat_model(**chat_model_args)
50 backend = RenderBackend.from_config(config)
51
52 return cls(
53 chat_model=chat_model,
54 image_generator=backend.image_generator,
55 video_generator=backend.video_generator,
56 working_dir=config["working_dir"],
57 )
58
59 async def extract_characters(
60 self,
61 story: str,
62 quiet: bool = False,
63 ):
64 save_path = os.path.join(self.working_dir, "characters.json")
65
66 if os.path.exists(save_path):
67 with open(save_path, "r", encoding="utf-8") as f:
68 characters = json.load(f)
69 characters = [CharacterInScene.model_validate(
70 character) for character in characters]
71 _pipeline_print(quiet, f"🚀 Loaded {len(characters)} characters from existing file.")
72 else:
73 characters = await self.character_extractor.extract_characters(story)
74 with open(save_path, "w", encoding="utf-8") as f:
75 json.dump([character.model_dump()
76 for character in characters], f, ensure_ascii=False, indent=4)
77 _pipeline_print(quiet, f"✅ Extracted {len(characters)} characters from story and saved to {save_path}.")
78
79 return characters
80
81 async def generate_character_portraits(
82 self,
83 characters: List[CharacterInScene],
84 character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]],
85 style: str,
86 ):
87 character_portraits_registry_path = os.path.join(
88 self.working_dir, "character_portraits_registry.json")
89 if character_portraits_registry is None:
90 if os.path.exists(character_portraits_registry_path):
91 with open(character_portraits_registry_path, 'r', encoding='utf-8') as f:
92 character_portraits_registry = json.load(f)
93 else:
94 character_portraits_registry = {}
95
96 tasks = [
97 self.generate_portraits_for_single_character(character, style)
98 for character in characters
99 if character.identifier_in_scene not in character_portraits_registry
100 # Characters never shown on screen (e.g. a voice or chat-only
101 # character) have no physical description, so asking the image
102 # model for front/side/back portraits of them is nonsensical and
103 # fails repeatedly (finish_reason=IMAGE_OTHER, empty candidates).
104 and character.is_visible
105 ]
106 if tasks:
107 for future in asyncio.as_completed(tasks):
108 character_portraits_registry.update(await future)
109 with open(character_portraits_registry_path, 'w', encoding='utf-8') as f:
110 json.dump(character_portraits_registry,
111 f, ensure_ascii=False, indent=4)
112
113 print(
114 f"✅ Completed character portrait generation for {len(characters)} characters.")
115 else:
116 print(
117 "🚀 All characters already have portraits, skipping portrait generation.")
118
119 return character_portraits_registry
120
121 async def develop_story(
122 self,
123 idea: str,
124 user_requirement: str,
125 quiet: bool = False,
126 ):
127 save_path = os.path.join(self.working_dir, "story.txt")
128 if os.path.exists(save_path):
129 with open(save_path, "r", encoding="utf-8") as f:
130 story = f.read()
131 _pipeline_print(quiet, f"🚀 Loaded story from existing file.")
132 else:
133 _pipeline_print(quiet, "🧠 Developing story...")
134 story = await self.screenwriter.develop_story(idea=idea, user_requirement=user_requirement)
135 with open(save_path, "w", encoding="utf-8") as f:
136 f.write(story)
137 _pipeline_print(quiet, f"✅ Developed story and saved to {save_path}.")
138
139 return story
140
141 async def write_script_based_on_story(
142 self,
143 story: str,
144 user_requirement: str,
145 quiet: bool = False,
146 ):
147 save_path = os.path.join(self.working_dir, "script.json")
148 if os.path.exists(save_path):
149 with open(save_path, "r", encoding="utf-8") as f:
150 script = json.load(f)
151 _pipeline_print(quiet, f"🚀 Loaded script from existing file.")
152 else:
153 _pipeline_print(quiet, "🧠 Writing script based on story...")
154 script = await self.screenwriter.write_script_based_on_story(story=story, user_requirement=user_requirement)
155 with open(save_path, "w", encoding="utf-8") as f:
156 json.dump(script, f, ensure_ascii=False, indent=4)
157 _pipeline_print(quiet, f"✅ Written script based on story and saved to {save_path}.")
158 return script
159
160 async def generate_portraits_for_single_character(
161 self,
162 character: CharacterInScene,
163 style: str,
164 ):
165 character_dir = os.path.join(
166 self.working_dir, "character_portraits", f"{character.idx}_{safe_path_component(character.identifier_in_scene)}")
167 os.makedirs(character_dir, exist_ok=True)
168
169 front_portrait_path = os.path.join(character_dir, "front.png")
170 if os.path.exists(front_portrait_path):
171 pass
172 else:
173 front_portrait_output = await self.character_portraits_generator.generate_front_portrait(character, style)
174 front_portrait_output.save(front_portrait_path)
175
176 side_portrait_path = os.path.join(character_dir, "side.png")
177 if os.path.exists(side_portrait_path):
178 pass
179 else:
180 try:
181 side_portrait_output = await self.character_portraits_generator.generate_side_portrait(character, front_portrait_path)
182 side_portrait_output.save(side_portrait_path)
183 except Exception as e:
184 # gemini-2.5-flash-image intermittently (sometimes beyond
185 # the tenacity retry budget) fails this front->side
186 # re-angling edit with finish_reason=IMAGE_OTHER / empty
187 # content. Fall back to the front portrait rather than
188 # aborting the whole pipeline.
189 print(f"⚠️ Side portrait generation failed for {character.identifier_in_scene} after retries ({e}); reusing front portrait as fallback.")
190 shutil.copy(front_portrait_path, side_portrait_path)
191
192 back_portrait_path = os.path.join(character_dir, "back.png")
193 if os.path.exists(back_portrait_path):
194 pass
195 else:
196 try:
197 back_portrait_output = await self.character_portraits_generator.generate_back_portrait(character, front_portrait_path)
198 back_portrait_output.save(back_portrait_path)
199 except Exception as e:
200 print(f"⚠️ Back portrait generation failed for {character.identifier_in_scene} after retries ({e}); reusing front portrait as fallback.")
201 shutil.copy(front_portrait_path, back_portrait_path)
202
203 print(
204 f"☑️ Completed character portrait generation for {character.identifier_in_scene}.")
205
206 return {
207 character.identifier_in_scene: {
208 "front": {
209 "path": front_portrait_path,
210 "description": f"A front view portrait of {character.identifier_in_scene}.",
211 },
212 "side": {
213 "path": side_portrait_path,
214 "description": f"A side view portrait of {character.identifier_in_scene}.",
215 },
216 "back": {
217 "path": back_portrait_path,
218 "description": f"A back view portrait of {character.identifier_in_scene}.",
219 },
220 }
221 }
222
223 async def __call__(
224 self,
225 idea: str,
226 user_requirement: str,
227 style: str,
228 quiet: bool = False,
229 ):
230
231 story = await self.develop_story(idea=idea, user_requirement=user_requirement, quiet=quiet)
232
233 characters = await self.extract_characters(story=story, quiet=quiet)
234
235 character_portraits_registry = await self.generate_character_portraits(
236 characters=characters,
237 character_portraits_registry=None,
238 style=style,
239 )
240
241 scene_scripts = await self.write_script_based_on_story(story=story, user_requirement=user_requirement, quiet=quiet)
242
243 all_video_paths = []
244
245 for idx, scene_script in enumerate(scene_scripts):
246 scene_working_dir = os.path.join(self.working_dir, f"scene_{idx}")
247 os.makedirs(scene_working_dir, exist_ok=True)
248 script2video_pipeline = Script2VideoPipeline(
249 chat_model=self.chat_model,
250 image_generator=self.image_generator,
251 video_generator=self.video_generator,
252 working_dir=scene_working_dir,
253 )
254 final_video_path = await script2video_pipeline(
255 script=scene_script,
256 user_requirement=user_requirement,
257 style=style,
258 characters=characters,
259 character_portraits_registry=character_portraits_registry,
260 quiet=quiet,
261 )
262 all_video_paths.append(final_video_path)
263
264 final_video_path = os.path.join(self.working_dir, "final_video.mp4")
265 if os.path.exists(final_video_path):
266 _pipeline_print(quiet, f"🚀 Skipped concatenating videos, already exists.")
267 else:
268 _pipeline_print(quiet, f"🎬 Starting concatenating videos...")
269 concatenate_video_files(all_video_paths, final_video_path)
270 _pipeline_print(quiet, f"☑️ Concatenated videos, saved to {final_video_path}.")
271 return final_video_path
272
272 lines PYTHON