debug

app.py

@@ -118,6 +118,8 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     while os.path.exists(save_dir):
         save_dir = Path(f"a2v_output/{date_str}/{save_dir_name}_{np.random.randint(10000):04d}")
     save_dir.mkdir(exist_ok=True, parents=True)
+    
+    print('=====1======')
 
     ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
     ref_image_np = crop_face(ref_image_np, lmk_extractor)
@@ -127,16 +129,22 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     ref_image_np = cv2.resize(ref_image_np, (size, size))
     ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
     
+    print('=====2======')
+    
     face_result = lmk_extractor(ref_image_np)
     if face_result is None: 
         return None, ref_image_pil
     
+    print('=====3======')
+    
     lmks = face_result['lmks'].astype(np.float32)
     ref_pose = vis.draw_landmarks((ref_image_np.shape[1], ref_image_np.shape[0]), lmks, normed=True)
     
     sample = prepare_audio_feature(input_audio, wav2vec_model_path=audio_infer_config['a2m_model']['model_path'])
     sample['audio_feature'] = torch.from_numpy(sample['audio_feature']).float().cuda()
     sample['audio_feature'] = sample['audio_feature'].unsqueeze(0)
+    
+    print('=====4======')
 
     # inference
     pred = a2m_model.infer(sample['audio_feature'], sample['seq_len'])
@@ -144,6 +152,8 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     pred = pred.reshape(pred.shape[0], -1, 3)
     pred = pred + face_result['lmks3d']
     
+    print('=====5======')
+    
     if headpose_video is not None:
         pose_seq = get_headpose_temp(headpose_video)
     else:
@@ -158,6 +168,8 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     for i, verts in enumerate(projected_vertices):
         lmk_img = vis.draw_landmarks((width, height), verts, normed=False)
         pose_images.append(lmk_img)
+        
+    print('=====6======')
 
     pose_list = []
     # pose_tensor_list = []
@@ -176,6 +188,8 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     pose_list = np.array(pose_list)
     
     video_length = len(pose_list)
+    
+    print('=====7======')
 
     video = pipe(
         ref_image_pil,
@@ -383,7 +397,7 @@ with gr.Blocks() as demo:
                     a2v_headpose_video = gr.Video(label="Option: upload head pose reference video", sources="upload")
 
                 with gr.Row():
-                    a2v_size_slider = gr.Slider(minimum=256, maximum=512, step=8, value=512, label="Video size (-W & -H)")
+                    a2v_size_slider = gr.Slider(minimum=256, maximum=512, step=8, value=384, label="Video size (-W & -H)")
                     a2v_step_slider = gr.Slider(minimum=5, maximum=20, step=1, value=15, label="Steps (--steps)")
                 
                 with gr.Row():
@@ -411,7 +425,7 @@ with gr.Blocks() as demo:
                     v2v_source_video = gr.Video(label="Upload source video", sources="upload")
                 
                 with gr.Row():
-                    v2v_size_slider = gr.Slider(minimum=256, maximum=512, step=8, value=512, label="Video size (-W & -H)")
+                    v2v_size_slider = gr.Slider(minimum=256, maximum=512, step=8, value=384, label="Video size (-W & -H)")
                     v2v_step_slider = gr.Slider(minimum=5, maximum=20, step=1, value=15, label="Steps (--steps)")
                 
                 with gr.Row():