DDCM_blind_face_image_restoration.py

import os
from functools import partial

import cv2
import gradio as gr
import spaces
from util.file import generate_binary_file, load_numpy_from_binary_bitwise
import torch
import yaml
from util.basicsr_img_util import img2tensor, tensor2img
from facexlib.utils.face_restoration_helper import FaceRestoreHelper
from torchvision.transforms.functional import resize

from guided_diffusion.gaussian_diffusion import create_sampler
from guided_diffusion.swinir import SwinIR
from guided_diffusion.unet import create_model


def create_swinir_model(ckpt_path):
    cfg = {
        'in_channels': 3,
        'out_channels': 3,
        'embed_dim': 180,
        'depths': [6, 6, 6, 6, 6, 6, 6, 6],
        'num_heads': [6, 6, 6, 6, 6, 6, 6, 6],
        'resi_connection': '1conv',
        'sf': 8
    }
    mmse_model = SwinIR(
        img_size=64,
        patch_size=1,
        in_chans=cfg['in_channels'],
        num_out_ch=cfg['out_channels'],
        embed_dim=cfg['embed_dim'],
        depths=cfg['depths'],
        num_heads=cfg['num_heads'],
        window_size=8,
        mlp_ratio=2,
        sf=cfg['sf'],
        img_range=1.0,
        upsampler="nearest+conv",
        resi_connection=cfg['resi_connection'],
        unshuffle=True,
        unshuffle_scale=8
    )
    ckpt = torch.load(ckpt_path, map_location="cpu")

    if 'params_ema' in ckpt:
        mmse_model.load_state_dict(ckpt['params_ema'])
    else:
        state_dict = ckpt['state_dict']
        state_dict = {layer_name.replace('model.', ''): weights for layer_name, weights in
                      state_dict.items()}
        state_dict = {layer_name.replace('module.', ''): weights for layer_name, weights in
                      state_dict.items()}
        mmse_model.load_state_dict(state_dict)
    for param in mmse_model.parameters():
        param.requires_grad = False
    return mmse_model


ffhq_diffusion_model = "./guided_diffusion/iddpm_ffhq512_ema500000.pth"
mmse_model_ckpt = "./guided_diffusion/swinir_restoration512_L1.pth"

if not os.path.exists(ffhq_diffusion_model):
    os.system(
        "wget https://github.com/zsyOAOA/DifFace/releases/download/V1.0/iddpm_ffhq512_ema500000.pth -O ./guided_diffusion/iddpm_ffhq512_ema500000.pth"
    )
if not os.path.exists(mmse_model_ckpt):
    os.system(
        "wget https://github.com/zsyOAOA/DifFace/releases/download/V1.0/swinir_restoration512_L1.pth -O ./guided_diffusion/swinir_restoration512_L1.pth"
    )


def load_yaml(file_path: str) -> dict:
    with open(file_path) as f:
        config = yaml.load(f, Loader=yaml.FullLoader)
    return config


model_config = './guided_diffusion/ffhq512_model_config.yaml'
diffusion_config = './guided_diffusion/diffusion_config.yaml'
model_config = load_yaml(model_config)
diffusion_config = load_yaml(diffusion_config)

models = {
    'main_model': create_model(**model_config),
    'mmse_model': create_swinir_model('./guided_diffusion/swinir_restoration512_L1.pth')
}
models['main_model'].eval()
models['mmse_model'].eval()


@torch.no_grad()
@spaces.GPU(duration=80)
def generate_reconstruction(degraded_face_img, K, T, iqa_metric, iqa_coef, loaded_indices):
    assert iqa_metric in ['niqe', 'clipiqa+', 'topiq_nr-face']
    diffusion_config['timestep_respacing'] = T
    sampler = create_sampler(**diffusion_config)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = models['main_model'].to(device)
    mmse_model = models['mmse_model'].to(device)

    sample_fn = partial(sampler.p_sample_loop_blind_restoration, model=model, num_opt_noises=K,
                        eta=1.0, iqa_metric=iqa_metric, iqa_coef=iqa_coef)

    if degraded_face_img is not None:
        mmse_img = mmse_model(degraded_face_img).clip(0, 1) * 2 - 1
        x_start = torch.randn(mmse_img.shape, device=device)
    else:
        mmse_img = None
        x_start = torch.randn(1, 3, 512, 512, device=device)
    restored_face, indices = sample_fn(x_start=x_start, mmse_img=mmse_img, loaded_indices=loaded_indices)

    return restored_face, indices


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


@torch.no_grad()
@spaces.GPU(duration=80)
def enhance_faces(img, face_helper, has_aligned, K, T, iqa_metric, iqa_coef, loaded_indices):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    face_helper.clean_all()
    if has_aligned:  # The inputs are already aligned
        img = cv2.resize(img, (512, 512), interpolation=cv2.INTER_LINEAR)
        face_helper.cropped_faces = [img]
    else:
        face_helper.read_image(img)
        face_helper.input_img = resize(face_helper.input_img, 640)
        face_helper.get_face_landmarks_5(only_center_face=False, eye_dist_threshold=5)
        face_helper.align_warp_face()
    if len(face_helper.cropped_faces) == 0:
        raise gr.Error("Could not identify any face in the image.")
    if has_aligned and len(face_helper.cropped_faces) > 1:
        raise gr.Error(
            "You marked that the input image is aligned, but multiple faces were detected."
        )
    restored_faces = []
    generated_indices = []
    for i, cropped_face in 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)
        cur_loaded_indices = loaded_indices[i] if loaded_indices is not None else None

        output, indices = generate_reconstruction(
            cropped_face_t,
            K,
            T,
            iqa_metric,
            iqa_coef,
            cur_loaded_indices
        )

        restored_face = tensor2img(
            output.to(torch.float32).squeeze(0), rgb2bgr=False, min_max=(-1, 1)
        )

        restored_face = restored_face.astype("uint8")
        restored_faces.append(restored_face),
        generated_indices.append(indices)
    return restored_faces, generated_indices


@torch.no_grad()
@spaces.GPU()
def decompress_face(K, T, iqa_metric, iqa_coef, loaded_indices):
    assert loaded_indices is not None

    output, indices = generate_reconstruction(
        None,
        K,
        T,
        iqa_metric,
        iqa_coef,
        loaded_indices
    )

    restored_face = tensor2img(
        output.to(torch.float32).squeeze(0), rgb2bgr=False, min_max=(-1, 1)
    ).astype("uint8")

    return restored_face, loaded_indices

@torch.no_grad()
@spaces.GPU(duration=80)
def inference(
        img,
        T,
        K,
        iqa_metric,
        iqa_coef,
        aligned,
        bitstream=None,
        progress=gr.Progress(track_tqdm=True),
):
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    iqa_metric_to_pyiqa_name = {
        'NIQE': 'niqe',
        'TOPIQ': 'topiq_nr-face',
        'CLIP-IQA': 'clipiqa+'
    }
    iqa_metric = iqa_metric_to_pyiqa_name[iqa_metric]
    indices = load_numpy_from_binary_bitwise(bitstream, K, T, 'ffhq', T)
    if indices is not None:
        indices = indices.to(device)

    if img is not None:
        img = cv2.imread(img, cv2.IMREAD_COLOR)
        h, w = img.shape[0:2]
        if h > 4500 or w > 4500:
            raise gr.Error("Image size too large.")

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

        x, indices = enhance_faces(
            img, face_helper, aligned, K=K, T=T, iqa_metric=iqa_metric, iqa_coef=iqa_coef,
            loaded_indices=indices,
        )
    else:
        x, indices = decompress_face(
            K=K, T=T, iqa_metric=iqa_metric, iqa_coef=iqa_coef, loaded_indices=indices,
        )

    torch.cuda.empty_cache()

    if bitstream is None:
        indices = [generate_binary_file(index.numpy(), K, T, 'ffhq') for index in indices]
        return x, indices
    return x