Coverage for apps/streamlecture/streamlecture_job.py: 94%
153 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1"""
2StreamLecture job to generate a lecture video.
3"""
4import asyncio
5import json
6import sys
7import aiofiles
9from PIL import Image
11from typing import override
12from typing import Any
13from typing import Dict
14from typing import List
15from typing import Optional
17from lecture_prompts import IMG_PROMPT
18from lecture_prompts import IMG_NEG_PROMPT
19from lecture_prompts import VIDEO_PROMPT
20from lecture_prompts import VIDEO_NEG_PROMPT
21from lecture_prompts import LECTURE_STYLE_PROMPT
22from lecture_prompts import LECTURE_CUSTOM_PROMPT
24# Local relative imports
25sys.path.append("..") # noqa: E402
26sys.path.append("../..") # noqa: E402
28from streamwise_job import StreamWiseJob
29from streamwise_job import JobStatus
30from streamwise_job import OutputMode
31from streamwise_job import MAX_LOG_TEXT
33from lmm_service_manager import LMMServiceManager
35from gen_video_chunked import GenVideoChunked
37from console_utils import bytes_to_human
39from file_utils import base64_to_binary
40from file_utils import save_base64_as_binary
42from pdf_utils import parse_pdf
44from media_utils import get_audio_duration
45from media_utils import get_video_frames
46from media_utils import get_video_file_info
47from media_utils import save_video_audio
48from media_utils import concatenate_videos
50from video import MAX_FT_DURATION_SECS
51from video import FANTASYTALKING_FPS
54class StreamLectureJob(StreamWiseJob):
55 """A job to generate a lecture 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 "streamlecture",
65 job_id,
66 service_manager,
67 config)
68 self.image: Optional[Image.Image] = None
70 @override
71 async def generate(
72 self,
73 job_config: Dict[str, Any],
74 ) -> None:
75 pdf_base64: Optional[str] = job_config.get("pdf_base64")
76 await self.gen_lecture(pdf_base64)
78 async def gen_scene(
79 self,
80 scene_id: int,
81 text: str,
82 ) -> str:
83 """
84 Generate a scene with TTS audio and video from the classroom image.
85 Returns the path to the saved scene video.
86 """
87 log_text = text[:MAX_LOG_TEXT] + "..." if len(text) > MAX_LOG_TEXT else text
88 self.logger.info(f"[{scene_id}] Generating scene with text '{log_text}'.")
90 # Generate TTS audio
91 voice = "af_heart"
92 speed = self.get_config_float("speech_speed", 1.1)
93 audio_base64 = await self.gen.gen_audio(
94 text,
95 voice=voice,
96 speed=speed,
97 task_id=f"{scene_id:03d}",
98 deadline=self.get_scene_deadline(scene_id),
99 )
100 if not audio_base64:
101 raise ValueError(f"Cannot generate audio for scene {scene_id}")
102 audio_duration = get_audio_duration(audio_base64)
103 audio_path = f"{self.job_path}/{scene_id:03d}.wav"
104 await save_base64_as_binary(audio_path, audio_base64)
105 self.logger.info(
106 f"[{scene_id}] Generated audio with {bytes_to_human(len(audio_base64))} "
107 f"and {audio_duration:.3f} seconds.")
109 if self.image is None:
110 raise ValueError(f"[{scene_id}] Classroom image not available for video generation.")
112 output_mode = self.get_config_output_mode()
113 num_steps = self.get_num_steps()
114 width, height = self.image.size
116 # Generate video from the classroom image + audio
117 if output_mode is OutputMode.AUDIO_ONLY:
118 # Audio-only: no video generation needed, handled at concatenation
119 return audio_path
121 if audio_duration < MAX_FT_DURATION_SECS and output_mode == OutputMode.VIDEO_AUDIO_SYNCED:
122 scene_video_binary = await self.gen.gen_video_audio_from_img(
123 img=self.image,
124 audio_base64=audio_base64,
125 prompt=VIDEO_PROMPT,
126 neg_prompt=VIDEO_NEG_PROMPT,
127 width=width,
128 height=height,
129 steps=num_steps,
130 task_id=f"{scene_id:03d}",
131 deadline=self.get_scene_deadline(scene_id),
132 )
133 elif audio_duration >= MAX_FT_DURATION_SECS and output_mode == OutputMode.VIDEO_AUDIO_SYNCED:
134 gen_video_chunked = GenVideoChunked(
135 video_id=scene_id,
136 gen=self.gen,
137 job_path=self.job_path,
138 logger=self.logger)
139 scene_video_binary = await gen_video_chunked.gen_video_chunked(
140 audio_path=audio_path,
141 image=self.image,
142 prompt=VIDEO_PROMPT,
143 neg_prompt=VIDEO_NEG_PROMPT,
144 width=width,
145 height=height,
146 num_steps=num_steps,
147 upscaling=self.get_config_bool("upscaling"),
148 debug=self.get_config_bool("debug_image"),
149 deadline=self.get_scene_deadline(scene_id),
150 )
151 else:
152 # VIDEO_AUDIO_UNSYNCED: generate video and merge audio separately
153 scene_video_binary = await self.gen.gen_video(
154 img=self.image,
155 prompt=VIDEO_PROMPT,
156 neg_prompt=VIDEO_NEG_PROMPT,
157 width=width,
158 height=height,
159 video_seconds=audio_duration,
160 steps=num_steps,
161 task_id=f"{scene_id:03d}",
162 wait_request=True,
163 deadline=self.get_scene_deadline(scene_id),
164 )
166 video_frames = await get_video_frames(scene_video_binary)
167 video_file_info = get_video_file_info(scene_video_binary)
168 video_info = video_file_info.get("video", {})
169 _video_fps = video_info.get("fps")
170 video_fps: float = _video_fps if _video_fps is not None else FANTASYTALKING_FPS
172 video_audio_path = f"{self.job_path}/{scene_id:03d}.mp4"
173 video_audio_path = await save_video_audio(
174 video_content=video_frames,
175 audio_path=audio_path,
176 out_video_path=video_audio_path,
177 fps=video_fps)
179 self.logger.info(f"[{scene_id}] Scene saved to '{video_audio_path}'.")
180 return video_audio_path
182 async def gen_lecture(
183 self,
184 pdf_base64: Optional[str],
185 ) -> None:
186 """
187 Generate a lecture video from a PDF.
188 """
189 async with self.job_status_handler():
190 if not pdf_base64:
191 self.logger.error("Document is required.")
192 await self.save_status(JobStatus.FAILED)
193 raise ValueError("Missing 'pdf_base64' in request")
194 self.logger.info(f"Generating lecture video for document with {bytes_to_human(len(pdf_base64))}.")
196 # Save PDF for processing and debugging
197 self.logger.info(f"Document base64 with {bytes_to_human(len(pdf_base64))}.")
198 pdf_path = f"{self.job_path}/document.pdf"
199 async with aiofiles.open(pdf_path, "wb") as file:
200 pdf_binary = base64_to_binary(pdf_base64)
201 await file.write(pdf_binary)
203 # Extract the materials from the PDF
204 pdf_text, pdf_images = parse_pdf(pdf_path)
205 num_pages = len(pdf_text)
206 self.logger.info(f"Extracted {num_pages} pages and {len(pdf_images)} images from PDF.")
208 # Save extracted text for debugging
209 txt_path = f"{self.job_path}/document.txt"
210 async with aiofiles.open(txt_path, "w", encoding="utf-8") as txt_file:
211 for page_ix, page_text in enumerate(pdf_text):
212 await txt_file.write(f"--- Page {page_ix + 1} ---\n{page_text}\n")
214 # Generate the main classroom image in parallel with transcript generation
215 image_task: asyncio.Task[Optional[Image.Image]] = asyncio.create_task(
216 self.gen.gen_image(
217 prompt=IMG_PROMPT,
218 neg_prompt=IMG_NEG_PROMPT,
219 width=self.width,
220 height=self.height,
221 task_id="main_image",
222 deadline=self.get_submission_time(),
223 )
224 )
226 await self.save_status(JobStatus.RUNNING)
228 # Generate the lecture transcript using the podcast transcript service (1 professor character)
229 output_mode = self.get_config_output_mode()
230 scene_texts: List[Dict] = []
231 async with aiofiles.open(f"{self.job_path}/lecture_transcript.jsonl", "wb") as file:
232 async for line_json in self.gen.gen_podcast_transcript(
233 task_id="transcript",
234 pdf_base64=pdf_base64,
235 num_characters=1,
236 style_prompt=LECTURE_STYLE_PROMPT,
237 custom_prompt=LECTURE_CUSTOM_PROMPT,
238 max_tokens=self.get_config_int("max_tokens"),
239 max_dialogues=self.get_config_int("max_dialogues"),
240 max_words_per_dialogue=self.get_config_int("max_words_per_dialogue"),
241 ):
242 line = (json.dumps(line_json) + "\n").encode("utf-8")
243 await file.write(line)
244 await file.flush()
246 line_type = line_json.get("type", "")
247 if line_type == "dialogue":
248 log_text = line_json.get("content", "")[:MAX_LOG_TEXT]
249 self.logger.info(f"Scene {len(scene_texts)}: {log_text}")
250 scene_texts.append(line_json)
251 self.logger.info(f"Lecture transcript generated with {len(scene_texts)} scenes.")
253 await self.save_status(JobStatus.RUNNING)
255 # Wait for classroom image
256 self.image = await image_task
257 if self.image is not None:
258 width, height = self.image.size
259 image_path = f"{self.job_path}/main_image.png"
260 self.image.save(image_path)
261 self.logger.info(f"Classroom image {width}x{height} saved to '{image_path}'.")
262 elif output_mode is not OutputMode.AUDIO_ONLY:
263 raise ValueError("Cannot generate classroom image for lecture video.")
265 # Generate each scene concurrently
266 scene_video_tasks: Dict[int, asyncio.Task] = {}
267 for scene_id, scene_json in enumerate(scene_texts):
268 scene_text = scene_json.get("content", "")
269 if not scene_text:
270 self.logger.warning(f"[{scene_id}] Empty scene text. Skipping.")
271 continue
272 scene_task = asyncio.create_task(
273 self.gen_scene(scene_id=scene_id, text=scene_text)
274 )
275 scene_video_tasks[scene_id] = scene_task
277 await self.save_status(JobStatus.RUNNING)
279 # Collect scene results
280 scene_video_paths: List[str] = []
281 results = await asyncio.gather(*scene_video_tasks.values(), return_exceptions=True)
282 for scene_id, result in enumerate(results):
283 if isinstance(result, Exception):
284 self._handle_scene_exception(scene_id, result)
285 elif isinstance(result, str):
286 scene_video_paths.append(result)
287 self.logger.info(f"[{scene_id}] Scene saved to '{result}'.")
288 else:
289 self.logger.warning(f"[{scene_id}] No video generated. Skipping.")
291 if not scene_video_paths:
292 raise ValueError("No scene videos generated. Cannot create final lecture video.")
294 await self.save_status(JobStatus.RUNNING)
296 # Concatenate all scene videos into a final lecture video
297 self.logger.info(f"Concatenating {len(scene_video_paths)} scenes into final lecture video...")
298 scene_binaries: List[bytes] = []
299 for scene_video_path in scene_video_paths:
300 async with aiofiles.open(scene_video_path, "rb") as vfile:
301 scene_binaries.append(await vfile.read())
303 video_binary = await concatenate_videos(scene_binaries, fast_copy=False)
304 if not video_binary:
305 raise ValueError("Cannot concatenate scenes into final lecture video.")
307 video_path = f"{self.job_path}/{self.job_id}.mp4"
308 async with aiofiles.open(video_path, "wb") as vfile:
309 await vfile.write(video_binary)
311 self.logger.info(f"Generated lecture video with {bytes_to_human(len(video_binary))} at '{video_path}'.")