Coverage for tests/test_wrapper_hunyuanavatar.py: 100%
216 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import os
4import sys
5import tempfile
6import numpy as np
7import pytest
9from unittest.mock import patch
10from unittest.mock import MagicMock
11from unittest.mock import AsyncMock
12from tests.torch_mock import TorchMock
14from PIL import Image
16mock_torch = TorchMock()
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 'sample_inference_audio': MagicMock(),
31 'transformers': MagicMock(),
32 'encode_data': MagicMock(),
33 'hymm_sp.config': MagicMock(),
34 'hymm_sp.data_kits.face_align': MagicMock(),
35 'hymm_sp.modules.parallel_states': MagicMock(),
36}
37mock_modules.update(mock_torch.get_sub_modules())
39sys.path.append("wrapper")
40sys.path.append("wrapper/hunyuanavatar")
42with patch.dict(sys.modules, mock_modules):
43 from image_utils import img_to_base64
44 from hunyuanavatar.wrapper_hunyuanavatar import HunyuanAvatarGeneration
45 from model_timing import GenTimer
46 _wrapper_module = sys.modules['hunyuanavatar.wrapper_hunyuanavatar']
49@pytest.mark.asyncio
50async def test_basic() -> None:
51 model = HunyuanAvatarGeneration()
52 assert model is not None
53 assert model.model_name == "hunyuanavatar"
54 assert model.status == "initializing"
56 with pytest.raises(TypeError):
57 model.init()
58 assert model.status == "failed"
59 health = model.get_health()
60 assert health is not None
61 timestamps = model.get_timestamps()
62 assert timestamps is not None
64 with pytest.raises(ValueError):
65 await model.get_rest_args({})
66 img = Image.new("RGB", (40, 30))
67 img_base64 = img_to_base64(img)
68 await model.get_rest_args({
69 "audio": "test",
70 "img": img_base64,
71 "prompt": "test prompt",
72 })
74 with pytest.raises(AssertionError):
75 # TODO improve the mocking
76 await model.warmup()
78 with pytest.raises(AssertionError):
79 # TODO improve the mocking
80 await model.generate(
81 img=Image.new('RGB', (100, 100)),
82 audio_path="test_audio.wav",
83 prompt="test prompt")
84 # assert video_frames is not None
86 del model
89@pytest.mark.asyncio
90async def test_get_rest_args_validation() -> None:
91 model = HunyuanAvatarGeneration()
93 with pytest.raises(ValueError):
94 await model.get_rest_args(None)
96 # Missing img
97 with pytest.raises(ValueError):
98 await model.get_rest_args({})
100 # Missing audio
101 img = Image.new("RGB", (40, 30))
102 img_base64 = img_to_base64(img)
103 with pytest.raises(ValueError):
104 await model.get_rest_args({"img": img_base64})
106 # Missing prompt - use valid base64 audio ("test" decodes cleanly)
107 with pytest.raises(ValueError):
108 await model.get_rest_args({"img": img_base64, "audio": "test"})
110 # All required params succeeds
111 result = await model.get_rest_args({
112 "img": img_base64,
113 "audio": "test",
114 "prompt": "test prompt",
115 })
116 assert result is not None
117 assert "args" in result
120@pytest.mark.asyncio
121async def test_get_rest_args_extra_params() -> None:
122 """Test get_rest_args with all optional parameters."""
123 model = HunyuanAvatarGeneration()
124 img = Image.new("RGB", (40, 30))
125 img_base64 = img_to_base64(img)
127 result = await model.get_rest_args({
128 "img": img_base64,
129 "audio": "test",
130 "prompt": "test prompt",
131 "height": 720,
132 "width": 1280,
133 "sampling_steps": 5,
134 "audio_scale": 0.8,
135 "cfg_scale": 3.0,
136 "audio_cfg_scale": 2.0,
137 "job_id": "test_job_001",
138 })
139 assert result["task"] == "hunyuanavatar"
140 assert result["args"]["height"] == 720
141 assert result["args"]["width"] == 1280
142 assert result["args"]["sampling_steps"] == 5
143 assert result["args"]["audio_scale"] == 0.8
144 assert result["args"]["cfg_scale"] == 3.0
145 assert result["args"]["audio_cfg_scale"] == 2.0
146 assert result["args"]["job_id"] == "test_job_001"
149def test_assert_model_init() -> None:
150 """Test _assert_model_init raises before model components are loaded."""
151 model = HunyuanAvatarGeneration()
152 # Components are None, so raises AssertionError
153 with pytest.raises(AssertionError, match="HunyuanVideoSampler is not initialized"):
154 model._assert_model_init()
157def test_assert_model_init_partial() -> None:
158 """Test _assert_model_init raises when only some components are initialized."""
159 model = HunyuanAvatarGeneration()
160 model.hunyuan_video_sampler = MagicMock()
161 with pytest.raises(AssertionError, match="Wav2Vec model is not initialized"):
162 model._assert_model_init()
164 model.wav2vec = MagicMock()
165 with pytest.raises(AssertionError, match="AlignImage instance is not initialized"):
166 model._assert_model_init()
168 model.align_instance = MagicMock()
169 with pytest.raises(AssertionError, match="Feature extractor is not initialized"):
170 model._assert_model_init()
172 model.feature_extractor = MagicMock()
173 with pytest.raises(AssertionError, match="Text encoder is not initialized"):
174 model._assert_model_init()
176 model.text_encoder = MagicMock()
177 with pytest.raises(AssertionError, match="Text encoder 2 is not initialized"):
178 model._assert_model_init()
180 model.text_encoder_2 = MagicMock()
181 # All initialized - should not raise
182 model._assert_model_init()
185def test_del_with_components() -> None:
186 """Test __del__ properly cleans up initialized components."""
187 model = HunyuanAvatarGeneration()
188 model.hunyuan_video_sampler = MagicMock()
189 model.wav2vec = MagicMock()
190 model.align_instance = MagicMock()
191 model.feature_extractor = MagicMock()
192 model.text_encoder = MagicMock()
193 model.text_encoder_2 = MagicMock()
194 # Should not raise
195 model.__del__()
196 assert model.hunyuan_video_sampler is None
197 assert model.wav2vec is None
198 assert model.align_instance is None
199 assert model.feature_extractor is None
200 assert model.text_encoder is None
201 assert model.text_encoder_2 is None
204def test_init_parallelism_no_master_addr() -> None:
205 """Test init_parallelism returns early when MASTER_ADDR is not set."""
206 model = HunyuanAvatarGeneration()
207 env_without_master = {k: v for k, v in os.environ.items() if k != "MASTER_ADDR"}
208 with patch.dict(os.environ, env_without_master, clear=True):
209 # Should return early without error
210 model.init_parallelism()
213def test_init_parallelism_world_size_one() -> None:
214 """Test init_parallelism with world_size=1 returns after setting device."""
215 model = HunyuanAvatarGeneration()
216 env = {"MASTER_ADDR": "localhost", "RANK": "0", "LOCAL_RANK": "0", "WORLD_SIZE": "1"}
217 with patch.dict(os.environ, env, clear=False):
218 model.init_parallelism()
219 assert model.world_size == 1
222@pytest.mark.asyncio
223async def test_output_video_pil() -> None:
224 """Test _output_video returns frames directly for pil output type."""
225 model = HunyuanAvatarGeneration()
226 gen_timer = GenTimer()
227 video_frames = np.zeros((5, 64, 64, 3), dtype=np.uint8)
228 result = await model._output_video(
229 job_id=None,
230 gen_timer=gen_timer,
231 audio_path="/tmp/test.wav",
232 video_frames=video_frames,
233 output_type="pil",
234 )
235 assert result is video_frames
238@pytest.mark.asyncio
239async def test_output_video_unknown_type() -> None:
240 """Test _output_video returns None for unknown output type."""
241 model = HunyuanAvatarGeneration()
242 gen_timer = GenTimer()
243 video_frames = np.zeros((5, 64, 64, 3), dtype=np.uint8)
244 result = await model._output_video(
245 job_id=None,
246 gen_timer=gen_timer,
247 audio_path="/tmp/test.wav",
248 video_frames=video_frames,
249 output_type="unknown_type",
250 )
251 assert result is None
254@pytest.mark.asyncio
255async def test_output_video_video_path() -> None:
256 """Test _output_video with video_path output type."""
257 model = HunyuanAvatarGeneration()
258 gen_timer = GenTimer()
259 video_frames = np.zeros((5, 64, 64, 3), dtype=np.uint8)
260 with patch.object(_wrapper_module, 'save_video_audio', new_callable=AsyncMock) as mock_save:
261 mock_save.return_value = "/tmp/test_output.mp4"
262 result = await model._output_video(
263 job_id="test_job",
264 gen_timer=gen_timer,
265 audio_path="/tmp/test.wav",
266 video_frames=video_frames,
267 output_type="video_path",
268 )
269 assert result == "/tmp/test_output.mp4"
272@pytest.mark.asyncio
273async def test_output_video_video_binary() -> None:
274 """Test _output_video with video_binary output type."""
275 model = HunyuanAvatarGeneration()
276 gen_timer = GenTimer()
277 video_frames = np.zeros((5, 64, 64, 3), dtype=np.uint8)
278 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as f:
279 tmp_path = f.name
280 f.write(b"fake_video_data")
282 try:
283 with patch.object(_wrapper_module, 'save_video_audio', new_callable=AsyncMock) as mock_save:
284 mock_save.return_value = tmp_path
285 result = await model._output_video(
286 job_id=None,
287 gen_timer=gen_timer,
288 audio_path="/tmp/test.wav",
289 video_frames=video_frames,
290 output_type="video_binary",
291 )
292 assert result == b"fake_video_data"
293 finally:
294 if os.path.exists(tmp_path):
295 os.remove(tmp_path)
298@pytest.mark.asyncio
299async def test_output_video_no_job_id() -> None:
300 """Test _output_video with video_path output type and no job_id creates temp file."""
301 model = HunyuanAvatarGeneration()
302 gen_timer = GenTimer()
303 video_frames = np.zeros((5, 64, 64, 3), dtype=np.uint8)
304 with patch.object(_wrapper_module, 'save_video_audio', new_callable=AsyncMock) as mock_save:
305 mock_save.return_value = "/tmp/generated_output.mp4"
306 result = await model._output_video(
307 job_id=None,
308 gen_timer=gen_timer,
309 audio_path="/tmp/test.wav",
310 video_frames=video_frames,
311 output_type="video_path",
312 )
313 assert result == "/tmp/generated_output.mp4"
314 # Verify save_video_audio was called with a temp path (no job_id)
315 call_kwargs = mock_save.call_args
316 assert call_kwargs is not None
319@pytest.mark.asyncio
320async def test_generate_assert_model_not_init() -> None:
321 """Test generate raises AssertionError when model is not initialized."""
322 model = HunyuanAvatarGeneration()
323 with pytest.raises(AssertionError):
324 await model.generate(
325 img=Image.new("RGB", (100, 100)),
326 audio_path="/tmp/test.wav",
327 prompt="test prompt",
328 )
331@pytest.mark.asyncio
332async def test_generate_with_mocked_model() -> None:
333 """Test generate with mocked model components."""
334 model = HunyuanAvatarGeneration()
336 # Set up all required mocked components
337 model.hunyuan_video_sampler = MagicMock()
338 model.wav2vec = MagicMock()
339 model.wav2vec.dtype = mock_torch.bfloat16
340 model.align_instance = MagicMock()
341 model.feature_extractor = MagicMock()
342 model.text_encoder = MagicMock()
343 model.text_encoder_2 = MagicMock()
344 model.data_loader = MagicMock()
346 # Mock encode_data result
347 model.data_loader.encode_data.return_value = {
348 "audio_len": 5,
349 }
351 # Mock hunyuan_video_sampler.predict
352 fake_sample = MagicMock()
353 fake_sample.unsqueeze.return_value = fake_sample
354 fake_sample.__getitem__ = MagicMock(return_value=fake_sample)
355 model.hunyuan_video_sampler.predict.return_value = {"samples": [fake_sample]}
357 # Set up the mock chain for video output:
358 # video = einops.rearrange(sample[0], ...) → mock_rearranged
359 # video = (video * 255.).data.cpu().numpy().astype(np.uint8) → fake_video (numpy)
360 fake_video = np.zeros((5, 64, 64, 3), dtype=np.uint8)
361 mock_rearranged = MagicMock()
362 mul_result = MagicMock()
363 mul_result.data.cpu.return_value.numpy.return_value.astype.return_value = fake_video
364 mock_rearranged.__mul__ = MagicMock(return_value=mul_result)
366 with patch.object(_wrapper_module, 'librosa') as mock_librosa, \
367 patch.object(_wrapper_module, 'einops') as mock_einops:
369 mock_librosa.get_duration.return_value = 2.0
370 mock_einops.rearrange.return_value = mock_rearranged
372 result = await model.generate(
373 img=Image.new("RGB", (100, 100)),
374 audio_path="/tmp/test.wav",
375 prompt="test prompt",
376 output_type="pil",
377 )
378 assert result is not None
381@pytest.mark.asyncio
382async def test_generate_audio_too_long() -> None:
383 """Test generate raises ValueError when audio exceeds MAX_FRAMES."""
384 model = HunyuanAvatarGeneration()
385 model.hunyuan_video_sampler = MagicMock()
386 model.wav2vec = MagicMock()
387 model.align_instance = MagicMock()
388 model.feature_extractor = MagicMock()
389 model.text_encoder = MagicMock()
390 model.text_encoder_2 = MagicMock()
391 model.data_loader = MagicMock()
393 # 7 seconds at 12.5 FPS gives num_frames = int(87.5 // 4) * 4 + 5 = 89 > MAX_FRAMES=81
394 with patch.object(_wrapper_module, 'librosa') as mock_librosa:
395 mock_librosa.get_duration.return_value = 7.0
396 with pytest.raises(ValueError, match="exceeds"):
397 await model.generate(
398 img=Image.new("RGB", (100, 100)),
399 audio_path="/tmp/long_audio.wav",
400 prompt="test prompt",
401 )