Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

detailer support for face restorer models

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2025-01-18 14:29:38 -05:00
parent bdf8ca7e27
commit e26de8cdba
11 changed files with 126 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/postprocess/aurasr_model.py
+1 -1
View File
@@ -3,7 +3,7 @@ import diffusers
from PIL import Image
from modules import shared, devices
from modules.upscaler import Upscaler, UpscalerData
from installer import install
class UpscalerAuraSR(Upscaler):
def __init__(self, dirname): # pylint: disable=super-init-not-called
modules/postprocess/codeformer_model.py
+1 -1
View File
@@ -4,7 +4,7 @@ import torch
import modules.detailer
from modules import shared, devices, modelloader, errors
from modules.paths import models_path
from installer import install
# codeformer people made a choice to include modified basicsr library to their project which makes
# it utterly impossible to use it alongside with other libraries that also use basicsr, like GFPGAN.
modules/postprocess/restorer.py
+61
View File
@@ -0,0 +1,61 @@
import time
import cv2
import numpy as np
from modules import shared, devices
face_helper = None
def restore(np_image, name, session, strength): # pylint: disable=unused-argument
t0 = time.time()
global face_helper # pylint: disable=global-statement
try:
from facelib.utils.face_restoration_helper import FaceRestoreHelper
from facelib.detection.retinaface import retinaface
except Exception as e:
shared.log.error(f"FaceRestorer error: {e}")
return np_image
if hasattr(retinaface, 'device'):
retinaface.device = devices.device
if face_helper is None:
face_helper = FaceRestoreHelper(1, face_size=512, crop_ratio=(1, 1), det_model='retinaface_resnet50', save_ext='png', use_parse=True, device=devices.device)
np_image = np_image[:, :, ::-1]
original_resolution = np_image.shape[0:2]
resolution = session.get_inputs()[0].shape[-2:]
if face_helper is None or session is None:
return np_image
face_helper.clean_all()
face_helper.read_image(np_image)
face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
face_helper.align_warp_face()
detected_faces = len(face_helper.cropped_faces)
for cropped_face in face_helper.cropped_faces:
cropped_face = cv2.resize(cropped_face, resolution, interpolation=cv2.INTER_LINEAR)
cropped_face = cropped_face.astype(np.float16)[:,:,::-1] / 255.0
cropped_face = cropped_face.transpose((2, 0, 1))
cropped_face = (cropped_face - 0.5) / 0.5
cropped_face = np.expand_dims(cropped_face, axis=0).astype(np.float16)
w = np.array([strength], dtype=np.double)
if 'codeformer' in name:
restored_face = session.run(None, {'x':cropped_face, 'w':w})[0][0]
else:
restored_face = session.run(None, {'input':cropped_face})[0][0]
restored_face = (restored_face.transpose(1,2,0).clip(-1,1) + 1) * 0.5
restored_face = (restored_face * 255)[:,:,::-1]
restored_face = restored_face.clip(0, 255).astype('uint8')
face_helper.add_restored_face(restored_face)
face_helper.get_inverse_affine(None)
restored_img = face_helper.paste_faces_to_input_image()
restored_img = restored_img[:, :, ::-1]
if original_resolution != restored_img.shape[0:2]:
restored_img = cv2.resize(restored_img, (0, 0), fx=original_resolution[1]/restored_img.shape[1], fy=original_resolution[0]/restored_img.shape[0], interpolation=cv2.INTER_LINEAR)
face_helper.clean_all()
t1 = time.time()
shared.log.info(f'Detailer: model="{name}" faces={detected_faces} strength={strength} time={t1-t0:.3f}')
return restored_img
modules/postprocess/yolo.py
+28 -10
View File
@@ -9,13 +9,17 @@ from modules import shared, processing, devices, processing_class, ui_common
from modules.detailer import Detailer
PREDEFINED = [ # <https://huggingface.co/vladmandic/yolo-detailers/tree/main>
predefined = [ # <https://huggingface.co/vladmandic/yolo-detailers/tree/main>
'https://github.com/ultralytics/assets/releases/download/v8.3.0/yolo11m.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/face-yolo8n.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/hand_yolov8n.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/person_yolov8n-seg.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/eyes-v1.pt',
'https://huggingface.co/vladmandic/yolo-detailers/resolve/main/eyes-full-v1.pt',
'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/codeformer.fp16.onnx',
'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/restoreformer.fp16.onnx',
'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GFPGANv1.4.fp16.onnx',
'https://huggingface.co/netrunner-exe/Face-Upscalers-onnx/resolve/main/GPEN-BFR-512.fp16.onnx',
]
load_lock = threading.Lock()
@@ -50,7 +54,7 @@ class YoloRestorer(Detailer):
self.list.clear()
files = []
downloaded = 0
for m in PREDEFINED:
for m in predefined:
name = os.path.splitext(os.path.basename(m))[0]
self.list[name] = m
files.append(name)
@@ -61,7 +65,7 @@ class YoloRestorer(Detailer):
name = os.path.splitext(os.path.basename(f))[0]
if name not in files:
self.list[name] = os.path.join(shared.opts.yolo_dir, f)
shared.log.info(f'Available Yolo: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}')
shared.log.info(f'Available Detailer: path="{shared.opts.yolo_dir}" items={len(list(self.list))} downloaded={downloaded}')
return self.list
def dependencies(self):
@@ -156,18 +160,30 @@ class YoloRestorer(Detailer):
with load_lock:
from modules import modelloader
model = None
self.dependencies()
if model_name is None:
model_name = list(self.list)[0]
if model_name in self.models:
return model_name, self.models[model_name]
else:
model_url = self.list.get(model_name)
model_url = self.list.get(model_name, None)
if model_url is None:
shared.log.error(f'Load: type=Detailer name="{model_name}" error="model not found"')
return None, None
file_name = os.path.basename(model_url)
model_file = None
try:
model_file = modelloader.load_file_from_url(url=model_url, model_dir=shared.opts.yolo_dir, file_name=file_name)
if model_file is not None:
if model_file is None:
shared.log.error(f'Load: type=Detailer name="{model_name}" url="{model_url}" error="failed to fetch model"')
elif model_file.endswith('.onnx'):
import onnxruntime as ort
options = ort.SessionOptions()
# options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession(model_file, sess_options=options, providers=devices.onnx)
self.models[model_name] = session
return model_name, session
else:
self.dependencies()
import ultralytics
model = ultralytics.YOLO(model_file)
classes = list(model.names.values())
@@ -200,6 +216,11 @@ class YoloRestorer(Detailer):
shared.log.warning(f'Detailer: model="{name}" not loaded')
continue
if name.endswith('.fp16'):
from modules.postprocess import restorer
np_image = restorer.restore(np_image, name, model, p.detailer_strength)
continue
image = Image.fromarray(np_image)
items = self.predict(model, image)
if len(items) == 0:
@@ -262,8 +283,7 @@ class YoloRestorer(Detailer):
p.steps = orig_p.get('steps', 0)
report = [{'label': i.label, 'score': i.score, 'size': f'{i.width}x{i.height}' } for i in items]
shared.log.info(f'Detailer: model="{name}" items={report} args={items[0].args} denoise={p.denoising_strength} blur={p.mask_blur} width={p.width} height={p.height} padding={p.inpaint_full_res_padding}')
# shared.log.debug(f'Detailer: prompt="{prompt}" negative="{negative}"')
shared.log.info(f'Detailer: model="{name}" items={report} args={items[0].args} strength={p.detailer_strength} blur={p.mask_blur} width={p.width} height={p.height} padding={p.inpaint_full_res_padding}')
models_used.append(name)
mask_all = []
@@ -304,8 +324,6 @@ class YoloRestorer(Detailer):
p.image_mask = blend([np.array(m) for m in mask_all])
p.image_mask = Image.fromarray(p.image_mask)
# if len(models_used) > 0:
# shared.log.debug(f'Detailer processed: models={models_used}')
return np_image
def ui(self, tab: str):