Refactor and Optimize Cross-Platform Support, Error Handling, and UI Enhancements

This commit is contained in:
Zohar Babin
2024-08-10 22:36:12 -04:00
parent 8d3072d906
commit 6b0cc74957
12 changed files with 453 additions and 387 deletions
+25 -26
View File
@@ -17,57 +17,56 @@ FRAME_PROCESSORS_INTERFACE = [
'process_video'
]
def load_frame_processor_module(frame_processor: str) -> Any:
def load_frame_processor_module(frame_processor: str) -> ModuleType:
try:
frame_processor_module = importlib.import_module(f'modules.processors.frame.{frame_processor}')
# Ensure all required methods are present
for method_name in FRAME_PROCESSORS_INTERFACE:
if not hasattr(frame_processor_module, method_name):
sys.exit()
raise AttributeError(f"Missing required method {method_name} in {frame_processor} module.")
except ImportError:
print(f"Frame processor {frame_processor} not found")
sys.exit()
print(f"Error: Frame processor '{frame_processor}' not found.")
sys.exit(1)
except AttributeError as e:
print(e)
sys.exit(1)
return frame_processor_module
def get_frame_processors_modules(frame_processors: List[str]) -> List[ModuleType]:
global FRAME_PROCESSORS_MODULES
if not FRAME_PROCESSORS_MODULES:
for frame_processor in frame_processors:
frame_processor_module = load_frame_processor_module(frame_processor)
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
FRAME_PROCESSORS_MODULES = [load_frame_processor_module(fp) for fp in frame_processors]
set_frame_processors_modules_from_ui(frame_processors)
return FRAME_PROCESSORS_MODULES
def set_frame_processors_modules_from_ui(frame_processors: List[str]) -> None:
global FRAME_PROCESSORS_MODULES
for frame_processor, state in modules.globals.fp_ui.items():
if state == True and frame_processor not in frame_processors:
frame_processor_module = load_frame_processor_module(frame_processor)
FRAME_PROCESSORS_MODULES.append(frame_processor_module)
if state and frame_processor not in frame_processors:
module = load_frame_processor_module(frame_processor)
FRAME_PROCESSORS_MODULES.append(module)
modules.globals.frame_processors.append(frame_processor)
if state == False:
try:
frame_processor_module = load_frame_processor_module(frame_processor)
FRAME_PROCESSORS_MODULES.remove(frame_processor_module)
modules.globals.frame_processors.remove(frame_processor)
except:
pass
elif not state and frame_processor in frame_processors:
module = load_frame_processor_module(frame_processor)
FRAME_PROCESSORS_MODULES.remove(module)
modules.globals.frame_processors.remove(frame_processor)
def multi_process_frame(source_path: str, temp_frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None], progress: Any = None) -> None:
with ThreadPoolExecutor(max_workers=modules.globals.execution_threads) as executor:
futures = []
for path in temp_frame_paths:
future = executor.submit(process_frames, source_path, [path], progress)
futures.append(future)
futures = [executor.submit(process_frames, source_path, [path], progress) for path in temp_frame_paths]
for future in futures:
future.result()
def process_video(source_path: str, frame_paths: list[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
def process_video(source_path: str, frame_paths: List[str], process_frames: Callable[[str, List[str], Any], None]) -> None:
progress_bar_format = '{l_bar}{bar}| {n_fmt}/{total_fmt} [{elapsed}<{remaining}, {rate_fmt}{postfix}]'
total = len(frame_paths)
with tqdm(total=total, desc='Processing', unit='frame', dynamic_ncols=True, bar_format=progress_bar_format) as progress:
progress.set_postfix({'execution_providers': modules.globals.execution_providers, 'execution_threads': modules.globals.execution_threads, 'max_memory': modules.globals.max_memory})
progress.set_postfix({
'execution_providers': modules.globals.execution_providers,
'execution_threads': modules.globals.execution_threads,
'max_memory': modules.globals.max_memory
})
multi_process_frame(source_path, frame_paths, process_frames, progress)
+3 -15
View File
@@ -8,7 +8,7 @@ import modules.globals
import modules.processors.frame.core
from modules.core import update_status
from modules.face_analyser import get_one_face
from modules.typing import Frame, Face
from modules.typing import Frame, Face # Ensure these are imported
from modules.utilities import conditional_download, resolve_relative_path, is_image, is_video
FACE_ENHANCER = None
@@ -16,34 +16,26 @@ THREAD_SEMAPHORE = threading.Semaphore()
THREAD_LOCK = threading.Lock()
NAME = 'DLC.FACE-ENHANCER'
def pre_check() -> bool:
download_directory_path = resolve_relative_path('..\models')
conditional_download(download_directory_path, ['https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth'])
return True
def pre_start() -> bool:
if not is_image(modules.globals.target_path) and not is_video(modules.globals.target_path):
update_status('Select an image or video for target path.', NAME)
return False
return True
def get_face_enhancer() -> Any:
global FACE_ENHANCER
with THREAD_LOCK:
if FACE_ENHANCER is None:
if os.name == 'nt':
model_path = resolve_relative_path('..\models\GFPGANv1.4.pth')
# todo: set models path https://github.com/TencentARC/GFPGAN/issues/399
else:
model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1) # type: ignore[attr-defined]
model_path = resolve_relative_path('../models/GFPGANv1.4.pth')
FACE_ENHANCER = gfpgan.GFPGANer(model_path=model_path, upscale=1) # type: ignore[attr-defined]
return FACE_ENHANCER
def enhance_face(temp_frame: Frame) -> Frame:
with THREAD_SEMAPHORE:
_, _, temp_frame = get_face_enhancer().enhance(
@@ -52,14 +44,12 @@ def enhance_face(temp_frame: Frame) -> Frame:
)
return temp_frame
def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
target_face = get_one_face(temp_frame)
if target_face:
temp_frame = enhance_face(temp_frame)
return temp_frame
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
for temp_frame_path in temp_frame_paths:
temp_frame = cv2.imread(temp_frame_path)
@@ -68,12 +58,10 @@ def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any
if progress:
progress.update(1)
def process_image(source_path: str, target_path: str, output_path: str) -> None:
target_frame = cv2.imread(target_path)
result = process_frame(None, target_frame)
cv2.imwrite(output_path, result)
def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
modules.processors.frame.core.process_video(None, temp_frame_paths, process_frames)
+6 -12
View File
@@ -2,6 +2,7 @@ from typing import Any, List
import cv2
import insightface
import threading
import os
import modules.globals
import modules.processors.frame.core
@@ -14,26 +15,25 @@ FACE_SWAPPER = None
THREAD_LOCK = threading.Lock()
NAME = 'DLC.FACE-SWAPPER'
def pre_check() -> bool:
download_directory_path = resolve_relative_path('../models')
conditional_download(download_directory_path, ['https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128_fp16.onnx'])
conditional_download(download_directory_path, [
'https://huggingface.co/hacksider/deep-live-cam/blob/main/inswapper_128_fp16.onnx'
])
return True
def pre_start() -> bool:
if not is_image(modules.globals.source_path):
update_status('Select an image for source path.', NAME)
return False
elif not get_one_face(cv2.imread(modules.globals.source_path)):
update_status('No face in source path detected.', NAME)
update_status('No face detected in the source path.', NAME)
return False
if not is_image(modules.globals.target_path) and not is_video(modules.globals.target_path):
update_status('Select an image or video for target path.', NAME)
return False
return True
def get_face_swapper() -> Any:
global FACE_SWAPPER
@@ -43,11 +43,9 @@ def get_face_swapper() -> Any:
FACE_SWAPPER = insightface.model_zoo.get_model(model_path, providers=modules.globals.execution_providers)
return FACE_SWAPPER
def swap_face(source_face: Face, target_face: Face, temp_frame: Frame) -> Frame:
return get_face_swapper().get(temp_frame, target_face, source_face, paste_back=True)
def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
if modules.globals.many_faces:
many_faces = get_many_faces(temp_frame)
@@ -60,7 +58,6 @@ def process_frame(source_face: Face, temp_frame: Frame) -> Frame:
temp_frame = swap_face(source_face, target_face, temp_frame)
return temp_frame
def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any = None) -> None:
source_face = get_one_face(cv2.imread(source_path))
for temp_frame_path in temp_frame_paths:
@@ -69,18 +66,15 @@ def process_frames(source_path: str, temp_frame_paths: List[str], progress: Any
result = process_frame(source_face, temp_frame)
cv2.imwrite(temp_frame_path, result)
except Exception as exception:
print(exception)
pass
print(f"Error processing frame {temp_frame_path}: {exception}")
if progress:
progress.update(1)
def process_image(source_path: str, target_path: str, output_path: str) -> None:
source_face = get_one_face(cv2.imread(source_path))
target_frame = cv2.imread(target_path)
result = process_frame(source_face, target_frame)
cv2.imwrite(output_path, result)
def process_video(source_path: str, temp_frame_paths: List[str]) -> None:
modules.processors.frame.core.process_video(source_path, temp_frame_paths, process_frames)