refactor the app file

app.py

@@ -1,13 +1,14 @@
 import os
 import platform
-import torch
-import gradio as gr
-from huggingface_hub import snapshot_download
 import uuid
 import shutil
 from pydub import AudioSegment
 import spaces
+import torch
+import gradio as gr
+from huggingface_hub import snapshot_download
 
+from examples.get_examples import get_examples
 from src.facerender.pirender_animate import AnimateFromCoeff_PIRender
 from src.utils.preprocess import CropAndExtract
 from src.test_audio2coeff import Audio2Coeff
@@ -16,179 +17,113 @@ from src.generate_batch import get_data
 from src.generate_facerender_batch import get_facerender_data
 from src.utils.init_path import init_path
 
-
-def get_source_image(image):
-    return image
-
-
-def toggle_audio_file(choice):
-    if choice == False:
-        return gr.update(visible=True), gr.update(visible=False)
-    else:
-        return gr.update(visible=False), gr.update(visible=True)
-
-
-def ref_video_fn(path_of_ref_video):
-    if path_of_ref_video is not None:
-        return gr.update(value=True)
-    else:
-        return gr.update(value=False)
-
-
-if torch.cuda.is_available():
-    device = "cuda"
-elif platform.system() == 'Darwin':  # macos
-    device = "mps"
-else:
-    device = "cpu"
-
-os.environ['TORCH_HOME'] = 'checkpoints'
-
 checkpoint_path = 'checkpoints'
 config_path = 'src/config'
+device = "cuda" if torch.cuda.is_available(
+) else "mps" if platform.system() == 'Darwin' else "cpu"
 
+os.environ['TORCH_HOME'] = checkpoint_path
 snapshot_download(repo_id='vinthony/SadTalker-V002rc',
-                  local_dir='./checkpoints', local_dir_use_symlinks=True)
+                  local_dir=checkpoint_path, local_dir_use_symlinks=True)
 
 
 def mp3_to_wav(mp3_filename, wav_filename, frame_rate):
-    mp3_file = AudioSegment.from_file(file=mp3_filename)
-    mp3_file.set_frame_rate(frame_rate).export(wav_filename, format="wav")
+    AudioSegment.from_file(file=mp3_filename).set_frame_rate(
+        frame_rate).export(wav_filename, format="wav")
 
 
 @spaces.GPU()
-def test(source_image, driven_audio, preprocess='crop',
-         still_mode=False,  use_enhancer=False, batch_size=1, size=256,
-         pose_style=0,
-         facerender='facevid2vid',
-         exp_scale=1.0,
-         use_ref_video=False,
-         ref_video=None,
-         ref_info=None,
-         use_idle_mode=False,
-         length_of_audio=0, use_blink=True,
-         result_dir='./results/'):
-
+def generate_video(source_image, driven_audio, preprocess='crop', still_mode=False, use_enhancer=False,
+                   batch_size=1, size=256, pose_style=0, facerender='facevid2vid', exp_scale=1.0,
+                   use_ref_video=False, ref_video=None, ref_info=None, use_idle_mode=False,
+                   length_of_audio=0, use_blink=True, result_dir='./results/'):
+    # Initialize models and paths
     sadtalker_paths = init_path(
         checkpoint_path, config_path, size, False, preprocess)
-
     audio_to_coeff = Audio2Coeff(sadtalker_paths, device)
     preprocess_model = CropAndExtract(sadtalker_paths, device)
+    animate_from_coeff = AnimateFromCoeff(sadtalker_paths, device) if facerender == 'facevid2vid' and device != 'mps' \
+        else AnimateFromCoeff_PIRender(sadtalker_paths, device)
 
-    if facerender == 'facevid2vid' and device != 'mps':
-        animate_from_coeff = AnimateFromCoeff(
-            sadtalker_paths, device)
-    elif facerender == 'pirender' or device == 'mps':
-        animate_from_coeff = AnimateFromCoeff_PIRender(
-            sadtalker_paths, device)
-        facerender = 'pirender'
-    else:
-        raise (RuntimeError('Unknown model: {}'.format(facerender)))
-
+    # Create directories for saving results
     time_tag = str(uuid.uuid4())
     save_dir = os.path.join(result_dir, time_tag)
     os.makedirs(save_dir, exist_ok=True)
-
     input_dir = os.path.join(save_dir, 'input')
     os.makedirs(input_dir, exist_ok=True)
 
-    print(source_image)
+    # Process source image
     pic_path = os.path.join(input_dir, os.path.basename(source_image))
     shutil.move(source_image, input_dir)
 
-    if driven_audio is not None and os.path.isfile(driven_audio):
+    # Process driven audio
+    if driven_audio and os.path.isfile(driven_audio):
         audio_path = os.path.join(input_dir, os.path.basename(driven_audio))
-
-        # mp3 to wav
         if '.mp3' in audio_path:
             mp3_to_wav(driven_audio, audio_path.replace('.mp3', '.wav'), 16000)
             audio_path = audio_path.replace('.mp3', '.wav')
         else:
             shutil.move(driven_audio, input_dir)
-
     elif use_idle_mode:
-        # generate audio from this new audio_path
         audio_path = os.path.join(
             input_dir, 'idlemode_'+str(length_of_audio)+'.wav')
-        from pydub import AudioSegment
-        one_sec_segment = AudioSegment.silent(
-            duration=1000*length_of_audio)  # duration in milliseconds
-        one_sec_segment.export(audio_path, format="wav")
+        AudioSegment.silent(
+            duration=1000*length_of_audio).export(audio_path, format="wav")
     else:
-        print(use_ref_video, ref_info)
-        assert use_ref_video == True and ref_info == 'all'
+        assert use_ref_video and ref_info == 'all'
 
-    if use_ref_video and ref_info == 'all':  # full ref mode
-        ref_video_videoname = os.path.basename(ref_video)
+    # Process reference video
+    if use_ref_video and ref_info == 'all':
+        ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
         audio_path = os.path.join(save_dir, ref_video_videoname+'.wav')
-        print('new audiopath:', audio_path)
-        # if ref_video contains audio, set the audio from ref_video.
-        cmd = r"ffmpeg -y -hide_banner -loglevel error -i %s %s" % (
-            ref_video, audio_path)
-        os.system(cmd)
-
-    os.makedirs(save_dir, exist_ok=True)
+        os.system(
+            f"ffmpeg -y -hide_banner -loglevel error -i {ref_video} {audio_path}")
+        ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
+        os.makedirs(ref_video_frame_dir, exist_ok=True)
+        ref_video_coeff_path, _, _ = preprocess_model.generate(
+            ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
+    else:
+        ref_video_coeff_path = None
 
-    # crop image and extract 3dmm from image
+    # Preprocess source image
     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, preprocess, True, size)
-
     if first_coeff_path is None:
         raise AttributeError("No face is detected")
 
+    # Determine reference coefficients
     if use_ref_video:
-        print('using ref video for genreation')
-        ref_video_videoname = os.path.splitext(os.path.split(ref_video)[-1])[0]
-        ref_video_frame_dir = os.path.join(save_dir, ref_video_videoname)
-        os.makedirs(ref_video_frame_dir, exist_ok=True)
-        print('3DMM Extraction for the reference video providing pose')
-        ref_video_coeff_path, _, _ = preprocess_model.generate(
-            ref_video, ref_video_frame_dir, preprocess, source_image_flag=False)
-    else:
-        ref_video_coeff_path = None
-
-    if use_ref_video:
+        ref_pose_coeff_path, ref_eyeblink_coeff_path = None, None
         if ref_info == 'pose':
             ref_pose_coeff_path = ref_video_coeff_path
-            ref_eyeblink_coeff_path = None
         elif ref_info == 'blink':
-            ref_pose_coeff_path = None
             ref_eyeblink_coeff_path = ref_video_coeff_path
         elif ref_info == 'pose+blink':
-            ref_pose_coeff_path = ref_video_coeff_path
-            ref_eyeblink_coeff_path = ref_video_coeff_path
-        elif ref_info == 'all':
-            ref_pose_coeff_path = None
-            ref_eyeblink_coeff_path = None
-        else:
-            raise ('error in refinfo')
-    else:
-        ref_pose_coeff_path = None
-        ref_eyeblink_coeff_path = None
+            ref_pose_coeff_path = ref_eyeblink_coeff_path = ref_video_coeff_path
 
-    # audio2ceoff
+    # Generate coefficients from audio or reference video
     if use_ref_video and ref_info == 'all':
-        # audio_to_coeff.generate(batch, save_dir, pose_style, ref_pose_coeff_path)
         coeff_path = ref_video_coeff_path
     else:
-        batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path, still=still_mode,
-                         idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink)  # longer audio?
+        batch = get_data(first_coeff_path, audio_path, device, ref_eyeblink_coeff_path=ref_eyeblink_coeff_path,
+                         still=still_mode, idlemode=use_idle_mode, length_of_audio=length_of_audio, use_blink=use_blink)
         coeff_path = audio_to_coeff.generate(
             batch, save_dir, pose_style, ref_pose_coeff_path)
 
-    # coeff2video
+    # Generate video from coefficients
     data = get_facerender_data(coeff_path, crop_pic_path, first_coeff_path, audio_path, batch_size, still_mode=still_mode,
                                preprocess=preprocess, size=size, expression_scale=exp_scale, facemodel=facerender)
-    return_path = animate_from_coeff.generate(
-        data, save_dir,  pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None, preprocess=preprocess, img_size=size)
+    return_path = animate_from_coeff.generate(data, save_dir, pic_path, crop_info, enhancer='gfpgan' if use_enhancer else None,
+                                              preprocess=preprocess, img_size=size)
     video_name = data['video_name']
     print(f'The generated video is named {video_name} in {save_dir}')
 
     return return_path
 
 
+# Gradio UI
 with gr.Blocks(analytics_enabled=False) as demo:
     with gr.Row():
         with gr.Column(variant='panel'):
@@ -214,8 +149,8 @@ with gr.Blocks(analytics_enabled=False) as demo:
                                 label="Use Idle Animation")
                             length_of_audio = gr.Number(
                                 value=5, label="The length(seconds) of the generated video.")
-                            use_idle_mode.change(toggle_audio_file, inputs=use_idle_mode, outputs=[
-                                driven_audio, driven_audio_no])  # todo
+                            use_idle_mode.change(lambda choice: (gr.update(visible=not choice), gr.update(visible=choice)),
+                                                 inputs=use_idle_mode, outputs=[driven_audio, driven_audio_no])
 
                     with gr.Row():
                         ref_video = gr.Video(
@@ -227,15 +162,13 @@ with gr.Blocks(analytics_enabled=False) as demo:
                             ref_info = gr.Radio(['pose', 'blink', 'pose+blink', 'all'], value='pose', label='Reference Video',
                                                 info="How to borrow from reference Video?((fully transfer, aka, video driving mode))")
 
-                        ref_video.change(ref_video_fn, inputs=ref_video, outputs=[
-                            use_ref_video])  # todo
+                        ref_video.change(lambda path: gr.update(
+                            value=path is not None), inputs=ref_video, outputs=use_ref_video)
 
         with gr.Column(variant='panel'):
             with gr.Tabs(elem_id="sadtalker_checkbox"):
                 with gr.TabItem('Settings'):
                     with gr.Column(variant='panel'):
-                        # width = gr.Slider(minimum=64, elem_id="img2img_width", maximum=2048, step=8, label="Manually Crop Width", value=512) # img2img_width
-                        # height = gr.Slider(minimum=64, elem_id="img2img_height", maximum=2048, step=8, label="Manually Crop Height", value=512) # img2img_width
                         with gr.Row():
                             pose_style = gr.Slider(
                                 minimum=0, maximum=45, step=1, label="Pose style", value=0)
@@ -265,104 +198,18 @@ with gr.Blocks(analytics_enabled=False) as demo:
                         submit = gr.Button(
                             'Generate', elem_id="sadtalker_generate", variant='primary')
 
-            with gr.Tabs(elem_id="sadtalker_genearted"):
+            with gr.Tabs(elem_id="sadtalker_generated"):
                 gen_video = gr.Video(label="Generated video")
 
     submit.click(
-        fn=test,
-        inputs=[source_image,
-                driven_audio,
-                preprocess_type,
-                is_still_mode,
-                enhancer,
-                batch_size,
-                size_of_image,
-                pose_style,
-                facerender,
-                exp_weight,
-                use_ref_video,
-                ref_video,
-                ref_info,
-                use_idle_mode,
-                length_of_audio,
-                blink_every
-                ],
+        fn=generate_video,
+        inputs=[source_image, driven_audio, preprocess_type, is_still_mode, enhancer, batch_size, size_of_image,
+                pose_style, facerender, exp_weight, use_ref_video, ref_video, ref_info, use_idle_mode, length_of_audio, blink_every],
         outputs=[gen_video],
     )
 
     with gr.Row():
-        gr.Examples(examples=[
-            [
-                'examples/source_image/full_body_1.png',
-                'examples/driven_audio/bus_chinese.wav',
-                'crop',
-                True,
-                False
-            ],
-            [
-                'examples/source_image/full_body_2.png',
-                'examples/driven_audio/japanese.wav',
-                'crop',
-                False,
-                False
-            ],
-            [
-                'examples/source_image/full3.png',
-                'examples/driven_audio/deyu.wav',
-                'crop',
-                False,
-                True
-            ],
-            [
-                'examples/source_image/full4.jpeg',
-                'examples/driven_audio/eluosi.wav',
-                'full',
-                False,
-                True
-            ],
-            [
-                'examples/source_image/full4.jpeg',
-                'examples/driven_audio/imagine.wav',
-                'full',
-                True,
-                True
-            ],
-            [
-                'examples/source_image/full_body_1.png',
-                'examples/driven_audio/bus_chinese.wav',
-                'full',
-                True,
-                False
-            ],
-            [
-                'examples/source_image/art_13.png',
-                'examples/driven_audio/fayu.wav',
-                'resize',
-                True,
-                False
-            ],
-            [
-                'examples/source_image/art_5.png',
-                'examples/driven_audio/chinese_news.wav',
-                'resize',
-                False,
-                False
-            ],
-            [
-                'examples/source_image/art_5.png',
-                'examples/driven_audio/RD_Radio31_000.wav',
-                'resize',
-                True,
-                True
-            ],
-        ],
-            inputs=[
-            source_image,
-            driven_audio,
-            preprocess_type,
-            is_still_mode,
-            enhancer],
-            outputs=[gen_video],
-            fn=test)
+        gr.Examples(examples=get_examples(), inputs=[source_image, driven_audio, preprocess_type, is_still_mode, enhancer],
+                    outputs=[gen_video], fn=generate_video)
 
 demo.launch(debug=True)

examples/get_examples.py

@@ -0,0 +1,66 @@
+def get_examples():
+    return [
+        [
+            'examples/source_image/full_body_1.png',
+            'examples/driven_audio/bus_chinese.wav',
+            'crop',
+            True,
+            False
+        ],
+        [
+            'examples/source_image/full_body_2.png',
+            'examples/driven_audio/japanese.wav',
+            'crop',
+            False,
+            False
+        ],
+        [
+            'examples/source_image/full3.png',
+            'examples/driven_audio/deyu.wav',
+            'crop',
+            False,
+            True
+        ],
+        [
+            'examples/source_image/full4.jpeg',
+            'examples/driven_audio/eluosi.wav',
+            'full',
+            False,
+            True
+        ],
+        [
+            'examples/source_image/full4.jpeg',
+            'examples/driven_audio/imagine.wav',
+            'full',
+            True,
+            True
+        ],
+        [
+            'examples/source_image/full_body_1.png',
+            'examples/driven_audio/bus_chinese.wav',
+            'full',
+            True,
+            False
+        ],
+        [
+            'examples/source_image/art_13.png',
+            'examples/driven_audio/fayu.wav',
+            'resize',
+            True,
+            False
+        ],
+        [
+            'examples/source_image/art_5.png',
+            'examples/driven_audio/chinese_news.wav',
+            'resize',
+            False,
+            False
+        ],
+        [
+            'examples/source_image/art_5.png',
+            'examples/driven_audio/RD_Radio31_000.wav',
+            'resize',
+            True,
+            True
+        ],
+    ]