Coverage for apps/streamdub/streamdub_job.py: 62%
265 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamDub job to generate a dubbed video.
3It coordinates the execution of the different models.
4"""
6import re
7import sys
8import json
9import aiofiles
11from dataclasses import asdict
13from typing import override
14from typing import Dict
15from typing import Any
16from typing import List
17from typing import Optional
19from dub_prompts import DUB_PROMPT
20from dub_prompts import VIDEO_DUB_PROMPT
21from dub_prompts import VIDEO_DUB_NEG_PROMPT
23from video import MAX_FT_DURATION_SECS
24from video import FANTASYTALKING_FPS
26from scenedetect import open_video
27from scenedetect import SceneManager
28from scenedetect.detectors import ContentDetector
29from scenedetect.stats_manager import StatsManager
31# Local relative imports
32sys.path.append("..") # noqa: E402
33sys.path.append("../..") # noqa: E402
35from streamwise_job import StreamWiseJob
36from streamwise_job import JobStatus
38from lmm_service_manager import LMMServiceManager
40from scene import SceneSegment
42from console_utils import bytes_to_human
44from file_utils import read_file_base64
45from file_utils import save_base64_as_binary
46from file_utils import read_file_bytes
48from media_utils import concatenate_videos
49from media_utils import fit_audio_to_duration
50from media_utils import chunk_video_binary
51from media_utils import get_video_file_info
52from media_utils import get_video_frames_at_fps
53from media_utils import get_audio_duration
54from media_utils import extract_audio_from_video
55from media_utils import get_video_frames
56from media_utils import chunk_audio_base64
57from media_utils import save_video_audio
59from language_utils import to_language
62def _is_empty_transcript(text: str) -> bool:
63 """Return True when the transcript contains no translatable text.
65 Transcripts such as ``""``, ``"-"``, ``"- -"``, and ``"♪ ♪"`` have no
66 actual words and should not be sent to the translation service.
67 """
68 return not bool(re.search(r'\w', text))
71class StreamDubJob(StreamWiseJob):
72 """A job to generate a dubbed video."""
74 def __init__(
75 self,
76 job_id: str,
77 service_manager: LMMServiceManager,
78 config: Dict[str, Any] = {},
79 ) -> None:
80 super().__init__(
81 "streamdub",
82 job_id,
83 service_manager,
84 config)
86 self.service_manager = service_manager
87 self.scenes: List[SceneSegment] = []
89 @override
90 async def generate(
91 self,
92 job_config: Dict[str, Any],
93 ) -> None:
94 video_base64 = job_config.get("video_base64", None)
95 output_language = job_config.get("output_language", None)
96 await self.gen_dub(
97 video_base64,
98 output_language)
100 async def save_scenes(self) -> None:
101 """Persist current scene metadata (including transcript/translation) to scenes.json."""
102 scenes_path = f"{self.job_path}/scenes.json"
103 async with aiofiles.open(scenes_path, "w") as scene_file:
104 scenes_dict_list = [asdict(scene) for scene in self.scenes]
105 scenes_json = json.dumps(scenes_dict_list, indent=2)
106 await scene_file.write(scenes_json)
108 async def gen_dub(
109 self,
110 video_base64: Optional[str] = None,
111 output_language: Optional[str] = "e", # Spanish
112 ) -> None:
113 """Generate a dubbed video."""
114 async with self.job_status_handler():
115 if not video_base64:
116 self.logger.error("Video is required.")
117 await self.save_status(JobStatus.FAILED)
118 raise ValueError("Missing 'video_base64' in request")
119 self.logger.info(f"Generating dubbed version for video with {bytes_to_human(len(video_base64))}.")
121 await self.save_status(JobStatus.RUNNING)
123 # Save as video for processing
124 self.logger.info(f"Saving input video with {bytes_to_human(len(video_base64))}.")
125 video_path = f"{self.job_path}/video.mp4"
126 await save_base64_as_binary(video_path, video_base64)
128 await self.save_status(JobStatus.RUNNING)
130 # Detect scenes
131 self.scenes = await self.detect_scenes()
132 if not self.scenes:
133 raise ValueError("No scenes detected in video.")
134 self.logger.info(f"Detected {len(self.scenes)} scenes.")
136 await self.save_status(JobStatus.RUNNING)
138 # Extract audio from each scene
139 await self.chunk_audio_into_scenes()
141 # Persist scene metadata so the UI can display scenes with original audio
142 self.logger.info("Scenes:")
143 for idx, scene in enumerate(self.scenes):
144 self.logger.info(f" Scene {idx}: {scene}")
145 await self.save_scenes()
147 scene_binaries = []
148 for scene in self.scenes:
149 try:
150 video_binary = await self.gen_dub_scene(
151 scene,
152 output_language)
153 scene_binaries.append(video_binary)
154 except Exception as ex:
155 self.logger.error(f"[{scene.scene_id}] Cannot generate dubbed scene: {ex}")
157 await self.save_status(JobStatus.RUNNING)
159 if not scene_binaries:
160 raise ValueError("No scenes generated.")
162 # Update scenes.json so the UI reflects the dubbed audio paths
163 await self.save_scenes()
165 # Concatenate scenes
166 video_binary = await concatenate_videos(
167 scene_binaries,
168 fast_copy=False)
169 if not video_binary:
170 raise ValueError("Cannot concatenate scenes into final dubbed video.")
171 video_info = get_video_file_info(video_binary)
172 video_duration = video_info["video"]["duration_seconds"]
173 video_fps = video_info["video"]["fps"]
175 video_path = f"{self.job_path}/{self.job_id}.mp4"
176 async with aiofiles.open(video_path, "wb") as file:
177 await file.write(video_binary)
179 self.logger.info(
180 f"Generated dubbed video with {video_duration:.3f} seconds, "
181 f"{video_fps} FPS, "
182 f"{bytes_to_human(len(video_binary))}, and "
183 f"'{video_path}'")
185 async def detect_scenes(
186 self,
187 threshold: float = 27.0,
188 min_scene_len: int = 15,
189 ) -> List[SceneSegment]:
190 """
191 Return list of (start_frame, end_frame, start_sec, end_sec).
192 TODO refactor with streamshort_job.py
193 """
194 video_path = f"{self.job_path}/video.mp4"
195 if not await aiofiles.os.path.exists(video_path):
196 raise FileNotFoundError(f"Video file not found: {video_path}")
197 video = open_video(video_path)
199 stats_manager = StatsManager()
200 scene_manager = SceneManager(stats_manager)
201 content_detector = ContentDetector(
202 threshold=threshold,
203 min_scene_len=min_scene_len)
204 scene_manager.add_detector(content_detector)
206 scene_manager.detect_scenes(video)
207 scene_list = scene_manager.get_scene_list()
209 scenes = []
210 for scene_id, (start_tc, end_tc) in enumerate(scene_list):
211 scene = SceneSegment(
212 scene_id,
213 start_tc.get_frames(), end_tc.get_frames(),
214 start_tc.get_seconds(), end_tc.get_seconds()
215 )
216 scenes.append(scene)
217 return scenes
219 async def chunk_audio_into_scenes(self) -> List[str]:
220 """
221 Chunk the audio of the video into scenes.
222 """
223 chunks = []
224 try:
225 video_path = f"{self.job_path}/video.mp4"
226 audio_path = f"{self.job_path}/audio.wav"
227 audio_path = await extract_audio_from_video(video_path, audio_path)
228 audio_base64 = await read_file_base64(audio_path)
229 self.logger.info(f"Extracted audio with {bytes_to_human(len(audio_base64))}.")
231 for scene in self.scenes:
232 scene_audio_base64 = chunk_audio_base64(
233 audio_base64=audio_base64,
234 start_seconds=scene.start_sec,
235 end_seconds=scene.end_sec)
236 scene.audio_path = f"scene_{scene.scene_id:03d}.wav"
237 scene_audio_path = f"{self.job_path}/{scene.audio_path}"
238 await save_base64_as_binary(scene_audio_path, scene_audio_base64)
239 chunks.append(scene_audio_base64)
240 except Exception as ex:
241 self.logger.error(f"Error during audio chunking: {ex} [{type(ex)}]")
242 return chunks
244 async def gen_dub_scene(
245 self,
246 scene: SceneSegment,
247 lang_code: Optional[str] = "e", # Spanish
248 ) -> bytes:
249 """
250 Generate scene with dubbed audio.
251 """
252 if lang_code is None:
253 lang_code = "e" # Default to Spanish (matches the parameter default above)
254 scene_id = scene.scene_id
255 self.logger.info(
256 f"[{scene_id}] Generating dubbed scene into {to_language(lang_code)} ('{lang_code}').")
258 # Transcribe audio
259 scene_transcript = await self.transcribe_audio(scene)
260 scene.transcript = scene_transcript
262 # TODO make this async into tasks
264 if "♪" in scene.transcript:
265 self.logger.info(f"[{scene_id}] Scene contains music ({scene.transcript}), skip dubbing.")
266 return await self.get_video_scene(scene)
267 if _is_empty_transcript(scene.transcript):
268 self.logger.info(f"[{scene_id}] Scene has no translatable text ({scene.transcript!r}), skip dubbing.")
269 return await self.get_video_scene(scene)
271 self.logger.info(f"[{scene_id}] Transcript: {scene.transcript[0:80]}...")
272 transcript_path = f"{self.job_path}/scene_{scene_id:03d}.txt"
273 async with aiofiles.open(transcript_path, "w") as file:
274 await file.write(scene.transcript)
275 await self.save_scenes()
277 # Translate transcription
278 scene.translation = await self.translate_scene(
279 scene,
280 output_lang_code=lang_code)
282 self.logger.info(f"[{scene_id}] Translation: {scene.translation[0:80]}...")
283 transcript_path = f"{self.job_path}/scene_{scene_id:03d}_translation.txt"
284 async with aiofiles.open(transcript_path, "w") as file:
285 await file.write(scene.translation)
286 await self.save_scenes()
288 if not scene.translation.strip():
289 self.logger.info(
290 f"[{scene_id}] No translation available, using original audio and skipping lip sync.")
291 return await self.get_video_scene(scene)
293 await self.save_status(JobStatus.RUNNING)
295 # Generate dubbed audio using the original scene audio as the voice reference
296 # so that the speaker's voice identity is preserved in the dubbed output.
297 deadline = self.get_submission_time() + scene.start_sec
298 original_audio_path = f"{self.job_path}/scene_{scene_id:03d}.wav"
299 voice_sample: Optional[str] = None
300 if self.config.get("voice_cloning", True):
301 try:
302 voice_sample = await read_file_base64(original_audio_path)
303 self.logger.info(
304 f"[{scene_id}] Using original scene audio for voice cloning "
305 f"({bytes_to_human(len(voice_sample))}).")
306 except Exception as ex:
307 self.logger.warning(
308 f"[{scene_id}] Could not read original scene audio for voice cloning: {ex}. "
309 "Falling back to default voice.")
310 else:
311 self.logger.info(f"[{scene_id}] Voice cloning disabled; using default voice.")
312 if voice_sample is not None:
313 try:
314 audio_base64 = await self.gen.gen_audio(
315 text=scene.translation,
316 lang_code=lang_code,
317 voice_sample=voice_sample,
318 task_id=f"{scene_id:03d}",
319 deadline=deadline,
320 )
321 except Exception as ex:
322 self.logger.warning(
323 f"[{scene_id}] Voice cloning failed: {ex}. Falling back to default voice.")
324 audio_base64 = await self.gen.gen_audio(
325 text=scene.translation,
326 lang_code=lang_code,
327 task_id=f"{scene_id:03d}",
328 deadline=deadline,
329 )
330 else:
331 audio_base64 = await self.gen.gen_audio(
332 text=scene.translation,
333 lang_code=lang_code,
334 task_id=f"{scene_id:03d}",
335 deadline=deadline,
336 )
337 scene.audio_path = f"scene_{scene_id:03d}_dubbed.wav"
338 scene_audio_path = f"{self.job_path}/{scene.audio_path}"
339 await save_base64_as_binary(scene_audio_path, audio_base64)
341 # Lip sync video scenes
342 # TODO if video too long, chunk it
343 if scene.duration_sec > MAX_FT_DURATION_SECS:
344 self.logger.warning(
345 f"[{scene_id}] Scene too long: "
346 f"{scene.duration_sec:.3f} > {MAX_FT_DURATION_SECS:.3f} seconds.")
347 scene_dubbed_video_binary = await self.gen_video_lip_synced(scene)
349 # Add subtitles (on by default)
350 if self.config.get("add_subtitles", True):
351 scene_dubbed_video_binary = await self._add_subtitles_to_video(
352 scene, scene_dubbed_video_binary)
354 # Save scene video
355 scene_dubbed_video_path = f"{self.job_path}/scene_{scene_id:03d}_dubbed.mp4"
356 async with aiofiles.open(scene_dubbed_video_path, "wb") as file:
357 await file.write(scene_dubbed_video_binary)
359 return scene_dubbed_video_binary
361 async def get_video_scene(
362 self,
363 scene: SceneSegment
364 ) -> bytes:
365 """
366 Get the video bytes for a scene.
367 """
368 video_path = f"{self.job_path}/video.mp4"
369 video_binary = await read_file_bytes(video_path)
370 scene_video_binary = chunk_video_binary(
371 video_binary,
372 start_seconds=scene.start_sec,
373 end_seconds=scene.end_sec
374 )
376 scene_video_path = f"{self.job_path}/scene_{scene.scene_id:03d}.mp4"
377 async with aiofiles.open(scene_video_path, "wb") as file:
378 await file.write(scene_video_binary)
380 return scene_video_binary
382 async def _add_subtitles_to_video(
383 self,
384 scene: SceneSegment,
385 video_binary: bytes,
386 ) -> bytes:
387 """
388 Overlay the translated subtitle text onto every frame of the dubbed video.
389 """
390 if not scene.translation:
391 return video_binary
393 scene_id = scene.scene_id
394 video_frames = await self._overlay_subtitles_on_frames(video_binary, scene.translation)
395 _fps = get_video_file_info(video_binary)["video"].get("fps")
396 video_fps: float = _fps if _fps is not None else FANTASYTALKING_FPS
398 scene_audio_path = f"{self.job_path}/{scene.audio_path}"
399 subtitled_path = f"{self.job_path}/scene_{scene_id:03d}_dubbed_subtitled.mp4"
400 subtitled_path = await save_video_audio(
401 video_content=video_frames,
402 audio_path=scene_audio_path,
403 fps=video_fps,
404 out_video_path=subtitled_path,
405 )
406 async with aiofiles.open(subtitled_path, "rb") as subtitle_file:
407 return await subtitle_file.read()
409 async def transcribe_audio(
410 self,
411 scene: SceneSegment
412 ) -> str:
413 """
414 Transcribe the audio of the video.
415 TODO refactor with streamshort_job.py
416 """
417 if not scene.audio_path:
418 return ""
420 audio_path = f"{self.job_path}/{scene.audio_path}"
421 audio_transcript, lang_code = await self.gen.gen_audio_transcript(
422 audio_path,
423 task_id=f"{scene.scene_id:03d}",
424 )
425 if not audio_transcript:
426 return ""
427 scene.language = lang_code
428 audio_transcript = audio_transcript.strip()
429 return audio_transcript
431 async def translate_scene(
432 self,
433 scene: SceneSegment,
434 input_lang_code: str = "a", # American English
435 output_lang_code: str = "e", # Spanish
436 ) -> str:
437 """
438 Translate text using LLM.
439 """
440 text = scene.transcript
441 if not text or _is_empty_transcript(text):
442 return ""
444 dub_prompt = DUB_PROMPT.format(
445 input_language=to_language(input_lang_code),
446 output_language=to_language(output_lang_code)
447 )
448 messages = [
449 {"role": "system", "content": dub_prompt},
450 {"role": "user", "content": text}
451 ]
452 prompt_path = f"{self.job_path}/translate_prompt_{scene.scene_id:03d}.txt"
453 async with aiofiles.open(prompt_path, "w") as prompt_file:
454 messages_json = json.dumps(messages, indent=2)
455 await prompt_file.write(messages_json)
457 translated_text = await self.gen.gen_text(
458 messages,
459 task_id=f"translate{scene.scene_id:03d}",
460 )
462 return translated_text
464 async def gen_video_lip_synced(
465 self,
466 scene: SceneSegment
467 ) -> bytes:
468 """
469 Generate lip synced video for a scene.
470 """
471 scene_id = scene.scene_id
473 # Chunk video for the scene
474 scene_video_binary = await self.get_video_scene(scene)
476 # Convert frames into the right sampling rate
477 scene_video_frames = await get_video_frames(scene_video_binary)
478 video_info = get_video_file_info(scene_video_binary)
479 _scene_fps = video_info["video"].get("fps")
480 scene_video_fps: float = _scene_fps if _scene_fps is not None else FANTASYTALKING_FPS
481 if scene_video_fps != FANTASYTALKING_FPS:
482 self.logger.info(
483 f"[{scene_id}] Resampling scene video from "
484 f"{scene_video_fps:.2f} FPS to {FANTASYTALKING_FPS:.2f} FPS.")
485 scene_video_frames = get_video_frames_at_fps(
486 scene_video_frames,
487 scene_video_fps,
488 FANTASYTALKING_FPS,
489 )
491 # Merge video with dubbed audio
492 scene_audio_path = f"{self.job_path}/{scene.audio_path}"
493 scene_audio_base64 = await read_file_base64(scene_audio_path)
494 scene_audio_seconds = get_audio_duration(scene_audio_base64)
496 scene_video_frame = scene_video_frames[0]
497 width = scene_video_frame.width
498 height = scene_video_frame.height
499 scene_video_seconds = len(scene_video_frames) / FANTASYTALKING_FPS
500 num_steps = self.get_num_steps()
502 self.logger.info(
503 f"[{scene_id}] Generating lip synced video for scene "
504 f"with video with {len(scene_video_frames)} frames, "
505 f"{scene_video_seconds:.2f} seconds, "
506 f"resolution {width}x{height}, "
507 f"and audio with {scene_audio_seconds:.2f} seconds, "
508 f"in steps {num_steps}.")
510 # Pad video or audio if needed
511 if scene_audio_seconds > scene_video_seconds + 0.1:
512 extra_seconds = scene_audio_seconds - scene_video_seconds
513 extra_frames = int(extra_seconds * FANTASYTALKING_FPS)
514 last_frame = scene_video_frames[-1]
515 for _ in range(extra_frames):
516 scene_video_frames.append(last_frame)
517 self.logger.warning(
518 f"[{scene_id}] Audio is longer than video: "
519 f"{scene_audio_seconds:.3f} > {scene_video_seconds:.3f} seconds. "
520 f"Added {extra_frames} extra frames ({extra_seconds:.3f} seconds).")
521 # TODO we may want to resample instead of extending the end
522 elif scene_audio_seconds < scene_video_seconds - 0.1:
523 self.logger.warning(
524 f"[{scene_id}] Audio is shorter than video: "
525 f"{scene_audio_seconds:.3f} < {scene_video_seconds:.3f} seconds.")
526 fit_audio_to_duration(
527 scene_audio_path,
528 target_duration=scene_video_seconds,
529 output_path=scene_audio_path) # TODO new output path?
531 deadline = self.get_submission_time() + scene.start_sec
532 scene_dubbed_video_binary = await self.gen.gen_video_audio_from_video(
533 video=scene_video_frames,
534 audio_base64=scene_audio_base64,
535 prompt=VIDEO_DUB_PROMPT,
536 neg_prompt=VIDEO_DUB_NEG_PROMPT,
537 width=width,
538 height=height,
539 steps=num_steps,
540 task_id=f"{scene.scene_id:03d}",
541 deadline=deadline,
542 )
544 return scene_dubbed_video_binary