Coverage for tests/test_wrapper_fantasytalking.py: 100%
92 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import math
4import sys
5import pytest
7from unittest.mock import patch
8from unittest.mock import MagicMock
9from tests.torch_mock import TorchMock
11from PIL import Image
13mock_torch = TorchMock()
15sys.path.append("wrapper")
16sys.path.append("wrapper/fantasytalking")
18mock_modules = {
19 'nvidia_smi': MagicMock(),
20 'imageio': MagicMock(),
21 'cv2': MagicMock(),
22 'torch': mock_torch,
23 'torchvision': MagicMock(),
24 'torchvision.transforms': MagicMock(),
25 'torchvision.transforms.functional': MagicMock(),
26 'xfuser': MagicMock(),
27 'xfuser.config': MagicMock(),
28 'xfuser.core': MagicMock(),
29 'xfuser.core.distributed': MagicMock(),
30 'diffsynth': MagicMock(),
31 'diffsynth.models': MagicMock(),
32 'diffsynth.models.wan_video_dit': MagicMock(),
33 'wan.distributed': MagicMock(),
34 'wan.distributed.xdit_context_parallel': MagicMock(),
35 'wan.distributed.fsdp': MagicMock(),
36 'transformers': MagicMock(),
37 'model': MagicMock(),
38 'utils': MagicMock(),
39}
40mock_modules.update(mock_torch.get_sub_modules())
42with patch.dict(sys.modules, mock_modules):
43 from fantasytalking.wrapper_fantasytalking import FantasyTalking
44 from fantasytalking.wrapper_fantasytalking import resample_frames
45 from fantasytalking.wrapper_fantasytalking import resample_and_normalize_frames
48@pytest.mark.asyncio
49async def test_fantasytalking_e2e() -> None:
50 model = FantasyTalking()
51 assert model is not None
52 assert model.model_name == "fantasytalking"
53 assert model.status == "initializing"
55 model.init()
56 assert model.status == "ok"
57 health = model.get_health()
58 assert health is not None
59 timestamps = model.get_timestamps()
60 assert timestamps is not None
62 with pytest.raises(ValueError, match="Missing JSON body"):
63 await model.get_rest_args(None)
65 with pytest.raises(ValueError, match="Missing 'audio' parameter"):
66 await model.get_rest_args({})
68 with pytest.raises(ValueError, match="Missing 'prompt' parameter"):
69 await model.get_rest_args({"audio": "test"})
71 rest_args = await model.get_rest_args({
72 "audio": "test",
73 "prompt": "test prompt",
74 })
75 assert "args" in rest_args
76 assert rest_args["args"]["audio_path"].startswith("/tmp/tmp")
77 assert rest_args["args"]["audio_path"].endswith(".wav")
79 await model.get_rest_args({
80 "audio": "test",
81 "prompt": "test prompt",
82 "job_id": "testjob",
83 })
84 assert "args" in rest_args
85 assert rest_args["args"]["audio_path"].endswith(".wav")
86 assert rest_args["args"]["audio_cfg_scale"] == 5.0
88 with pytest.raises((TypeError, ValueError)):
89 await model.warmup()
91 with pytest.raises(TypeError, match="required positional arguments"):
92 await model.generate()
94 with pytest.raises(ValueError, match="Audio file 'nonexisting.wav' does not exist"):
95 await model.generate(
96 img=Image.new('RGB', (100, 100)),
97 audio_path="nonexisting.wav",
98 prompt="test prompt",
99 video=None)
101 with pytest.raises((TypeError, ValueError)):
102 # TODO improve mocking
103 await model.generate(
104 img=Image.new('RGB', (100, 100)),
105 audio_path="tests/data/audio_4675.wav",
106 prompt="test prompt",
107 video=None)
109 with pytest.raises(TypeError):
110 # TODO improve mocking
111 await model.generate(
112 video=[
113 Image.new('RGB', (100, 100))
114 for _ in range(3)
115 ],
116 audio_path="tests/data/audio_4675.wav",
117 prompt="test prompt")
119 with pytest.raises((TypeError, ValueError)):
120 # TODO improve mocking
121 await model.generate(
122 img=None,
123 video=None,
124 audio_path="tests/data/audio_4675.wav",
125 prompt="test prompt")
127 del model
130def test_resample_frames() -> None:
131 resampled_frames = resample_frames([], 30, 23)
132 assert resampled_frames == []
134 resampled_frames = resample_frames([
135 Image.new("RGB", (100, 100))
136 for _ in range(3)
137 ], 30, 23)
138 assert len(resampled_frames) == 2
140 resampled_frames = resample_frames([
141 Image.new("RGB", (100, 100))
142 for _ in range(24)
143 ], 24, 23, 0.5)
144 assert len(resampled_frames) == 12
147def test_resample_frames_may_differ_from_num_frames() -> None:
148 """Reproduce the bug scenario: 73 frames at 30 FPS with ~2.4s audio.
150 Before the fix, resample_frames returns 55 frames while num_frames (computed
151 from audio at 23 FPS aligned to 1+4n) is 57, causing a latent/noise tensor
152 shape mismatch (14 vs 15 frames) inside CustomWanVideoPipeline.__call__.
153 This test documents that resample_frames can return fewer frames than
154 num_frames; the fix in generate() normalises the result via
155 resample_and_normalize_frames() so both tensors share the same latent dimension.
156 """
158 FPS = 23.0
159 SRC_FPS = 30.0
160 vae_stride = 4
162 audio_duration = 2.4 # seconds — representative value from the bug report
163 audio_num_frames = int(math.ceil(FPS * audio_duration))
164 num_frames = int(1 + math.ceil((audio_num_frames - 1) / vae_stride) * vae_stride)
166 src_frames = [Image.new("RGB", (100, 100)) for _ in range(73)]
168 # resample_frames alone produces fewer frames than num_frames due to rounding
169 resampled_raw = resample_frames(src_frames, SRC_FPS, FPS, audio_duration)
170 assert len(resampled_raw) < num_frames, (
171 f"Expected resampled ({len(resampled_raw)}) < num_frames ({num_frames})"
172 )
174 # resample_and_normalize_frames must return exactly num_frames
175 normalized = resample_and_normalize_frames(
176 src_frames, SRC_FPS, FPS, num_frames, audio_duration)
177 assert len(normalized) == num_frames, (
178 f"After normalization, expected {num_frames} frames but got {len(normalized)}"
179 )
181 # Verify that the latent dimensions now match
182 lat_expected = (num_frames - 1) // vae_stride + 1
183 lat_actual = (len(normalized) - 1) // vae_stride + 1
184 assert lat_expected == lat_actual, (
185 f"Latent frame count mismatch: noise={lat_expected}, latents={lat_actual}"
186 )
189def test_normalize_frames_trimming() -> None:
190 """Test that resample_and_normalize_frames trims oversized resampled video to num_frames."""
192 FPS = 23.0
193 SRC_FPS = 30.0
194 vae_stride = 4
196 # Audio duration just long enough that audio_num_frames rounds up such that
197 # num_frames (1+4n aligned) is smaller than what resample_frames would produce
198 # without the truncation step. We achieve this by supplying no audio_duration
199 # to resample_frames so the truncation branch is skipped, giving more dst frames.
200 audio_duration = 1.0 # 1 second
201 audio_num_frames = int(math.ceil(FPS * audio_duration)) # 23
202 num_frames = int(1 + math.ceil((audio_num_frames - 1) / vae_stride) * vae_stride) # 1+24=25
204 # Provide a video that, when resampled WITHOUT the audio-duration truncation,
205 # produces more frames than num_frames.
206 # 30 frames at 30 FPS = 1.0 s → round(1.0 * 23) = 23 frames < 25, so use a
207 # slightly longer source to get more dest frames.
208 # 35 frames at 30 FPS = 1.167 s → round(1.167 * 23) = round(26.83) = 27 frames > 25
209 src_frames = [Image.new("RGB", (100, 100)) for _ in range(35)]
211 # Without normalization, resample_frames produces more frames than num_frames
212 resampled_raw = resample_frames(src_frames, SRC_FPS, FPS)
213 assert len(resampled_raw) > num_frames, (
214 f"Expected resampled ({len(resampled_raw)}) > num_frames ({num_frames})"
215 )
217 # resample_and_normalize_frames must return exactly num_frames
218 normalized = resample_and_normalize_frames(src_frames, SRC_FPS, FPS, num_frames)
219 assert len(normalized) == num_frames
221 lat_expected = (num_frames - 1) // vae_stride + 1
222 lat_actual = (len(normalized) - 1) // vae_stride + 1
223 assert lat_expected == lat_actual