Coverage for tests/streamwise_app/lmm_generator_mock.py: 99%
119 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1import math
2import sys
3import os
4import asyncio
5import tempfile
6import aiofiles
8from PIL import Image
10from typing import Any
11from typing import Optional
12from typing import Dict
13from typing import List
14from typing import AsyncGenerator
15from typing import cast
17from unittest.mock import patch
18from unittest.mock import MagicMock
20# Add current path
21sys.path.append(os.getcwd())
23from video import FANTASYTALKING_FPS
24from video import HUNYUANFRAMEPACK_FPS
26from file_utils import read_file_base64
27from file_utils import save_base64_as_binary
29from tests.torch_mock import TorchMock
30from tests.k8s_mock import K8sMock
31from tests.fantasytalking_mock import FantasyTalkingMock
33mock_torch = TorchMock()
34mock_k8s = K8sMock()
35mock_ft = FantasyTalkingMock()
37mock_modules = {
38 "imageio": MagicMock(),
39 "tabulate": MagicMock(),
40 "soundfile": MagicMock(),
41}
42mock_modules.update(mock_torch.get_sub_modules())
43mock_modules.update(mock_k8s.get_sub_modules())
44mock_modules.update(mock_ft.get_sub_modules())
46with patch.dict(sys.modules, mock_modules):
47 from media_utils import save_video_frames
48 from media_utils import get_video_frames
49 from media_utils import get_frame_with_text
50 from media_utils import get_video_file_info
51 from media_utils import save_video_audio
52 from media_utils import get_audio_duration
54 from apps.lmm_generator import LMMGenerator
56 from apps.client import ServiceRequest
57 from apps.client import RequestStatus
60# Each mock function outputs a color for easy asserting
61MOCK_COLORS = {
62 "default": "pink",
63 "gen_image": "white",
64 "gen_edit_image": "lightgray",
65 "gen_extract_characters": "gray",
66 "gen_video": "blue",
67 "gen_video_audio_from_img": "green",
68 "gen_video_audio_from_video": "red",
69 "gen_video_upscale": "yellow",
70 "gen_video_from_latents": "purple",
71 "gen_intermediate_video_frames": "orange",
72}
73MOCK_COLORS_RGB = {
74 "default": (255, 192, 203),
75 "gen_image": (255, 255, 255),
76 "gen_edit_image": (211, 211, 211),
77 "gen_extract_characters": (128, 128, 128),
78 "gen_video": (0, 0, 255),
79 "gen_video_audio_from_img": (0, 128, 0),
80 "gen_video_audio_from_video": (253, 0, 0),
81 "gen_video_upscale": (255, 255, 0),
82 "gen_video_from_latents": (128, 0, 128),
83 "gen_intermediate_video_frames": (255, 165, 0),
84}
87class LMMGeneratorMock(LMMGenerator):
88 """Mock LMMGenerator for testing."""
90 def __init__(
91 self,
92 *_: Any,
93 **__: Any,
94 ) -> None:
95 self.app_name = "mock"
96 self.job_id = "mock"
98 self.request_executor = MagicMock()
99 self.service_manager = MagicMock()
100 # The mock uses a dict for O(1) lookup, while the base class uses a list for ordered
101 # request tracking. Both serve the same role (request registry) but with different APIs.
102 self.requests: Dict[str, ServiceRequest] = {} # type: ignore[assignment]
104 async def gen_podcast_transcript(
105 self, *args: Any, **kwargs: Any
106 ) -> AsyncGenerator[Dict[str, str], None]:
107 yield {"type": "image", "content": "Image prompt."}
108 yield {"type": "character", "name": "Jane", "gender": "Female"}
109 yield {"type": "character", "name": "Joe", "gender": "Male"}
110 yield {"type": "dialogue", "character": "Jane", "content": "Hello"}
111 yield {"type": "dialogue", "character": "Joe", "content": "World"}
112 yield {"type": "dialogue", "character": "Unknown", "content": "Unknown character"}
113 yield {"type": "dialogue", "content": "No character"}
114 yield {"type": "dialogue"}
116 async def gen_image( # type: ignore[override]
117 self,
118 *_: Any,
119 width: int,
120 height: int,
121 **__: Any,
122 ) -> Image.Image:
123 return Image.new("RGB", (width, height), color=MOCK_COLORS["gen_image"])
125 async def gen_edit_image( # type: ignore[override]
126 self,
127 *_: Any,
128 width: int,
129 height: int,
130 **__: Any,
131 ) -> Image.Image:
132 return Image.new("RGB", (width, height), color=MOCK_COLORS["gen_edit_image"])
134 async def gen_extract_characters(
135 self,
136 image: Image.Image,
137 *_: Any,
138 num_characters: int = 2,
139 **__: Any,
140 ) -> List[Image.Image]:
141 width, height = image.size
142 return [
143 Image.new("RGB", (width, height), color=MOCK_COLORS["gen_extract_characters"])
144 for _ in range(num_characters)
145 ]
147 async def gen_audio(self, *args: Any, **kwargs: Any) -> str:
148 mock_audio_path = "tests/data/audio_4675.wav"
149 # TODO cut it to a size?
150 return await read_file_base64(mock_audio_path)
152 async def _gen_synthetic_video(
153 self,
154 num_frames: int = 81,
155 width: int = 160,
156 height: int = 100,
157 fps: int = FANTASYTALKING_FPS,
158 audio_path: Optional[str] = None,
159 color: str = MOCK_COLORS["default"],
160 ) -> bytes:
161 """Synthetic test video."""
162 frames: List[Image.Image] = [
163 cast(Image.Image, get_frame_with_text(
164 width,
165 height,
166 f"frame{idx:02d}",
167 output_type="pil",
168 background_color=color,
169 ))
170 for idx in range(num_frames)
171 ]
173 if audio_path and await aiofiles.os.path.exists(audio_path):
174 video_path = await save_video_audio(
175 frames,
176 audio_path=audio_path,
177 fps=fps
178 )
179 else:
180 video_path = await save_video_frames(
181 frames,
182 fps=fps
183 )
185 async with aiofiles.open(video_path, "rb") as file:
186 video_binary = await file.read()
188 return video_binary
190 async def gen_video( # type: ignore[override]
191 self,
192 img: Image.Image,
193 prompt: str,
194 neg_prompt: str = "",
195 width: int = 640,
196 height: int = 400,
197 num_frames: int = 81,
198 wait_request: bool = True,
199 task_id: Optional[str] = None,
200 **_: Any,
201 ) -> ServiceRequest | bytes:
202 video_binary = await self._gen_synthetic_video(
203 width=width,
204 height=height,
205 num_frames=num_frames,
206 fps=HUNYUANFRAMEPACK_FPS,
207 color=MOCK_COLORS["gen_video"],
208 )
210 if wait_request is True:
211 return video_binary
213 service_name = "hunyuanframepackf1"
214 request = ServiceRequest(
215 request_id=f"{self.job_id}_{task_id}_{service_name}",
216 service_name=service_name,
217 payload_json={
218 "job_id": self.job_id,
219 "img": b"AAAAA", # This would come from img
220 "width": width,
221 "height": height,
222 "num_frames": num_frames,
223 },
224 )
225 request.status = RequestStatus.COMPLETED
226 request.future = asyncio.Future()
227 request.future.set_result(("video/mp4", video_binary))
228 return request
230 async def gen_video_audio_from_img( # type: ignore[override]
231 self,
232 *,
233 audio_base64: str,
234 width: int = 640,
235 height: int = 400,
236 **_: Any,
237 ) -> bytes:
238 audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
239 audio_path = await save_base64_as_binary(audio_path, audio_base64)
240 audio_duration = get_audio_duration(audio_base64)
241 num_frames = (((math.ceil(audio_duration * FANTASYTALKING_FPS) - 1) // 4) * 4) + 1
242 return await self._gen_synthetic_video(
243 width=width,
244 height=height,
245 num_frames=num_frames,
246 fps=FANTASYTALKING_FPS,
247 audio_path=audio_path,
248 color=MOCK_COLORS["gen_video_audio_from_img"],
249 )
251 async def gen_video_audio_from_video( # type: ignore[override]
252 self,
253 video: List[Image.Image],
254 audio_base64: str,
255 prompt: str,
256 neg_prompt: str = "",
257 width: int = 640,
258 height: int = 400,
259 **_: Any,
260 ) -> bytes:
261 audio_path = tempfile.NamedTemporaryFile(delete=False, suffix=".wav").name
262 audio_path = await save_base64_as_binary(audio_path, audio_base64)
263 return await self._gen_synthetic_video(
264 width=width,
265 height=height,
266 num_frames=len(video),
267 fps=FANTASYTALKING_FPS,
268 audio_path=audio_path,
269 color=MOCK_COLORS["gen_video_audio_from_video"],
270 )
272 async def gen_video_upscale( # type: ignore[override]
273 self,
274 *,
275 video_binary: bytes,
276 width: int = 640,
277 height: int = 400,
278 **_: Any,
279 ) -> bytes:
280 video_file_info = get_video_file_info(video_binary)
281 video_info = video_file_info["video"]
282 num_frames = video_info["num_frames"] or 81
283 fps = int(video_info["fps"] or FANTASYTALKING_FPS)
284 return await self._gen_synthetic_video(
285 width=width,
286 height=height,
287 num_frames=num_frames,
288 fps=fps,
289 color=MOCK_COLORS["gen_video_upscale"],
290 )
292 async def gen_video_from_latents(
293 self,
294 *_: Any,
295 # latents: torch.Tensor,
296 **__: Any,
297 ) -> bytes:
298 # TODO guess parameters from latents
299 return await self._gen_synthetic_video(
300 width=320,
301 height=200,
302 num_frames=81,
303 # fps=FANTASYTALKING_FPS,
304 color=MOCK_COLORS["gen_video_from_latents"],
305 )
307 async def gen_intermediate_video_frames( # type: ignore[override]
308 self,
309 *_: Any,
310 video_gen_request: ServiceRequest,
311 **__: Any,
312 ) -> AsyncGenerator[Image.Image, None]:
313 assert video_gen_request.payload_json is not None
314 video_binary = await self._gen_synthetic_video(
315 width=video_gen_request.payload_json["width"],
316 height=video_gen_request.payload_json["height"],
317 num_frames=video_gen_request.payload_json["num_frames"],
318 color=MOCK_COLORS["gen_intermediate_video_frames"],
319 )
320 frames = await get_video_frames(video_binary)
321 for frame in frames:
322 yield frame
324 def get_requests(self) -> Dict[str, ServiceRequest]:
325 return self.requests
327 def get_queued_requests(self) -> List[str]:
328 return [] # TODO
330 async def gen_text(
331 self,
332 messages: List[Dict],
333 *_: Any,
334 **__: Any,
335 ) -> str:
336 """Return a canned text reply for testing."""
337 return "This is a mock response."
339 async def gen_audio_transcript(
340 self,
341 audio_path: str,
342 *_: Any,
343 **__: Any,
344 ) -> tuple[str, str]:
345 """Return a canned transcript for testing."""
346 return ("Mock transcript text.", "en")
349"""
350async def _mock_generation() -> MagicMock:
351 "" "
352 Mock multi-modal generation components.
353 "" "
354 gen = MagicMock()
356 gen.job_id = "mock_chunked_job_id"
358 gen.job_path = f"/tmp/{gen.job_id}"
359 os.makedirs(gen.job_path, exist_ok=True)
361 loop = asyncio.get_running_loop()
362 gen.image_task = loop.create_future()
363 gen.image_task.set_result(Image.new("RGB", (640, 480), color="blue"))
365 video_frames = [
366 Image.new("RGB", (640, 480), color=color)
367 for color in ["blue", "green", "red", "yellow"]
368 ]
369 video_path = await save_video_frames(video_frames)
370 video_binary = await read_file_bytes(video_path)
371 gen.gen_video_audio_from_img = AsyncMock(return_value=video_binary)
373 gen.gen_video_audio_from_video = AsyncMock(return_value=video_binary)
375 async def gen_video_mock(*args, **kwargs) -> Any:
376 if kwargs.get("wait_request", None):
377 return video_binary
378 req = ServiceRequest(
379 request_id="mock_request_id",
380 service_name="hunyuanframepackf1",
381 payload_json={"job_id": gen.job_id},
382 )
383 req.status = RequestStatus.COMPLETED
384 req.future = asyncio.Future()
385 req.future.set_result(("video/mp4", video_binary))
386 return req
388 gen.gen_video = AsyncMock(side_effect=gen_video_mock)
389 return gen
390"""