Coverage for utils / pipeline_utils.py: 96.67%

1from diffusers import StableDiffusionPipeline, TextToVideoSDPipeline, StableVideoDiffusionPipeline 

2from typing import Any, Dict, List, Optional, Union 

3 

4# Pipeline type constants 

5PIPELINE_TYPE_IMAGE = "image" 

6PIPELINE_TYPE_TEXT_TO_VIDEO = "t2v" 

7PIPELINE_TYPE_IMAGE_TO_VIDEO = "i2v" 

8 

9def get_pipeline_type(pipeline) -> Optional[str]: 

10 """ 

11 Determine the type of diffusion pipeline. 

12  

13 Args: 

14 pipeline: The diffusion pipeline object 

15  

16 Returns: 

17 str: One of the pipeline type constants or None if not recognized 

18 """ 

19 if isinstance(pipeline, StableDiffusionPipeline): 

20 return PIPELINE_TYPE_IMAGE 

21 elif isinstance(pipeline, TextToVideoSDPipeline): 

22 return PIPELINE_TYPE_TEXT_TO_VIDEO 

23 elif isinstance(pipeline, StableVideoDiffusionPipeline): 

24 return PIPELINE_TYPE_IMAGE_TO_VIDEO 

25 else: 

26 return None 

27 

28def is_video_pipeline(pipeline) -> bool: 

29 """ 

30 Check if the pipeline is a video generation pipeline. 

31  

32 Args: 

33 pipeline: The diffusion pipeline object 

34  

35 Returns: 

36 bool: True if the pipeline is a video generation pipeline, False otherwise 

37 """ 

38 pipeline_type = get_pipeline_type(pipeline) 

39 return pipeline_type in [PIPELINE_TYPE_TEXT_TO_VIDEO, PIPELINE_TYPE_IMAGE_TO_VIDEO] 

40 

41def is_image_pipeline(pipeline) -> bool: 

42 """ 

43 Check if the pipeline is an image generation pipeline. 

44  

45 Args: 

46 pipeline: The diffusion pipeline object 

47  

48 Returns: 

49 bool: True if the pipeline is an image generation pipeline, False otherwise 

50 """ 

51 return get_pipeline_type(pipeline) == PIPELINE_TYPE_IMAGE 

52 

53def is_t2v_pipeline(pipeline) -> bool: 

54 """ 

55 Check if the pipeline is a text-to-video pipeline. 

56  

57 Args: 

58 pipeline: The diffusion pipeline object 

59  

60 Returns: 

61 bool: True if the pipeline is a text-to-video pipeline, False otherwise 

62 """ 

63 return get_pipeline_type(pipeline) == PIPELINE_TYPE_TEXT_TO_VIDEO 

64 

65def is_i2v_pipeline(pipeline) -> bool: 

66 """ 

67 Check if the pipeline is an image-to-video pipeline. 

68  

69 Args: 

70 pipeline: The diffusion pipeline object 

71  

72 Returns: 

73 bool: True if the pipeline is an image-to-video pipeline, False otherwise 

74 """ 

75 return get_pipeline_type(pipeline) == PIPELINE_TYPE_IMAGE_TO_VIDEO 

76 

77def get_pipeline_requirements(pipeline_type: str) -> Dict[str, Any]: 

78 """ 

79 Get the requirements for a specific pipeline type (required parameters, etc.) 

80  

81 Args: 

82 pipeline_type: The pipeline type string 

83  

84 Returns: 

85 Dict: A dictionary containing the pipeline requirements 

86 """ 

87 if pipeline_type == PIPELINE_TYPE_IMAGE: 

88 return { 

89 "required_params": [], 

90 "optional_params": ["height", "width", "num_images_per_prompt"] 

91 } 

92 elif pipeline_type == PIPELINE_TYPE_TEXT_TO_VIDEO: 

93 return { 

94 "required_params": ["num_frames"], 

95 "optional_params": ["height", "width", "fps"] 

96 } 

97 elif pipeline_type == PIPELINE_TYPE_IMAGE_TO_VIDEO: 

98 return { 

99 "required_params": ["input_image", "num_frames"], 

100 "optional_params": ["height", "width", "fps"] 

101 } 

102 else: 

103 return {"required_params": [], "optional_params": []}