Update gradio_tabs/vid_edit.py

gradio_tabs/vid_edit.py

@@ -143,11 +143,13 @@ def vid_edit(gen, chunk_size, device):
 		video_target_tensor = video_target_tensor.to(device)
 
 		edited_video_tensor = gen.edit_vid_batch(video_target_tensor, labels_v, selected_s, chunk_size)
+		edited_image_tensor = edited_video_tensor[:,:,0,:,:]       
 
 		# de-norm
 		animated_video, animated_all_video = vid_all_save(video_target_tensor, edited_video_tensor, fps)
+		edited_image = img_postprocessing(edited_image_tensor)
 
-		return animated_video, animated_all_video
+		return edited_image, animated_video, animated_all_video
 
 
 	def clear_media():
@@ -246,7 +248,7 @@ def vid_edit(gen, chunk_size, device):
 		animate_btn.click(
 			fn=edit_vid,
 			inputs=[video_input] + inputs_s,  # [image_input, video_input] + inputs_s,
-			outputs=[video_output, video_all_output],
+			outputs=[image_output, video_output, video_all_output],
 		)
 
 		clear_btn.click(
@@ -269,7 +271,7 @@ def vid_edit(gen, chunk_size, device):
 			],
             fn=edit_vid,
 			inputs=[video_input] + inputs_s,
-            outputs=[video_output, video_all_output],
+            outputs=[image_output, video_output, video_all_output],
             cache_examples=True,
 		)