Update app.py

app.py

@@ -31,7 +31,6 @@ from src.utils.crop_face_single import crop_face
 from src.audio2vid import get_headpose_temp, smooth_pose_seq
 from src.utils.frame_interpolation import init_frame_interpolation_model, batch_images_interpolation_tool
 
-
 config = OmegaConf.load('./configs/prompts/animation_audio.yaml')
 if config.weight_dtype == "fp16":
     weight_dtype = torch.float16
@@ -99,7 +98,7 @@ pipe = pipe.to("cuda", dtype=weight_dtype)
 frame_inter_model = init_frame_interpolation_model()
 
 @spaces.GPU
-def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, length=60, seed=42):   
+def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, length=60, seed=42):
     fps = 30
     cfg = 3.5
     fi_step = 3
@@ -161,16 +160,9 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
         pose_images.append(lmk_img)
 
     pose_list = []
-    # pose_tensor_list = []
-
-    # pose_transform = transforms.Compose(
-    #     [transforms.Resize((height, width)), transforms.ToTensor()]
-    # )
     args_L = len(pose_images) if length==0 or length > len(pose_images) else length
     args_L = min(args_L, 90)
     for pose_image_np in pose_images[: args_L : fi_step]:
-        # pose_image_pil = Image.fromarray(cv2.cvtColor(pose_image_np, cv2.COLOR_BGR2RGB))
-        # pose_tensor_list.append(pose_transform(pose_image_pil))
         pose_image_np = cv2.resize(pose_image_np,  (width, height))
         pose_list.append(pose_image_np)
     
@@ -200,11 +192,6 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
         fps=fps,
     )
     
-    # save_path = f"{save_dir}/{size}x{size}_{time_str}_noaudio"
-    # save_pil_imgs(video, save_path)
-    
-    # save_path = batch_images_interpolation_tool(save_path, frame_inter_model, int(fps))
-    
     stream = ffmpeg.input(save_path)
     audio = ffmpeg.input(input_audio)
     ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac', shortest=None).run()
@@ -212,149 +199,6 @@ def audio2video(input_audio, ref_img, headpose_video=None, size=512, steps=25, l
     
     return save_path.replace('_noaudio.mp4', '.mp4'), ref_image_pil
 
-@spaces.GPU
-def video2video(ref_img, source_video, size=512, steps=25, length=60, seed=42):
-    cfg = 3.5
-    fi_step = 3
-    
-    generator = torch.manual_seed(seed)
-    
-    lmk_extractor = LMKExtractor()
-    vis = FaceMeshVisualizer()
-
-    width, height = size, size
-
-    date_str = datetime.now().strftime("%Y%m%d")
-    time_str = datetime.now().strftime("%H%M")
-    save_dir_name = f"{time_str}--seed_{seed}-{size}x{size}"
-
-    save_dir = Path(f"v2v_output/{date_str}/{save_dir_name}")
-    while os.path.exists(save_dir):
-        save_dir = Path(f"v2v_output/{date_str}/{save_dir_name}_{np.random.randint(10000):04d}")
-    save_dir.mkdir(exist_ok=True, parents=True)
-    
-    ref_image_np = cv2.cvtColor(ref_img, cv2.COLOR_RGB2BGR)
-    ref_image_np = crop_face(ref_image_np, lmk_extractor)
-    if ref_image_np is None:
-        return None, Image.fromarray(ref_img)
-    
-    ref_image_np = cv2.resize(ref_image_np, (size, size))
-    ref_image_pil = Image.fromarray(cv2.cvtColor(ref_image_np, cv2.COLOR_BGR2RGB))
-    
-    face_result = lmk_extractor(ref_image_np)
-    if face_result is None: 
-        return None, ref_image_pil
-    
-    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)
-
-    source_images = read_frames(source_video)
-    src_fps = get_fps(source_video)
-    pose_transform = transforms.Compose(
-        [transforms.Resize((height, width)), transforms.ToTensor()]
-    )
-    
-    step = 1
-    if src_fps == 60:
-        src_fps = 30
-        step = 2
-    
-    pose_trans_list = []
-    verts_list = []
-    bs_list = []
-    args_L = len(source_images) if length==0 or length*step > len(source_images) else length*step
-    args_L = min(args_L, 90*step)
-    for src_image_pil in source_images[: args_L : step*fi_step]:
-        src_img_np = cv2.cvtColor(np.array(src_image_pil), cv2.COLOR_RGB2BGR)
-        frame_height, frame_width, _ = src_img_np.shape
-        src_img_result = lmk_extractor(src_img_np)
-        if src_img_result is None:
-            break
-        pose_trans_list.append(src_img_result['trans_mat'])
-        verts_list.append(src_img_result['lmks3d'])
-        bs_list.append(src_img_result['bs'])
-
-    trans_mat_arr = np.array(pose_trans_list)
-    verts_arr = np.array(verts_list)
-    bs_arr = np.array(bs_list)
-    min_bs_idx = np.argmin(bs_arr.sum(1))
-    
-    # compute delta pose
-    pose_arr = np.zeros([trans_mat_arr.shape[0], 6])
-
-    for i in range(pose_arr.shape[0]):
-        euler_angles, translation_vector = matrix_to_euler_and_translation(trans_mat_arr[i]) # real pose of source
-        pose_arr[i, :3] =  euler_angles
-        pose_arr[i, 3:6] =  translation_vector
-    
-    init_tran_vec = face_result['trans_mat'][:3, 3] # init translation of tgt
-    pose_arr[:, 3:6] = pose_arr[:, 3:6] - pose_arr[0, 3:6] + init_tran_vec # (relative translation of source) + (init translation of tgt)
-
-    pose_arr_smooth = smooth_pose_seq(pose_arr, window_size=3)
-    pose_mat_smooth = [euler_and_translation_to_matrix(pose_arr_smooth[i][:3], pose_arr_smooth[i][3:6]) for i in range(pose_arr_smooth.shape[0])]    
-    pose_mat_smooth = np.array(pose_mat_smooth)   
-
-    # face retarget
-    verts_arr = verts_arr - verts_arr[min_bs_idx] + face_result['lmks3d']
-    # project 3D mesh to 2D landmark
-    projected_vertices = project_points_with_trans(verts_arr, pose_mat_smooth, [frame_height, frame_width])
-    
-    pose_list = []
-    for i, verts in enumerate(projected_vertices):
-        lmk_img = vis.draw_landmarks((frame_width, frame_height), verts, normed=False)
-        pose_image_np = cv2.resize(lmk_img,  (width, height))
-        pose_list.append(pose_image_np)
-    
-    pose_list = np.array(pose_list)
-    
-    video_length = len(pose_list)
-
-    video = pipe(
-        ref_image_pil,
-        pose_list,
-        ref_pose,
-        width,
-        height,
-        video_length,
-        steps,
-        cfg,
-        generator=generator,
-    ).videos
-    
-    video = batch_images_interpolation_tool(video, frame_inter_model, inter_frames=fi_step-1)
-
-    save_path = f"{save_dir}/{size}x{size}_{time_str}_noaudio.mp4"
-    save_videos_grid(
-        video,
-        save_path,
-        n_rows=1,
-        fps=src_fps,
-    )
-    
-    # save_path = f"{save_dir}/{size}x{size}_{time_str}_noaudio"
-    # save_pil_imgs(video, save_path)
-    
-    # save_path = batch_images_interpolation_tool(save_path, frame_inter_model, int(src_fps))
-    
-    audio_output = f'{save_dir}/audio_from_video.aac'
-    # extract audio
-    try:
-        ffmpeg.input(source_video).output(audio_output, acodec='copy').run()
-        # merge audio and video
-        stream = ffmpeg.input(save_path)
-        audio = ffmpeg.input(audio_output)
-        ffmpeg.output(stream.video, audio.audio, save_path.replace('_noaudio.mp4', '.mp4'), vcodec='copy', acodec='aac', shortest=None).run()
-    
-        os.remove(save_path)
-        os.remove(audio_output)
-    except:
-        shutil.move(
-            save_path,
-            save_path.replace('_noaudio.mp4', '.mp4')
-        )
-    
-    return save_path.replace('_noaudio.mp4', '.mp4'), ref_image_pil
-
 
 ################# GUI ################
 
@@ -405,32 +249,23 @@ with gr.Blocks() as demo:
         )
             
     
-    with gr.Tab("Video2video"):
+    with gr.Tab("TTS"):
         with gr.Row():
             with gr.Column():
                 with gr.Row():
-                    v2v_ref_img = gr.Image(label="Upload reference image", sources="upload")
-                    v2v_source_video = gr.Video(label="Upload source video", sources="upload")
+                    tts_text_input = gr.Textbox(lines=5, label="Input text", placeholder="Enter text to synthesize...")
+                    tts_ref_img = gr.Image(label="Upload reference image", sources="upload")
                 
                 with gr.Row():
-                    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)")
+                    tts_size_slider = gr.Slider(minimum=256, maximum=512, step=8, value=384, label="Video size (-W & -H)")
+                    tts_step_slider = gr.Slider(minimum=5, maximum=20, step=1, value=15, label="Steps (--steps)")
                 
                 with gr.Row():
-                    v2v_length = gr.Slider(minimum=0, maximum=90, step=1, value=30, label="Length (-L)")
-                    v2v_seed = gr.Number(value=42, label="Seed (--seed)")
+                    tts_length = gr.Slider(minimum=0, maximum=90, step=1, value=30, label="Length (-L)")
+                    tts_seed = gr.Number(value=42, label="Seed (--seed)")
                 
-                v2v_botton = gr.Button("Generate", variant="primary")
-            v2v_output_video = gr.PlayableVideo(label="Result", interactive=False)
-            
-        gr.Examples(
-            examples=[
-                ["configs/inference/ref_images/Aragaki.png", "configs/inference/video/Aragaki_song.mp4"],
-                ["configs/inference/ref_images/solo.png", "configs/inference/video/Aragaki_song.mp4"],
-                ["configs/inference/ref_images/lyl.png", "configs/inference/head_pose_temp/pose_ref_video.mp4"],
-                ],
-            inputs=[v2v_ref_img, v2v_source_video, a2v_headpose_video],
-        )
+                tts_button = gr.Button("Generate", variant="primary")
+            tts_output_video = gr.PlayableVideo(label="Result", interactive=False)
             
     a2v_botton.click(
         fn=audio2video,
@@ -438,11 +273,11 @@ with gr.Blocks() as demo:
                 a2v_size_slider, a2v_step_slider, a2v_length, a2v_seed], 
         outputs=[a2v_output_video, a2v_ref_img]
     )
-    v2v_botton.click(
-        fn=video2video,
-        inputs=[v2v_ref_img, v2v_source_video,
-                v2v_size_slider, v2v_step_slider, v2v_length, v2v_seed], 
-        outputs=[v2v_output_video, v2v_ref_img]
+    tts_button.click(
+        fn=audio2video,  # 추후 TTS 관련 함수로 대체 필요
+        inputs=[tts_text_input, tts_ref_img, None,
+                tts_size_slider, tts_step_slider, tts_length, tts_seed], 
+        outputs=[tts_output_video, tts_ref_img]
     )
     
 demo.launch()