Coverage for tests/test_wrapper_wan.py: 99%
350 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-09 04:47 +0000
1#!/usr/bin/env python3
3import sys
4import gc
5import pytest
7from unittest.mock import patch
8from unittest.mock import MagicMock
10from PIL import Image
12from tests.torch_mock import TorchMock
14mock_torch = TorchMock()
16sys.path.append("wrapper")
17sys.path.append("wrapper/wan")
19mock_modules = {
20 'nvidia_smi': MagicMock(),
21 'imageio': MagicMock(),
22 'cv2': MagicMock(),
23 'torch': mock_torch,
24 'torchvision': MagicMock(),
25 'torchvision.transforms': MagicMock(),
26 'torchvision.transforms.functional': MagicMock(),
27 'xfuser': MagicMock(),
28 'xfuser.config': MagicMock(),
29 'xfuser.core': MagicMock(),
30 'xfuser.core.distributed': MagicMock(),
31 'transformers': MagicMock(),
32 'wan.modules': MagicMock(),
33 'wan.modules.t5': MagicMock(),
34 'wan.modules.clip': MagicMock(),
35 'wan.modules.vae': MagicMock(),
36 'wan.modules.model': MagicMock(),
37 'wan.utils': MagicMock(),
38 'wan.utils.utils': MagicMock(),
39 'wan.utils.fm_solvers_unipc': MagicMock(),
40 'wan.distributed': MagicMock(),
41 'wan.distributed.fsdp': MagicMock(),
42 'wan.distributed.xdit_context_parallel': MagicMock(),
43}
44mock_modules.update(mock_torch.get_sub_modules())
46with patch.dict(sys.modules, mock_modules):
47 from image_utils import img_to_base64
48 from wan.wrapper_wan21 import Wan21VideoGeneration
49 from wan.vae import WanVAE
52@pytest.mark.asyncio
53async def test_init() -> None:
54 model = Wan21VideoGeneration()
55 assert model is not None
56 assert model.model_name == "wan"
58 model.init()
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(None)
66 with pytest.raises(ValueError):
67 await model.get_rest_args({})
68 img = Image.new("RGB", (40, 30))
69 img_base64 = img_to_base64(img)
70 await model.get_rest_args({
71 "img": img_base64,
72 "prompt": "test prompt",
73 })
75 with pytest.raises(ValueError):
76 # TODO implement fixture for torch tensor
77 # TF.to_tensor(img_resized).sub_(0.5).div_(0.5).to(self.device)
78 await model.warmup()
80 with pytest.raises(ValueError, match="not enough values to unpack"):
81 # TODO implement fixture for torch tensor
82 # TF.to_tensor(img_resized).sub_(0.5).div_(0.5).to(self.device)
83 await model.generate(
84 img=img,
85 prompt="test prompt",
86 output_type="video_frames")
88 del model
89 gc.collect()
92@pytest.mark.asyncio
93async def test_vae() -> None:
94 """Test the VAE wrapper."""
95 vae = WanVAE()
96 assert vae is not None
98 mock_tensor = mock_torch.randn(1, 3, 4, 64, 64)
100 encoded_ret = vae.encode(
101 videos=[mock_tensor],
102 start_frames=1,
103 end_frames=0)
104 assert encoded_ret is not None
105 assert len(encoded_ret) == 1
107 decoded_ret = vae.decode(zs=[mock_tensor])
108 assert decoded_ret is not None
109 assert len(decoded_ret) == 1
111 for ret in vae.decode_stream(z=mock_tensor):
112 assert ret is not None
115@pytest.mark.asyncio
116async def test_assert_args() -> None:
117 with patch.dict(sys.modules, mock_modules):
118 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
120 model = _Wan21()
121 model.init()
122 model.vae_stride = (4, 8, 8)
124 # Valid args: 480%8=0, 640%8=0, (5-1)%4=0
125 model._assert_args(height=480, width=640, num_frames=5)
127 # Height not divisible by vae_stride[1]=8
128 with pytest.raises(ValueError, match="Height"):
129 model._assert_args(height=481, width=640, num_frames=5)
131 # Width not divisible by vae_stride[2]=8
132 with pytest.raises(ValueError, match="Width"):
133 model._assert_args(height=480, width=641, num_frames=5)
135 # num_frames: (2-1)%4 != 0
136 with pytest.raises(ValueError):
137 model._assert_args(height=480, width=640, num_frames=2)
139 # Too many frames: > 1+80
140 with pytest.raises(ValueError):
141 model._assert_args(height=480, width=640, num_frames=100)
143 # vae_stride not set
144 model.vae_stride = None
145 with pytest.raises(ValueError, match="VAE stride"):
146 model._assert_args(height=480, width=640, num_frames=5)
149@pytest.mark.asyncio
150async def test_get_rest_args_full() -> None:
151 with patch.dict(sys.modules, mock_modules):
152 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
153 from image_utils import img_to_base64 as _img_to_base64
155 model = _Wan21()
156 model.init()
158 img = Image.new("RGB", (40, 30))
159 img_base64 = _img_to_base64(img)
161 # Missing img still raises ValueError
162 with pytest.raises(ValueError):
163 await model.get_rest_args({"prompt": "test"})
165 # With num_frames, width, height, seed
166 result = await model.get_rest_args({
167 "img": img_base64,
168 "prompt": "test prompt",
169 "num_frames": 17,
170 "width": 640,
171 "height": 480,
172 })
173 assert result["args"]["num_frames"] == 17
174 assert result["args"]["width"] == 640
175 assert result["args"]["height"] == 480
176 assert result["args"]["prompt"] == "test prompt"
179@pytest.mark.asyncio
180async def test_get_rest_args_extra_params() -> None:
181 """Test get_rest_args with neg_prompt, sampling_steps, output_type and video_seconds."""
182 with patch.dict(sys.modules, mock_modules):
183 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
184 from image_utils import img_to_base64 as _img_to_base64
186 model = _Wan21()
187 model.init()
189 img = Image.new("RGB", (40, 30))
190 img_base64 = _img_to_base64(img)
192 # neg_prompt, sampling_steps, output_type
193 result = await model.get_rest_args({
194 "img": img_base64,
195 "prompt": "test prompt",
196 "neg_prompt": "bad quality",
197 "sampling_steps": 20,
198 "output_type": "video_path",
199 })
200 assert result["args"]["neg_prompt"] == "bad quality"
201 assert result["args"]["sampling_steps"] == 20
202 assert result["args"]["output_type"] == "video_path"
204 # video_seconds branch: converts seconds to num_frames using FPS and vae_stride
205 result_secs = await model.get_rest_args({
206 "img": img_base64,
207 "prompt": "test prompt",
208 "video_seconds": 2.0,
209 })
210 # FPS=16, vae_stride[0]=4: num_frames = 1 + ((int(2.0*16)-1) // 4) * 4 = 29
211 assert result_secs["args"]["num_frames"] == 29
214def test_assert_model_init() -> None:
215 """Test that _assert_model_init raises when model components are not loaded."""
216 with patch.dict(sys.modules, mock_modules):
217 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
219 model = _Wan21()
220 # Status is "initializing" (init() not called yet), base raises ValueError
221 with pytest.raises(ValueError, match="Model not initialized"):
222 model._assert_model_init()
225def test_model_compile_no_op() -> None:
226 """Test model_compile returns immediately when torch_compile=False."""
227 with patch.dict(sys.modules, mock_modules):
228 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
230 model = _Wan21()
231 model.init()
232 model.torch_compile = False
233 model.model_compile() # should not raise
236def test_model_compile_with_torch_compile() -> None:
237 """Test model_compile calls torch.compile when torch_compile=True."""
238 with patch.dict(sys.modules, mock_modules):
239 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
241 model = _Wan21()
242 model.init()
243 model.torch_compile = True
244 model.model_compile() # should call torch.compile (mocked, no-op)
247@pytest.mark.asyncio
248async def test_assert_model_init_text_encoder_none() -> None:
249 """Test _assert_model_init raises when text_encoder is None."""
250 with patch.dict(sys.modules, mock_modules):
251 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
253 model = _Wan21()
254 model.init()
255 model.text_encoder = None
256 with pytest.raises(ValueError, match="Text encoder"):
257 model._assert_model_init()
260@pytest.mark.asyncio
261async def test_assert_model_init_vae_none() -> None:
262 """Test _assert_model_init raises when vae is None."""
263 with patch.dict(sys.modules, mock_modules):
264 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
266 model = _Wan21()
267 model.init()
268 model.vae = None
269 with pytest.raises(ValueError, match="VAE"):
270 model._assert_model_init()
273@pytest.mark.asyncio
274async def test_wan21_assert_model_init_image_encoder_none() -> None:
275 """Test Wan21 _assert_model_init raises when image_encoder is None."""
276 with patch.dict(sys.modules, mock_modules):
277 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
279 model = _Wan21()
280 model.init()
281 model.image_encoder = None
282 with pytest.raises(ValueError, match="Image encoder"):
283 model._assert_model_init()
286@pytest.mark.asyncio
287async def test_wan21_assert_model_init_image_encoder_model_none() -> None:
288 """Test Wan21 _assert_model_init raises when image_encoder.model is None."""
289 with patch.dict(sys.modules, mock_modules):
290 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
292 model = _Wan21()
293 model.init()
294 # Use a fresh MagicMock to avoid polluting the shared mock state
295 fresh_encoder = MagicMock()
296 fresh_encoder.model = None
297 model.image_encoder = fresh_encoder
298 with pytest.raises(ValueError, match="Image encoder"):
299 model._assert_model_init()
302@pytest.mark.asyncio
303async def test_wan21_assert_model_init_dit_model_none() -> None:
304 """Test Wan21 _assert_model_init raises when dit model is None."""
305 with patch.dict(sys.modules, mock_modules):
306 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
308 model = _Wan21()
309 model.init()
310 # Ensure image_encoder and its .model remain truthy; only null out the DiT model
311 model.image_encoder = MagicMock()
312 model.image_encoder.model = MagicMock()
313 model.model = None
314 with pytest.raises(ValueError, match="DiT model"):
315 model._assert_model_init()
318@pytest.mark.asyncio
319async def test_output_video_tensor() -> None:
320 """Test _output_video returns tensor directly for output_type='tensor'."""
321 with patch.dict(sys.modules, mock_modules):
322 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
324 model = _Wan21()
325 model.init()
326 gen_timer = model._new_gen_timer(None)
327 mock_video = MagicMock()
329 result = await model._output_video(None, gen_timer, mock_video, "tensor")
330 assert result is mock_video
333@pytest.mark.asyncio
334async def test_output_video_unknown_type() -> None:
335 """Test _output_video returns None for an unknown output_type."""
336 with patch.dict(sys.modules, mock_modules):
337 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
339 model = _Wan21()
340 model.init()
341 gen_timer = model._new_gen_timer(None)
342 mock_video = MagicMock()
344 result = await model._output_video(None, gen_timer, mock_video, "not_a_real_type")
345 assert result is None
348@pytest.mark.asyncio
349async def test_get_rest_args_steps_fallback() -> None:
350 """Test get_rest_args uses 'steps' when sampling_steps=0."""
351 with patch.dict(sys.modules, mock_modules):
352 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
353 from image_utils import img_to_base64 as _img_to_base64
355 model = _Wan21()
356 model.init()
358 img = Image.new("RGB", (40, 30))
359 img_base64 = _img_to_base64(img)
361 result = await model.get_rest_args({
362 "img": img_base64,
363 "prompt": "test",
364 "sampling_steps": 0,
365 "steps": 15,
366 })
367 assert result["args"]["sampling_steps"] == 15
370@pytest.mark.asyncio
371async def test_get_rest_args_video_seconds_no_vae_stride() -> None:
372 """Test get_rest_args raises when video_seconds is set but vae_stride is None."""
373 with patch.dict(sys.modules, mock_modules):
374 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
375 from image_utils import img_to_base64 as _img_to_base64
377 model = _Wan21()
378 model.init()
379 model.vae_stride = None
381 img = Image.new("RGB", (40, 30))
382 img_base64 = _img_to_base64(img)
384 with pytest.raises(ValueError, match="VAE stride"):
385 await model.get_rest_args({
386 "img": img_base64,
387 "prompt": "test",
388 "video_seconds": 2.0,
389 })
392@pytest.mark.asyncio
393async def test_assert_args_num_frames_too_large() -> None:
394 """Test _assert_args raises when num_frames is too large (passes modulo check)."""
395 with patch.dict(sys.modules, mock_modules):
396 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
398 model = _Wan21()
399 model.init()
400 model.vae_stride = (4, 8, 8)
402 # num_frames=85: (85-1)%4=0 and 85 > 1+80=81 → hits the upper-bound error
403 with pytest.raises(ValueError, match="num_frames"):
404 model._assert_args(height=480, width=640, num_frames=85)
407@pytest.mark.asyncio
408async def test_get_rest_args_img_not_string() -> None:
409 """Test get_rest_args raises when img is truthy but not a string."""
410 with patch.dict(sys.modules, mock_modules):
411 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
413 model = _Wan21()
414 model.init()
416 with pytest.raises(ValueError, match="'img' parameter must be a base64-encoded string"):
417 await model.get_rest_args({"img": 12345, "prompt": "test"})
420@pytest.mark.asyncio
421async def test_get_rest_args_missing_prompt() -> None:
422 """Test get_rest_args raises when prompt is missing."""
423 with patch.dict(sys.modules, mock_modules):
424 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
425 from image_utils import img_to_base64 as _img_to_base64
427 model = _Wan21()
428 model.init()
430 img = Image.new("RGB", (40, 30))
431 img_base64 = _img_to_base64(img)
433 with pytest.raises(ValueError, match="Missing 'prompt' parameter"):
434 await model.get_rest_args({"img": img_base64})
437@pytest.mark.asyncio
438async def test_output_video_pil() -> None:
439 """Test _output_video calls _tensor_to_pil for output_type='pil'."""
440 with patch.dict(sys.modules, mock_modules):
441 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
443 model = _Wan21()
444 model.init()
445 gen_timer = model._new_gen_timer(None)
446 mock_video = MagicMock()
447 mock_pil_frames = [MagicMock()]
449 with patch.object(model, '_tensor_to_pil', return_value=mock_pil_frames) as mock_pil:
450 result = await model._output_video(None, gen_timer, mock_video, "pil")
451 mock_pil.assert_called_once_with(mock_video)
452 assert result is mock_pil_frames
455@pytest.mark.asyncio
456async def test_output_video_video_path() -> None:
457 """Test _output_video returns a video path for output_type='video_path'."""
458 with patch.dict(sys.modules, mock_modules):
459 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
461 model = _Wan21()
462 model.init()
463 gen_timer = model._new_gen_timer(None)
464 mock_video = MagicMock()
466 with patch.object(model, '_save_video', return_value=None):
467 result = await model._output_video("test_job", gen_timer, mock_video, "video_path")
468 assert result == "/tmp/test_job.mp4"
471@pytest.mark.asyncio
472async def test_output_video_video_binary() -> None:
473 """Test _output_video returns video bytes for output_type='video_binary'."""
474 import tempfile
475 import os
476 with patch.dict(sys.modules, mock_modules):
477 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
479 model = _Wan21()
480 model.init()
481 gen_timer = model._new_gen_timer(None)
482 mock_video = MagicMock()
484 # Create a temp file so aiofiles.open can read it
485 with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as tmp:
486 tmp.write(b"fake video data")
487 tmp_path = tmp.name
489 try:
490 with patch.object(model, '_save_video', return_value=None):
491 # Patch NamedTemporaryFile so _output_video uses our pre-created file
492 mock_ntf_instance = MagicMock()
493 mock_ntf_instance.name = tmp_path
494 with patch('tempfile.NamedTemporaryFile', return_value=mock_ntf_instance):
495 result = await model._output_video(None, gen_timer, mock_video, "video_binary")
496 assert isinstance(result, bytes)
497 assert result == b"fake video data"
498 finally:
499 os.unlink(tmp_path)
502def test_vae_decode() -> None:
503 """Test vae_decode passes through correctly with a properly mocked tensor."""
504 with patch.dict(sys.modules, mock_modules):
505 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
507 model = _Wan21()
508 model.init()
510 # Create a fake tensor that passes isinstance(x, torch.Tensor)
511 class FakeLatent(mock_torch.Tensor): # type: ignore[name-defined]
512 def to(self, *args: object, **kwargs: object) -> 'FakeLatent':
513 return self
515 def dim(self) -> int:
516 return 4
518 @property
519 def shape(self) -> tuple:
520 return (20, 21, 68, 90)
522 fake_lat = FakeLatent()
523 mock_pixel = MagicMock()
524 model.vae.decode = MagicMock(return_value=[mock_pixel])
526 result = model.vae_decode(fake_lat)
527 assert result is mock_pixel
530def test_vae_encode() -> None:
531 """Test vae_encode passes through correctly with a properly mocked tensor."""
532 with patch.dict(sys.modules, mock_modules):
533 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
535 model = _Wan21()
536 model.init()
538 # Create a fake tensor that passes isinstance(x, torch.Tensor)
539 class FakePixels(mock_torch.Tensor): # type: ignore[name-defined]
540 def to(self, *args: object, **kwargs: object) -> 'FakePixels':
541 return self
543 def dim(self) -> int:
544 return 4
546 @property
547 def shape(self) -> tuple:
548 return (21, 3, 68, 90) # shape[1] == 3 (RGB channels)
550 fake_pix = FakePixels()
551 mock_latent = MagicMock()
552 model.vae.encode = MagicMock(return_value=[mock_latent])
554 result = model.vae_encode(fake_pix)
555 assert result is mock_latent
558@pytest.mark.asyncio
559async def test_get_rest_args_negative_values() -> None:
560 """get_rest_args raises ValueError for non-positive numeric parameters."""
561 with patch.dict(sys.modules, mock_modules):
562 from wan.wrapper_wan21 import Wan21VideoGeneration as _Wan21
563 from image_utils import img_to_base64 as _img_to_base64
565 model = _Wan21()
566 model.init()
568 img = Image.new("RGB", (40, 30))
569 img_base64 = _img_to_base64(img)
570 base = {"img": img_base64, "prompt": "test prompt"}
572 with pytest.raises(ValueError, match="num_frames"):
573 await model.get_rest_args({**base, "num_frames": -3})
575 with pytest.raises(ValueError, match="num_frames"):
576 await model.get_rest_args({**base, "num_frames": 0})
578 with pytest.raises(ValueError, match="height"):
579 await model.get_rest_args({**base, "height": -480})
581 with pytest.raises(ValueError, match="height"):
582 await model.get_rest_args({**base, "height": 0})
584 with pytest.raises(ValueError, match="width"):
585 await model.get_rest_args({**base, "width": -640})
587 with pytest.raises(ValueError, match="width"):
588 await model.get_rest_args({**base, "width": 0})
590 with pytest.raises(ValueError, match="sampling_steps"):
591 await model.get_rest_args({**base, "sampling_steps": -10, "steps": -5})
593 with pytest.raises(ValueError, match="video_seconds"):
594 await model.get_rest_args({**base, "video_seconds": -1.0})