Coverage for tests/test_wrapper_fluxupscaler.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 

11from PIL import Image 

12 

13mock_torch = TorchMock() 

14mock_diffusers = DiffusersMock() 

15 

16sys.path.append("wrapper") 

17sys.path.append("wrapper/flux") 

18sys.path.append("wrapper/fluxupscaler") 

19 

20mock_modules = { 

21 'nvidia_smi': MagicMock(), 

22 'imageio': MagicMock(), 

23 'cv2': MagicMock(), 

24 'torch': mock_torch, 

25 'xfuser': MagicMock(), 

26 'xfuser.config': MagicMock(), 

27 'xfuser.core': MagicMock(), 

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

29 'xfuser.model_executor': MagicMock(), 

30 'xfuser.model_executor.models': MagicMock(), 

31 'xfuser.model_executor.models.transformers.transformer_flux': MagicMock(), 

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

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

34} 

35mock_modules.update(mock_torch.get_sub_modules()) 

36mock_modules.update(mock_diffusers.get_sub_modules()) 

37 

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

39 from image_utils import img_to_base64 

40 from fluxupscaler.wrapper_fluxupscaler import FluxUpscalerGeneration 

41 

42 

43@pytest.mark.asyncio 

44async def test_wrapper_fluxupscaler() -> None: 

45 model = FluxUpscalerGeneration() 

46 assert model is not None 

47 assert model.model_name == "fluxupscaler" 

48 assert model.status == "initializing" 

49 

50 img = Image.new("RGB", (40, 30)) 

51 img_base64 = img_to_base64(img) 

52 

53 with pytest.raises(ValueError, match="Model not initialized"): 

54 await model.generate( 

55 img=img, 

56 width=128, 

57 height=80, 

58 prompt="test prompt") 

59 

60 model.init() 

61 assert model.status == "ok" 

62 

63 health = model.get_health() 

64 assert health is not None 

65 timestamps = model.get_timestamps() 

66 assert timestamps is not None 

67 

68 await model.get_rest_args({}) 

69 

70 await model.get_rest_args({ 

71 "img": img_base64, 

72 }) 

73 

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

75 await model.get_rest_args(None) 

76 

77 # Success case 

78 await model.get_rest_args({ 

79 "job_id": "unittest", 

80 "prompt": "Test prompt", 

81 "width": 80, 

82 "height": 60, 

83 "img": img_base64 

84 }) 

85 

86 await model.warmup() 

87 

88 image = await model.generate( 

89 img=img, 

90 width=256, 

91 height=160, 

92 prompt="Test prompt", 

93 ) 

94 assert image is not None 

95 assert isinstance(image, Image.Image) 

96 

97 # Bad cases: _assert_args only rejects when world_size > 1 

98 model.world_size = 2 

99 with pytest.raises(ValueError, match="48x48 not supported for 2 GPUs"): 

100 await model.generate( 

101 img=img, 

102 width=48, 

103 height=48, 

104 prompt="Test prompt", 

105 ) 

106 with pytest.raises(ValueError, match="208x116 not supported for 2 GPUs"): 

107 await model.generate( 

108 img=img, 

109 width=208, 

110 height=116, 

111 prompt="Test prompt", 

112 ) 

113 

114 del model