Coverage for tests/test_wrapper_hunyuanframepackvae.py: 96%

1#!/usr/bin/env python3 

2 

3import sys 

4import pytest 

5 

6from unittest.mock import patch 

7from unittest.mock import MagicMock 

8from tests.torch_mock import TorchMock 

9from tests.diffusers_mock import DiffusersMock 

10 

11mock_torch = TorchMock() 

12mock_diffusers = DiffusersMock() 

13 

14sys.path.append("wrapper") 

15sys.path.append("wrapper/hunyuanframepackvae") 

16sys.path.append("wrapper/hunyuanframepack") 

17 

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.envs': MagicMock(), 

28 'xfuser.config': MagicMock(), 

29 'xfuser.core': MagicMock(), 

30 'xfuser.core.distributed': MagicMock(), 

31 'xfuser.core.cache_manager': MagicMock(), 

32 'xfuser.core.cache_manager.cache_manager': MagicMock(), 

33 'xfuser.core.long_ctx_attention': MagicMock(), 

34 'xfuser.model_executor': MagicMock(), 

35 'xfuser.model_executor.layers': MagicMock(), 

36 'xfuser.model_executor.layers.attention_processor': MagicMock(), 

37 'transformers': MagicMock(), 

38 'einops': MagicMock(), 

39 'diffusers_helper': MagicMock(), 

40 'diffusers_helper.hunyuan': MagicMock(), 

41 'diffusers_helper.utils': MagicMock(), 

42 'diffusers_helper.clip_vision': MagicMock(), 

43 'diffusers_helper.models': MagicMock(), 

44 'diffusers_helper.models.hunyuan_video_packed': MagicMock(), 

45 'diffusers_helper.pipelines': MagicMock(), 

46 'diffusers_helper.pipelines.k_diffusion_hunyuan': MagicMock(), 

47 'diffusers_helper.k_diffusion': MagicMock(), 

48 'diffusers_helper.k_diffusion.uni_pc_fm': MagicMock(), 

49 'diffusers_helper.k_diffusion.wrapper': MagicMock(), 

50 'flash_attn': MagicMock(), 

51} 

52mock_modules.update(mock_torch.get_sub_modules()) 

53mock_modules.update(mock_diffusers.get_sub_modules()) 

54 

55with patch.dict(sys.modules, mock_modules): 

56 from hunyuanframepackvae.wrapper_hunyuanframepackvae import HunyuanFramepackVAEGeneration 

57 

58 

59@pytest.mark.asyncio 

60async def test_wrapper_hunyuanframepackvae() -> None: 

61 model = HunyuanFramepackVAEGeneration() 

62 assert model is not None 

63 assert model.model_name == "hunyuanframepackvae" 

64 assert model.status == "initializing" 

65 

66 model.init() 

67 assert model.status == "ok" 

68 health = model.get_health() 

69 assert health is not None 

70 assert len(health) >= 10 

71 timestamps = model.get_timestamps() 

72 assert timestamps is not None 

73 assert len(timestamps) == 4 

74 

75 with pytest.raises(ValueError, match="Missing JSON body"): 

76 await model.get_rest_args(None) 

77 with pytest.raises(ValueError, match="Missing 'latents' parameter"): 

78 await model.get_rest_args({}) 

79 

80 try: 

81 args = await model.get_rest_args({"latents": ""}) 

82 assert args is not None 

83 assert args["task"] == "hunyuanframepackvae" 

84 except EOFError: 

85 pass # This is expected due to the mocks not providing actual latents 

86 

87 with pytest.raises(TypeError, match="Expected latents to be a torch.Tensor"): 

88 # TODO fix mocks 

89 await model.warmup() 

90 

91 with pytest.raises(ValueError, match="Latents cannot be None"): 

92 # TODO implement fixtures 

93 await model.generate(latents=None) 

94 # assert video_frames is not None 

95 

96 del model