Next (#122)

* Clear VRAM of face analyser on post process

* Mark as NEXT

* Reduce tensorflow memory to 512 MB

* Cosmetics on installer

* Add is_download_done to pre_process() hook to prevent errors

* Use latest onnxruntime

* Testing for download methods, Make get_download_size more robust

* Testing for download methods

* Introduce --skip-download argument

* Catch exception causes by a firewall

* Looks stable to me

README.md

@@ -56,6 +56,7 @@ python run.py [options]
 --execution-providers {cpu} [{cpu} ...]                                                          choose from the available execution providers (choices: cpu, ...)
 --execution-thread-count EXECUTION_THREAD_COUNT                                                  specify the number of execution threads
 --execution-queue-count EXECUTION_QUEUE_COUNT                                                    specify the number of execution queries
+--skip-download                                                                                  omit automate downloads and lookups
 --headless                                                                                       run the program in headless mode
 -v, --version                                                                                    show program's version number and exit
 ```

facefusion/core.py

@@ -55,6 +55,7 @@ def parse_args() -> None:
 	program.add_argument('--execution-providers', help = wording.get('execution_providers_help').format(choices = 'cpu'), dest = 'execution_providers', default = ['cpu'], choices = suggest_execution_providers_choices(), nargs = '+')
 	program.add_argument('--execution-thread-count', help = wording.get('execution_thread_count_help'), dest = 'execution_thread_count', type = int, default = suggest_execution_thread_count_default())
 	program.add_argument('--execution-queue-count', help = wording.get('execution_queue_count_help'), dest = 'execution_queue_count', type = int, default = 1)
+	program.add_argument('--skip-download', help = wording.get('skip_download_help'), dest = 'skip_download', action = 'store_true')
 	program.add_argument('--headless', help = wording.get('headless_help'), dest = 'headless', action = 'store_true')
 	program.add_argument('-v', '--version', version = metadata.get('name') + ' ' + metadata.get('version'), action = 'version')
 
@@ -86,6 +87,7 @@ def parse_args() -> None:
 	facefusion.globals.execution_providers = decode_execution_providers(args.execution_providers)
 	facefusion.globals.execution_thread_count = args.execution_thread_count
 	facefusion.globals.execution_queue_count = args.execution_queue_count
+	facefusion.globals.skip_download = args.skip_download
 	facefusion.globals.headless = args.headless
 
 
@@ -104,7 +106,7 @@ def limit_resources() -> None:
 	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)
+			tensorflow.config.experimental.VirtualDeviceConfiguration(memory_limit = 512)
 		])
 	# limit memory usage
 	if facefusion.globals.max_memory:

facefusion/globals.py

@@ -5,7 +5,6 @@ from facefusion.typing import FaceRecognition, FaceAnalyserDirection, FaceAnalys
 source_path : Optional[str] = None
 target_path : Optional[str] = None
 output_path : Optional[str] = None
-headless : Optional[bool] = None
 frame_processors : List[str] = []
 ui_layouts : List[str] = []
 keep_fps : Optional[bool] = None
@@ -29,3 +28,5 @@ max_memory : Optional[int] = None
 execution_providers : List[str] = []
 execution_thread_count : Optional[int] = None
 execution_queue_count : Optional[int] = None
+skip_download : Optional[bool] = None
+headless : Optional[bool] = None

facefusion/installer.py

@@ -13,11 +13,11 @@ from facefusion import metadata, wording
 
 ONNXRUNTIMES : Dict[str, Tuple[str, str]] =\
 {
-	'cpu': ('onnxruntime', '1.15.1'),
-	'cuda': ('onnxruntime-gpu', '1.15.1'),
+	'cpu': ('onnxruntime', '1.16.0 '),
+	'cuda': ('onnxruntime-gpu', '1.16.0'),
 	'coreml-legacy': ('onnxruntime-coreml', '1.13.1'),
 	'coreml-silicon': ('onnxruntime-silicon', '1.14.2'),
-	'directml': ('onnxruntime-directml', '1.15.1'),
+	'directml': ('onnxruntime-directml', '1.16.0'),
 	'openvino': ('onnxruntime-openvino', '1.15.0')
 }
 
@@ -60,6 +60,6 @@ def run() -> None:
 			wheel_name = '-'.join([ 'onnxruntime_silicon', onnxruntime_version, python_id, python_id, 'macosx_12_0_arm64.whl' ])
 			wheel_path = os.path.join(tempfile.gettempdir(), wheel_name)
 			wheel_url = 'https://github.com/cansik/onnxruntime-silicon/releases/download/v' + onnxruntime_version + '/' + wheel_name
-			subprocess.call([ 'curl', wheel_url, '-o', wheel_path, '-L' ])
+			subprocess.call([ 'curl', '--silent', '--location', '--continue-at', '-', '--output', wheel_path, wheel_url ])
 			subprocess.call([ 'pip', 'install', wheel_path ])
 			os.remove(wheel_path)

facefusion/metadata.py

@@ -2,7 +2,7 @@ METADATA =\
 {
 	'name': 'FaceFusion',
 	'description': 'Next generation face swapper and enhancer',
-	'version': '1.2.0',
+	'version': '1.2.1',
 	'license': 'MIT',
 	'author': 'Henry Ruhs',
 	'url': 'https://facefusion.io'

facefusion/processors/frame/modules/face_enhancer.py

@@ -5,15 +5,17 @@ from gfpgan.utils import GFPGANer
 import facefusion.globals
 from facefusion import wording, utilities
 from facefusion.core import update_status
-from facefusion.face_analyser import get_many_faces
+from facefusion.face_analyser import get_many_faces, clear_face_analyser
 from facefusion.typing import Frame, Face, ProcessMode
-from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video
+from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video, is_file, is_download_done
 from facefusion.vision import read_image, read_static_image, write_image
 
 FRAME_PROCESSOR = None
 THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore()
 THREAD_LOCK : threading.Lock = threading.Lock()
 NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_ENHANCER'
+MODEL_URL = 'https://github.com/facefusion/facefusion-assets/releases/download/models/GFPGANv1.4.pth'
+MODEL_PATH = resolve_relative_path('../.assets/models/GFPGANv1.4.pth')
 
 
 def get_frame_processor() -> Any:
@@ -21,9 +23,8 @@ def get_frame_processor() -> Any:
 
 	with THREAD_LOCK:
 		if FRAME_PROCESSOR is None:
-			model_path = resolve_relative_path('../.assets/models/GFPGANv1.4.pth')
 			FRAME_PROCESSOR = GFPGANer(
-				model_path = model_path,
+				model_path = MODEL_PATH,
 				upscale = 1,
 				device = utilities.get_device(facefusion.globals.execution_providers)
 			)
@@ -37,12 +38,19 @@ def clear_frame_processor() -> None:
 
 
 def pre_check() -> bool:
+	if not facefusion.globals.skip_download:
 		download_directory_path = resolve_relative_path('../.assets/models')
-	conditional_download(download_directory_path, [ 'https://github.com/facefusion/facefusion-assets/releases/download/models/GFPGANv1.4.pth' ])
+		conditional_download(download_directory_path, [ MODEL_URL ])
 	return True
 
 
 def pre_process(mode : ProcessMode) -> bool:
+	if not is_download_done(MODEL_URL, MODEL_PATH):
+		update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME)
+		return False
+	elif not is_file(MODEL_PATH):
+		update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME)
+		return False
 	if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path):
 		update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
 		return False
@@ -54,6 +62,7 @@ def pre_process(mode : ProcessMode) -> bool:
 
 def post_process() -> None:
 	clear_frame_processor()
+	clear_face_analyser()
 	read_static_image.cache_clear()
 
 

facefusion/processors/frame/modules/face_swapper.py

@@ -6,15 +6,17 @@ import facefusion.globals
 import facefusion.processors.frame.core as frame_processors
 from facefusion import wording
 from facefusion.core import update_status
-from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces
+from facefusion.face_analyser import get_one_face, get_many_faces, find_similar_faces, clear_face_analyser
 from facefusion.face_reference import get_face_reference, set_face_reference
 from facefusion.typing import Face, Frame, ProcessMode
-from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video
+from facefusion.utilities import conditional_download, resolve_relative_path, is_image, is_video, is_file, is_download_done
 from facefusion.vision import read_image, read_static_image, write_image
 
 FRAME_PROCESSOR = None
 THREAD_LOCK : threading.Lock = threading.Lock()
 NAME = 'FACEFUSION.FRAME_PROCESSOR.FACE_SWAPPER'
+MODEL_URL = 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx'
+MODEL_PATH = resolve_relative_path('../.assets/models/inswapper_128.onnx')
 
 
 def get_frame_processor() -> Any:
@@ -22,8 +24,7 @@ def get_frame_processor() -> Any:
 
 	with THREAD_LOCK:
 		if FRAME_PROCESSOR is None:
-			model_path = resolve_relative_path('../.assets/models/inswapper_128.onnx')
-			FRAME_PROCESSOR = insightface.model_zoo.get_model(model_path, providers = facefusion.globals.execution_providers)
+			FRAME_PROCESSOR = insightface.model_zoo.get_model(MODEL_PATH, providers = facefusion.globals.execution_providers)
 	return FRAME_PROCESSOR
 
 
@@ -34,12 +35,19 @@ def clear_frame_processor() -> None:
 
 
 def pre_check() -> bool:
+	if not facefusion.globals.skip_download:
 		download_directory_path = resolve_relative_path('../.assets/models')
-	conditional_download(download_directory_path, [ 'https://github.com/facefusion/facefusion-assets/releases/download/models/inswapper_128.onnx' ])
+		conditional_download(download_directory_path, [ MODEL_URL ])
 	return True
 
 
 def pre_process(mode : ProcessMode) -> bool:
+	if not facefusion.globals.skip_download and not is_download_done(MODEL_URL, MODEL_PATH):
+		update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME)
+		return False
+	elif not is_file(MODEL_PATH):
+		update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME)
+		return False
 	if not is_image(facefusion.globals.source_path):
 		update_status(wording.get('select_image_source') + wording.get('exclamation_mark'), NAME)
 		return False
@@ -48,6 +56,7 @@ def pre_process(mode : ProcessMode) -> bool:
 		return False
 	if mode in [ 'output', 'preview' ] and not is_image(facefusion.globals.target_path) and not is_video(facefusion.globals.target_path):
 		update_status(wording.get('select_image_or_video_target') + wording.get('exclamation_mark'), NAME)
+		return False
 	if mode == 'output' and not facefusion.globals.output_path:
 		update_status(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME)
 		return False
@@ -56,6 +65,7 @@ def pre_process(mode : ProcessMode) -> bool:
 
 def post_process() -> None:
 	clear_frame_processor()
+	clear_face_analyser()
 	read_static_image.cache_clear()
 
 

facefusion/processors/frame/modules/frame_enhancer.py

@@ -3,18 +3,21 @@ import threading
 from basicsr.archs.rrdbnet_arch import RRDBNet
 from realesrgan import RealESRGANer
 
-import facefusion
+import facefusion.globals
 import facefusion.processors.frame.core as frame_processors
 from facefusion import wording, utilities
 from facefusion.core import update_status
+from facefusion.face_analyser import clear_face_analyser
 from facefusion.typing import Frame, Face, ProcessMode
-from facefusion.utilities import conditional_download, resolve_relative_path
+from facefusion.utilities import conditional_download, resolve_relative_path, is_file, is_download_done
 from facefusion.vision import read_image, read_static_image, write_image
 
 FRAME_PROCESSOR = None
 THREAD_SEMAPHORE : threading.Semaphore = threading.Semaphore()
 THREAD_LOCK : threading.Lock = threading.Lock()
 NAME = 'FACEFUSION.FRAME_PROCESSOR.FRAME_ENHANCER'
+MODEL_URL = 'https://github.com/facefusion/facefusion-assets/releases/download/models/RealESRGAN_x4plus.pth'
+MODEL_PATH = resolve_relative_path('../.assets/models/RealESRGAN_x4plus.pth')
 
 
 def get_frame_processor() -> Any:
@@ -22,9 +25,8 @@ def get_frame_processor() -> Any:
 
 	with THREAD_LOCK:
 		if FRAME_PROCESSOR is None:
-			model_path = resolve_relative_path('../.assets/models/RealESRGAN_x4plus.pth')
 			FRAME_PROCESSOR = RealESRGANer(
-				model_path = model_path,
+				model_path = MODEL_PATH,
 				model = RRDBNet(
 					num_in_ch = 3,
 					num_out_ch = 3,
@@ -49,12 +51,19 @@ def clear_frame_processor() -> None:
 
 
 def pre_check() -> bool:
+	if not facefusion.globals.skip_download:
 		download_directory_path = resolve_relative_path('../.assets/models')
-	conditional_download(download_directory_path, [ 'https://github.com/facefusion/facefusion-assets/releases/download/models/RealESRGAN_x4plus.pth' ])
+		conditional_download(download_directory_path, [ MODEL_URL ])
 	return True
 
 
 def pre_process(mode : ProcessMode) -> bool:
+	if not facefusion.globals.skip_download and not facefusion.globals.skip_download and not is_download_done(MODEL_URL, MODEL_PATH):
+		update_status(wording.get('model_download_not_done') + wording.get('exclamation_mark'), NAME)
+		return False
+	elif not is_file(MODEL_PATH):
+		update_status(wording.get('model_file_not_present') + wording.get('exclamation_mark'), NAME)
+		return False
 	if mode == 'output' and not facefusion.globals.output_path:
 		update_status(wording.get('select_file_or_directory_output') + wording.get('exclamation_mark'), NAME)
 		return False
@@ -63,6 +72,7 @@ def pre_process(mode : ProcessMode) -> bool:
 
 def post_process() -> None:
 	clear_frame_processor()
+	clear_face_analyser()
 	read_static_image.cache_clear()
 
 

facefusion/uis/choices.py

@@ -2,6 +2,6 @@ from typing import List
 
 from facefusion.uis.typing import WebcamMode
 
-settings : List[str] = [ 'keep-fps', 'keep-temp', 'skip-audio' ]
+settings : List[str] = [ 'keep-fps', 'keep-temp', 'skip-audio', 'skip-download' ]
 webcam_mode : List[WebcamMode] = [ 'inline', 'stream_udp', 'stream_v4l2' ]
 webcam_resolution : List[str] = [ '320x240', '640x480', '1280x720', '1920x1080', '2560x1440', '3840x2160' ]

facefusion/uis/components/processors.py

@@ -29,9 +29,10 @@ def listen() -> None:
 def update_frame_processors(frame_processors : List[str]) -> Update:
 	clear_frame_processors_modules()
 	facefusion.globals.frame_processors = frame_processors
-	for frame_processor in facefusion.globals.frame_processors:
+	for frame_processor in frame_processors:
 		frame_processor_module = load_frame_processor_module(frame_processor)
-		frame_processor_module.pre_check()
+		if not frame_processor_module.pre_check():
+			return gradio.update()
 	return gradio.update(value = frame_processors, choices = sort_frame_processors(frame_processors))
 
 

facefusion/uis/components/settings.py

@@ -19,6 +19,8 @@ def render() -> None:
 		value.append('keep-temp')
 	if facefusion.globals.skip_audio:
 		value.append('skip-audio')
+	if facefusion.globals.skip_download:
+		value.append('skip-download')
 	SETTINGS_CHECKBOX_GROUP = gradio.Checkboxgroup(
 		label = wording.get('settings_checkbox_group_label'),
 		choices = choices.settings,
@@ -34,4 +36,5 @@ def update(settings : List[str]) -> Update:
 	facefusion.globals.keep_fps = 'keep-fps' in settings
 	facefusion.globals.keep_temp = 'keep-temp' in settings
 	facefusion.globals.skip_audio = 'skip-audio' in settings
+	facefusion.globals.skip_download = 'skip-download' in settings
 	return gradio.update(value = settings)

facefusion/uis/layouts/benchmark.py

@@ -1,10 +1,12 @@
 import gradio
 
+import facefusion.globals
 from facefusion.uis.components import about, processors, execution, execution_thread_count, execution_queue_count, limit_resources, benchmark_settings, benchmark
 from facefusion.utilities import conditional_download
 
 
 def pre_check() -> bool:
+	if not facefusion.globals.skip_download:
 		conditional_download('.assets/examples',
 		[
 			'https://github.com/facefusion/facefusion-assets/releases/download/examples/source.jpg',
@@ -17,6 +19,7 @@ def pre_check() -> bool:
 			'https://github.com/facefusion/facefusion-assets/releases/download/examples/target-2160p.mp4'
 		])
 		return True
+	return False
 
 
 def pre_render() -> bool:

facefusion/utilities.py

@@ -1,4 +1,5 @@
 from typing import List, Optional
+from functools import lru_cache
 from pathlib import Path
 from tqdm import tqdm
 import glob
@@ -194,12 +195,19 @@ def conditional_download(download_directory_path : str, urls : List[str]) -> Non
 						progress.update(current - progress.n)
 
 
-def get_download_size(url : str) -> int:
+@lru_cache(maxsize = None)
+def get_download_size(url : str) -> Optional[int]:
+	try:
 		response = urllib.request.urlopen(url) # type: ignore[attr-defined]
-	content_length = response.getheader('Content-Length')
-	if content_length:
-		return int(content_length)
-	return 0
+		return int(response.getheader('Content-Length'))
+	except (OSError, ValueError):
+		return None
+
+
+def is_download_done(url : str, file_path : str) -> bool:
+	if is_file(file_path):
+		return get_download_size(url) == os.path.getsize(file_path)
+	return False
 
 
 def resolve_relative_path(path : str) -> str:

facefusion/wording.py

@@ -29,6 +29,7 @@ WORDING =\
 	'execution_providers_help': 'choose from the available execution providers (choices: {choices}, ...)',
 	'execution_thread_count_help': 'specify the number of execution threads',
 	'execution_queue_count_help': 'specify the number of execution queries',
+	'skip_download_help': 'omit automate downloads and lookups',
 	'headless_help': 'run the program in headless mode',
 	'creating_temp': 'Creating temporary resources',
 	'extracting_frames_fps': 'Extracting frames with {fps} FPS',
@@ -47,6 +48,8 @@ WORDING =\
 	'processing_image_failed': 'Processing to image failed',
 	'processing_video_succeed': 'Processing to video succeed',
 	'processing_video_failed': 'Processing to video failed',
+	'model_download_not_done': 'Download of the model is not done',
+	'model_file_not_present': 'File of the model is not present',
 	'select_image_source': 'Select an image for source path',
 	'select_image_or_video_target': 'Select an image or video for target path',
 	'select_file_or_directory_output': 'Select an file or directory for output path',

tests/test_utilities.py

@@ -4,7 +4,7 @@ import subprocess
 import pytest
 
 import facefusion.globals
-from facefusion.utilities import conditional_download, extract_frames, create_temp, get_temp_directory_path, clear_temp, normalize_output_path, is_file, is_directory, is_image, is_video, encode_execution_providers, decode_execution_providers
+from facefusion.utilities import conditional_download, extract_frames, create_temp, get_temp_directory_path, clear_temp, normalize_output_path, is_file, is_directory, is_image, is_video, get_download_size, is_download_done, encode_execution_providers, decode_execution_providers
 
 
 @pytest.fixture(scope = 'module', autouse = True)
@@ -140,6 +140,18 @@ def test_is_video() -> None:
 	assert is_video('invalid') is False
 
 
+def test_get_download_size() -> None:
+	assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4') == 191675
+	assert get_download_size('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-360p.mp4') == 370732
+	assert get_download_size('invalid') is None
+
+
+def test_is_download_done() -> None:
+	assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4', '.assets/examples/target-240p.mp4') is True
+	assert is_download_done('https://github.com/facefusion/facefusion-assets/releases/download/examples/target-240p.mp4','invalid') is False
+	assert is_download_done('invalid', 'invalid') is False
+
+
 def test_encode_execution_providers() -> None:
 	assert encode_execution_providers([ 'CPUExecutionProvider' ]) == [ 'cpu' ]