Fix dependency issues and add some updates

README.md

@@ -4,7 +4,7 @@ emoji: 👁
 colorFrom: indigo
 colorTo: indigo
 sdk: gradio
-sdk_version: 3.0.19
+sdk_version: 3.48.0
 app_file: app.py
 pinned: false
 ---

app.py

@@ -1,15 +1,18 @@
-import gradio as gr
 import os
-import shutil
+import pathlib
+
+import gradio as gr
 import torch
 from PIL import Image
-import argparse
-import pathlib
 
-os.system("git clone https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model")
-os.chdir("Thin-Plate-Spline-Motion-Model")
-os.system("mkdir checkpoints")
-os.system("gdown 1-CKOjv_y_TzNe-dwQsjjeVxJUuyBAb5X -O checkpoints/vox.pth.tar")
+repo_dir = pathlib.Path("Thin-Plate-Spline-Motion-Model").absolute()
+if not repo_dir.exists():
+    os.system("git clone https://github.com/yoyo-nb/Thin-Plate-Spline-Motion-Model")
+os.chdir(repo_dir.name)
+if not (repo_dir / "checkpoints").exists():
+    os.system("mkdir checkpoints")
+if not (repo_dir / "checkpoints/vox.pth.tar").exists():
+    os.system("gdown 1-CKOjv_y_TzNe-dwQsjjeVxJUuyBAb5X -O checkpoints/vox.pth.tar")
 
 
 
@@ -40,24 +43,21 @@ def update_style_image(style_name: str) -> dict:
     return gr.Markdown.update(value=text)
 
 
-def set_example_image(example: list) -> dict:
-    return gr.Image.update(value=example[0])
+def inference(img, vid):
+    if not os.path.exists('temp'):
+        os.system('mkdir temp')
 
-def set_example_video(example: list) -> dict:
-    return gr.Video.update(value=example[0])
+    img.save("temp/image.jpg", "JPEG")
+    if torch.cuda.is_available():
+        os.system(f"python demo.py --config config/vox-256.yaml --checkpoint ./checkpoints/vox.pth.tar --source_image 'temp/image.jpg' --driving_video {vid} --result_video './temp/result.mp4'")
+    else:
+        os.system(f"python demo.py --config config/vox-256.yaml --checkpoint ./checkpoints/vox.pth.tar --source_image 'temp/image.jpg' --driving_video {vid} --result_video './temp/result.mp4' --cpu")
+    return './temp/result.mp4'
 
-def inference(img,vid):
-  if not os.path.exists('temp'):
-    os.system('mkdir temp')
-  
-  img.save("temp/image.jpg", "JPEG")
-  os.system(f"python demo.py --config config/vox-256.yaml --checkpoint ./checkpoints/vox.pth.tar --source_image 'temp/image.jpg' --driving_video {vid} --result_video './temp/result.mp4' --cpu")
-  return './temp/result.mp4'
-  
 
 
 def main():
-    with gr.Blocks(theme="huggingface", css='style.css') as demo:
+    with gr.Blocks(css='style.css') as demo:
         gr.Markdown(title)
         gr.Markdown(DESCRIPTION)
 
@@ -71,12 +71,11 @@ def main():
                     with gr.Row():
                         input_image = gr.Image(label='Input Image',
                                                type="pil")
-                        
+
             with gr.Row():
                 paths = sorted(pathlib.Path('assets').glob('*.png'))
-                example_images = gr.Dataset(components=[input_image],
-                                            samples=[[path.as_posix()]
-                                                     for path in paths])
+                gr.Examples(inputs=[input_image],
+                            examples=[[path.as_posix()] for path in paths])
 
         with gr.Box():
             gr.Markdown('''## Step 2 (Select Driving Video)
@@ -86,17 +85,16 @@ def main():
                 with gr.Column():
                     with gr.Row():
                         driving_video = gr.Video(label='Driving Video',
-                                               format="mp4")
+                                                 format="mp4")
 
             with gr.Row():
                 paths = sorted(pathlib.Path('assets').glob('*.mp4'))
-                example_video = gr.Dataset(components=[driving_video],
-                                            samples=[[path.as_posix()]
-                                                     for path in paths])
+                gr.Examples(inputs=[driving_video],
+                            examples=[[path.as_posix()] for path in paths])
 
         with gr.Box():
             gr.Markdown('''## Step 3 (Generate Animated Image based on the Video)
-- Hit the **Generate** button. (Note: As it runs on the CPU, it takes ~ 3 minutes to generate final results.)
+- Hit the **Generate** button. (Note: On cpu-basic, it takes ~ 10 minutes to generate final results.)
 ''')
             with gr.Row():
                 with gr.Column():
@@ -104,7 +102,7 @@ def main():
                         generate_button = gr.Button('Generate')
 
                 with gr.Column():
-                    result = gr.Video(type="file", label="Output")
+                    result = gr.Video(label="Output")
         gr.Markdown(FOOTER)
         generate_button.click(fn=inference,
                               inputs=[
@@ -112,17 +110,8 @@ def main():
                                   driving_video
                               ],
                               outputs=result)
-        example_images.click(fn=set_example_image,
-                             inputs=example_images,
-                             outputs=example_images.components)
-        example_video.click(fn=set_example_video,
-                             inputs=example_video,
-                             outputs=example_video.components)
-
-    demo.launch(
-        enable_queue=True,
-        debug=True
-    )
+
+    demo.queue(max_size=10).launch()
 
 if __name__ == '__main__':
-    main()
+    main()

requirements.txt

@@ -1,27 +1,78 @@
-cffi==1.14.6
-cycler==0.10.0
-decorator==5.1.0
-face-alignment==1.3.5
-imageio==2.9.0
-imageio-ffmpeg==0.4.5
-kiwisolver==1.3.2
-matplotlib==3.4.3
-networkx==2.6.3
-numpy==1.20.3
-pandas==1.3.3
-Pillow
-pycparser==2.20
-pyparsing==2.4.7
+aiofiles==23.2.1
+altair==5.1.2
+annotated-types==0.6.0
+anyio==3.7.1
+attrs==23.1.0
+beautifulsoup4==4.12.2
+certifi==2023.7.22
+charset-normalizer==3.3.0
+click==8.1.7
+cmake==3.27.7
+contourpy==1.1.1
+cycler==0.12.1
+exceptiongroup==1.1.3
+fastapi==0.103.2
+ffmpy==0.3.1
+filelock==3.12.4
+fonttools==4.43.1
+fsspec==2023.9.2
+gdown==4.7.1
+gradio==3.48.0
+gradio_client==0.6.1
+h11==0.14.0
+httpcore==0.18.0
+httpx==0.25.0
+huggingface-hub==0.18.0
+idna==3.4
+imageio==2.31.5
+imageio-ffmpeg==0.4.9
+importlib-resources==6.1.0
+Jinja2==3.1.2
+jsonschema==4.19.1
+jsonschema-specifications==2023.7.1
+kiwisolver==1.4.5
+lazy_loader==0.3
+lit==17.0.3
+MarkupSafe==2.1.3
+matplotlib==3.8.0
+mpmath==1.3.0
+networkx==3.1
+numpy==1.25.2
+orjson==3.9.9
+packaging==23.2
+pandas==2.1.1
+Pillow==10.1.0
+psutil==5.9.6
+pydantic==2.4.2
+pydantic_core==2.10.1
+pydub==0.25.1
+pyparsing==3.1.1
+PySocks==1.7.1
 python-dateutil==2.8.2
-pytz==2021.1
-PyWavelets==1.1.1
-PyYAML
-scikit-image
-scikit-learn
-scipy
+python-multipart==0.0.6
+pytz==2023.3.post1
+PyYAML==6.0.1
+referencing==0.30.2
+requests==2.31.0
+rpds-py==0.10.6
+scikit-image==0.22.0
+scipy==1.11.3
+semantic-version==2.10.0
+setuptools-scm==8.0.4
 six==1.16.0
-torch==1.11.0
-torchvision==0.12.0
-tqdm==4.62.3
-gradio
-gdown
+sniffio==1.3.0
+soupsieve==2.5
+starlette==0.27.0
+sympy==1.12
+tifffile==2023.9.26
+tomli==2.0.1
+toolz==0.12.0
+torch==2.0.0
+torchvision==0.15.1
+tqdm==4.66.1
+triton==2.0.0
+typing_extensions==4.8.0
+tzdata==2023.3
+urllib3==2.0.7
+uvicorn==0.23.2
+websockets==11.0.3