From 89676189b9afe0e8b3a01ceceef0ccfa70f06f93 Mon Sep 17 00:00:00 2001
From: Tom Nijhof <nijhof.dns@gmail.com>
Date: Fri, 12 Jul 2024 08:25:16 +0200
Subject: [PATCH] fix: fast check for ffmpeg (#101)

* add ffmpeg to readme

* check if it installed

* formating
---
 inference.py | 26 +++++++++++++++++++++-----
 readme.md    |  2 ++
 2 files changed, 23 insertions(+), 5 deletions(-)

diff --git a/inference.py b/inference.py
index 24a9fd7..894fecb 100644
--- a/inference.py
+++ b/inference.py
@@ -6,12 +6,21 @@ from src.config.argument_config import ArgumentConfig
 from src.config.inference_config import InferenceConfig
 from src.config.crop_config import CropConfig
 from src.live_portrait_pipeline import LivePortraitPipeline
+import subprocess
 
 
 def partial_fields(target_class, kwargs):
     return target_class(**{k: v for k, v in kwargs.items() if hasattr(target_class, k)})
 
 
+def fast_check_ffmpeg():
+    try:
+        subprocess.run(["ffmpeg", "-version"], capture_output=True, check=True)
+        return True
+    except:
+        return False
+
+
 def fast_check_args(args: ArgumentConfig):
     if not osp.exists(args.source_image):
         raise FileNotFoundError(f"source image not found: {args.source_image}")
@@ -24,21 +33,28 @@ def main():
     tyro.extras.set_accent_color("bright_cyan")
     args = tyro.cli(ArgumentConfig)
 
+    if not fast_check_ffmpeg():
+        raise ImportError(
+            "FFmpeg is not installed. Please install FFmpeg before running this script. https://ffmpeg.org/download.html"
+        )
     # fast check the args
     fast_check_args(args)
 
     # specify configs for inference
-    inference_cfg = partial_fields(InferenceConfig, args.__dict__)  # use attribute of args to initial InferenceConfig
-    crop_cfg = partial_fields(CropConfig, args.__dict__)  # use attribute of args to initial CropConfig
+    inference_cfg = partial_fields(
+        InferenceConfig, args.__dict__
+    )  # use attribute of args to initial InferenceConfig
+    crop_cfg = partial_fields(
+        CropConfig, args.__dict__
+    )  # use attribute of args to initial CropConfig
 
     live_portrait_pipeline = LivePortraitPipeline(
-        inference_cfg=inference_cfg,
-        crop_cfg=crop_cfg
+        inference_cfg=inference_cfg, crop_cfg=crop_cfg
     )
 
     # run
     live_portrait_pipeline.execute(args)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()
diff --git a/readme.md b/readme.md
index 607637d..deb8c90 100644
--- a/readme.md
+++ b/readme.md
@@ -59,6 +59,8 @@ conda activate LivePortrait
 pip install -r requirements.txt
 ```
 
+Make sure your system has [FFmpeg](https://ffmpeg.org/)
+
 ### 2. Download pretrained weights
 
 Download the pretrained weights from HuggingFace: