Lauri Gates
e57116de68
Add two new face enhancement processors using GPEN-BFR ONNX models at 256x256 and 512x512 resolutions. Models auto-download on first use from GitHub releases. Integrates into existing frame processor pipeline alongside GFPGAN enhancer with UI toggle switches. - modules/paths.py: Shared path constants module - modules/processors/frame/_onnx_enhancer.py: ONNX enhancement utilities - modules/processors/frame/face_enhancer_gpen256.py: GPEN-BFR 256 processor - modules/processors/frame/face_enhancer_gpen512.py: GPEN-BFR 512 processor - modules/core.py: Add GPEN choices to --frame-processor CLI arg - modules/globals.py: Add GPEN entries to fp_ui toggle dict - modules/ui.py: Add GPEN toggle switches and processing integration Closes #1663 Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
126 lines
3.9 KiB
Python
126 lines
3.9 KiB
Python
"""GPEN-BFR-512 face enhancer — ONNX-based face restoration at 512x512."""
|
|
|
|
from typing import Any, List
|
|
import os
|
|
import threading
|
|
|
|
import cv2
|
|
import numpy as np
|
|
|
|
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.utilities import (
|
|
is_image,
|
|
is_video,
|
|
)
|
|
from modules.processors.frame._onnx_enhancer import (
|
|
create_onnx_session,
|
|
warmup_session,
|
|
enhance_face_onnx,
|
|
)
|
|
|
|
NAME = "DLC.FACE-ENHANCER-GPEN512"
|
|
INPUT_SIZE = 512
|
|
MODEL_URL = "https://github.com/harisreedhar/Face-Upscalers-ONNX/releases/download/GPEN-BFR/GPEN-BFR-512.onnx"
|
|
MODEL_FILE = "GPEN-BFR-512.onnx"
|
|
|
|
ENHANCER = None
|
|
THREAD_LOCK = threading.Lock()
|
|
|
|
abs_dir = os.path.dirname(os.path.abspath(__file__))
|
|
models_dir = os.path.join(
|
|
os.path.dirname(os.path.dirname(os.path.dirname(abs_dir))), "models"
|
|
)
|
|
|
|
|
|
def pre_check() -> bool:
|
|
model_path = os.path.join(models_dir, MODEL_FILE)
|
|
if not os.path.exists(model_path):
|
|
update_status(f"Downloading {MODEL_FILE}...", NAME)
|
|
from modules.utilities import conditional_download
|
|
conditional_download(models_dir, [MODEL_URL])
|
|
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_enhancer() -> Any:
|
|
global ENHANCER
|
|
with THREAD_LOCK:
|
|
if ENHANCER is None:
|
|
model_path = os.path.join(models_dir, MODEL_FILE)
|
|
if not os.path.exists(model_path):
|
|
from modules.utilities import conditional_download
|
|
conditional_download(models_dir, [MODEL_URL])
|
|
if not os.path.exists(model_path):
|
|
raise FileNotFoundError(f"Model file not found: {model_path}")
|
|
print(f"{NAME}: Loading ONNX model from {model_path}")
|
|
ENHANCER = create_onnx_session(model_path)
|
|
warmup_session(ENHANCER)
|
|
print(f"{NAME}: Model loaded successfully.")
|
|
return ENHANCER
|
|
|
|
|
|
def enhance_face(temp_frame: Frame, face: Face) -> Frame:
|
|
try:
|
|
session = get_enhancer()
|
|
except Exception as e:
|
|
print(f"{NAME}: {e}")
|
|
return temp_frame
|
|
try:
|
|
return enhance_face_onnx(temp_frame, face, session, INPUT_SIZE)
|
|
except Exception as e:
|
|
print(f"{NAME}: Error during face enhancement: {e}")
|
|
return temp_frame
|
|
|
|
|
|
def process_frame(source_face: Face | None, temp_frame: Frame) -> Frame:
|
|
target_face = get_one_face(temp_frame)
|
|
if target_face is None:
|
|
return temp_frame
|
|
return enhance_face(temp_frame, target_face)
|
|
|
|
|
|
def process_frame_v2(temp_frame: Frame) -> Frame:
|
|
target_face = get_one_face(temp_frame)
|
|
if target_face:
|
|
temp_frame = enhance_face(temp_frame, target_face)
|
|
return temp_frame
|
|
|
|
|
|
def process_frames(
|
|
source_path: str | None, temp_frame_paths: List[str], progress: Any = None
|
|
) -> None:
|
|
for temp_frame_path in temp_frame_paths:
|
|
temp_frame = cv2.imread(temp_frame_path)
|
|
if temp_frame is None:
|
|
if progress:
|
|
progress.update(1)
|
|
continue
|
|
result = process_frame(None, temp_frame)
|
|
cv2.imwrite(temp_frame_path, result)
|
|
if progress:
|
|
progress.update(1)
|
|
|
|
|
|
def process_image(source_path: str | None, target_path: str, output_path: str) -> None:
|
|
target_frame = cv2.imread(target_path)
|
|
if target_frame is None:
|
|
print(f"{NAME}: Error: Failed to read target image {target_path}")
|
|
return
|
|
result_frame = process_frame(None, target_frame)
|
|
cv2.imwrite(output_path, result_frame)
|
|
print(f"{NAME}: Enhanced image saved to {output_path}")
|
|
|
|
|
|
def process_video(source_path: str | None, temp_frame_paths: List[str]) -> None:
|
|
modules.processors.frame.core.process_video(source_path, temp_frame_paths, process_frames)
|