From e26de8cdba321c840dc8a56292f5a05f9f48b970 Mon Sep 17 00:00:00 2001 From: Vladimir Mandic Date: Sat, 18 Jan 2025 14:29:38 -0500 Subject: [PATCH] detailer support for face restorer models Signed-off-by: Vladimir Mandic --- CHANGELOG.md | 12 +++-- installer.py | 12 +++++ launch.py | 7 +-- modules/detailer.py | 4 +- modules/postprocess/aurasr_model.py | 2 +- modules/postprocess/codeformer_model.py | 2 +- modules/postprocess/restorer.py | 61 +++++++++++++++++++++++++ modules/postprocess/yolo.py | 38 +++++++++++---- modules/shared.py | 2 +- webui.sh | 7 ++- wiki | 2 +- 11 files changed, 126 insertions(+), 23 deletions(-) create mode 100644 modules/postprocess/restorer.py diff --git a/CHANGELOG.md b/CHANGELOG.md index de7142e81..defa805f1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,19 +1,23 @@ # Change Log for SD.Next -## Update for 2025-01-17 +## Update for 2025-01-18 -- **Other**: +- **Detailer**: + - in addition as standard behavior of detect & run-generate, it can now also run face-restore models + - included models are: *CodeFormer, RestoreFormer, GFPGan, GPEN-BFR* +- **Other**: + - **ipex**: update supported torch versions - **gallery**: add http fallback for slow/unreliable links - **upscale**: code refactor to unify latent, resize and model based upscalers - **splash**: add legacy mode indicator on splash screen -- **Fixes**: +- **Fixes**: - non-full vae decode - send-to image transfer - sana vae tiling - increase gallery timeouts - update ui element ids - modernui use local font - - unique font family registration + - unique font family registration ## Update for 2025-01-15 diff --git a/installer.py b/installer.py index 88b3b30eb..ea37f55f0 100644 --- a/installer.py +++ b/installer.py @@ -734,6 +734,17 @@ def install_torch_addons(): ts('addons', t_start) +# check cudnn +def check_cudnn(): + import site + site_packages = site.getsitepackages() + cuda_path = os.environ.get('CUDA_PATH', '') + for site_package in site_packages: + folder = os.path.join(site_package, 'nvidia', 'cudnn', 'lib') + if os.path.exists(folder) and folder not in cuda_path: + os.environ['CUDA_PATH'] = f"{cuda_path}:{folder}" + + # check torch version def check_torch(): t_start = time.time() @@ -845,6 +856,7 @@ def check_torch(): return if not args.skip_all: install_torch_addons() + check_cudnn() if args.profile: pr.disable() print_profile(pr, 'Torch') diff --git a/launch.py b/launch.py index 07fa93697..f92e9bfda 100755 --- a/launch.py +++ b/launch.py @@ -66,9 +66,10 @@ def get_custom_args(): installer.log.trace(f'Environment: {installer.print_dict(env)}') env = [f'{k}={v}' for k, v in os.environ.items() if k.startswith('SD_')] installer.log.debug(f'Env flags: {env}') - ldd = os.environ.get('LD_PRELOAD', None) - if ldd is not None: - installer.log.debug(f'Linker flags: "{ldd}"') + ldpreload = os.environ.get('LD_PRELOAD', None) + ldpath = os.environ.get('LD_LIBRARY_PATH', None) + if ldpreload is not None or ldpath is not None: + installer.log.debug(f'Linker flags: preload="{ldpreload}" path="{ldpath}"') rec('args') diff --git a/modules/detailer.py b/modules/detailer.py index 31908c295..6c31371aa 100644 --- a/modules/detailer.py +++ b/modules/detailer.py @@ -1,3 +1,4 @@ +from abc import abstractmethod from modules import shared @@ -5,6 +6,7 @@ class Detailer: # abstract class used for postprocessing def name(self): return "None" + @abstractmethod def restore(self, np_image): return np_image @@ -13,5 +15,5 @@ def detail(np_image, p=None): # postprocesses the image detailers = [x for x in shared.detailers if x.name() == shared.opts.detailer_model or shared.opts.detailer_model is None] if len(detailers) == 0: return np_image - detailer = detailers[0] + detailer: Detailer = detailers[0] return detailer.restore(np_image, p) diff --git a/modules/postprocess/aurasr_model.py b/modules/postprocess/aurasr_model.py index ab5844d1a..546adf35b 100644 --- a/modules/postprocess/aurasr_model.py +++ b/modules/postprocess/aurasr_model.py @@ -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 diff --git a/modules/postprocess/codeformer_model.py b/modules/postprocess/codeformer_model.py index c601f2b40..d0509a120 100644 --- a/modules/postprocess/codeformer_model.py +++ b/modules/postprocess/codeformer_model.py @@ -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. diff --git a/modules/postprocess/restorer.py b/modules/postprocess/restorer.py new file mode 100644 index 000000000..827a08336 --- /dev/null +++ b/modules/postprocess/restorer.py @@ -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 diff --git a/modules/postprocess/yolo.py b/modules/postprocess/yolo.py index 85a50cc2a..5e6d5ed05 100644 --- a/modules/postprocess/yolo.py +++ b/modules/postprocess/yolo.py @@ -9,13 +9,17 @@ from modules import shared, processing, devices, processing_class, ui_common from modules.detailer import Detailer -PREDEFINED = [ # +predefined = [ # '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): diff --git a/modules/shared.py b/modules/shared.py index 720ca7c55..54204c201 100644 --- a/modules/shared.py +++ b/modules/shared.py @@ -1172,7 +1172,7 @@ opts.data['uni_pc_lower_order_final'] = opts.schedulers_use_loworder # compatibi opts.data['uni_pc_order'] = max(2, opts.schedulers_solver_order) # compatibility log.info(f'Engine: backend={backend} compute={devices.backend} device={devices.get_optimal_device_name()} attention="{opts.cross_attention_optimization}" mode={devices.inference_context.__name__}') if not native: - log.warning('Backend=original is in maintainance-only mode') + log.warning('Backend=original: legacy mode / maintainance-only') opts.data['diffusers_offload_mode'] = 'none' prompt_styles = modules.styles.StyleDatabase(opts) diff --git a/webui.sh b/webui.sh index b6ae67847..dc814713b 100755 --- a/webui.sh +++ b/webui.sh @@ -84,7 +84,12 @@ fi # Add venv lib folder to PATH if [ -d "$(realpath "$venv_dir")/lib/" ] && [[ -z "${DISABLE_VENV_LIBS}" ]] then - export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(realpath "$venv_dir")/lib/ + if [[ -v LD_LIBRARY_PATH ]] + then + export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:$(realpath "$venv_dir")/lib/ + else + export LD_LIBRARY_PATH=$(realpath "$venv_dir")/lib/ + fi fi # Add ROCm to PATH if it's not already diff --git a/wiki b/wiki index 3dcc0808d..1e3ebab7d 160000 --- a/wiki +++ b/wiki @@ -1 +1 @@ -Subproject commit 3dcc0808db1e9c351a845184a3bf0fe7ca783190 +Subproject commit 1e3ebab7dcb772a6ad4e4c03be97973c72c36fa4