| 1 | import os |
| 2 | import shutil |
| 3 | import json |
| 4 | import logging |
| 5 | import asyncio |
| 6 | import time |
| 7 | import hashlib |
| 8 | from pathlib import Path |
| 9 | from uuid import uuid4 |
| 10 | from typing import Any, Callable, Optional, Dict, List, Tuple, Literal, Type, TypeVar |
| 11 | from moviepy import VideoFileClip, concatenate_videoclips |
| 12 | from PIL import Image |
| 13 | from agents import * |
| 14 | from agents.best_image_selector import BestImageSelector |
| 15 | import yaml |
| 16 | from interfaces import * |
| 17 | from langchain.chat_models import init_chat_model |
| 18 | from tools.render_backend import RenderBackend |
| 19 | from utils.provider_presets import resolve_chat_model_config |
| 20 | from utils.image_selection import DEFAULT_IMAGE_CANDIDATES, image_candidate_count_from_config, validate_image_candidate_count |
| 21 | |
| 22 | |
| 23 | |
| 24 | |
| 25 | TModel = TypeVar("TModel") |
| 26 | |
| 27 | |
| 28 | def _normalize_model_list(items: Any, model_cls: Type[TModel], field_name: str) -> List[TModel]: |
| 29 | if items is None: |
| 30 | return [] |
| 31 | if not isinstance(items, list): |
| 32 | raise TypeError(f"{field_name} must be a list, got {type(items).__name__}") |
| 33 | normalized: List[TModel] = [] |
| 34 | for idx, item in enumerate(items): |
| 35 | if isinstance(item, model_cls): |
| 36 | normalized.append(item) |
| 37 | elif isinstance(item, dict): |
| 38 | normalized.append(model_cls.model_validate(item)) |
| 39 | else: |
| 40 | raise TypeError(f"{field_name}[{idx}] must be {model_cls.__name__} or dict, got {type(item).__name__}") |
| 41 | return normalized |
| 42 | |
| 43 | |
| 44 | def _group_shots_into_cameras(shot_descriptions: List[ShotDescription]) -> List[Camera]: |
| 45 | cameras_by_idx: Dict[int, Camera] = {} |
| 46 | for shot_description in shot_descriptions: |
| 47 | camera = cameras_by_idx.get(shot_description.cam_idx) |
| 48 | if camera is None: |
| 49 | camera = Camera(idx=shot_description.cam_idx, active_shot_idxs=[]) |
| 50 | cameras_by_idx[shot_description.cam_idx] = camera |
| 51 | camera.active_shot_idxs.append(shot_description.idx) |
| 52 | return list(cameras_by_idx.values()) |
| 53 | |
| 54 | def _collect_priority_shot_idxs(camera_tree: List[Camera]) -> List[int]: |
| 55 | """Shot indices that other cameras depend on.""" |
| 56 | return [camera.parent_shot_idx for camera in camera_tree if camera.parent_shot_idx is not None] |
| 57 | |
| 58 | |
| 59 | def _pipeline_print(quiet: bool, message: str) -> None: |
| 60 | if not quiet: |
| 61 | print(message) |
| 62 | |
| 63 | |
| 64 | def _emit_text_plan_progress(progress, stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None: |
| 65 | if progress is not None: |
| 66 | progress(stage, message, metadata or {}) |
| 67 | |
| 68 | |
| 69 | def _emit_render_progress(progress, stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None: |
| 70 | if progress is not None: |
| 71 | progress(stage, message, metadata or {}) |
| 72 | |
| 73 | |
| 74 | def _scoped_progress(progress, **scope): |
| 75 | if progress is None: |
| 76 | return None |
| 77 | |
| 78 | def emit(stage: str, message: str, metadata: Dict[str, Any] | None = None) -> None: |
| 79 | payload = dict(scope) |
| 80 | payload.update(metadata or {}) |
| 81 | _emit_render_progress(progress, stage, message, payload) |
| 82 | |
| 83 | return emit |
| 84 | |
| 85 | |
| 86 | class Script2VideoPipeline: |
| 87 | |
| 88 | def __init__( |
| 89 | self, |
| 90 | chat_model: str, |
| 91 | image_generator, |
| 92 | video_generator, |
| 93 | working_dir: str, |
| 94 | num_image_candidates: int = DEFAULT_IMAGE_CANDIDATES, |
| 95 | ): |
| 96 | |
| 97 | self.chat_model = chat_model |
| 98 | self.image_generator = image_generator |
| 99 | self.video_generator = video_generator |
| 100 | self.num_image_candidates = validate_image_candidate_count(num_image_candidates) |
| 101 | self.best_image_selector = BestImageSelector(chat_model=self.chat_model) |
| 102 | |
| 103 | self.character_extractor = CharacterExtractor(chat_model=self.chat_model) |
| 104 | self.character_portraits_generator = CharacterPortraitsGenerator(image_generator=self.image_generator) |
| 105 | self.storyboard_artist = StoryboardArtist(chat_model=self.chat_model) |
| 106 | self.camera_image_generator = CameraImageGenerator(chat_model=self.chat_model, image_generator=self.image_generator, video_generator=self.video_generator) |
| 107 | self.reference_image_selector = ReferenceImageSelector(chat_model=self.chat_model) |
| 108 | |
| 109 | self.working_dir = working_dir |
| 110 | os.makedirs(self.working_dir, exist_ok=True) |
| 111 | self.character_portrait_events = {} |
| 112 | self.shot_desc_events = {} |
| 113 | self.frame_events = {} |
| 114 | |
| 115 | |
| 116 | async def plan_text_artifacts( |
| 117 | self, |
| 118 | script: str, |
| 119 | user_requirement: str, |
| 120 | style: str, |
| 121 | characters: List[CharacterInScene] = None, |
| 122 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 123 | quiet: bool = False, |
| 124 | ): |
| 125 | """Generate only structured text artifacts required before rendering. |
| 126 | |
| 127 | This helper intentionally stops before character portraits, frame generation, |
| 128 | video generation, and final concatenation so an agent loop can pause for |
| 129 | user review after narrative planning. |
| 130 | """ |
| 131 | self.character_portrait_events = {} |
| 132 | self.shot_desc_events = {} |
| 133 | self.frame_events = {} |
| 134 | |
| 135 | if characters is None: |
| 136 | _emit_text_plan_progress(progress, "extract_characters", "Extracting characters from script") |
| 137 | characters = await self.extract_characters(script=script, quiet=quiet) |
| 138 | else: |
| 139 | characters = _normalize_model_list(characters, CharacterInScene, "characters") |
| 140 | _emit_text_plan_progress(progress, "extract_characters", "Using provided characters", {"provided": True, "count": len(characters)}) |
| 141 | characters_path = os.path.join(self.working_dir, "characters.json") |
| 142 | if not os.path.exists(characters_path): |
| 143 | with open(characters_path, "w", encoding="utf-8") as f: |
| 144 | json.dump([character.model_dump() for character in characters], f, ensure_ascii=False, indent=4) |
| 145 | for character in characters: |
| 146 | self.character_portrait_events[character.idx] = asyncio.Event() |
| 147 | |
| 148 | _emit_text_plan_progress(progress, "design_storyboard", "Designing storyboard") |
| 149 | storyboard = await self.design_storyboard( |
| 150 | script=script, |
| 151 | characters=characters, |
| 152 | user_requirement=user_requirement, |
| 153 | quiet=quiet, |
| 154 | ) |
| 155 | _emit_text_plan_progress(progress, "decompose_shots", "Decomposing shot visual descriptions", {"shot_count": len(storyboard)}) |
| 156 | shot_descriptions = await self.decompose_visual_descriptions( |
| 157 | shot_brief_descriptions=storyboard, |
| 158 | characters=characters, |
| 159 | quiet=quiet, |
| 160 | ) |
| 161 | camera_tree = None |
| 162 | for attempt in range(2): |
| 163 | try: |
| 164 | stage = "construct_camera_tree" if attempt == 0 else "construct_camera_tree_retry" |
| 165 | message = "Constructing camera tree" if attempt == 0 else "Retrying camera tree construction after schema/type failure" |
| 166 | _emit_text_plan_progress(progress, stage, message, {"shot_count": len(shot_descriptions), "attempt": attempt + 1}) |
| 167 | camera_tree = await self.construct_camera_tree( |
| 168 | shot_descriptions=shot_descriptions, |
| 169 | quiet=quiet, |
| 170 | ) |
| 171 | break |
| 172 | except Exception: |
| 173 | camera_tree_path = os.path.join(self.working_dir, "camera_tree.json") |
| 174 | if os.path.exists(camera_tree_path): |
| 175 | os.remove(camera_tree_path) |
| 176 | if attempt == 1: |
| 177 | raise |
| 178 | assert camera_tree is not None |
| 179 | return { |
| 180 | "characters": characters, |
| 181 | "storyboard": storyboard, |
| 182 | "shot_descriptions": shot_descriptions, |
| 183 | "camera_tree": camera_tree, |
| 184 | } |
| 185 | |
| 186 | |
| 187 | @classmethod |
| 188 | def init_from_config(cls, config_path: str): |
| 189 | with open(config_path, "r") as f: |
| 190 | config = yaml.safe_load(f) |
| 191 | |
| 192 | chat_model_args = resolve_chat_model_config(config["chat_model"]["init_args"]) |
| 193 | chat_model = init_chat_model(**chat_model_args) |
| 194 | backend = RenderBackend.from_config(config) |
| 195 | |
| 196 | return cls( |
| 197 | chat_model=chat_model, |
| 198 | image_generator=backend.image_generator, |
| 199 | video_generator=backend.video_generator, |
| 200 | working_dir=config["working_dir"], |
| 201 | num_image_candidates=image_candidate_count_from_config(config), |
| 202 | ) |
| 203 | |
| 204 | async def __call__( |
| 205 | self, |
| 206 | script: str, |
| 207 | user_requirement: str, |
| 208 | style: str, |
| 209 | characters: List[CharacterInScene] = None, |
| 210 | character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]] = None, |
| 211 | quiet: bool = False, |
| 212 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 213 | ): |
| 214 | _emit_render_progress(progress, "render_start", "Starting script2video render") |
| 215 | if characters is None: |
| 216 | _emit_render_progress(progress, "extract_characters", "Extracting characters before render") |
| 217 | characters = await self.extract_characters(script=script, quiet=quiet) |
| 218 | |
| 219 | # characters_path = os.path.join(self.working_dir, "characters.json") |
| 220 | # if os.path.exists(characters_path): |
| 221 | # with open(characters_path, "r", encoding="utf-8") as f: |
| 222 | # characters = [CharacterInScene.model_validate(c) for c in json.load(f)] |
| 223 | # print(f"🚀 Loaded {len(characters)} characters from existing file.") |
| 224 | # else: |
| 225 | # print(f"🔍 Extracting characters from script...") |
| 226 | # characters = await self.extract_characters(script=script) |
| 227 | # with open(characters_path, "w", encoding="utf-8") as f: |
| 228 | # json.dump([c.model_dump() for c in characters], f, ensure_ascii=False, indent=4) |
| 229 | # print(f"☑️ Extracted {len(characters)} characters from script and saved to {characters_path}.") |
| 230 | else: |
| 231 | characters = _normalize_model_list(characters, CharacterInScene, "characters") |
| 232 | _emit_render_progress(progress, "extract_characters", "Using provided characters for render", {"provided": True, "count": len(characters)}) |
| 233 | for character in characters: |
| 234 | self.character_portrait_events[character.idx] = asyncio.Event() |
| 235 | |
| 236 | if character_portraits_registry is None: |
| 237 | character_portraits_registry_path = os.path.join(self.working_dir, "character_portraits_registry.json") |
| 238 | if os.path.exists(character_portraits_registry_path): |
| 239 | with open(character_portraits_registry_path, "r", encoding="utf-8") as f: |
| 240 | character_portraits_registry = json.load(f) |
| 241 | print(f"🚀 Loaded {len(character_portraits_registry)} character portraits from existing file.") |
| 242 | _emit_render_progress(progress, "character_portraits_loaded", "Loaded existing character portraits", {"count": len(character_portraits_registry)}) |
| 243 | else: |
| 244 | print(f"🔍 Generating character portraits...") |
| 245 | _emit_render_progress(progress, "character_portraits_start", "Generating character portraits", {"character_count": len(characters)}) |
| 246 | character_portraits_registry = await self.generate_character_portraits( |
| 247 | characters=characters, |
| 248 | character_portraits_registry=None, |
| 249 | style=style, |
| 250 | progress=progress, |
| 251 | ) |
| 252 | |
| 253 | with open(character_portraits_registry_path, "w", encoding="utf-8") as f: |
| 254 | json.dump(character_portraits_registry, f, ensure_ascii=False, indent=4) |
| 255 | print(f"☑️ Generated {len(character_portraits_registry)} character portraits and saved to {character_portraits_registry_path}.") |
| 256 | _emit_render_progress(progress, "character_portraits_done", "Character portraits ready", {"count": len(character_portraits_registry)}) |
| 257 | |
| 258 | |
| 259 | |
| 260 | # design shots |
| 261 | _emit_render_progress(progress, "load_storyboard", "Loading or designing storyboard") |
| 262 | storyboard = await self.design_storyboard( |
| 263 | script=script, |
| 264 | characters=characters, |
| 265 | user_requirement=user_requirement, |
| 266 | quiet=quiet, |
| 267 | ) |
| 268 | _emit_render_progress(progress, "storyboard_ready", "Storyboard ready", {"shot_count": len(storyboard)}) |
| 269 | |
| 270 | # decompose visual descriptions of shots |
| 271 | _emit_render_progress(progress, "load_shot_descriptions", "Loading or decomposing shot descriptions", {"shot_count": len(storyboard)}) |
| 272 | shot_descriptions = await self.decompose_visual_descriptions( |
| 273 | shot_brief_descriptions=storyboard, |
| 274 | characters=characters, |
| 275 | quiet=quiet, |
| 276 | ) |
| 277 | _emit_render_progress(progress, "shot_descriptions_ready", "Shot descriptions ready", {"shot_count": len(shot_descriptions)}) |
| 278 | |
| 279 | # construct camera tree |
| 280 | _emit_render_progress(progress, "load_camera_tree", "Loading or constructing camera tree", {"shot_count": len(shot_descriptions)}) |
| 281 | camera_tree = await self.construct_camera_tree( |
| 282 | shot_descriptions=shot_descriptions, |
| 283 | quiet=quiet, |
| 284 | ) |
| 285 | _emit_render_progress(progress, "camera_tree_ready", "Camera tree ready", {"camera_count": len(camera_tree)}) |
| 286 | |
| 287 | priority_shot_idxs = [camera.parent_cam_idx for camera in camera_tree if camera.parent_cam_idx is not None] |
| 288 | _emit_render_progress(progress, "frames_start", "Generating frames for cameras", {"camera_count": len(camera_tree), "shot_count": len(shot_descriptions)}) |
| 289 | tasks = [ |
| 290 | self.generate_frames_for_single_camera( |
| 291 | camera=camera, |
| 292 | shot_descriptions=shot_descriptions, |
| 293 | characters=characters, |
| 294 | character_portraits_registry=character_portraits_registry, |
| 295 | priority_shot_idxs=priority_shot_idxs, |
| 296 | progress=progress, |
| 297 | ) |
| 298 | for camera in camera_tree |
| 299 | ] |
| 300 | |
| 301 | _emit_render_progress(progress, "video_clips_start", "Generating video clips for shots", {"shot_count": len(shot_descriptions)}) |
| 302 | video_tasks = [ |
| 303 | self.generate_video_for_single_shot( |
| 304 | shot_description=shot_description, |
| 305 | progress=progress, |
| 306 | ) |
| 307 | for shot_description in shot_descriptions |
| 308 | ] |
| 309 | tasks.extend(video_tasks) |
| 310 | await asyncio.gather(*tasks) |
| 311 | |
| 312 | final_video_path = os.path.join(self.working_dir, "final_video.mp4") |
| 313 | if os.path.exists(final_video_path): |
| 314 | print(f"🚀 Skipped concatenating videos, already exists.") |
| 315 | _emit_render_progress(progress, "final_video_exists", "Final video already exists", {"path": final_video_path}) |
| 316 | else: |
| 317 | print(f"🎬 Starting concatenating videos...") |
| 318 | _emit_render_progress(progress, "concat_start", "Concatenating video clips", {"shot_count": len(shot_descriptions)}) |
| 319 | video_clips = [ |
| 320 | VideoFileClip(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "video.mp4")) |
| 321 | for shot_description in shot_descriptions |
| 322 | ] |
| 323 | final_video = concatenate_videoclips(video_clips) |
| 324 | final_video.write_videofile(final_video_path, codec="libx264", preset="medium") |
| 325 | print(f"☑️ Concatenated videos, saved to {final_video_path}.") |
| 326 | _emit_render_progress(progress, "concat_done", "Final video concatenated", {"path": final_video_path}) |
| 327 | |
| 328 | _emit_render_progress(progress, "render_done", "Script2video render complete", {"final_video_path": final_video_path}) |
| 329 | return final_video_path |
| 330 | |
| 331 | |
| 332 | async def generate_frames_for_single_camera( |
| 333 | self, |
| 334 | camera: Camera, |
| 335 | shot_descriptions: List[ShotDescription], |
| 336 | characters: List[CharacterInScene], |
| 337 | character_portraits_registry: Dict[str, Dict[str, Dict[str, str]]], |
| 338 | priority_shot_idxs: List[int], |
| 339 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 340 | ): |
| 341 | # 1. generate the first_frame of the first shot of the camera |
| 342 | first_shot_idx = camera.active_shot_idxs[0] |
| 343 | first_shot_ff_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", "first_frame.png") |
| 344 | _emit_render_progress(progress, "camera_frames_start", f"Generating frames for camera {camera.idx}", {"camera_idx": camera.idx, "active_shot_idxs": camera.active_shot_idxs}) |
| 345 | |
| 346 | if os.path.exists(first_shot_ff_path): |
| 347 | print(f"🚀 Skipped generating first_frame for shot {first_shot_idx}, already exists.") |
| 348 | self.frame_events[first_shot_idx]["first_frame"].set() |
| 349 | _emit_render_progress(progress, "frame_exists", f"First frame for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path}) |
| 350 | |
| 351 | else: |
| 352 | print(f"🖼️ Starting first_frame generation for shot {first_shot_idx}...") |
| 353 | _emit_render_progress(progress, "frame_start", f"Generating first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame"}) |
| 354 | available_image_path_and_text_pairs = [] |
| 355 | |
| 356 | for character_idx in shot_descriptions[first_shot_idx].ff_vis_char_idxs: |
| 357 | identifier_in_scene = characters[character_idx].identifier_in_scene |
| 358 | registry_item = character_portraits_registry[identifier_in_scene] |
| 359 | for view, item in registry_item.items(): |
| 360 | available_image_path_and_text_pairs.append((item["path"], item["description"])) |
| 361 | |
| 362 | # generate the first_frame based on the shot_description.ff_desc |
| 363 | if camera.parent_shot_idx is not None: |
| 364 | # generate the first_frame based on the transition video |
| 365 | parent_shot_idx = camera.parent_shot_idx |
| 366 | await self.frame_events[parent_shot_idx]["first_frame"].wait() |
| 367 | parent_shot_ff_path = os.path.join(self.working_dir, "shots", f"{parent_shot_idx}", "first_frame.png") |
| 368 | transition_video_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", f"transition_video_from_shot_{parent_shot_idx}.mp4") |
| 369 | |
| 370 | if os.path.exists(transition_video_path): |
| 371 | print(f"🚀 Skipped generating transition video for shot {first_shot_idx} from shot {parent_shot_idx}, already exists.") |
| 372 | _emit_render_progress(progress, "transition_video_exists", f"Transition video for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx, "path": transition_video_path}) |
| 373 | else: |
| 374 | print(f"🖼️ Starting transition video generation for shot {first_shot_idx} from shot {parent_shot_idx}...") |
| 375 | _emit_render_progress(progress, "transition_video_start", f"Generating transition video for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx}) |
| 376 | transition_video_output = await self.camera_image_generator.generate_transition_video( |
| 377 | first_shot_visual_desc=shot_descriptions[parent_shot_idx].visual_desc, |
| 378 | second_shot_visual_desc=shot_descriptions[first_shot_idx].visual_desc, |
| 379 | first_shot_ff_path=parent_shot_ff_path, |
| 380 | progress=_scoped_progress(progress, camera_idx=camera.idx, shot_idx=first_shot_idx, parent_shot_idx=parent_shot_idx, artifact="transition_video"), |
| 381 | ) |
| 382 | transition_video_output.save(transition_video_path) |
| 383 | print(f"☑️ Generated transition video for shot {first_shot_idx} from shot {parent_shot_idx}, saved to {transition_video_path}.") |
| 384 | _emit_render_progress(progress, "transition_video_done", f"Transition video for shot {first_shot_idx} generated", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "parent_shot_idx": parent_shot_idx, "path": transition_video_path}) |
| 385 | |
| 386 | new_camera_image_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", f"new_camera_{camera.idx}.png") |
| 387 | if os.path.exists(new_camera_image_path): |
| 388 | print(f"🚀 Skipped generating new camera image for shot {first_shot_idx}, already exists.") |
| 389 | _emit_render_progress(progress, "new_camera_image_exists", f"New camera image for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "path": new_camera_image_path}) |
| 390 | else: |
| 391 | print(f"🖼️ Starting new camera image generation for shot {first_shot_idx}...") |
| 392 | _emit_render_progress(progress, "new_camera_image_start", f"Extracting new camera image for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx}) |
| 393 | new_camera_image = self.camera_image_generator.get_new_camera_image(transition_video_path) |
| 394 | new_camera_image.save(new_camera_image_path) |
| 395 | print(f"☑️ Generated new camera image for shot {first_shot_idx} (not completed), saved to {new_camera_image_path}.") |
| 396 | _emit_render_progress(progress, "new_camera_image_done", f"New camera image for shot {first_shot_idx} extracted", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "path": new_camera_image_path}) |
| 397 | |
| 398 | available_image_path_and_text_pairs.append( |
| 399 | ( |
| 400 | new_camera_image_path, |
| 401 | f"The composition and background are correct but some elements may be wrong. The wrong elements should be replaced.\nWrong elements: {camera.missing_info}.\nYou must select this image as the main reference and replace the characters in the image with the provided character portraits. Don't change the background." |
| 402 | ) |
| 403 | ) |
| 404 | |
| 405 | |
| 406 | # 如果子镜头缺少信息，则需要选择参考图像生成 |
| 407 | if camera.parent_shot_idx is None or camera.missing_info is not None: |
| 408 | ff_selector_output_path = os.path.join(self.working_dir, "shots", f"{first_shot_idx}", "first_frame_selector_output.json") |
| 409 | if os.path.exists(ff_selector_output_path): |
| 410 | with open(ff_selector_output_path, 'r', encoding='utf-8') as f: |
| 411 | ff_selector_output = json.load(f) |
| 412 | print(f"🚀 Loaded existing reference image selection and prompt for first_frame of shot {first_shot_idx} from {ff_selector_output_path}.") |
| 413 | _emit_render_progress(progress, "frame_prompt_exists", f"First frame prompt for shot {first_shot_idx} already exists", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": ff_selector_output_path}) |
| 414 | else: |
| 415 | print(f"🔍 Selecting reference images and generating prompt for first_frame of shot {first_shot_idx}...") |
| 416 | _emit_render_progress(progress, "frame_prompt_start", f"Selecting references for first frame of shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame"}) |
| 417 | ff_selector_output = await self.reference_image_selector.select_reference_images_and_generate_prompt( |
| 418 | available_image_path_and_text_pairs=available_image_path_and_text_pairs, |
| 419 | frame_description=shot_descriptions[first_shot_idx].ff_desc |
| 420 | ) |
| 421 | with open(ff_selector_output_path, 'w', encoding='utf-8') as f: |
| 422 | json.dump(ff_selector_output, f, ensure_ascii=False, indent=4) |
| 423 | |
| 424 | print(f"☑️ Selected reference images and generated prompt for first_frame of shot {first_shot_idx}, saved to {ff_selector_output_path}.") |
| 425 | _emit_render_progress(progress, "frame_prompt_done", f"Selected references for first frame of shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": ff_selector_output_path}) |
| 426 | |
| 427 | reference_image_path_and_text_pairs, prompt = ff_selector_output["reference_image_path_and_text_pairs"], ff_selector_output["text_prompt"] |
| 428 | prefix_prompt = "" |
| 429 | for i, (image_path, text) in enumerate(reference_image_path_and_text_pairs): |
| 430 | prefix_prompt += f"Image {i}: {text}\n" |
| 431 | prompt = f"{prefix_prompt}\n{prompt}" |
| 432 | reference_image_paths = [item[0] for item in reference_image_path_and_text_pairs] |
| 433 | ff_image: ImageOutput = await self.generate_and_select_best_image( |
| 434 | prompt=prompt, |
| 435 | reference_image_paths=reference_image_paths, |
| 436 | reference_image_path_and_text_pairs=reference_image_path_and_text_pairs, |
| 437 | target_description=shot_descriptions[first_shot_idx].ff_desc, |
| 438 | size="1600x900", |
| 439 | candidates_save_dir=os.path.join(self.working_dir, "shots", str(first_shot_idx), "first_frame_candidates"), |
| 440 | progress=_scoped_progress(progress, camera_idx=camera.idx, shot_idx=first_shot_idx, frame_type="first_frame"), |
| 441 | ) |
| 442 | ff_image.save(first_shot_ff_path) |
| 443 | self.frame_events[first_shot_idx]["first_frame"].set() |
| 444 | print(f"☑️ Generated first_frame for shot {first_shot_idx}, saved to {first_shot_ff_path}.") |
| 445 | _emit_render_progress(progress, "frame_done", f"Generated first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path}) |
| 446 | else: |
| 447 | shutil.copy(new_camera_image_path, first_shot_ff_path) |
| 448 | self.frame_events[first_shot_idx]["first_frame"].set() |
| 449 | print(f"☑️ Generated first_frame for shot {first_shot_idx}, saved to {first_shot_ff_path}.") |
| 450 | _emit_render_progress(progress, "frame_done", f"Generated first frame for shot {first_shot_idx}", {"camera_idx": camera.idx, "shot_idx": first_shot_idx, "frame_type": "first_frame", "path": first_shot_ff_path}) |
| 451 | |
| 452 | |
| 453 | # 2. generate the following frames of the camera |
| 454 | priority_tasks = [] |
| 455 | normal_tasks = [] |
| 456 | |
| 457 | if shot_descriptions[first_shot_idx].variation_type in ["medium", "large"]: |
| 458 | task = self.generate_frame_for_single_shot( |
| 459 | shot_idx=first_shot_idx, |
| 460 | frame_type="last_frame", |
| 461 | first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc), |
| 462 | frame_desc=shot_descriptions[first_shot_idx].lf_desc, |
| 463 | visible_characters=[characters[idx] for idx in shot_descriptions[first_shot_idx].lf_vis_char_idxs], |
| 464 | character_portraits_registry=character_portraits_registry, |
| 465 | progress=progress, |
| 466 | ) |
| 467 | normal_tasks.append(task) |
| 468 | |
| 469 | for shot_idx in camera.active_shot_idxs[1:]: |
| 470 | first_frame_task = self.generate_frame_for_single_shot( |
| 471 | shot_idx=shot_idx, |
| 472 | frame_type="first_frame", |
| 473 | first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc), |
| 474 | frame_desc=shot_descriptions[shot_idx].ff_desc, |
| 475 | visible_characters=[characters[idx] for idx in shot_descriptions[shot_idx].ff_vis_char_idxs], |
| 476 | character_portraits_registry=character_portraits_registry, |
| 477 | progress=progress, |
| 478 | ) |
| 479 | if shot_idx in priority_shot_idxs: |
| 480 | priority_tasks.append(first_frame_task) |
| 481 | else: |
| 482 | normal_tasks.append(first_frame_task) |
| 483 | |
| 484 | |
| 485 | if shot_descriptions[shot_idx].variation_type in ["medium", "large"]: |
| 486 | last_frame_task = self.generate_frame_for_single_shot( |
| 487 | shot_idx=shot_idx, |
| 488 | frame_type="last_frame", |
| 489 | first_shot_ff_path_and_text_pair=(first_shot_ff_path, shot_descriptions[first_shot_idx].ff_desc), |
| 490 | frame_desc=shot_descriptions[shot_idx].lf_desc, |
| 491 | visible_characters=[characters[idx] for idx in shot_descriptions[shot_idx].lf_vis_char_idxs], |
| 492 | character_portraits_registry=character_portraits_registry, |
| 493 | progress=progress, |
| 494 | ) |
| 495 | normal_tasks.append(last_frame_task) |
| 496 | |
| 497 | |
| 498 | await asyncio.gather(*priority_tasks) |
| 499 | await asyncio.gather(*normal_tasks) |
| 500 | _emit_render_progress(progress, "camera_frames_done", f"Frames for camera {camera.idx} ready", {"camera_idx": camera.idx, "active_shot_idxs": camera.active_shot_idxs}) |
| 501 | |
| 502 | |
| 503 | |
| 504 | async def generate_video_for_single_shot( |
| 505 | self, |
| 506 | shot_description: ShotDescription, |
| 507 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 508 | ): |
| 509 | video_path = os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "video.mp4") |
| 510 | if os.path.exists(video_path): |
| 511 | print(f"🚀 Skipped generating video for shot {shot_description.idx}, already exists.") |
| 512 | _emit_render_progress(progress, "video_clip_exists", f"Video clip for shot {shot_description.idx} already exists", {"shot_idx": shot_description.idx, "path": video_path}) |
| 513 | else: |
| 514 | _emit_render_progress(progress, "video_clip_waiting_for_frames", f"Waiting for frames before video clip {shot_description.idx}", {"shot_idx": shot_description.idx}) |
| 515 | await self.frame_events[shot_description.idx]["first_frame"].wait() |
| 516 | if shot_description.variation_type in ["medium", "large"]: |
| 517 | await self.frame_events[shot_description.idx]["last_frame"].wait() |
| 518 | |
| 519 | frame_paths = [] |
| 520 | frame_paths.append(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "first_frame.png")) |
| 521 | if shot_description.variation_type in ["medium", "large"]: |
| 522 | frame_paths.append(os.path.join(self.working_dir, "shots", f"{shot_description.idx}", "last_frame.png")) |
| 523 | |
| 524 | print(f"🎬 Starting video generation for shot {shot_description.idx}...") |
| 525 | _emit_render_progress(progress, "video_clip_start", f"Generating video clip for shot {shot_description.idx}", {"shot_idx": shot_description.idx, "frame_count": len(frame_paths)}) |
| 526 | video_output = await self.video_generator.generate_single_video( |
| 527 | prompt=shot_description.motion_desc + "\n" + shot_description.audio_desc, |
| 528 | reference_image_paths=frame_paths, |
| 529 | progress=_scoped_progress(progress, shot_idx=shot_description.idx, artifact="video_clip"), |
| 530 | ) |
| 531 | video_output.save(video_path) |
| 532 | print(f"☑️ Generated video for shot {shot_description.idx}, saved to {video_path}.") |
| 533 | _emit_render_progress(progress, "video_clip_done", f"Generated video clip for shot {shot_description.idx}", {"shot_idx": shot_description.idx, "path": video_path}) |
| 534 | |
| 535 | async def generate_and_select_best_image( |
| 536 | self, |
| 537 | prompt: str, |
| 538 | reference_image_paths: List[str], |
| 539 | reference_image_path_and_text_pairs: List[Tuple[str, str]], |
| 540 | target_description: str, |
| 541 | size: str = "1600x900", |
| 542 | num_candidates: int | None = None, |
| 543 | candidates_save_dir: str | None = None, |
| 544 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 545 | ) -> ImageOutput: |
| 546 | count = validate_image_candidate_count(self.num_image_candidates if num_candidates is None else num_candidates) |
| 547 | if count == 1: |
| 548 | return await self.image_generator.generate_single_image( |
| 549 | prompt=prompt, reference_image_paths=reference_image_paths, size=size, |
| 550 | ) |
| 551 | |
| 552 | save_dir = Path(candidates_save_dir) if candidates_save_dir else Path(self.working_dir) / "image_candidates" / uuid4().hex |
| 553 | save_dir.mkdir(parents=True, exist_ok=True) |
| 554 | selection_path = save_dir / "selection.json" |
| 555 | references = [] |
| 556 | for path, description in reference_image_path_and_text_pairs: |
| 557 | stat = os.stat(path) |
| 558 | references.append((path, description, stat.st_size, stat.st_mtime_ns)) |
| 559 | fingerprint = hashlib.sha256(json.dumps({ |
| 560 | "prompt": prompt, "size": size, "references": references, |
| 561 | "reference_image_paths": reference_image_paths, "target_description": target_description, |
| 562 | "generator": type(self.image_generator).__name__, |
| 563 | "model": str(getattr(self.image_generator, "model", "")), |
| 564 | }, sort_keys=True).encode()).hexdigest() |
| 565 | previous = {} |
| 566 | if selection_path.exists(): |
| 567 | try: |
| 568 | previous = json.loads(selection_path.read_text(encoding="utf-8")) |
| 569 | except (OSError, ValueError): |
| 570 | pass |
| 571 | reuse = isinstance(previous, dict) and previous.get("request_fingerprint") == fingerprint |
| 572 | reusable_indices = { |
| 573 | item["index"] for item in previous.get("candidates", []) |
| 574 | if isinstance(item, dict) and item.get("status") == "ready" and "index" in item |
| 575 | } if reuse else set() |
| 576 | selection = {"request_fingerprint": fingerprint, "requested_candidates": count, "status": "generating", "candidates": []} |
| 577 | |
| 578 | def save_selection(): |
| 579 | temporary = selection_path.with_suffix(".json.tmp") |
| 580 | temporary.write_text(json.dumps(selection, ensure_ascii=False, indent=2), encoding="utf-8") |
| 581 | temporary.replace(selection_path) |
| 582 | |
| 583 | save_selection() |
| 584 | _emit_render_progress(progress, "image_candidates_start", f"Generating {count} image candidates", {"count": count}) |
| 585 | |
| 586 | async def generate_candidate(index): |
| 587 | path = save_dir / f"candidate_{index}.png" |
| 588 | record = {"index": index, "path": path.name, "status": "generating"} |
| 589 | selection["candidates"].append(record) |
| 590 | try: |
| 591 | if index in reusable_indices and path.is_file(): |
| 592 | record["reused"] = True |
| 593 | else: |
| 594 | output = await self.image_generator.generate_single_image( |
| 595 | prompt=prompt, reference_image_paths=reference_image_paths, size=size, |
| 596 | ) |
| 597 | if output is None: |
| 598 | raise RuntimeError("Image provider returned no image") |
| 599 | await asyncio.to_thread(output.save, str(path)) |
| 600 | with Image.open(path) as image: |
| 601 | width, height = image.size |
| 602 | output = ImageOutput(fmt="pil", ext="png", data=image.copy()) |
| 603 | record.update(width=width, height=height) |
| 604 | if width <= height: |
| 605 | raise ValueError("Candidate is not landscape") |
| 606 | record["status"] = "ready" |
| 607 | save_selection() |
| 608 | _emit_render_progress(progress, "image_candidate_done", f"Image candidate {index + 1}/{count} ready", {"index": index, "count": count}) |
| 609 | return index, output, str(path) |
| 610 | except Exception as exc: |
| 611 | record.update(status="error", error_type=type(exc).__name__) |
| 612 | save_selection() |
| 613 | _emit_render_progress(progress, "image_candidate_failed", f"Image candidate {index + 1}/{count} failed", {"index": index, "error_type": type(exc).__name__}) |
| 614 | raise |
| 615 | |
| 616 | tasks = [asyncio.create_task(generate_candidate(index)) for index in range(count)] |
| 617 | try: |
| 618 | results = await asyncio.gather(*tasks, return_exceptions=True) |
| 619 | cancelled = next((result for result in results if isinstance(result, asyncio.CancelledError)), None) |
| 620 | if cancelled is not None: |
| 621 | raise cancelled |
| 622 | valid = [result for result in results if not isinstance(result, Exception)] |
| 623 | if not valid: |
| 624 | selection["status"] = "generation_failed" |
| 625 | save_selection() |
| 626 | raise RuntimeError(f"All {count} image candidates failed: {results[0]}") from results[0] |
| 627 | |
| 628 | if len(valid) == 1: |
| 629 | selected = valid[0] |
| 630 | reason = "Only one valid landscape candidate; VLM comparison skipped." |
| 631 | selection["selection_method"] = "single_valid_candidate" |
| 632 | else: |
| 633 | selection["status"] = "selecting" |
| 634 | save_selection() |
| 635 | _emit_render_progress(progress, "image_selection_start", f"VLM selecting from {len(valid)} candidates", {"count": len(valid)}) |
| 636 | try: |
| 637 | response = await self.best_image_selector.select( |
| 638 | reference_image_path_and_text_pairs=reference_image_path_and_text_pairs, |
| 639 | target_description=target_description, |
| 640 | candidate_image_paths=[item[2] for item in valid], |
| 641 | ) |
| 642 | except Exception as exc: |
| 643 | selection.update(status="selection_failed", error_type=type(exc).__name__) |
| 644 | save_selection() |
| 645 | raise |
| 646 | selected = valid[response.best_image_index] |
| 647 | reason = response.reason |
| 648 | selection["selection_method"] = "vlm" |
| 649 | selection.update(status="selected", selected_candidate_index=selected[0], reason=reason) |
| 650 | save_selection() |
| 651 | _emit_render_progress(progress, "image_selection_done", f"Selected image candidate {selected[0] + 1}", {"index": selected[0], "count": len(valid), "selection_method": selection["selection_method"]}) |
| 652 | return selected[1] |
| 653 | except asyncio.CancelledError: |
| 654 | for task in tasks: |
| 655 | task.cancel() |
| 656 | await asyncio.gather(*tasks, return_exceptions=True) |
| 657 | selection["status"] = "cancelled" |
| 658 | save_selection() |
| 659 | raise |
| 660 | |
| 661 | |
| 662 | async def generate_frame_for_single_shot( |
| 663 | self, |
| 664 | shot_idx: int, |
| 665 | frame_type: Literal["first_frame", "last_frame"], |
| 666 | first_shot_ff_path_and_text_pair: Tuple[str, str], |
| 667 | frame_desc: str, |
| 668 | visible_characters: List[CharacterInScene], |
| 669 | character_portraits_registry: Dict[str, Dict[str, Dict[str, str]]], |
| 670 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 671 | ) -> ImageOutput: |
| 672 | |
| 673 | frame_image_path = os.path.join(self.working_dir, "shots", f"{shot_idx}", f"{frame_type}.png") |
| 674 | |
| 675 | if os.path.exists(frame_image_path): |
| 676 | print(f"🚀 Skipped generating {frame_type} for shot {shot_idx}, already exists.") |
| 677 | _emit_render_progress(progress, "frame_exists", f"{frame_type} for shot {shot_idx} already exists", {"shot_idx": shot_idx, "frame_type": frame_type, "path": frame_image_path}) |
| 678 | |
| 679 | else: |
| 680 | print(f"🖼️ Starting {frame_type} generation for shot {shot_idx}...") |
| 681 | _emit_render_progress(progress, "frame_start", f"Generating {frame_type} for shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type}) |
| 682 | available_image_path_and_text_pairs = [] |
| 683 | for visible_character in visible_characters: |
| 684 | identifier_in_scene = visible_character.identifier_in_scene |
| 685 | registry_item = character_portraits_registry[identifier_in_scene] |
| 686 | for view, item in registry_item.items(): |
| 687 | available_image_path_and_text_pairs.append((item["path"], item["description"])) |
| 688 | |
| 689 | available_image_path_and_text_pairs.append(first_shot_ff_path_and_text_pair) |
| 690 | |
| 691 | selector_output_path = os.path.join(self.working_dir, "shots", f"{shot_idx}", f"{frame_type}_selector_output.json") |
| 692 | if os.path.exists(selector_output_path): |
| 693 | with open(selector_output_path, 'r', encoding='utf-8') as f: |
| 694 | selector_output = json.load(f) |
| 695 | print(f"🚀 Loaded existing reference image selection and prompt for {frame_type} frame of shot {shot_idx} from {selector_output_path}.") |
| 696 | _emit_render_progress(progress, "frame_prompt_exists", f"Prompt for {frame_type} of shot {shot_idx} already exists", {"shot_idx": shot_idx, "frame_type": frame_type, "path": selector_output_path}) |
| 697 | else: |
| 698 | print(f"🔍 Selecting reference images and generating prompt for {frame_type} frame of shot {shot_idx}...") |
| 699 | _emit_render_progress(progress, "frame_prompt_start", f"Selecting references for {frame_type} of shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type}) |
| 700 | selector_output = await self.reference_image_selector.select_reference_images_and_generate_prompt( |
| 701 | available_image_path_and_text_pairs=available_image_path_and_text_pairs, |
| 702 | frame_description=frame_desc |
| 703 | ) |
| 704 | with open(selector_output_path, 'w', encoding='utf-8') as f: |
| 705 | json.dump(selector_output, f, ensure_ascii=False, indent=4) |
| 706 | print(f"☑️ Selected reference images and generated prompt for {frame_type} frame of shot {shot_idx}, saved to {selector_output_path}.") |
| 707 | _emit_render_progress(progress, "frame_prompt_done", f"Selected references for {frame_type} of shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type, "path": selector_output_path}) |
| 708 | |
| 709 | reference_image_path_and_text_pairs, prompt = selector_output["reference_image_path_and_text_pairs"], selector_output["text_prompt"] |
| 710 | prefix_prompt = "" |
| 711 | for i, (image_path, text) in enumerate(reference_image_path_and_text_pairs): |
| 712 | prefix_prompt += f"Image {i}: {text}\n" |
| 713 | prompt = f"{prefix_prompt}\n{prompt}" |
| 714 | reference_image_paths = [item[0] for item in reference_image_path_and_text_pairs] |
| 715 | |
| 716 | frame_image: ImageOutput = await self.generate_and_select_best_image( |
| 717 | prompt=prompt, |
| 718 | reference_image_paths=reference_image_paths, |
| 719 | reference_image_path_and_text_pairs=reference_image_path_and_text_pairs, |
| 720 | target_description=frame_desc, |
| 721 | size="1600x900", |
| 722 | candidates_save_dir=os.path.join(self.working_dir, "shots", str(shot_idx), f"{frame_type}_candidates"), |
| 723 | progress=_scoped_progress(progress, shot_idx=shot_idx, frame_type=frame_type), |
| 724 | ) |
| 725 | frame_image.save(frame_image_path) |
| 726 | print(f"☑️ Generated {frame_type} frame for shot {shot_idx}, saved to {frame_image_path}.") |
| 727 | _emit_render_progress(progress, "frame_done", f"Generated {frame_type} for shot {shot_idx}", {"shot_idx": shot_idx, "frame_type": frame_type, "path": frame_image_path}) |
| 728 | |
| 729 | |
| 730 | self.frame_events[shot_idx][frame_type].set() |
| 731 | return frame_image_path |
| 732 | |
| 733 | |
| 734 | async def construct_camera_tree( |
| 735 | self, |
| 736 | shot_descriptions: List[ShotDescription], |
| 737 | quiet: bool = False, |
| 738 | ): |
| 739 | camera_tree_path = os.path.join(self.working_dir, "camera_tree.json") |
| 740 | |
| 741 | if os.path.exists(camera_tree_path): |
| 742 | with open(camera_tree_path, "r", encoding="utf-8") as f: |
| 743 | camera_tree = json.load(f) |
| 744 | camera_tree = [Camera.model_validate(camera) for camera in camera_tree] |
| 745 | _pipeline_print(quiet, f"🚀 Loaded {len(camera_tree)} cameras from existing file.") |
| 746 | return camera_tree |
| 747 | |
| 748 | shot_descriptions = _normalize_model_list(shot_descriptions, ShotDescription, "shot_descriptions") |
| 749 | cameras = _group_shots_into_cameras(shot_descriptions) |
| 750 | |
| 751 | camera_tree = await self.camera_image_generator.construct_camera_tree(cameras=cameras, shot_descs=shot_descriptions) |
| 752 | camera_tree = _normalize_model_list(camera_tree, Camera, "camera_tree") |
| 753 | with open(camera_tree_path, "w", encoding="utf-8") as f: |
| 754 | json.dump([camera.model_dump() for camera in camera_tree], f, ensure_ascii=False, indent=4) |
| 755 | _pipeline_print(quiet, f"✅ Constructed camera tree and saved to {camera_tree_path}.") |
| 756 | return camera_tree |
| 757 | |
| 758 | |
| 759 | |
| 760 | |
| 761 | async def extract_characters( |
| 762 | self, |
| 763 | script: str, |
| 764 | quiet: bool = False, |
| 765 | ): |
| 766 | save_path = os.path.join(self.working_dir, "characters.json") |
| 767 | |
| 768 | if os.path.exists(save_path): |
| 769 | with open(save_path, "r", encoding="utf-8") as f: |
| 770 | characters = json.load(f) |
| 771 | characters = [CharacterInScene.model_validate(character) for character in characters] |
| 772 | _pipeline_print(quiet, f"🚀 Loaded {len(characters)} characters from existing file.") |
| 773 | else: |
| 774 | characters = await self.character_extractor.extract_characters(script) |
| 775 | with open(save_path, "w", encoding="utf-8") as f: |
| 776 | json.dump([character.model_dump() for character in characters], f, ensure_ascii=False, indent=4) |
| 777 | _pipeline_print(quiet, f"✅ Extracted {len(characters)} characters from script and saved to {save_path}.") |
| 778 | |
| 779 | for character in characters: |
| 780 | self.character_portrait_events[character.idx] = asyncio.Event() |
| 781 | |
| 782 | return characters |
| 783 | |
| 784 | |
| 785 | async def generate_character_portraits( |
| 786 | self, |
| 787 | characters: List[CharacterInScene], |
| 788 | character_portraits_registry: Optional[Dict[str, Dict[str, Dict[str, str]]]], |
| 789 | style: str, |
| 790 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 791 | ): |
| 792 | character_portraits_registry_path = os.path.join(self.working_dir, "character_portraits_registry.json") |
| 793 | if character_portraits_registry is None: |
| 794 | if os.path.exists(character_portraits_registry_path): |
| 795 | with open(character_portraits_registry_path, 'r', encoding='utf-8') as f: |
| 796 | character_portraits_registry = json.load(f) |
| 797 | else: |
| 798 | character_portraits_registry = {} |
| 799 | |
| 800 | |
| 801 | tasks = [ |
| 802 | self.generate_portraits_for_single_character(character, style, progress=progress) |
| 803 | for character in characters |
| 804 | if character.identifier_in_scene not in character_portraits_registry |
| 805 | ] |
| 806 | if tasks: |
| 807 | for future in asyncio.as_completed(tasks): |
| 808 | character_portraits_registry.update(await future) |
| 809 | with open(character_portraits_registry_path, 'w', encoding='utf-8') as f: |
| 810 | json.dump(character_portraits_registry, f, ensure_ascii=False, indent=4) |
| 811 | |
| 812 | print(f"✅ Completed character portrait generation for {len(characters)} characters.") |
| 813 | _emit_render_progress(progress, "character_portraits_done", "Completed character portrait generation", {"character_count": len(characters)}) |
| 814 | else: |
| 815 | print("🚀 All characters already have portraits, skipping portrait generation.") |
| 816 | _emit_render_progress(progress, "character_portraits_exist", "All character portraits already exist", {"character_count": len(characters)}) |
| 817 | return character_portraits_registry |
| 818 | |
| 819 | |
| 820 | async def generate_portraits_for_single_character( |
| 821 | self, |
| 822 | character: CharacterInScene, |
| 823 | style: str, |
| 824 | progress: Callable[[str, str, Dict[str, Any] | None], None] | None = None, |
| 825 | ): |
| 826 | character_dir = os.path.join(self.working_dir, "character_portraits", f"{character.idx}_{character.identifier_in_scene}") |
| 827 | os.makedirs(character_dir, exist_ok=True) |
| 828 | _emit_render_progress(progress, "character_portrait_start", f"Generating portraits for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene}) |
| 829 | |
| 830 | front_portrait_path = os.path.join(character_dir, "front.png") |
| 831 | if os.path.exists(front_portrait_path): |
| 832 | pass |
| 833 | else: |
| 834 | _emit_render_progress(progress, "character_portrait_front_start", f"Generating front portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene}) |
| 835 | front_portrait_output = await self.character_portraits_generator.generate_front_portrait(character, style) |
| 836 | front_portrait_output.save(front_portrait_path) |
| 837 | _emit_render_progress(progress, "character_portrait_front_done", f"Generated front portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": front_portrait_path}) |
| 838 | |
| 839 | |
| 840 | side_portrait_path = os.path.join(character_dir, "side.png") |
| 841 | if os.path.exists(side_portrait_path): |
| 842 | pass |
| 843 | else: |
| 844 | _emit_render_progress(progress, "character_portrait_side_start", f"Generating side portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene}) |
| 845 | side_portrait_output = await self.character_portraits_generator.generate_side_portrait(character, front_portrait_path) |
| 846 | side_portrait_output.save(side_portrait_path) |
| 847 | _emit_render_progress(progress, "character_portrait_side_done", f"Generated side portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": side_portrait_path}) |
| 848 | |
| 849 | back_portrait_path = os.path.join(character_dir, "back.png") |
| 850 | if os.path.exists(back_portrait_path): |
| 851 | pass |
| 852 | else: |
| 853 | _emit_render_progress(progress, "character_portrait_back_start", f"Generating back portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene}) |
| 854 | back_portrait_output = await self.character_portraits_generator.generate_back_portrait(character, front_portrait_path) |
| 855 | back_portrait_output.save(back_portrait_path) |
| 856 | _emit_render_progress(progress, "character_portrait_back_done", f"Generated back portrait for {character.identifier_in_scene}", {"character_idx": character.idx, "identifier": character.identifier_in_scene, "path": back_portrait_path}) |
| 857 | |
| 858 | self.character_portrait_events[character.idx].set() |
| 859 | |
| 860 | print(f"☑️ Completed character portrait generation for {character.identifier_in_scene}.") |
| 861 | _emit_render_progress(progress, "character_portrait_done", f"Portraits for {character.identifier_in_scene} ready", {"character_idx": character.idx, "identifier": character.identifier_in_scene}) |
| 862 | |
| 863 | return { |
| 864 | character.identifier_in_scene: { |
| 865 | "front": { |
| 866 | "path": front_portrait_path, |
| 867 | "description": f"A front view portrait of {character.identifier_in_scene}.", |
| 868 | }, |
| 869 | "side": { |
| 870 | "path": side_portrait_path, |
| 871 | "description": f"A side view portrait of {character.identifier_in_scene}.", |
| 872 | }, |
| 873 | "back": { |
| 874 | "path": back_portrait_path, |
| 875 | "description": f"A back view portrait of {character.identifier_in_scene}.", |
| 876 | }, |
| 877 | } |
| 878 | } |
| 879 | |
| 880 | |
| 881 | |
| 882 | async def design_storyboard( |
| 883 | self, |
| 884 | script: str, |
| 885 | characters: List[CharacterInScene], |
| 886 | user_requirement: str, |
| 887 | quiet: bool = False, |
| 888 | ): |
| 889 | storyboard_path = os.path.join(self.working_dir, "storyboard.json") |
| 890 | if os.path.exists(storyboard_path): |
| 891 | with open(storyboard_path, 'r', encoding='utf-8') as f: |
| 892 | storyboard = json.load(f) |
| 893 | storyboard = [ShotBriefDescription.model_validate(shot) for shot in storyboard] |
| 894 | _pipeline_print(quiet, f"🚀 Loaded {len(storyboard)} shot brief descriptions from existing file.") |
| 895 | else: |
| 896 | _pipeline_print(quiet, f"🔍 Designing storyboard...") |
| 897 | storyboard = await self.storyboard_artist.design_storyboard( |
| 898 | script=script, |
| 899 | characters=characters, |
| 900 | user_requirement=user_requirement, |
| 901 | retry_timeout=150, |
| 902 | ) |
| 903 | storyboard = _normalize_model_list(storyboard, ShotBriefDescription, "storyboard") |
| 904 | with open(storyboard_path, 'w', encoding='utf-8') as f: |
| 905 | json.dump([shot.model_dump() for shot in storyboard], f, ensure_ascii=False, indent=4) |
| 906 | _pipeline_print(quiet, f"✅ Designed storyboard and saved to {storyboard_path}.") |
| 907 | |
| 908 | for shot_brief_description in storyboard: |
| 909 | self.shot_desc_events[shot_brief_description.idx] = asyncio.Event() |
| 910 | |
| 911 | return storyboard |
| 912 | |
| 913 | |
| 914 | |
| 915 | async def decompose_visual_descriptions( |
| 916 | self, |
| 917 | shot_brief_descriptions: List[ShotBriefDescription], |
| 918 | characters: List[CharacterInScene], |
| 919 | quiet: bool = False, |
| 920 | ): |
| 921 | tasks = [ |
| 922 | self.decompose_visual_description_for_single_shot_brief_description(shot_brief_description, characters, quiet=quiet) |
| 923 | for shot_brief_description in shot_brief_descriptions |
| 924 | ] |
| 925 | |
| 926 | shot_descriptions = await asyncio.gather(*tasks) |
| 927 | return shot_descriptions |
| 928 | |
| 929 | |
| 930 | async def decompose_visual_description_for_single_shot_brief_description( |
| 931 | self, |
| 932 | shot_brief_description: ShotBriefDescription, |
| 933 | characters: List[CharacterInScene], |
| 934 | quiet: bool = False, |
| 935 | ): |
| 936 | shot_description_path = os.path.join(self.working_dir, "shots", f"{shot_brief_description.idx}", "shot_description.json") |
| 937 | os.makedirs(os.path.dirname(shot_description_path), exist_ok=True) |
| 938 | |
| 939 | if os.path.exists(shot_description_path): |
| 940 | with open(shot_description_path, 'r', encoding='utf-8') as f: |
| 941 | shot_description = ShotDescription.model_validate(json.load(f)) |
| 942 | _pipeline_print(quiet, f"🚀 Loaded shot {shot_brief_description.idx} description from existing file.") |
| 943 | else: |
| 944 | shot_description = await self.storyboard_artist.decompose_visual_description( |
| 945 | shot_brief_desc=shot_brief_description, |
| 946 | characters=characters, |
| 947 | retry_timeout=120, |
| 948 | ) |
| 949 | shot_description = _normalize_model_list([shot_description], ShotDescription, "shot_description")[0] |
| 950 | with open(shot_description_path, 'w', encoding='utf-8') as f: |
| 951 | json.dump(shot_description.model_dump(), f, ensure_ascii=False, indent=4) |
| 952 | _pipeline_print(quiet, f"✅ Decomposed visual description for shot {shot_brief_description.idx} and saved to {shot_description_path}.") |
| 953 | |
| 954 | self.shot_desc_events[shot_brief_description.idx].set() |
| 955 | |
| 956 | if shot_description.variation_type in ["medium", "large"]: |
| 957 | self.frame_events[shot_brief_description.idx] = { |
| 958 | "first_frame": asyncio.Event(), |
| 959 | "last_frame": asyncio.Event(), |
| 960 | } |
| 961 | else: |
| 962 | self.frame_events[shot_brief_description.idx] = { |
| 963 | "first_frame": asyncio.Event(), |
| 964 | } |
| 965 | |
| 966 | return shot_description |
| 967 |