nodes.py

# Code from https://huggingface.co/spaces/ohayonguy/PMRF/blob/main/app.py
# Some of the implementations below are adopted from
# https://huggingface.co/spaces/sczhou/CodeFormer and https://huggingface.co/spaces/wzhouxiff/RestoreFormerPlusPlus
import cv2
from tqdm import tqdm
import torch
from basicsr.archs.rrdbnet_arch import RRDBNet
from basicsr.utils import img2tensor, tensor2img
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
from realesrgan.utils import RealESRGANer
from .lightning_models.mmse_rectified_flow import MMSERectifiedFlow
import torchvision
import numpy as np
from PIL import Image
import os
import folder_paths

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

def set_realesrgan():
    use_half = False
    if torch.cuda.is_available():  # set False in CPU/MPS mode
        no_half_gpu_list = ["1650", "1660"]  # set False for GPUs that don't support f16
        if not True in [
            gpu in torch.cuda.get_device_name(0) for gpu in no_half_gpu_list
        ]:
            use_half = True

    model = RRDBNet(
        num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2,
    )
    upscale_models_path = os.path.join(folder_paths.models_dir, "upscale_models")
    realesrgan_path = os.path.join(upscale_models_path, "RealESRGAN_x2plus.pth")
    upsampler = RealESRGANer(
        scale=2,
        model_path=realesrgan_path,
        model=model,
        tile=400,
        tile_pad=40,
        pre_pad=0,
        half=use_half,
    )
    return upsampler

pmrf_path = os.path.join(folder_paths.models_dir, "pmrf")
upsampler = set_realesrgan()
pmrf = MMSERectifiedFlow.from_pretrained(pmrf_path).to(device=device)

def generate_reconstructions(pmrf_model, x, y, non_noisy_z0, num_flow_steps, device):
    source_dist_samples = pmrf_model.create_source_distribution_samples(
        x, y, non_noisy_z0
    )
    dt = (1.0 / num_flow_steps) * (1.0 - pmrf_model.hparams.eps)
    x_t_next = source_dist_samples.clone()
    t_one = torch.ones(x.shape[0], device=device)
    for i in tqdm(range(num_flow_steps)):
        num_t = (i / num_flow_steps) * (
            1.0 - pmrf_model.hparams.eps
        ) + pmrf_model.hparams.eps
        v_t_next = pmrf_model(x_t=x_t_next, t=t_one * num_t, y=y).to(x_t_next.dtype)
        x_t_next = x_t_next.clone() + v_t_next * dt
    return x_t_next.clip(0, 1)

def resize(img, size, interpolation):
    # From https://github.com/sczhou/CodeFormer/blob/master/facelib/utils/face_restoration_helper.py
    h, w = img.shape[0:2]
    scale = float(size) / float(min(h, w))
    h, w = int(round(h * scale)), int(round(w * scale))
    return cv2.resize(img, (w, h), interpolation=interpolation)

@torch.inference_mode()
def enhance_face(img, face_helper, num_flow_steps, scale=2, interpolation=cv2.INTER_LANCZOS4):
    face_helper.clean_all()
    face_helper.read_image(img)
    face_helper.input_img = resize(face_helper.input_img, 640, interpolation)
    face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
    face_helper.align_warp_face()

    # face restoration
    for i, cropped_face in tqdm(enumerate(face_helper.cropped_faces)):
        cropped_face_t = img2tensor(cropped_face / 255.0, bgr2rgb=True, float32=True)
        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)

        output = generate_reconstructions(
            pmrf,
            torch.zeros_like(cropped_face_t),
            cropped_face_t,
            None,
            num_flow_steps,
            device,
        )
        restored_face = tensor2img(
            output.to(torch.float32).squeeze(0), rgb2bgr=True, min_max=(0, 1)
        )
        restored_face = restored_face.astype("uint8")
        face_helper.add_restored_face(restored_face)

    # upsample the background
    # Now only support RealESRGAN for upsampling background
    bg_img = upsampler.enhance(img, outscale=scale)[0]
    face_helper.get_inverse_affine(None)
    # paste each restored face to the input image
    restored_img = face_helper.paste_faces_to_input_image(upsample_img=bg_img)
    return face_helper.cropped_faces, face_helper.restored_faces, restored_img

@torch.inference_mode()
def inference(
    imgs,
    scale,
    num_flow_steps,
    seed,
    interpolation,
):
    torch.manual_seed(seed)
    if interpolation == "lanczos4":
        interpolation = cv2.INTER_LANCZOS4
    elif interpolation == "nearest":
        interpolation = cv2.INTER_NEAREST
    elif interpolation == "linear":
        interpolation = cv2.INTER_LINEAR
    elif interpolation == "cubic":
        interpolation = cv2.INTER_CUBIC
    elif interpolation == "area":
        interpolation = cv2.INTER_AREA
    elif interpolation == "linear_exact":
        interpolation = cv2.INTER_LINEAR_EXACT
    elif interpolation == "nearest_exact":
        interpolation = cv2.INTER_NEAREST_EXACT
    imgs_output = []
    for img in imgs:
        img = img.permute(2, 0, 1)
        img = torchvision.transforms.functional.to_pil_image(img.clamp(0, 1)).convert("RGB")
        img = np.array(img)
        img = img[:, :, ::-1].copy()
        h, w = img.shape[0:2]
        size = min(h, w)
        
        face_scale = scale*(size/640)
        face_scale = face_scale if face_scale < scale else scale
        face_scale = face_scale if face_scale > 1.0 else 1.0

        face_helper = FaceRestoreHelper(
            face_scale,
            face_size=512,
            crop_ratio=(1, 1),
            det_model="retinaface_resnet50",
            save_ext="png",
            use_parse=True,
            device=device,
            model_rootpath=None,
        )

        cropped_face, restored_faces, restored_img = enhance_face(
            img,
            face_helper,
            num_flow_steps=num_flow_steps,
            scale=face_scale,
            interpolation=interpolation,
        )

        output = restored_img
        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
        output = resize(output, size*scale, interpolation)

        torch.cuda.empty_cache()
        output = torchvision.transforms.functional.pil_to_tensor(Image.fromarray(output)).to(torch.float32) / 255.0
        output = output.permute(1, 2, 0)
        imgs_output.append(output[None,])
    return (torch.cat(tuple(imgs_output), dim=0),)

class PMRF:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "images": ("IMAGE",),
                "scale": ("FLOAT", {"default": 1.0, "min": 1.0, "max": 40.0, "step": 0.1}),
                "num_steps": ("INT", {"default": 25, "min": 1, "max": 400, "step": 1}),
                "seed": ("INT", {"default": 123, "min": 0, "max": 2**32, "step": 1}),
                "interpolation": (["lanczos4", "nearest", "linear", "cubic", "area", "linear_exact", "nearest_exact"],)
            },
        }

    RETURN_TYPES = ("IMAGE",)
    RETURN_NAMES = ("images", )
    FUNCTION = "pmrf"
    CATEGORY = "PMRF"

    def pmrf(self, images, scale, num_steps, seed, interpolation):
        return inference(images, scale, num_steps, seed, interpolation)
    
NODE_CLASS_MAPPINGS = {
    "PMRF": PMRF,
}
NODE_DISPLAY_NAME_MAPPINGS = {
    "PMRF": "PMRF",
}