Coverage for apps/streamedit/streamedit_job.py: 77%
202 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamEdit job to generate an edited video.
3"""
4import sys
5import json
6import aiofiles
7import aiofiles.os
9from typing import override
10from typing import Any
11from typing import Dict
12from typing import List
13from typing import Optional
15from dataclasses import asdict
17from scenedetect import open_video
18from scenedetect import SceneManager
19from scenedetect.detectors import ContentDetector
20from scenedetect.stats_manager import StatsManager
22# Local relative imports
23sys.path.append("..") # noqa: E402
24sys.path.append("../..") # noqa: E402
26from streamwise_job import StreamWiseJob
27from streamwise_job import JobStatus
29from lmm_service_manager import LMMServiceManager
31from scene import SceneSegment
33from console_utils import bytes_to_human
35from file_utils import save_base64_as_binary
36from file_utils import read_file_bytes
37from file_utils import read_file_base64
38from media_utils import chunk_video_binary
39from media_utils import chunk_audio_base64
40from media_utils import get_video_frames
41from media_utils import extract_audio_from_video
42from media_utils import concatenate_videos
44from tts_utils import get_audio_chunks_by_silences
46from video import MAX_FT_DURATION_SECS
47from video import FANTASYTALKING_FPS
48from video import get_num_video_frames_from_duration
50from edit_prompts import build_edit_prompt
51from edit_prompts import EDIT_PROMPT
54class StreamEditJob(StreamWiseJob):
55 """A job to generate an edited video."""
57 def __init__(
58 self,
59 job_id: str,
60 service_manager: LMMServiceManager,
61 config: Dict[str, Any] = {},
62 ) -> None:
63 super().__init__(
64 "streamedit",
65 job_id,
66 service_manager,
67 config)
68 self.scenes: List[SceneSegment] = [] # Populated by detect_scenes() during gen_edit()
70 @override
71 async def generate(
72 self,
73 job_config: Dict[str, Any],
74 ) -> None:
75 video_base64 = job_config.get("video_base64", None)
76 assert video_base64 is not None
77 await self.gen_edit(video_base64)
79 def _save_frame_from_video(self, video_path: str, frame_path: str, frame_num: int = 0) -> bool:
80 """Extract a single frame from a video file and save it as a JPEG.
82 Returns True on success, False if the frame could not be read.
83 """
84 import cv2 # noqa: PLC0415 - deferred to avoid module-level dependency
85 cap = cv2.VideoCapture(video_path)
86 cap.set(cv2.CAP_PROP_POS_FRAMES, frame_num)
87 ok, frame = cap.read()
88 cap.release()
89 if ok:
90 cv2.imwrite(frame_path, frame)
91 return ok
93 async def extract_scene_frames(self) -> None:
94 """Extract one representative frame per scene from the input video and save as PNG.
96 Populates ``scene.frame_image_paths`` for each scene so the WebUI can
97 display a thumbnail alongside the scene details.
98 """
99 video_path = f"{self.job_path}/video.mp4"
100 for scene in self.scenes:
101 frame_file_name = f"scene_{scene.scene_id:03d}_frame.png"
102 frame_path = f"{self.job_path}/{frame_file_name}"
103 ok = self._save_frame_from_video(video_path, frame_path, scene.start_frame)
104 if ok:
105 scene.add_image_path(frame_file_name)
106 else:
107 self.logger.warning(
108 f"[{scene.scene_id}] Failed to extract frame at position {scene.start_frame}.")
110 async def _save_video_as_output(self, video_path: str) -> None:
111 """Copy a video file to the job output path."""
112 out_path = f"{self.job_path}/{self.job_id}.mp4"
113 video_binary = await read_file_bytes(video_path)
114 async with aiofiles.open(out_path, "wb") as file:
115 await file.write(video_binary)
116 self.logger.info(f"Video ({bytes_to_human(len(video_binary))}) saved to '{out_path}'.")
118 async def gen_edit(
119 self,
120 video_base64: str,
121 ) -> None:
122 """
123 Generate an edited video given an input video.
124 """
125 async with self.job_status_handler():
126 if not video_base64:
127 self.logger.error("Video is required.")
128 await self.save_status(JobStatus.FAILED)
129 raise ValueError("Missing 'video_base64' in request")
130 self.logger.info(f"Generating edit for video with {bytes_to_human(len(video_base64))}.")
132 # Save video locally for processing and debugging
133 self.logger.info(f"Saving input video with {bytes_to_human(len(video_base64))}.")
134 video_path = f"{self.job_path}/video.mp4"
135 await save_base64_as_binary(video_path, video_base64)
137 await self.save_status(JobStatus.RUNNING)
139 # Detect scenes first (this validates the video and may take time)
140 self.scenes = await self.detect_scenes(video_path)
141 self.logger.info(f"Detected {len(self.scenes)} scenes.")
143 await self.save_status(JobStatus.RUNNING)
145 # Build edit prompt from user instructions
146 edit_instructions = self.get_config_str("edit_instructions")
147 edit_prompt = build_edit_prompt(edit_instructions)
149 # If no scenes detected, fall back to returning the original video unchanged
150 if not self.scenes:
151 self.logger.warning("No scenes detected. Saving original video as output.")
152 await self._save_video_as_output(video_path)
153 return
155 # Chunk audio into per-scene files
156 await self.chunk_audio_into_scenes()
158 # Extract a thumbnail frame per scene for the WebUI
159 await self.extract_scene_frames()
161 # Write scenes debug file (audio_path and frame_image_paths are now populated)
162 scenes_path = f"{self.job_path}/scenes.json"
163 async with aiofiles.open(scenes_path, "w") as scene_file:
164 scenes_dict_list = [asdict(scene) for scene in self.scenes]
165 await scene_file.write(json.dumps(scenes_dict_list, indent=2))
167 await self.save_status(JobStatus.RUNNING)
169 # Edit each scene
170 scene_video_paths: List[Optional[str]] = []
171 for scene in self.scenes:
172 try:
173 scene_path = await self.gen_edit_scene(scene, edit_prompt)
174 scene_video_paths.append(scene_path)
175 self.logger.info(f"[{scene.scene_id}] Edited scene saved to '{scene_path}'.")
176 # Extract a thumbnail from the edited scene for the WebUI
177 edit_frame_path = f"{self.job_path}/scene_{scene.scene_id:03d}_edit_frame.png"
178 if not self._save_frame_from_video(scene_path, edit_frame_path):
179 self.logger.warning(f"[{scene.scene_id}] Could not extract edit frame.")
180 except Exception as ex:
181 self.logger.error(f"[{scene.scene_id}] Error editing scene: {ex}")
182 scene_video_paths.append(None)
184 await self.save_status(JobStatus.RUNNING)
186 # Combine the edited scenes into a final video
187 valid_paths = [p for p in scene_video_paths if p]
188 if not valid_paths:
189 # No scenes could be edited – fall back to the original video
190 self.logger.warning("No edited scenes produced. Saving original video as output.")
191 await self._save_video_as_output(video_path)
192 return
194 self.logger.info(f"Combining {len(valid_paths)} edited scenes into final video...")
195 scene_binaries: List[bytes] = []
196 for path in valid_paths:
197 scene_binary = await read_file_bytes(path)
198 scene_binaries.append(scene_binary)
200 video_binary = await concatenate_videos(scene_binaries)
201 if not video_binary:
202 raise ValueError("Cannot concatenate edited scenes into final video.")
204 out_path = f"{self.job_path}/{self.job_id}.mp4"
205 async with aiofiles.open(out_path, "wb") as file:
206 await file.write(video_binary)
208 self.logger.info(f"Generated edited video with {bytes_to_human(len(video_binary))} at '{out_path}'.")
210 async def detect_scenes(
211 self,
212 video_path: str,
213 threshold: float = 27.0,
214 min_scene_len: int = 15,
215 ) -> List[SceneSegment]:
216 """
217 Return list of scenes for a video.
218 TODO Move to a library later.
219 """
220 if not await aiofiles.os.path.exists(video_path):
221 raise FileNotFoundError(f"Video file not found: {video_path}")
222 video = open_video(video_path)
224 stats_manager = StatsManager()
225 scene_manager = SceneManager(stats_manager)
226 content_detector = ContentDetector(
227 threshold=threshold,
228 min_scene_len=min_scene_len)
229 scene_manager.add_detector(content_detector)
231 scene_manager.detect_scenes(video)
232 scene_list = scene_manager.get_scene_list()
234 scenes = []
235 for scene_id, (start_tc, end_tc) in enumerate(scene_list):
236 scene = SceneSegment(
237 scene_id,
238 start_tc.get_frames(), end_tc.get_frames(),
239 start_tc.get_seconds(), end_tc.get_seconds()
240 )
241 scenes.append(scene)
242 return scenes
244 async def chunk_audio_into_scenes(self) -> List[str]:
245 """
246 Chunk the audio of the video into per-scene WAV files.
247 Sets scene.audio_path on each SceneSegment to the relative filename
248 (e.g. scene_000.wav) and saves the file to the job directory.
249 Returns a list of base64-encoded audio chunks, one per scene.
250 """
251 chunks = []
252 try:
253 video_path = f"{self.job_path}/video.mp4"
254 audio_path = f"{self.job_path}/audio.wav"
255 audio_path = await extract_audio_from_video(video_path, audio_path)
256 audio_base64 = await read_file_base64(audio_path)
257 self.logger.info(f"Extracted audio with {bytes_to_human(len(audio_base64))}.")
259 for scene in self.scenes:
260 scene_audio_base64 = chunk_audio_base64(
261 audio_base64=audio_base64,
262 start_seconds=scene.start_sec,
263 end_seconds=scene.end_sec)
264 scene_audio_path = f"{self.job_path}/scene_{scene.scene_id:03d}.wav"
265 await save_base64_as_binary(scene_audio_path, scene_audio_base64)
266 scene.audio_path = f"scene_{scene.scene_id:03d}.wav"
267 chunks.append(scene_audio_base64)
268 except Exception as ex:
269 self.logger.error(f"Error during audio chunking: {ex} [{type(ex)}]")
270 return chunks
272 async def gen_edit_scene(
273 self,
274 scene: SceneSegment,
275 edit_prompt: str = EDIT_PROMPT,
276 ) -> str:
277 """
278 Generate edited version of a scene and save it to disk.
279 Returns the path to the saved file.
280 """
281 input_video_path = f"{self.job_path}/video.mp4"
282 input_video_binary = await read_file_bytes(input_video_path)
284 scene_binary = chunk_video_binary(
285 video_binary=input_video_binary,
286 start_seconds=scene.start_sec,
287 end_seconds=scene.end_sec,
288 )
290 scene_audio_path = f"{self.job_path}/{scene.audio_path}"
291 scene_audio_base64 = await read_file_base64(scene_audio_path)
293 scene_duration = scene.end_sec - scene.start_sec
294 if scene_duration > MAX_FT_DURATION_SECS:
295 scene_edit_binary = await self._gen_edit_scene_chunked(
296 scene=scene,
297 scene_binary=scene_binary,
298 scene_audio_path=scene_audio_path,
299 scene_audio_base64=scene_audio_base64,
300 edit_prompt=edit_prompt,
301 )
302 else:
303 scene_video_frames = await get_video_frames(scene_binary)
304 scene_edit_binary = await self.gen.gen_video_audio_from_video(
305 video=scene_video_frames,
306 audio_base64=scene_audio_base64,
307 prompt=edit_prompt,
308 task_id=f"{scene.scene_id:03d}",
309 deadline=self.get_submission_time() + scene.start_sec,
310 )
312 scene_path = f"{self.job_path}/scene_{scene.scene_id:03d}_edit.mp4"
313 async with aiofiles.open(scene_path, "wb") as file:
314 await file.write(scene_edit_binary)
315 return scene_path
317 async def _gen_edit_scene_chunked(
318 self,
319 scene: SceneSegment,
320 scene_binary: bytes,
321 scene_audio_path: str,
322 scene_audio_base64: str,
323 edit_prompt: str,
324 ) -> bytes:
325 """
326 Generate edited scene in sub-chunks when the scene exceeds MAX_FT_DURATION_SECS.
327 Splits the scene audio by silences, generates each sub-chunk independently,
328 then concatenates the results.
329 """
330 scene_duration = scene.end_sec - scene.start_sec
331 self.logger.info(
332 "[%d] Scene too long (%.3f > %.3f s), using chunked generation.",
333 scene.scene_id, scene_duration, MAX_FT_DURATION_SECS)
335 audio_splits = get_audio_chunks_by_silences(
336 scene_audio_path,
337 max_duration_seconds=MAX_FT_DURATION_SECS,
338 chunk_alignment_seconds=1.0 / FANTASYTALKING_FPS,
339 )
341 if not audio_splits:
342 raise ValueError(f"[{scene.scene_id}] No audio splits produced for long scene.")
344 self.logger.info(
345 "[%d] Split into %d sub-chunks.", scene.scene_id, len(audio_splits))
347 sub_binaries: List[bytes] = []
348 for sub_idx, (sub_start_s, sub_end_s) in enumerate(audio_splits):
349 sub_duration = sub_end_s - sub_start_s
350 self.logger.info(
351 "[%d.%d] Sub-chunk %.3f-%.3f s (%.3f s).",
352 scene.scene_id, sub_idx, sub_start_s, sub_end_s, sub_duration)
354 sub_scene_binary = chunk_video_binary(
355 video_binary=scene_binary,
356 start_seconds=sub_start_s,
357 end_seconds=sub_end_s,
358 )
359 sub_scene_frames = await get_video_frames(sub_scene_binary)
361 if not sub_scene_frames:
362 self.logger.warning(
363 "[%d.%d] No frames for sub-chunk, skipping.", scene.scene_id, sub_idx)
364 continue
366 # Align frame count to what FantasyTalking expects for the audio duration
367 expected_frames = get_num_video_frames_from_duration(sub_duration)
368 if len(sub_scene_frames) < expected_frames:
369 sub_scene_frames += [sub_scene_frames[-1]] * (expected_frames - len(sub_scene_frames))
370 elif len(sub_scene_frames) > expected_frames:
371 sub_scene_frames = sub_scene_frames[:expected_frames]
373 sub_audio_base64 = chunk_audio_base64(
374 audio_base64=scene_audio_base64,
375 start_seconds=sub_start_s,
376 end_seconds=sub_end_s,
377 )
379 sub_edit_binary = await self.gen.gen_video_audio_from_video(
380 video=sub_scene_frames,
381 audio_base64=sub_audio_base64,
382 prompt=edit_prompt,
383 task_id=f"{scene.scene_id:03d}_{sub_idx:03d}",
384 deadline=self.get_submission_time() + scene.start_sec + sub_start_s,
385 )
386 sub_binaries.append(sub_edit_binary)
388 if not sub_binaries:
389 raise ValueError(f"[{scene.scene_id}] All sub-chunks failed for long scene.")
391 scene_edit_binary = await concatenate_videos(sub_binaries)
392 if not scene_edit_binary:
393 raise ValueError(f"[{scene.scene_id}] Cannot concatenate sub-chunk results.")
394 return scene_edit_binary