Coverage for apps/streammovie/streammovie_job.py: 89%
188 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamMovie job to generate a movie.
3It coordinates the execution of the different models.
4"""
5import sys
6import json
7import asyncio
8import aiofiles
10from typing import override
11from typing import Dict
12from typing import Any
13from typing import List
14from typing import Optional
16from movie_prompts import SYSTEM_PROMPT
18# Local relative imports
19sys.path.append("..") # noqa: E402
20sys.path.append("../..") # noqa: E402
22from streamwise_job import StreamWiseJob
23from streamwise_job import JobStatus
24from streamwise_job import OutputMode
26from lmm_service_manager import LMMServiceManager
28from console_utils import bytes_to_human
30from media_utils import concatenate_videos
31from media_utils import save_video_audio
32from media_utils import save_video_frames
33from media_utils import get_audio_duration
35from file_utils import save_base64_as_binary
37from video import HUNYUANFRAMEPACK_FPS
40DEFAULT_SHOT_DURATION_SECS = 4.0
41DEFAULT_MAX_TOKENS = 8192
42DEFAULT_SPEECH_SPEED = 1.1
43MAX_LOG_TEXT = 80
44SHOT_DEADLINE_BUFFER_SECS = 120.0 # Extra buffer per shot on top of its timeline offset
47class StreamMovieJob(StreamWiseJob):
48 """Job class for StreamMovie movie generation."""
50 def __init__(
51 self,
52 job_id: str,
53 config: Dict[str, Any],
54 service_manager: LMMServiceManager
55 ) -> None:
56 super().__init__(
57 "streammovie",
58 job_id=job_id,
59 config=config,
60 service_manager=service_manager
61 )
63 @override
64 async def generate(
65 self,
66 job_config: Dict[str, Any],
67 ) -> None:
68 movie_description: Optional[str] = job_config.get("movie_description")
69 await self.gen_movie(movie_description)
71 @staticmethod
72 def build_movie_messages(movie_description: str, max_shots: int = -1) -> list:
73 """
74 Build LLM messages for movie planning using the system prompt.
75 Returns messages suitable for gen_text to generate a movie structure.
76 """
77 user_content = f"Generate a movie based on the following description:\n\n{movie_description}"
78 if max_shots > 0:
79 user_content += f"\n\nIMPORTANT: Generate EXACTLY {max_shots} shots. No more, no less."
80 return [
81 {"role": "system", "content": SYSTEM_PROMPT},
82 {"role": "user", "content": user_content},
83 ]
85 async def gen_movie(
86 self,
87 movie_description: Optional[str]
88 ) -> None:
89 """
90 Generate a movie based on the provided description.
92 Steps:
93 1. Stream a structured movie script (JSONL) from the LLM.
94 2. Collect shot_description objects from the script.
95 3. For each shot, generate an image and a video (with optional lip-synced audio).
96 4. Concatenate all shot videos into the final movie.
97 """
98 async with self.job_status_handler():
99 if not movie_description:
100 self.logger.error("Movie description is required.")
101 await self.save_status(JobStatus.FAILED)
102 raise ValueError("Missing 'movie_description' in request")
104 log_desc = movie_description[:MAX_LOG_TEXT]
105 self.logger.info(f"Generating movie for description: '{log_desc}...'")
107 max_tokens = self.get_config_int("max_tokens", DEFAULT_MAX_TOKENS)
108 max_shots = self.get_config_int("max_shots", -1)
110 # Stream the movie script from the LLM
111 shot_descriptions = await self._stream_movie_script(
112 movie_description=movie_description,
113 max_tokens=max_tokens,
114 max_shots=max_shots,
115 )
117 await self.save_status(JobStatus.RUNNING)
119 self.logger.info(f"Parsed {len(shot_descriptions)} shots from the movie script.")
120 if not shot_descriptions:
121 raise ValueError("No shots generated from the movie script.")
123 # Limit number of shots if configured (safety cap in case LLM ignores the instruction)
124 if max_shots > 0 and len(shot_descriptions) > max_shots:
125 self.logger.info(f"Limiting to {max_shots} shots (out of {len(shot_descriptions)}).")
126 shot_descriptions = shot_descriptions[:max_shots]
128 # Generate each shot: image -> video (+ audio if dialogue present)
129 shot_tasks: Dict[int, asyncio.Task] = {}
130 for idx, shot in enumerate(shot_descriptions):
131 task = asyncio.create_task(self._gen_shot(idx, shot))
132 shot_tasks[idx] = task
134 await self.save_status(JobStatus.RUNNING)
136 shot_video_paths: List[str] = []
137 results = await asyncio.gather(*shot_tasks.values(), return_exceptions=True)
138 for idx, result in enumerate(results):
139 if isinstance(result, Exception):
140 self.logger.error(f"[{idx}] Shot generation failed: {result}")
141 elif isinstance(result, str):
142 shot_video_paths.append(result)
143 self.logger.info(f"[{idx}] Shot video saved to '{result}'.")
144 else:
145 self.logger.warning(f"[{idx}] No video generated for shot.")
147 if not shot_video_paths:
148 raise ValueError("No shot videos generated. Cannot create final movie.")
150 await self.save_status(JobStatus.RUNNING)
152 # Concatenate all shot videos into the final movie
153 self.logger.info(f"Concatenating {len(shot_video_paths)} shots into final movie...")
154 scene_binaries: List[bytes] = []
155 for video_path in shot_video_paths:
156 async with aiofiles.open(video_path, "rb") as file:
157 scene_binary = await file.read()
158 scene_binaries.append(scene_binary)
160 video_binary = await concatenate_videos(scene_binaries, fast_copy=False)
161 if not video_binary:
162 raise ValueError("Cannot concatenate shots into final movie.")
164 video_path = f"{self.job_path}/{self.job_id}.mp4"
165 async with aiofiles.open(video_path, "wb") as file:
166 await file.write(video_binary)
168 self.logger.info(
169 f"Generated movie with {bytes_to_human(len(video_binary))} at '{video_path}'.")
171 async def _stream_movie_script(
172 self,
173 movie_description: str,
174 max_tokens: int = DEFAULT_MAX_TOKENS,
175 max_shots: int = -1,
176 ) -> List[Dict[str, Any]]:
177 """
178 Stream a structured movie script from the LLM in a single call.
180 Filters out any non-JSON lines (prose, markdown fences, etc.) and
181 returns a list of shot_description dicts collected from the stream.
182 """
183 messages: List[Dict[str, Any]] = self.build_movie_messages(
184 movie_description, max_shots=max_shots
185 )
187 script_path = f"{self.job_path}/movie_script.jsonl"
188 shot_descriptions: List[Dict[str, Any]] = []
189 total_line_count = 0
191 async with aiofiles.open(script_path, "w") as script_file:
192 buffer = ""
194 async for chunk in self.gen.gen_text_stream(
195 messages=messages,
196 max_tokens=max_tokens,
197 task_id="movie_script",
198 ):
199 if not chunk:
200 continue
201 buffer += chunk
202 # Parse complete JSONL lines as they arrive
203 while "\n" in buffer:
204 line, buffer = buffer.split("\n", 1)
205 line = line.strip()
206 if not line:
207 continue
208 # Skip markdown code fences
209 if line.startswith("```"):
210 continue
211 parsed = self._try_parse_json(line, total_line_count + 1)
212 if parsed is None:
213 continue
214 await script_file.write(line + "\n")
215 await script_file.flush()
216 total_line_count += 1
217 line_type = parsed.get("type", "")
218 if line_type == "shot_description":
219 shot_descriptions.append(parsed)
220 self.logger.info(
221 f"Shot {parsed.get('shot_id', total_line_count)}: "
222 f"{str(parsed.get('visual_prompt', ''))[:MAX_LOG_TEXT]}...")
224 # Flush any remaining buffer content
225 remainder = buffer.strip()
226 if remainder and not remainder.startswith("```"):
227 parsed = self._try_parse_json(remainder, total_line_count + 1)
228 if parsed is not None:
229 await script_file.write(remainder + "\n")
230 total_line_count += 1
231 if parsed.get("type") == "shot_description":
232 shot_descriptions.append(parsed)
234 self.logger.info(
235 f"Movie script complete: {total_line_count} JSONL lines, {len(shot_descriptions)} shots.")
236 return shot_descriptions
238 def _try_parse_json(
239 self,
240 line: str,
241 line_num: int,
242 ) -> Optional[Dict[str, Any]]:
243 """Try to parse a JSON line; return None on failure."""
244 try:
245 return json.loads(line)
246 except json.JSONDecodeError:
247 self.logger.debug(f"Skipping non-JSON line {line_num}: {line[:80]}")
248 return None
250 def _get_shot_deadline(
251 self,
252 shot_idx: int,
253 shot_duration: float,
254 ) -> float:
255 """
256 Compute the scheduling deadline for a shot.
257 Uses the shot's time offset in the final movie plus a generation buffer,
258 so earlier shots get higher priority while all shots have enough time.
259 """
260 return (
261 self.get_submission_time()
262 + (shot_idx * shot_duration)
263 + SHOT_DEADLINE_BUFFER_SECS
264 )
266 async def _gen_shot(
267 self,
268 shot_idx: int,
269 shot: Dict[str, Any],
270 ) -> Optional[str]:
271 """
272 Generate a single movie shot: image -> video (with optional lip-synced audio).
273 Returns the path to the saved shot video, or None on failure.
274 """
275 visual_prompt = shot.get("visual_prompt", "")
276 neg_prompt = shot.get("negative_prompt", "")
277 dialogue = shot.get("dialogue", None)
278 tech_specs = shot.get("technical_specs", {})
279 shot_duration = float(tech_specs.get("duration_seconds", DEFAULT_SHOT_DURATION_SECS))
281 deadline = self._get_shot_deadline(shot_idx, shot_duration)
283 self.logger.info(
284 f"[{shot_idx}] Generating shot with prompt: '{visual_prompt[:MAX_LOG_TEXT]}...'")
286 # Generate image for this shot
287 image = await self.gen.gen_image(
288 prompt=visual_prompt,
289 neg_prompt=neg_prompt,
290 width=self.width,
291 height=self.height,
292 task_id=f"shot_{shot_idx:03d}_img",
293 deadline=deadline,
294 )
296 image_path = f"{self.job_path}/shot_{shot_idx:03d}.png"
297 image.save(image_path)
298 self.logger.info(f"[{shot_idx}] Image saved to '{image_path}'.")
300 output_mode = self.get_config_output_mode()
301 video_binary: Optional[bytes] = None
303 if dialogue and dialogue.strip():
304 # Generate audio for the dialogue (for all output modes)
305 speech_speed = self.get_config_float("speech_speed", DEFAULT_SPEECH_SPEED)
306 audio_base64 = await self.gen.gen_audio(
307 text=dialogue.strip(),
308 speed=speech_speed,
309 task_id=f"shot_{shot_idx:03d}_audio",
310 deadline=deadline,
311 )
312 if audio_base64:
313 audio_path = f"{self.job_path}/shot_{shot_idx:03d}.wav"
314 await save_base64_as_binary(audio_path, audio_base64)
315 audio_duration = get_audio_duration(audio_base64)
316 self.logger.info(
317 f"[{shot_idx}] Generated audio with {bytes_to_human(len(audio_base64))} "
318 f"and {audio_duration:.2f} seconds.")
320 if output_mode is OutputMode.AUDIO_ONLY:
321 # Static image + audio: no generative video
322 num_frames = int(round(audio_duration * HUNYUANFRAMEPACK_FPS))
323 video_frames = [image] * num_frames
324 merged_path = f"{self.job_path}/shot_{shot_idx:03d}_merged.mp4"
325 await save_video_audio(
326 video_frames,
327 audio_path,
328 fps=HUNYUANFRAMEPACK_FPS,
329 out_video_path=merged_path,
330 )
331 async with aiofiles.open(merged_path, "rb") as fh:
332 video_binary = await fh.read()
333 elif output_mode is OutputMode.VIDEO_AUDIO_SYNCED:
334 # Lip-synced video
335 video_binary = await self.gen.gen_video_audio_from_img(
336 img=image,
337 audio_base64=audio_base64,
338 prompt=visual_prompt,
339 neg_prompt=neg_prompt,
340 width=self.width,
341 height=self.height,
342 steps=self.get_num_steps(),
343 task_id=f"shot_{shot_idx:03d}_video",
344 deadline=deadline,
345 )
346 else:
347 # Unsynced: generate video then merge audio track
348 raw_video_binary = await self.gen.gen_video(
349 img=image,
350 prompt=visual_prompt,
351 neg_prompt=neg_prompt,
352 width=self.width,
353 height=self.height,
354 video_seconds=audio_duration,
355 steps=self.get_num_steps(),
356 task_id=f"shot_{shot_idx:03d}_video",
357 deadline=deadline,
358 )
359 merged_path = f"{self.job_path}/shot_{shot_idx:03d}_merged.mp4"
360 await save_video_audio(
361 raw_video_binary,
362 audio_path,
363 out_video_path=merged_path,
364 )
365 async with aiofiles.open(merged_path, "rb") as fh:
366 video_binary = await fh.read()
368 if video_binary is None:
369 if output_mode is OutputMode.AUDIO_ONLY:
370 # No dialogue (or audio generation failed): static image video, no generative model
371 num_frames = int(round(shot_duration * HUNYUANFRAMEPACK_FPS))
372 video_frames_static = [image] * num_frames
373 static_path = f"{self.job_path}/shot_{shot_idx:03d}_static.mp4"
374 await save_video_frames(
375 video_frames_static,
376 fps=HUNYUANFRAMEPACK_FPS,
377 out_video_path=static_path)
378 async with aiofiles.open(static_path, "rb") as fh:
379 video_binary = await fh.read()
380 else:
381 # No dialogue, or audio generation failed: plain AI-generated video
382 video_binary = await self.gen.gen_video(
383 img=image,
384 prompt=visual_prompt,
385 neg_prompt=neg_prompt,
386 width=self.width,
387 height=self.height,
388 video_seconds=shot_duration,
389 steps=self.get_num_steps(),
390 task_id=f"shot_{shot_idx:03d}_video",
391 deadline=deadline,
392 )
394 if not video_binary:
395 self.logger.error(f"[{shot_idx}] No video binary produced for shot.")
396 return None
398 shot_path = f"{self.job_path}/shot_{shot_idx:03d}.mp4"
399 async with aiofiles.open(shot_path, "wb") as file:
400 await file.write(video_binary)
402 self._log_video_info(f"[{shot_idx}] Shot video", video_binary)
403 return shot_path