Update app.py

app.py

@@ -115,10 +115,85 @@ def main(args):
         mapping_checkpoint = os.path.join(current_root_path, args.checkpoint_dir, 'mapping_00229-model.pth.tar')
         facerender_yaml_path = os.path.join(current_root_path, 'src', 'config', 'facerender.yaml')
 
-    face_path = "/home/user/app/images/download_1.mp4"  # Replace with the path to your face image or video
+
+    # preprocess_model = CropAndExtract(sadtalker_paths, device)
+    #init model
+    print(path_of_net_recon_model)
+    preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device)
+
+    # audio_to_coeff = Audio2Coeff(sadtalker_paths,  device)
+    audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, 
+                                audio2exp_checkpoint, audio2exp_yaml_path, 
+                                wav2lip_checkpoint, device)
+    # animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device)
+    animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, 
+                                            facerender_yaml_path, device)
+
+    first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
+    os.makedirs(first_frame_dir, exist_ok=True)
+    # first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess,\
+                                                                            #  source_image_flag=True, pic_size=args.size)
+
+   
+    first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess, source_image_flag=True)
+    print('first_coeff_path ',first_coeff_path)
+    print('crop_pic_path ',crop_pic_path)
+
+    if first_coeff_path is None:
+        print("Can't get the coeffs of the input")
+        return
+
+    if ref_eyeblink is not None:
+        ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[0]
+        ref_eyeblink_frame_dir = os.path.join(save_dir, ref_eyeblink_videoname)
+        os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
+        # ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir, args.preprocess, source_image_flag=False)
+        ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir)
+    else:
+        ref_eyeblink_coeff_path=None
+        print('ref_eyeblink_coeff_path',ref_eyeblink_coeff_path)
+
+    if ref_pose is not None:
+        if ref_pose == ref_eyeblink:
+            ref_pose_coeff_path = ref_eyeblink_coeff_path
+        else:
+            ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
+            ref_pose_frame_dir = os.path.join(save_dir, ref_pose_videoname)
+            os.makedirs(ref_pose_frame_dir, exist_ok=True)
+            # ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir, args.preprocess, source_image_flag=False)
+            ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir)
+    else:
+        ref_pose_coeff_path=None
+        print('ref_eyeblink_coeff_path',ref_pose_coeff_path)
+
+    batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
+    coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
+
+
+    if args.face3dvis:
+        from src.face3d.visualize import gen_composed_video
+        gen_composed_video(args, device, first_coeff_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4'))
+  
+    # data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path,
+                                # batch_size, input_yaw_list, input_pitch_list, input_roll_list,
+                                # expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size)
+
+
+    data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, 
+                                batch_size, input_yaw_list, input_pitch_list, input_roll_list,
+                                expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess)
+
+    # result, base64_video,temp_file_path= animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
+                                # enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size)
+
+
+    result, base64_video,temp_file_path = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
+                                enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess)
+
+    face_path = temp_file_path
     audio_path = "/home/user/app/images/audio_1.mp3"  # Replace with the path to your audio file
-    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')  # You can change suffix based on your file type
-    temp_file_path = temp_file.name
+    # temp_file = tempfile.NamedTemporaryFile(delete=False, suffix='.mp4')
+    # temp_file_path = temp_file.name
     output_path = temp_file_path
 
     # Call the function
@@ -136,90 +211,16 @@ def main(args):
         LNet_batch_size=16,
         without_rl1=False
     )
-
-    # # preprocess_model = CropAndExtract(sadtalker_paths, device)
-    # #init model
-    # print(path_of_net_recon_model)
-    # preprocess_model = CropAndExtract(path_of_lm_croper, path_of_net_recon_model, dir_of_BFM_fitting, device)
-
-    # # audio_to_coeff = Audio2Coeff(sadtalker_paths,  device)
-    # audio_to_coeff = Audio2Coeff(audio2pose_checkpoint, audio2pose_yaml_path, 
-    #                             audio2exp_checkpoint, audio2exp_yaml_path, 
-    #                             wav2lip_checkpoint, device)
-    # # animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device)
-    # animate_from_coeff = AnimateFromCoeff(free_view_checkpoint, mapping_checkpoint, 
-    #                                         facerender_yaml_path, device)
-
-    # first_frame_dir = os.path.join(save_dir, 'first_frame_dir')
-    # os.makedirs(first_frame_dir, exist_ok=True)
-    # # first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess,\
-    #                                                                         #  source_image_flag=True, pic_size=args.size)
-
-   
-    # first_coeff_path, crop_pic_path, crop_info =  preprocess_model.generate(pic_path, first_frame_dir, args.preprocess, source_image_flag=True)
-    # print('first_coeff_path ',first_coeff_path)
-    # print('crop_pic_path ',crop_pic_path)
-
-    # if first_coeff_path is None:
-    #     print("Can't get the coeffs of the input")
-    #     return
-
-    # if ref_eyeblink is not None:
-    #     ref_eyeblink_videoname = os.path.splitext(os.path.split(ref_eyeblink)[-1])[0]
-    #     ref_eyeblink_frame_dir = os.path.join(save_dir, ref_eyeblink_videoname)
-    #     os.makedirs(ref_eyeblink_frame_dir, exist_ok=True)
-    #     # ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir, args.preprocess, source_image_flag=False)
-    #     ref_eyeblink_coeff_path, _, _ =  preprocess_model.generate(ref_eyeblink, ref_eyeblink_frame_dir)
-    # else:
-    #     ref_eyeblink_coeff_path=None
-    #     print('ref_eyeblink_coeff_path',ref_eyeblink_coeff_path)
-
-    # if ref_pose is not None:
-    #     if ref_pose == ref_eyeblink:
-    #         ref_pose_coeff_path = ref_eyeblink_coeff_path
-    #     else:
-    #         ref_pose_videoname = os.path.splitext(os.path.split(ref_pose)[-1])[0]
-    #         ref_pose_frame_dir = os.path.join(save_dir, ref_pose_videoname)
-    #         os.makedirs(ref_pose_frame_dir, exist_ok=True)
-    #         # ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir, args.preprocess, source_image_flag=False)
-    #         ref_pose_coeff_path, _, _ =  preprocess_model.generate(ref_pose, ref_pose_frame_dir)
-    # else:
-    #     ref_pose_coeff_path=None
-    #     print('ref_eyeblink_coeff_path',ref_pose_coeff_path)
-
-    # batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path, still=args.still)
-    # coeff_path = audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
-
-
-    # if args.face3dvis:
-    #     from src.face3d.visualize import gen_composed_video
-    #     gen_composed_video(args, device, first_coeff_path, coeff_path, audio_path, os.path.join(save_dir, '3dface.mp4'))
-  
-    # # data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path,
-    #                             # batch_size, input_yaw_list, input_pitch_list, input_roll_list,
-    #                             # expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess, size=args.size)
-
-
-    # data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, 
-    #                             batch_size, input_yaw_list, input_pitch_list, input_roll_list,
-    #                             expression_scale=args.expression_scale, still_mode=args.still, preprocess=args.preprocess)
-
-    # # result, base64_video,temp_file_path= animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
-    #                             # enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess, img_size=args.size)
-
-
-    # result, base64_video,temp_file_path = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, \
-    #                             enhancer=args.enhancer, background_enhancer=args.background_enhancer, preprocess=args.preprocess)
     
-    # print('The video is generated')
+    print('The video is generated')
 
-    # video_clip = mp.VideoFileClip(temp_file_path)
-    # duration = video_clip.duration
+    video_clip = mp.VideoFileClip(temp_file_path)
+    duration = video_clip.duration
     
-    # app.config['temp_response'] = base64_video
-    # app.config['final_video_path'] = temp_file_path
-    # app.config['final_video_duration'] = duration
-    # return base64_video, temp_file_path, duration
+    app.config['temp_response'] = base64_video
+    app.config['final_video_path'] = temp_file_path
+    app.config['final_video_duration'] = duration
+    return base64_video, temp_file_path, duration
 
     # shutil.move(result, save_dir+'.mp4')