Update gradio_app.py

gradio_app.py

@@ -3,6 +3,14 @@ import os
 os.environ['HYDRA_FULL_ERROR']='1'
 os.environ['CUDA_VISIBLE_DEVICES'] = '0'
 
+from huggingface_hub import snapshot_download 
+
+# Download weights
+snapshot_download(
+    repo_id = "ChaolongYang/KDTalker",
+    local_dir = "./"
+)
+
 import argparse
 import shutil
 import uuid
@@ -13,7 +21,7 @@ import cv2
 from rich.progress import track
 import tyro
 
-
+import gradio as gr 
 from PIL import Image
 import time
 import torch
@@ -133,7 +141,7 @@ class Inferencer(object):
 
         from model import get_model
         self.point_diffusion = get_model()
-        ckpt = torch.load('KDTalker.pth')
+        ckpt = torch.load('./KDTalker.pth')
 
         self.point_diffusion.load_state_dict(ckpt['model'])
         self.point_diffusion.eval()
@@ -368,16 +376,34 @@ class Inferencer(object):
         os.remove(path)
         os.remove(new_audio_path)
 
+def gradio_infer(source_image, driven_audio):
 
-if __name__ == '__main__':
-    parser = argparse.ArgumentParser()
-    parser.add_argument("-source_image", type=str, default="example/source_image/WDA_BenCardin1_000.png",
-                        help="source image")
-    parser.add_argument("-driven_audio", type=str, default="example/driven_audio/WDA_BenCardin1_000.wav",
-                        help="driving audio")
-    parser.add_argument("-output", type=str, default="results/output.mp4", help="output video file name", )
-
-    args = parser.parse_args()
+    output_path = "results/output.mp4"
 
     Infer = Inferencer()
-    Infer.generate_with_audio_img(args.source_image, args.driven_audio, args.output)
+    Infer.generate_with_audio_img(source_image, driven_audio, output)
+
+    return output_path
+
+with gr.Blocks() as demo:
+    with gr.Column():
+        gr.Markdown("# KDTalker")
+
+        with gr.Row():
+
+            with gr.Column():
+                source_image = gr.Image(label="Source Image", type="filepath")
+                driven_audio = gr.Audio(label="Driven Audio")
+                submit_btn = gr.Button("Submit")
+
+            with gr.Column():
+                output_video = gr.Video(label="Output Video")
+
+    submit_btn.click(
+        fn = gradio_infer,
+        inputs = [source_image, driven_audio],
+        outputs = [output_video]
+    )
+
+demo.launch()
+