diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index de2ea3e..c27c904 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -28,3 +28,5 @@ jobs: - run: pip install -r requirements.txt gdown - run: gdown 13QpWFWJ37EB-nHrEOY64CEtQWY-tz7DZ - run: python run.py -f=.github/examples/face.jpg -t=.github/examples/target.mp4 -o=.github/examples/output.mp4 + - run: ffmpeg -i .github/examples/snapshot.mp4 -i .github/examples/output.mp4 -filter_complex psnr -f null - + diff --git a/roop/core.py b/roop/core.py index 2bf3ebd..aaa9571 100755 --- a/roop/core.py +++ b/roop/core.py @@ -2,11 +2,13 @@ import os import sys -from typing import List os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2' # single thread doubles performance of gpu-mode - needs to be set before torch import if any(arg.startswith('--gpu-vendor') for arg in sys.argv): os.environ['OMP_NUM_THREADS'] = '1' +import warnings +warnings.simplefilter(action='ignore', category=FutureWarning) +from typing import List import platform import signal import shutil @@ -20,8 +22,7 @@ import cv2 import roop.globals from roop.swapper import process_video, process_img, process_faces -from roop.utilities import has_image_extention, is_image, detect_fps, create_video, extract_frames, \ - get_temp_frames_paths, restore_audio, create_temp, clean_temp, is_video +from roop.utilities import has_image_extention, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frames_paths, restore_audio, create_temp, move_temp, clean_temp from roop.analyser import get_face_single import roop.ui as ui @@ -216,7 +217,9 @@ def start(preview_callback = None) -> None: else: status('Restoring audio might cause issues as fps are not kept...') restore_audio(args.target_path, args.output_path) - clean_temp(args.target_path, args.output_path) + else: + move_temp(args.target_path, args.output_path) + clean_temp(args.target_path) if is_video(args.target_path): status('Swapping to video succeed!') else: @@ -265,8 +268,9 @@ def run() -> None: handle_parse() pre_check() limit_resources() - start() - if not roop.globals.headless: + if roop.globals.headless: + start() + else: window = ui.init( { 'all_faces': args.all_faces, diff --git a/roop/utilities.py b/roop/utilities.py index acf1632..487cbb7 100644 --- a/roop/utilities.py +++ b/roop/utilities.py @@ -14,7 +14,7 @@ def run_ffmpeg(args: List) -> None: commands.extend(args) try: subprocess.check_output(commands, stderr=subprocess.STDOUT) - except Exception: + except Exception as exception: pass @@ -29,18 +29,17 @@ def detect_fps(source_path: str) -> int: def extract_frames(target_path: str) -> None: - temp_directory_path = get_temp_directory_path(target_path) - run_ffmpeg(['-i', target_path, temp_directory_path + os.sep + '%04d.png']) + run_ffmpeg(['-i', target_path, get_temp_directory_path(target_path) + os.sep + '%04d.png']) def create_video(target_path: str, fps: int) -> None: - temp_directory_path = get_temp_directory_path(target_path) - temp_file_path = get_temp_file_path(target_path) - run_ffmpeg(['-i', temp_directory_path + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', temp_file_path]) + run_ffmpeg(['-i', get_temp_directory_path(target_path) + os.sep + '%04d.png', '-framerate', str(fps), '-c:v', 'libx264', '-crf', '7', '-pix_fmt', 'yuv420p', '-y', get_temp_file_path(target_path)]) def restore_audio(target_path: str, output_path: str) -> None: run_ffmpeg(['-i', get_temp_file_path(target_path), '-i', target_path, '-c:v', 'copy', '-map', '0:v:0', '-map', '1:a:0', '-y', output_path]) + if not os.path.isfile(output_path): + move_temp(target_path, output_path) def get_temp_frames_paths(target_path: str) -> List: @@ -48,8 +47,7 @@ def get_temp_frames_paths(target_path: str) -> List: def get_temp_directory_path(target_path: str) -> str: - target_directory_path = os.path.dirname(target_path) - return target_directory_path + os.sep + 'temp' + return os.path.dirname(target_path) + os.sep + 'temp' def get_temp_file_path(target_path: str) -> str: @@ -60,10 +58,13 @@ def create_temp(target_path: str) -> None: Path(get_temp_directory_path(target_path)).mkdir(exist_ok=True) -def clean_temp(target_path: str, output_path: str) -> None: +def move_temp(target_path: str, output_path: str) -> None: temp_file_path = get_temp_file_path(target_path) - if not roop.globals.keep_audio and os.path.isfile(temp_file_path): + if os.path.isfile(temp_file_path): shutil.move(temp_file_path, output_path) + + +def clean_temp(target_path: str) -> None: if not roop.globals.keep_frames: shutil.rmtree(get_temp_directory_path(target_path))