Coverage for wrapper/imageresize/wrapper_imageresize.py: 100%

1import logging 

2 

3from PIL import Image 

4 

5from typing import override 

6from typing import List 

7from typing import Optional 

8from typing import Dict 

9from typing import Any 

10from typing import Union 

11 

12from wrapper_model import ModelGeneration 

13 

14from image_utils import base64_to_img 

15from media_utils import base64_to_video_frames 

16 

17 

18class ImageResize(ModelGeneration): 

19 """Image resizing model generation.""" 

20 

21 def __init__(self) -> None: 

22 super().__init__("imageresize") 

23 

24 async def warmup(self) -> None: 

25 logging.info("Warmup for ImageResize generation") 

26 dummy_image = Image.new('RGB', (128, 128), color="red") 

27 await self.generate(image=dummy_image) 

28 

29 @override 

30 async def generate( 

31 self, 

32 image: Image.Image, 

33 video: Optional[List[Image.Image]] = None, 

34 height: int = 1024, 

35 width: int = 1024, 

36 job_id: Optional[str] = None, 

37 ) -> Union[Image.Image, List[Image.Image]]: 

38 gen_timer = self._new_gen_timer(job_id) 

39 

40 self.running = True 

41 

42 try: 

43 # Video 

44 if video is not None: 

45 return [ 

46 frame.resize((width, height), Image.Resampling.LANCZOS) 

47 for frame in video 

48 ] 

49 # Image 

50 if image is not None: 

51 return image.resize((width, height), Image.Resampling.LANCZOS) 

52 # Missing inputs 

53 raise ValueError("Image is required for resizing generation.") 

54 finally: 

55 gen_timer.end("total") 

56 self.running = False 

57 

58 async def get_rest_args( 

59 self, 

60 data_json: Dict[str, Union[str, int, float]], 

61 ) -> Dict[str, Any]: 

62 if data_json is None: 

63 raise ValueError("Missing JSON body") 

64 

65 img_base64 = data_json.get("img", None) 

66 img = None 

67 if img_base64 is not None: 

68 assert isinstance(img_base64, str) 

69 img = base64_to_img(img_base64) 

70 

71 video_base64 = data_json.get("video", None) 

72 video = None 

73 if video_base64 is not None: 

74 assert isinstance(video_base64, str) 

75 video = base64_to_video_frames(video_base64) 

76 

77 return { 

78 "task": self.model_name, 

79 "args": { 

80 "image": img, 

81 "video": video, 

82 "height": data_json.get("height", 1024), 

83 "width": data_json.get("width", 1024), 

84 } 

85 }