detailer support for face restorer models

Signed-off-by: Vladimir Mandic <mandic00@live.com>

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
 

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')

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')
 
 

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)

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

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.

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

modules/postprocess/yolo.py

@@ -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):

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)

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