Coverage for apps/streamanimate/streamanimate_job.py: 88%

1""" 

2StreamAnimate job to generate an animated video. 

3""" 

4import sys 

5import aiofiles 

6 

7from PIL import Image 

8 

9from io import BytesIO 

10from base64 import b64decode 

11 

12from typing import override 

13from typing import Any 

14from typing import Dict 

15from typing import Optional 

16 

17# Local relative imports 

18sys.path.append("..") # noqa: E402 

19sys.path.append("../..") # noqa: E402 

20 

21from streamwise_job import StreamWiseJob 

22from streamwise_job import JobStatus 

23from streamwise_job import OutputMode 

24 

25from lmm_service_manager import LMMServiceManager 

26 

27from animate_prompts import IMG_PROMPT 

28from animate_prompts import IMG_NEG_PROMPT 

29from animate_prompts import VIDEO_PROMPT 

30from animate_prompts import VIDEO_NEG_PROMPT 

31 

32from console_utils import bytes_to_human 

33 

34from file_utils import save_base64_as_binary 

35 

36from media_utils import get_audio_duration 

37from media_utils import get_video_frames 

38from media_utils import get_video_file_info 

39from media_utils import save_video_audio 

40 

41from video import MAX_FT_DURATION_SECS 

42from video import FANTASYTALKING_FPS 

43 

44 

45def _append_base_prompt(user_text: str, base_prompt: str) -> str: 

46 """Combine an optional user text prefix with a base prompt string.""" 

47 if user_text: 

48 return f"{user_text.rstrip('.!?,;')}. {base_prompt}" 

49 return base_prompt 

50 

51 

52class StreamAnimateJob(StreamWiseJob): 

53 """A job to generate an animated video.""" 

54 

55 def __init__( 

56 self, 

57 job_id: str, 

58 service_manager: LMMServiceManager, 

59 config: Dict[str, Any] = {}, 

60 ) -> None: 

61 super().__init__( 

62 "streamanimate", 

63 job_id, 

64 service_manager, 

65 config) 

66 

67 @override 

68 async def generate( 

69 self, 

70 job_config: Dict[str, Any], 

71 ) -> None: 

72 image_base64: Optional[str] = job_config.get("image_base64", None) 

73 text_prompt: str = job_config.get("text_prompt", "") 

74 narration_text: str = job_config.get("narration_text", "") 

75 await self.gen_animate( 

76 image_base64=image_base64, 

77 text_prompt=text_prompt, 

78 narration_text=narration_text, 

79 ) 

80 

81 async def gen_animate( 

82 self, 

83 image_base64: Optional[str], 

84 text_prompt: str = "", 

85 narration_text: str = "", 

86 ) -> None: 

87 """ 

88 Generate an animated video from an image (or a text-to-image prompt). 

89 

90 Pipeline: 

91 1. If no image is provided, generate one from text_prompt. 

92 2. Generate a video animation from the image. 

93 3. If narration_text is provided, generate TTS audio and merge it with the video. 

94 4. Save the final video. 

95 """ 

96 async with self.job_status_handler(): 

97 output_mode = self.get_config_output_mode() 

98 num_steps = self.get_num_steps() 

99 

100 # --- Step 1: Obtain source image --- 

101 image: Optional[Image.Image] = None 

102 if image_base64: 

103 self.logger.info(f"Using uploaded image with {bytes_to_human(len(image_base64))}.") 

104 image_bytes = b64decode(image_base64) 

105 image = Image.open(BytesIO(image_bytes)).convert("RGB") 

106 image_path = f"{self.job_path}/input_image.png" 

107 image.save(image_path) 

108 self.logger.info(f"Input image {image.width}x{image.height} saved to '{image_path}'.") 

109 else: 

110 if not text_prompt: 

111 raise ValueError("Either 'image_base64' or 'text_prompt' must be provided.") 

112 img_prompt = _append_base_prompt(text_prompt, IMG_PROMPT) 

113 self.logger.info(f"Generating image for prompt: '{img_prompt[:80]}'.") 

114 image = await self.gen.gen_image( 

115 prompt=img_prompt, 

116 neg_prompt=IMG_NEG_PROMPT, 

117 width=self.width, 

118 height=self.height, 

119 task_id="input_image", 

120 deadline=self.get_submission_time(), 

121 ) 

122 if image is None: 

123 raise ValueError("Failed to generate image from text prompt.") 

124 image_path = f"{self.job_path}/input_image.png" 

125 image.save(image_path) 

126 self.logger.info(f"Generated image {image.width}x{image.height} saved to '{image_path}'.") 

127 

128 await self.save_status(JobStatus.RUNNING) 

129 

130 width, height = image.size 

131 video_prompt = _append_base_prompt(text_prompt, VIDEO_PROMPT) 

132 video_duration_seconds = self.get_config_float("video_duration_seconds", 5.0) 

133 

134 # --- Step 2: Generate TTS narration (if requested) --- 

135 audio_base64: Optional[str] = None 

136 audio_path: Optional[str] = None 

137 if narration_text and output_mode is not OutputMode.AUDIO_ONLY: 

138 voice = "af_heart" 

139 speed = self.get_config_float("speech_speed", 1.1) 

140 self.logger.info(f"Generating narration audio for: '{narration_text[:80]}'.") 

141 audio_base64 = await self.gen.gen_audio( 

142 narration_text, 

143 voice=voice, 

144 speed=speed, 

145 task_id="narration", 

146 deadline=self.get_submission_time(), 

147 ) 

148 if audio_base64: 

149 audio_duration = get_audio_duration(audio_base64) 

150 audio_path = f"{self.job_path}/narration.wav" 

151 await save_base64_as_binary(audio_path, audio_base64) 

152 self.logger.info( 

153 f"Narration audio: {bytes_to_human(len(audio_base64))}, " 

154 f"{audio_duration:.3f} seconds.") 

155 video_duration_seconds = audio_duration # Match video length to narration 

156 

157 await self.save_status(JobStatus.RUNNING) 

158 

159 # --- Step 3: Generate animation video --- 

160 scene_video_binary: Optional[bytes] = None 

161 if output_mode is OutputMode.AUDIO_ONLY: 

162 # Audio-only mode: narration over a static image slide 

163 self.logger.info("Audio-only mode: skipping video generation.") 

164 elif ( 

165 audio_base64 

166 and audio_path 

167 and output_mode == OutputMode.VIDEO_AUDIO_SYNCED 

168 and video_duration_seconds <= MAX_FT_DURATION_SECS 

169 ): 

170 # Lip-synced animation (Fantasy Talking) 

171 self.logger.info( 

172 f"Generating lip-synced animation {width}x{height} " 

173 f"for {video_duration_seconds:.1f} seconds.") 

174 scene_video_binary = await self.gen.gen_video_audio_from_img( 

175 img=image, 

176 audio_base64=audio_base64, 

177 prompt=video_prompt, 

178 neg_prompt=VIDEO_NEG_PROMPT, 

179 width=width, 

180 height=height, 

181 steps=num_steps, 

182 task_id="animate", 

183 deadline=self.get_submission_time(), 

184 ) 

185 else: 

186 # Standard animation (FramePack / HunyuanVideo) 

187 self.logger.info( 

188 f"Generating animation {width}x{height} " 

189 f"for {video_duration_seconds:.1f} seconds.") 

190 scene_video_binary = await self.gen.gen_video( 

191 img=image, 

192 prompt=video_prompt, 

193 neg_prompt=VIDEO_NEG_PROMPT, 

194 width=width, 

195 height=height, 

196 video_seconds=video_duration_seconds, 

197 steps=num_steps, 

198 task_id="animate", 

199 wait_request=True, 

200 deadline=self.get_submission_time(), 

201 ) 

202 

203 await self.save_status(JobStatus.RUNNING) 

204 

205 # --- Step 4: Merge audio and video, then save --- 

206 out_path = f"{self.job_path}/{self.job_id}.mp4" 

207 if scene_video_binary and audio_path: 

208 # Merge narration audio with generated video 

209 video_frames = await get_video_frames(scene_video_binary) 

210 video_file_info = get_video_file_info(scene_video_binary) 

211 video_info = video_file_info.get("video", {}) 

212 _video_fps = video_info.get("fps") 

213 video_fps: float = _video_fps if _video_fps is not None else FANTASYTALKING_FPS 

214 

215 out_path = await save_video_audio( 

216 video_content=video_frames, 

217 audio_path=audio_path, 

218 out_video_path=out_path, 

219 fps=video_fps) 

220 elif scene_video_binary: 

221 # Video without narration 

222 async with aiofiles.open(out_path, "wb") as file: 

223 await file.write(scene_video_binary) 

224 elif audio_path: 

225 # Audio-only: no video was generated, save narration as final output 

226 self.logger.info("Audio-only mode: final output is the narration audio.") 

227 out_path = audio_path 

228 else: 

229 raise ValueError("Neither video nor audio was generated.") 

230 

231 self.logger.info( 

232 f"Generated animation saved to '{out_path}'.")