Update gradio_tabs/animation.py

gradio_tabs/animation.py

@@ -137,11 +137,12 @@ def animation(gen, chunk_size, device):
 		video_target_tensor = vid_target_tensor.to(device)
 
 		animated_video = gen.animate_batch(image_tensor, video_target_tensor, labels_v, selected_s, chunk_size)
+		edited_image = animated_video[:,:,0,:,:]
 
 		# postprocessing
 		animated_video = vid_postprocessing(animated_video, fps)
-
-		return animated_video
+		edited_image = img_postprocessing(edited_image)
+		return edited_image, animated_video
 
 
 	def clear_media():
@@ -249,8 +250,8 @@ def animation(gen, chunk_size, device):
 
 		animate_btn.click(
 			fn=animate_media,
-			inputs=[image_input, video_input] + inputs_s,  # [image_input, video_input] + inputs_s,
-			outputs=[video_output],
+			inputs=[image_input, video_input] + inputs_s,
+			outputs=[image_output, video_output],
             show_progress=True
 		)
 
@@ -274,7 +275,7 @@ def animation(gen, chunk_size, device):
 			],
             fn=animate_media,
 			inputs=[image_input, video_input] + inputs_s,
-            outputs=[video_output],
+            outputs=[image_output, video_output],
             cache_examples=True,
 		)