Create app.py

app.py

@@ -0,0 +1,404 @@
+#!/usr/bin/env python3
+
+import os
+import sys
+import shutil
+# single thread doubles cuda performance - needs to be set before torch import
+if any(arg.startswith('--execution-provider') for arg in sys.argv):
+    os.environ['OMP_NUM_THREADS'] = '1'
+
+import warnings
+from typing import List
+import platform
+import signal
+import torch
+import onnxruntime
+import pathlib
+import argparse
+
+from time import time
+
+import roop.globals
+import roop.metadata
+import roop.utilities as util
+import roop.util_ffmpeg as ffmpeg
+import ui.main as main
+from settings import Settings
+from roop.face_util import extract_face_images
+from roop.ProcessEntry import ProcessEntry
+from roop.ProcessMgr import ProcessMgr
+from roop.ProcessOptions import ProcessOptions
+from roop.capturer import get_video_frame_total, release_video
+
+
+clip_text = None
+
+call_display_ui = None
+
+process_mgr = None
+
+
+if 'ROCMExecutionProvider' in roop.globals.execution_providers:
+    del torch
+
+warnings.filterwarnings('ignore', category=FutureWarning, module='insightface')
+warnings.filterwarnings('ignore', category=UserWarning, module='torchvision')
+
+
+def parse_args() -> None:
+    signal.signal(signal.SIGINT, lambda signal_number, frame: destroy())
+    roop.globals.headless = False
+
+    program = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=100))
+    program.add_argument('--server_share', help='Public server', dest='server_share', action='store_true', default=False)
+    program.add_argument('--cuda_device_id', help='Index of the cuda gpu to use', dest='cuda_device_id', type=int, default=0)
+    roop.globals.startup_args = program.parse_args()
+    # Always enable all processors when using GUI
+    roop.globals.frame_processors = ['face_swapper', 'face_enhancer']
+
+
+def encode_execution_providers(execution_providers: List[str]) -> List[str]:
+    return [execution_provider.replace('ExecutionProvider', '').lower() for execution_provider in execution_providers]
+
+
+def decode_execution_providers(execution_providers: List[str]) -> List[str]:
+    list_providers = [provider for provider, encoded_execution_provider in zip(onnxruntime.get_available_providers(), encode_execution_providers(onnxruntime.get_available_providers()))
+            if any(execution_provider in encoded_execution_provider for execution_provider in execution_providers)]
+    
+    try:
+        for i in range(len(list_providers)):
+            if list_providers[i] == 'CUDAExecutionProvider':
+                list_providers[i] = ('CUDAExecutionProvider', {'device_id': roop.globals.cuda_device_id})
+                torch.cuda.set_device(roop.globals.cuda_device_id)
+                break
+    except:
+        pass
+
+    return list_providers
+    
+
+
+def suggest_max_memory() -> int:
+    if platform.system().lower() == 'darwin':
+        return 4
+    return 16
+
+
+def suggest_execution_providers() -> List[str]:
+    return encode_execution_providers(onnxruntime.get_available_providers())
+
+
+def suggest_execution_threads() -> int:
+    if 'DmlExecutionProvider' in roop.globals.execution_providers:
+        return 1
+    if 'ROCMExecutionProvider' in roop.globals.execution_providers:
+        return 1
+    return 8
+
+
+def limit_resources() -> None:
+    # limit memory usage
+    if roop.globals.max_memory:
+        memory = roop.globals.max_memory * 1024 ** 3
+        if platform.system().lower() == 'darwin':
+            memory = roop.globals.max_memory * 1024 ** 6
+        if platform.system().lower() == 'windows':
+            import ctypes
+            kernel32 = ctypes.windll.kernel32  # type: ignore[attr-defined]
+            kernel32.SetProcessWorkingSetSize(-1, ctypes.c_size_t(memory), ctypes.c_size_t(memory))
+        else:
+            import resource
+            resource.setrlimit(resource.RLIMIT_DATA, (memory, memory))
+
+
+
+def release_resources() -> None:
+    import gc
+    global process_mgr
+
+    if process_mgr is not None:
+        process_mgr.release_resources()
+        process_mgr = None
+
+    gc.collect()
+    # if 'CUDAExecutionProvider' in roop.globals.execution_providers and torch.cuda.is_available():
+    #     with torch.cuda.device('cuda'):
+    #         torch.cuda.empty_cache()
+    #         torch.cuda.ipc_collect()
+
+
+def pre_check() -> bool:
+    if sys.version_info < (3, 9):
+        update_status('Python version is not supported - please upgrade to 3.9 or higher.')
+        return False
+    
+    download_directory_path = util.resolve_relative_path('../models')
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/inswapper_128.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GFPGANv1.4.onnx'])
+    util.conditional_download(download_directory_path, ['https://github.com/csxmli2016/DMDNet/releases/download/v1/DMDNet.pth'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/GPEN-BFR-512.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/restoreformer_plus_plus.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/xseg.onnx'])
+    download_directory_path = util.resolve_relative_path('../models/CLIP')
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/rd64-uni-refined.pth'])
+    download_directory_path = util.resolve_relative_path('../models/CodeFormer')
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/CodeFormerv0.1.onnx'])
+    download_directory_path = util.resolve_relative_path('../models/Frame')
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_artistic.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/deoldify_stable.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/isnet-general-use.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x4.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/real_esrgan_x2.onnx'])
+    util.conditional_download(download_directory_path, ['https://huggingface.co/countfloyd/deepfake/resolve/main/lsdir_x4.onnx'])
+
+    if not shutil.which('ffmpeg'):
+       update_status('ffmpeg is not installed.')
+    return True
+
+def set_display_ui(function):
+    global call_display_ui
+
+    call_display_ui = function
+
+
+def update_status(message: str) -> None:
+    global call_display_ui
+
+    print(message)
+    if call_display_ui is not None:
+        call_display_ui(message)
+
+
+
+
+def start() -> None:
+    if roop.globals.headless:
+        print('Headless mode currently unsupported - starting UI!')
+        # faces = extract_face_images(roop.globals.source_path,  (False, 0))
+        # roop.globals.INPUT_FACES.append(faces[roop.globals.source_face_index])
+        # faces = extract_face_images(roop.globals.target_path,  (False, util.has_image_extension(roop.globals.target_path)))
+        # roop.globals.TARGET_FACES.append(faces[roop.globals.target_face_index])
+        # if 'face_enhancer' in roop.globals.frame_processors:
+        #     roop.globals.selected_enhancer = 'GFPGAN'
+       
+    batch_process_regular(None, False, None)
+
+
+def get_processing_plugins(masking_engine):
+    processors = {  "faceswap": {}}
+    if masking_engine is not None:
+        processors.update({masking_engine: {}})
+    
+    if roop.globals.selected_enhancer == 'GFPGAN':
+        processors.update({"gfpgan": {}})
+    elif roop.globals.selected_enhancer == 'Codeformer':
+        processors.update({"codeformer": {}})
+    elif roop.globals.selected_enhancer == 'DMDNet':
+        processors.update({"dmdnet": {}})
+    elif roop.globals.selected_enhancer == 'GPEN':
+        processors.update({"gpen": {}})
+    elif roop.globals.selected_enhancer == 'Restoreformer++':
+        processors.update({"restoreformer++": {}})
+    return processors
+
+
+def live_swap(frame, options):
+    global process_mgr
+
+    if frame is None:
+        return frame
+
+    if process_mgr is None:
+        process_mgr = ProcessMgr(None)
+    
+#    if len(roop.globals.INPUT_FACESETS) <= selected_index:
+#        selected_index = 0
+    process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
+    newframe = process_mgr.process_frame(frame)
+    if newframe is None:
+        return frame
+    return newframe
+
+
+def batch_process_regular(output_method, files:list[ProcessEntry], masking_engine:str, new_clip_text:str, use_new_method, imagemask, restore_original_mouth, num_swap_steps, progress, selected_index = 0) -> None:
+    global clip_text, process_mgr
+
+    release_resources()
+    limit_resources()
+    if process_mgr is None:
+        process_mgr = ProcessMgr(progress)
+    mask = imagemask["layers"][0] if imagemask is not None else None
+    if len(roop.globals.INPUT_FACESETS) <= selected_index:
+        selected_index = 0
+    options = ProcessOptions(get_processing_plugins(masking_engine), roop.globals.distance_threshold, roop.globals.blend_ratio,
+                              roop.globals.face_swap_mode, selected_index, new_clip_text, mask, num_swap_steps,
+                              roop.globals.subsample_size, False, restore_original_mouth)
+    process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
+    batch_process(output_method, files, use_new_method)
+    return
+
+def batch_process_with_options(files:list[ProcessEntry], options, progress):
+    global clip_text, process_mgr
+
+    release_resources()
+    limit_resources()
+    if process_mgr is None:
+        process_mgr = ProcessMgr(progress)
+    process_mgr.initialize(roop.globals.INPUT_FACESETS, roop.globals.TARGET_FACES, options)
+    roop.globals.keep_frames = False
+    roop.globals.wait_after_extraction = False
+    roop.globals.skip_audio = False
+    batch_process("Files", files, True)
+
+
+
+def batch_process(output_method, files:list[ProcessEntry], use_new_method) -> None:
+    global clip_text, process_mgr
+
+    roop.globals.processing = True
+
+    # limit threads for some providers
+    max_threads = suggest_execution_threads()
+    if max_threads == 1:
+        roop.globals.execution_threads = 1
+
+    imagefiles:list[ProcessEntry] = []
+    videofiles:list[ProcessEntry] = []
+           
+    update_status('Sorting videos/images')
+
+
+    for index, f in enumerate(files):
+        fullname = f.filename
+        if util.has_image_extension(fullname):
+            destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'.{roop.globals.CFG.output_image_format}')
+            destination = util.replace_template(destination, index=index)
+            pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
+            f.finalname = destination
+            imagefiles.append(f)
+
+        elif util.is_video(fullname) or util.has_extension(fullname, ['gif']):
+            destination = util.get_destfilename_from_path(fullname, roop.globals.output_path, f'__temp.{roop.globals.CFG.output_video_format}')
+            f.finalname = destination
+            videofiles.append(f)
+
+
+
+    if(len(imagefiles) > 0):
+        update_status('Processing image(s)')
+        origimages = []
+        fakeimages = []
+        for f in imagefiles:
+            origimages.append(f.filename)
+            fakeimages.append(f.finalname)
+
+        process_mgr.run_batch(origimages, fakeimages, roop.globals.execution_threads)
+        origimages.clear()
+        fakeimages.clear()
+
+    if(len(videofiles) > 0):
+        for index,v in enumerate(videofiles):
+            if not roop.globals.processing:
+                end_processing('Processing stopped!')
+                return
+            fps = v.fps if v.fps > 0 else util.detect_fps(v.filename)
+            if v.endframe == 0:
+                v.endframe = get_video_frame_total(v.filename)
+
+            is_streaming_only = output_method == "Virtual Camera"
+            if is_streaming_only == False:
+                update_status(f'Creating {os.path.basename(v.finalname)} with {fps} FPS...')
+
+            start_processing = time()
+            if is_streaming_only == False and roop.globals.keep_frames or not use_new_method:
+                util.create_temp(v.filename)
+                update_status('Extracting frames...')
+                ffmpeg.extract_frames(v.filename,v.startframe,v.endframe, fps)
+                if not roop.globals.processing:
+                    end_processing('Processing stopped!')
+                    return
+
+                temp_frame_paths = util.get_temp_frame_paths(v.filename)
+                process_mgr.run_batch(temp_frame_paths, temp_frame_paths, roop.globals.execution_threads)
+                if not roop.globals.processing:
+                    end_processing('Processing stopped!')
+                    return
+                if roop.globals.wait_after_extraction:
+                    extract_path = os.path.dirname(temp_frame_paths[0])
+                    util.open_folder(extract_path)
+                    input("Press any key to continue...")
+                    print("Resorting frames to create video")
+                    util.sort_rename_frames(extract_path)                                    
+                
+                ffmpeg.create_video(v.filename, v.finalname, fps)
+                if not roop.globals.keep_frames:
+                    util.delete_temp_frames(temp_frame_paths[0])
+            else:
+                if util.has_extension(v.filename, ['gif']):
+                    skip_audio = True
+                else:
+                    skip_audio = roop.globals.skip_audio
+                process_mgr.run_batch_inmem(output_method, v.filename, v.finalname, v.startframe, v.endframe, fps,roop.globals.execution_threads)
+                
+            if not roop.globals.processing:
+                end_processing('Processing stopped!')
+                return
+            
+            video_file_name = v.finalname
+            if os.path.isfile(video_file_name):
+                destination = ''
+                if util.has_extension(v.filename, ['gif']):
+                    gifname = util.get_destfilename_from_path(v.filename, roop.globals.output_path, '.gif')
+                    destination = util.replace_template(gifname, index=index)
+                    pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
+
+                    update_status('Creating final GIF')
+                    ffmpeg.create_gif_from_video(video_file_name, destination)
+                    if os.path.isfile(destination):
+                        os.remove(video_file_name)
+                else:
+                    skip_audio = roop.globals.skip_audio
+                    destination = util.replace_template(video_file_name, index=index)
+                    pathlib.Path(os.path.dirname(destination)).mkdir(parents=True, exist_ok=True)
+
+                    if not skip_audio:
+                        ffmpeg.restore_audio(video_file_name, v.filename, v.startframe, v.endframe, destination)
+                        if os.path.isfile(destination):
+                            os.remove(video_file_name)
+                    else:
+                        shutil.move(video_file_name, destination)
+
+            elif is_streaming_only == False:
+                update_status(f'Failed processing {os.path.basename(v.finalname)}!')
+            elapsed_time = time() - start_processing
+            average_fps = (v.endframe - v.startframe) / elapsed_time
+            update_status(f'\nProcessing {os.path.basename(destination)} took {elapsed_time:.2f} secs, {average_fps:.2f} frames/s')
+    end_processing('Finished')
+
+
+def end_processing(msg:str):
+    update_status(msg)
+    roop.globals.target_folder_path = None
+    release_resources()
+
+
+def destroy() -> None:
+    if roop.globals.target_path:
+        util.clean_temp(roop.globals.target_path)
+    release_resources()        
+    sys.exit()
+
+
+def run() -> None:
+    parse_args()
+    if not pre_check():
+        return
+    roop.globals.CFG = Settings('config.yaml')
+    roop.globals.cuda_device_id = roop.globals.startup_args.cuda_device_id
+    roop.globals.execution_threads = roop.globals.CFG.max_threads
+    roop.globals.video_encoder = roop.globals.CFG.output_video_codec
+    roop.globals.video_quality = roop.globals.CFG.video_quality
+    roop.globals.max_memory = roop.globals.CFG.memory_limit if roop.globals.CFG.memory_limit > 0 else None
+    if roop.globals.startup_args.server_share:
+        roop.globals.CFG.server_share = True
+    main.run()