Coverage for tests/test_wrapper_hunyuanimage.py: 100%

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("hunyuanimage") 

15 

16mock_modules = { 

17 "torch": mock_torch, 

18 "diffusers": mock_diffusers, 

19 "transformers": MagicMock(), 

20 "xfuser": MagicMock(), 

21 "xfuser.config": MagicMock(), 

22} 

23mock_modules.update(mock_torch.get_sub_modules()) 

24mock_modules.update(mock_diffusers.get_sub_modules()) 

25 

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

27 from hunyuanimage.wrapper_hunyuanimage import HunyuanImageGeneration 

28 

29 

30@pytest.mark.asyncio 

31async def test_wrapper_hunyuanimage() -> None: 

32 model = HunyuanImageGeneration() 

33 assert model is not None 

34 assert model.model_name == "hunyuanimage" 

35 assert model.status == "initializing" 

36 

37 with pytest.raises(ValueError, match="Model not initialized. Current status: initializing."): 

38 await model.generate(64, 48, "test prompt") 

39 

40 model.status = "ok" 

41 with pytest.raises(ValueError, match="HunyuanImage model not loaded."): 

42 await model.generate(64, 48, "test prompt") 

43 model.status == "initializing" 

44 

45 model.init() 

46 assert model.status == "ok" 

47 

48 # Mock model init() 

49 model.model.vae.config.ffactor_spatial = 8 

50 

51 health = model.get_health() 

52 assert health is not None 

53 

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

55 await model.get_rest_args(None) 

56 with pytest.raises(ValueError, match="Missing 'prompt' parameter"): 

57 await model.get_rest_args({}) 

58 

59 rest_args = await model.get_rest_args({ 

60 "prompt": "test prompt", 

61 }) 

62 assert rest_args == { 

63 "task": "hunyuanimage", 

64 "args": { 

65 "width": 640, 

66 "height": 480, 

67 "prompt": "test prompt", 

68 "sampling_steps": 25, 

69 "seed": None, 

70 } 

71 } 

72 

73 await model.warmup() 

74 

75 image = await model.generate( 

76 width=1024, 

77 height=1024, 

78 prompt="Test prompt") 

79 assert image is not None 

80 

81 with pytest.raises(ValueError, match="2048x1024 too large. Max is 1024 x 1024."): 

82 await model.generate(width=2048, height=1024, prompt="Prompt") 

83 

84 with pytest.raises(ValueError, match="Width 1027 not supported. Must be multiple of 8."): 

85 await model.generate(width=1027, height=512, prompt="Prompt") 

86 

87 with pytest.raises(ValueError, match="Height 513 not supported. Must be multiple of 8."): 

88 await model.generate(width=512, height=513, prompt="Prompt") 

89 

90 

91@pytest.mark.asyncio 

92async def test_wrapper_hunyuanimage_parallel() -> None: 

93 model = HunyuanImageGeneration() 

94 model.init() 

95 assert model.status == "ok" 

96 model.rank = 1 # Simulate distributed rank 

97 image = await model.generate( 

98 width=128, 

99 height=128, 

100 prompt="Test prompt") 

101 assert image is None # Only rank 0 generates images