Restore Face node

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "extern/SadTalker"]
 	path = extern/SadTalker
 	url = https://github.com/OpenTalker/SadTalker.git
+[submodule "extern/GFPGAN"]
+	path = extern/GFPGAN
+	url = https://github.com/TencentARC/GFPGAN.git

__init__.py

@@ -1,10 +1,11 @@
 import traceback
-from .log import log, blue_text, get_summary, get_label
+from .log import log, blue_text, cyan_text, get_summary, get_label
 from .utils import here
 import importlib
 import os
 
 NODE_CLASS_MAPPINGS = {}
+NODE_DISPLAY_NAME_MAPPINGS = {}
 NODE_CLASS_MAPPINGS_DEBUG = {}
 
 
@@ -64,16 +65,17 @@ def load_nodes():
 nodes = load_nodes()
 for node_class in nodes:
     class_name = node_class.__name__
-    class_name = node_class.__name__
-    node_name = f"{get_label(class_name)} (mtb)"
-    NODE_CLASS_MAPPINGS[node_name] = node_class
-    NODE_CLASS_MAPPINGS_DEBUG[node_name] = node_class.__doc__
-
+    node_label = f"{get_label(class_name)} (mtb)"
+    NODE_CLASS_MAPPINGS[node_label] = node_class
+    NODE_DISPLAY_NAME_MAPPINGS[class_name] = node_label
+    NODE_CLASS_MAPPINGS_DEBUG[node_label] = node_class.__doc__
+    # TODO: I removed this, I find it more convenient to write without spaces, but it breaks every of my workflows
+    # TODO (cont): and until I find a way to automate the conversion, I'll leave it like this
 
-log.debug(
+log.info(
     f"Loaded the following nodes:\n\t"
     + "\n\t".join(
-        f"{k}: {blue_text(get_summary(doc)) if doc else '-'}"
+        f"{cyan_text(k)}: {blue_text(get_summary(doc)) if doc else '-'}"
         for k, doc in NODE_CLASS_MAPPINGS_DEBUG.items()
     )
 )

log.py

@@ -1,9 +1,21 @@
 import logging
 import re
+import os
+
+base_log_level = logging.DEBUG if os.environ.get("MTB_DEBUG") else logging.INFO
+print(f"Log level: {base_log_level}")
+
+
+# Custom object that discards the output
+class NullWriter:
+    def write(self, text):
+        pass
 
 
 class Formatter(logging.Formatter):
     grey = "\x1b[38;20m"
+    cyan = "\x1b[36;20m"
+    purple = "\x1b[35;20m"
     yellow = "\x1b[33;20m"
     red = "\x1b[31;20m"
     bold_red = "\x1b[31;1m"
@@ -12,8 +24,8 @@ class Formatter(logging.Formatter):
     format = "[%(name)s] | %(levelname)s -> %(message)s"
 
     FORMATS = {
-        logging.DEBUG: grey + format + reset,
-        logging.INFO: grey + format + reset,
+        logging.DEBUG: purple + format + reset,
+        logging.INFO: cyan + format + reset,
         logging.WARNING: yellow + format + reset,
         logging.ERROR: red + format + reset,
         logging.CRITICAL: bold_red + format + reset,
@@ -25,21 +37,26 @@ def format(self, record):
         return formatter.format(record)
 
 
-def mklog(name, level=logging.DEBUG):
+def mklog(name, level=base_log_level):
     logger = logging.getLogger(name)
     logger.setLevel(level)
-    # create console handler with a higher log level
-    ch = logging.StreamHandler()
-    ch.setLevel(logging.DEBUG)
 
-    ch.setFormatter(Formatter())
+    for handler in logger.handlers:
+        logger.removeHandler(handler)
 
+    ch = logging.StreamHandler()
+    ch.setLevel(level)
+    ch.setFormatter(Formatter())
     logger.addHandler(ch)
+
+    # Disable log propagation
+    logger.propagate = False
+
     return logger
 
 
 # - The main app logger
-log = mklog(__package__)
+log = mklog(__package__, base_log_level)
 
 
 def log_user(arg):
@@ -54,6 +71,10 @@ def blue_text(text):
     return f"\033[94m{text}\033[0m"
 
 
+def cyan_text(text):
+    return f"\033[96m{text}\033[0m"
+
+
 def get_label(label):
     words = re.findall(r"(?:^|[A-Z])[a-z]*", label)
     return " ".join(words).strip()

nodes/faceenhance.py

@@ -0,0 +1,235 @@
+import logging
+from gfpgan import GFPGANer
+import cv2
+import numpy as np
+import os
+from pathlib import Path
+import folder_paths
+from basicsr.utils import imwrite
+from PIL import Image
+from ..utils import pil2tensor, tensor2pil, np2tensor, tensor2np
+import torch
+from munch import Munch
+from ..log import NullWriter, log
+from comfy import model_management
+import comfy
+from typing import Tuple
+
+
+class LoadFaceEnhanceModel:
+    def __init__(self) -> None:
+        pass
+
+    @classmethod
+    def get_models_root(cls):
+        return Path(folder_paths.models_dir) / "upscale_models"
+
+    @classmethod
+    def get_models(cls):
+        models_path = cls.get_models_root()
+
+        return [
+            x
+            for x in models_path.iterdir()
+            if x.name.endswith(".pth")
+            and ("GFPGAN" in x.name or "RestoreFormer" in x.name)
+        ]
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "model_name": (
+                    [x.name for x in cls.get_models()],
+                    {"default": "None"},
+                ),
+                "upscale": ("INT", {"default": 2}),
+            },
+            "optional": {"bg_upsampler": ("UPSCALE_MODEL", {"default": None})},
+        }
+
+    RETURN_TYPES = ("FACEENHANCE_MODEL",)
+    RETURN_NAMES = ("model",)
+    FUNCTION = "load_model"
+    CATEGORY = "face"
+
+    def load_model(self, model_name, upscale=2, bg_upsampler=None):
+        basic = "RestoreFormer" not in model_name
+
+        root = self.get_models_root()
+
+        if bg_upsampler is not None:
+            log.warning(
+                f"Upscale value overridden to {bg_upsampler.scale} from bg_upsampler"
+            )
+            upscale = bg_upsampler.scale
+            bg_upsampler = BGUpscaleWrapper(bg_upsampler)
+
+        sys.stdout = NullWriter()
+        model = GFPGANer(
+            model_path=(root / model_name).as_posix(),
+            upscale=upscale,
+            arch="clean" if basic else "RestoreFormer",  # or original for v1.0 only
+            channel_multiplier=2,  # 1 for v1.0 only
+            bg_upsampler=bg_upsampler,
+        )
+
+        sys.stdout = sys.__stdout__
+        return (model,)
+
+
+class BGUpscaleWrapper:
+    def __init__(self, upscale_model) -> None:
+        self.upscale_model = upscale_model
+
+    def enhance(self, img: Image.Image, outscale=2):
+        device = model_management.get_torch_device()
+        self.upscale_model.to(device)
+
+        tile = 128 + 64
+        overlap = 8
+
+        imgt = np2tensor(img)
+        imgt = imgt.movedim(-1, -3).to(device)
+
+        steps = imgt.shape[0] * comfy.utils.get_tiled_scale_steps(
+            imgt.shape[3], imgt.shape[2], tile_x=tile, tile_y=tile, overlap=overlap
+        )
+
+        log.debug(f"Steps: {steps}")
+
+        pbar = comfy.utils.ProgressBar(steps)
+
+        s = comfy.utils.tiled_scale(
+            imgt,
+            lambda a: self.upscale_model(a),
+            tile_x=tile,
+            tile_y=tile,
+            overlap=overlap,
+            upscale_amount=self.upscale_model.scale,
+            pbar=pbar,
+        )
+
+        self.upscale_model.cpu()
+        s = torch.clamp(s.movedim(-3, -1), min=0, max=1.0)
+        return (tensor2np(s),)
+
+
+import sys
+
+
+class RestoreFace:
+    def __init__(self) -> None:
+        pass
+
+    RETURN_TYPES = ("IMAGE",)
+    FUNCTION = "restore"
+    CATEGORY = "face"
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "image": ("IMAGE",),
+                "model": ("FACEENHANCE_MODEL",),
+                # Input are aligned faces
+                "aligned": (["true", "false"], {"default": "false"}),
+                # Only restore the center face
+                "only_center_face": (["true", "false"], {"default": "false"}),
+                # Adjustable weights
+                "weight": ("FLOAT", {"default": 0.5}),
+                "save_tmp_steps": (["true", "false"], {"default": "true"}),
+            }
+        }
+
+    def do_restore(
+        self,
+        image: torch.Tensor,
+        model: GFPGANer,
+        aligned,
+        only_center_face,
+        weight,
+        save_tmp_steps,
+    ) -> torch.Tensor:
+        pimage = tensor2pil(image)
+        width, height = pimage.size
+
+        source_img = cv2.cvtColor(np.array(pimage), cv2.COLOR_RGB2BGR)
+
+        sys.stdout = NullWriter()
+        cropped_faces, restored_faces, restored_img = model.enhance(
+            source_img,
+            has_aligned=aligned,
+            only_center_face=only_center_face,
+            paste_back=True,
+            # TODO: weight has no effect in 1.3 and 1.4 (only tested these for now...)
+            weight=weight,
+        )
+        sys.stdout = sys.__stdout__
+        log.warning(f"Weight value has no effect for now. (value: {weight})")
+
+        if save_tmp_steps:
+            self.save_intermediate_images(cropped_faces, restored_faces, height, width)
+        output = None
+        if restored_img is not None:
+            output = Image.fromarray(cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB))
+            # imwrite(restored_img, save_restore_path)
+
+        return pil2tensor(output)
+
+    def restore(
+        self,
+        image: torch.Tensor,
+        model: GFPGANer,
+        aligned="false",
+        only_center_face="false",
+        weight=0.5,
+        save_tmp_steps="true",
+    ) -> Tuple[torch.Tensor]:
+        save_tmp_steps = save_tmp_steps == "true"
+        aligned = aligned == "true"
+        only_center_face = only_center_face == "true"
+
+        out = [
+            self.do_restore(
+                image[i], model, aligned, only_center_face, weight, save_tmp_steps
+            )
+            for i in range(image.size(0))
+        ]
+
+        return (torch.cat(out, dim=0),)
+
+    def get_step_image_path(self, step, idx):
+        (
+            full_output_folder,
+            filename,
+            counter,
+            _subfolder,
+            _filename_prefix,
+        ) = folder_paths.get_save_image_path(
+            f"{step}_{idx:03}",
+            folder_paths.temp_directory,
+        )
+        file = f"{filename}_{counter:05}_.png"
+
+        return os.path.join(full_output_folder, file)
+
+    def save_intermediate_images(self, cropped_faces, restored_faces, height, width):
+        for idx, (cropped_face, restored_face) in enumerate(
+            zip(cropped_faces, restored_faces)
+        ):
+            face_id = idx + 1
+            file = self.get_step_image_path("cropped_faces", face_id)
+            imwrite(cropped_face, file)
+
+            file = self.get_step_image_path("cropped_faces_restored", face_id)
+            imwrite(restored_face, file)
+
+            file = self.get_step_image_path("cropped_faces_compare", face_id)
+
+            # save comparison image
+            cmp_img = np.concatenate((cropped_face, restored_face), axis=1)
+            imwrite(cmp_img, file)
+
+
+__nodes__ = [RestoreFace, LoadFaceEnhanceModel]