utils/inference.py

import os
from typing import overload, Generator, Dict
from argparse import Namespace

import numpy as np
import torch
from PIL import Image
from omegaconf import OmegaConf

from model.cldm import ControlLDM
from model.gaussian_diffusion import Diffusion
from model.bsrnet import RRDBNet
from model.scunet import SCUNet
from model.swinir import SwinIR
from utils.common import instantiate_from_config, load_file_from_url, count_vram_usage
from utils.face_restoration_helper import FaceRestoreHelper
from utils.helpers import (
    Pipeline,
    BSRNetPipeline, SwinIRPipeline, SCUNetPipeline,
    bicubic_resize
)
from utils.cond_fn import MSEGuidance, WeightedMSEGuidance


MODELS = {
    ### stage_1 model weights
    "bsrnet": "https://github.com/cszn/KAIR/releases/download/v1.0/BSRNet.pth",
    # the following checkpoint is up-to-date, but we use the old version in our paper
    # "swinir_face": "https://github.com/zsyOAOA/DifFace/releases/download/V1.0/General_Face_ffhq512.pth",
    "swinir_face": "https://huggingface.co/lxq007/DiffBIR/resolve/main/face_swinir_v1.ckpt",
    "scunet_psnr": "https://github.com/cszn/KAIR/releases/download/v1.0/scunet_color_real_psnr.pth",
    "swinir_general": "https://huggingface.co/lxq007/DiffBIR/resolve/main/general_swinir_v1.ckpt",
    ### stage_2 model weights
    "sd_v21": "https://huggingface.co/stabilityai/stable-diffusion-2-1-base/resolve/main/v2-1_512-ema-pruned.ckpt",
    "v1_face": "https://huggingface.co/lxq007/DiffBIR-v2/resolve/main/v1_face.pth",
    "v1_general": "https://huggingface.co/lxq007/DiffBIR-v2/resolve/main/v1_general.pth",
    "v2": "https://huggingface.co/lxq007/DiffBIR-v2/resolve/main/v2.pth"
}


def load_model_from_url(url: str) -> Dict[str, torch.Tensor]:
    sd_path = load_file_from_url(url, model_dir="weights")
    sd = torch.load(sd_path, map_location="cpu")
    if "state_dict" in sd:
        sd = sd["state_dict"]
    if list(sd.keys())[0].startswith("module"):
        sd = {k[len("module."):]: v for k, v in sd.items()}
    return sd


class InferenceLoop:

    def __init__(self, args: Namespace) -> "InferenceLoop":
        self.args = args
        self.loop_ctx = {}
        self.pipeline: Pipeline = None
        self.init_stage1_model()
        self.init_stage2_model()
        self.init_cond_fn()
        self.init_pipeline()

    @overload
    def init_stage1_model(self) -> None:
        ...

    @count_vram_usage
    def init_stage2_model(self) -> None:
        ### load uent, vae, clip
        self.cldm: ControlLDM = instantiate_from_config(OmegaConf.load("configs/inference/my_cldm.yaml"))
        sd = load_model_from_url(MODELS["sd_v21"])
        unused = self.cldm.load_pretrained_sd(sd)
        print(f"strictly load pretrained sd_v2.1, unused weights: {unused}")
        ### load controlnet
        if self.args.version == "v1":
            if self.args.task == "fr":
                control_sd = load_model_from_url(MODELS["v1_face"])
            elif self.args.task == "sr":
                control_sd = load_model_from_url(MODELS["v1_general"])
            else:
                raise ValueError(f"DiffBIR v1 doesn't support task: {self.args.task}, please use v2 by passsing '--version v2'")
        else:
            control_sd = load_model_from_url(MODELS["v2"])
        self.cldm.load_controlnet_from_ckpt(control_sd)
        print(f"strictly load controlnet weight")
        self.cldm.eval().to(self.args.device)
        ### load diffusion
        self.diffusion: Diffusion = instantiate_from_config(OmegaConf.load("configs/inference/diffusion.yaml"))
        self.diffusion.to(self.args.device)

    def init_cond_fn(self) -> None:
        if not self.args.guidance:
            self.cond_fn = None
            return
        if self.args.g_loss == "mse":
            cond_fn_cls = MSEGuidance
        elif self.args.g_loss == "w_mse":
            cond_fn_cls = WeightedMSEGuidance
        else:
            raise ValueError(self.args.g_loss)
        self.cond_fn = cond_fn_cls(
            scale=self.args.g_scale, t_start=self.args.g_start, t_stop=self.args.g_stop,
            space=self.args.g_space, repeat=self.args.g_repeat
        )

    @overload
    def init_pipeline(self) -> None:
        ...

    def setup(self) -> None:
        self.output_dir = self.args.output
        os.makedirs(self.output_dir, exist_ok=True)

    def lq_loader(self) -> Generator[np.ndarray, None, None]:
        img_exts = [".png", ".jpg", ".jpeg"]
        if os.path.isdir(self.args.input):
            file_names = sorted([
                file_name for file_name in os.listdir(self.args.input) if os.path.splitext(file_name)[-1] in img_exts
            ])
            file_paths = [os.path.join(self.args.input, file_name) for file_name in file_names]
        else:
            assert os.path.splitext(self.args.input)[-1] in img_exts
            file_paths = [self.args.input]

        def _loader() -> Generator[np.ndarray, None, None]:
            for file_path in file_paths:
                ### load lq
                lq = np.array(Image.open(file_path).convert("RGB"))
                print(f"load lq: {file_path}")
                ### set context for saving results
                self.loop_ctx["file_stem"] = os.path.splitext(os.path.basename(file_path))[0]
                for i in range(self.args.n_samples):
                    self.loop_ctx["repeat_idx"] = i
                    yield lq

        return _loader
    
    def after_load_lq(self, lq: np.ndarray) -> np.ndarray:
        return lq

    @torch.no_grad()
    def run(self) -> None:
        self.setup()
        # We don't support batch processing since input images may have different size
        loader = self.lq_loader()
        for i, lq in enumerate(loader()):
            lq = self.after_load_lq(lq)
            sample = self.pipeline.run(
                lq[None], self.args.steps, 1.0, self.args.tiled,
                self.args.tile_size, self.args.tile_stride,
                self.args.pos_prompt, self.args.neg_prompt, self.args.cfg_scale,
                self.args.better_start,
                index=i, input=self.args.input
            )[0]
            self.save(sample)

    def save(self, sample: np.ndarray) -> None:
        file_stem, repeat_idx = self.loop_ctx["file_stem"], self.loop_ctx["repeat_idx"]
        file_name = f"{file_stem}_{repeat_idx}.png" if self.args.n_samples > 1 else f"{file_stem}.png"
        save_path = os.path.join(self.args.output, file_name)
        Image.fromarray(sample).save(save_path)
        print(f"save result to {save_path}")


class BSRInferenceLoop(InferenceLoop):

    @count_vram_usage
    def init_stage1_model(self) -> None:
        self.bsrnet: RRDBNet = instantiate_from_config(OmegaConf.load("configs/inference/bsrnet.yaml"))
        sd = load_model_from_url(MODELS["bsrnet"])
        self.bsrnet.load_state_dict(sd, strict=True)
        self.bsrnet.eval().to(self.args.device)

    def init_pipeline(self) -> None:
        self.pipeline = BSRNetPipeline(self.bsrnet, self.cldm, self.diffusion, self.cond_fn, self.args.device, self.args.upscale)


class BFRInferenceLoop(InferenceLoop):

    @count_vram_usage
    def init_stage1_model(self) -> None:
        self.swinir_face: SwinIR = instantiate_from_config(OmegaConf.load("configs/inference/swinir.yaml"))
        sd = load_model_from_url(MODELS["swinir_face"])
        self.swinir_face.load_state_dict(sd, strict=True)
        self.swinir_face.eval().to(self.args.device)

    def init_pipeline(self) -> None:
        self.pipeline = SwinIRPipeline(self.swinir_face, self.cldm, self.diffusion, self.cond_fn, self.args.device)

    def after_load_lq(self, lq: np.ndarray) -> np.ndarray:
        # For BFR task, super resolution is achieved by directly upscaling lq
        return bicubic_resize(lq, self.args.upscale)


class BIDInferenceLoop(InferenceLoop):

    @count_vram_usage
    def init_stage1_model(self) -> None:
        self.scunet_psnr: SCUNet = instantiate_from_config(OmegaConf.load("configs/inference/scunet.yaml"))
        sd = load_model_from_url(MODELS["scunet_psnr"])
        self.scunet_psnr.load_state_dict(sd, strict=True)
        self.scunet_psnr.eval().to(self.args.device)

    def init_pipeline(self) -> None:
        self.pipeline = SCUNetPipeline(self.scunet_psnr, self.cldm, self.diffusion, self.cond_fn, self.args.device)

    def after_load_lq(self, lq: np.ndarray) -> np.ndarray:
        # For BID task, super resolution is achieved by directly upscaling lq
        return bicubic_resize(lq, self.args.upscale)


class V1InferenceLoop(InferenceLoop):

    @count_vram_usage
    def init_stage1_model(self) -> None:
        self.swinir: SwinIR = instantiate_from_config(OmegaConf.load("configs/inference/swinir.yaml"))
        if self.args.task == "fr":
            sd = load_model_from_url(MODELS["swinir_face"])
        elif self.args.task == "sr":
            sd = load_model_from_url(MODELS["swinir_general"])
        else:
            raise ValueError(f"DiffBIR v1 doesn't support task: {self.args.task}, please use v2 by passsing '--version v2'")
        self.swinir.load_state_dict(sd, strict=True)
        self.swinir.eval().to(self.args.device)

    def init_pipeline(self) -> None:
        self.pipeline = SwinIRPipeline(self.swinir, self.cldm, self.diffusion, self.cond_fn, self.args.device)

    def after_load_lq(self, lq: np.ndarray) -> np.ndarray:
        # For BFR task, super resolution is achieved by directly upscaling lq
        return bicubic_resize(lq, self.args.upscale)


class UnAlignedBFRInferenceLoop(InferenceLoop):

    @count_vram_usage
    def init_stage1_model(self) -> None:
        self.bsrnet: RRDBNet = instantiate_from_config(OmegaConf.load("configs/inference/bsrnet.yaml"))
        sd = load_model_from_url(MODELS["bsrnet"])
        self.bsrnet.load_state_dict(sd, strict=True)
        self.bsrnet.eval().to(self.args.device)

        self.swinir_face: SwinIR = instantiate_from_config(OmegaConf.load("configs/inference/swinir.yaml"))
        sd = load_model_from_url(MODELS["swinir_face"])
        self.swinir_face.load_state_dict(sd, strict=True)
        self.swinir_face.eval().to(self.args.device)

    def init_pipeline(self) -> None:
        self.pipes = {
            "bg": BSRNetPipeline(self.bsrnet, self.cldm, self.diffusion, self.cond_fn, self.args.device, self.args.upscale),
            "face": SwinIRPipeline(self.swinir_face, self.cldm, self.diffusion, self.cond_fn, self.args.device)
        }
        self.pipeline = self.pipes["face"]

    def setup(self) -> None:
        super().setup()
        self.cropped_face_dir = os.path.join(self.args.output, "cropped_faces")
        os.makedirs(self.cropped_face_dir, exist_ok=True)
        self.restored_face_dir = os.path.join(self.args.output, "restored_faces")
        os.makedirs(self.restored_face_dir, exist_ok=True)
        self.restored_bg_dir = os.path.join(self.args.output, "restored_backgrounds")
        os.makedirs(self.restored_bg_dir, exist_ok=True)

    def lq_loader(self) -> Generator[np.ndarray, None, None]:
        base_loader = super().lq_loader()
        self.face_helper = FaceRestoreHelper(
            device=self.args.device,
            upscale_factor=1,
            face_size=512,
            use_parse=True,
            det_model="retinaface_resnet50"
        )
        
        def _loader() -> Generator[np.ndarray, None, None]:
            for lq in base_loader():
                ### set input image
                self.face_helper.clean_all()
                upscaled_bg = bicubic_resize(lq, self.args.upscale)
                self.face_helper.read_image(upscaled_bg)
                ### get face landmarks for each face
                self.face_helper.get_face_landmarks_5(resize=640, eye_dist_threshold=5)
                self.face_helper.align_warp_face()
                print(f"detect {len(self.face_helper.cropped_faces)} faces")
                ### restore each face (has been upscaeled)
                for i, lq_face in enumerate(self.face_helper.cropped_faces):
                    self.loop_ctx["is_face"] = True
                    self.loop_ctx["face_idx"] = i
                    self.loop_ctx["cropped_face"] = lq_face
                    yield lq_face
                ### restore background (hasn't been upscaled)
                self.loop_ctx["is_face"] = False
                yield lq
        
        return _loader

    def after_load_lq(self, lq: np.ndarray) -> np.ndarray:
        if self.loop_ctx["is_face"]:
            self.pipeline = self.pipes["face"]
        else:
            self.pipeline = self.pipes["bg"]
        return lq

    def save(self, sample: np.ndarray) -> None:
        file_stem, repeat_idx = self.loop_ctx["file_stem"], self.loop_ctx["repeat_idx"]
        if self.loop_ctx["is_face"]:
            face_idx = self.loop_ctx["face_idx"]
            file_name = f"{file_stem}_{repeat_idx}_face_{face_idx}.png"
            Image.fromarray(sample).save(os.path.join(self.restored_face_dir, file_name))

            cropped_face = self.loop_ctx["cropped_face"]
            Image.fromarray(cropped_face).save(os.path.join(self.cropped_face_dir, file_name))

            self.face_helper.add_restored_face(sample)
        else:
            self.face_helper.get_inverse_affine()
            # paste each restored face to the input image
            restored_img = self.face_helper.paste_faces_to_input_image(
                upsample_img=sample
            )
            file_name = f"{file_stem}_{repeat_idx}.png"
            Image.fromarray(sample).save(os.path.join(self.restored_bg_dir, file_name))
            Image.fromarray(restored_img).save(os.path.join(self.output_dir, file_name))