first commit

app.py

@@ -0,0 +1,170 @@
+import os
+
+import cv2
+import gradio as gr
+import torch
+from basicsr.archs.srvgg_arch import SRVGGNetCompact
+from basicsr.utils import img2tensor, tensor2img
+from facexlib.utils.face_restoration_helper import FaceRestoreHelper
+from realesrgan.utils import RealESRGANer
+import spaces
+
+from lightning_models.mmse_rectified_flow import MMSERectifiedFlow
+
+torch.set_grad_enabled(False)
+
+if os.getenv('SPACES_ZERO_GPU') == "true":
+    os.environ['SPACES_ZERO_GPU'] = "1"
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+if not os.path.exists('pretrained_models'):
+    os.makedirs('pretrained_models')
+realesr_model_path = 'pretrained_models/RealESRGAN_x4plus.pth'
+if not os.path.exists(realesr_model_path):
+    os.system(
+        "wget https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.5.0/realesr-general-x4v3.pth -O experiments/pretrained_models/RealESRGAN_x4plus.pth")
+
+pmrf_model_path = 'blind_face_restoration_pmrf.ckpt'
+
+# background enhancer with RealESRGAN
+model = SRVGGNetCompact(num_in_ch=3, num_out_ch=3, num_feat=64, num_conv=32, upscale=4, act_type='prelu')
+half = True if torch.cuda.is_available() else False
+upsampler = RealESRGANer(scale=4, model_path=realesr_model_path, model=model, tile=0, tile_pad=10, pre_pad=0, half=half)
+
+pmrf = MMSERectifiedFlow.load_from_checkpoint('./blind_face_restoration_pmrf.ckpt',
+                                              mmse_model_arch='swinir_L',
+                                              mmse_model_ckpt_path=None,
+                                              map_location='cpu').to(device)
+
+os.makedirs('output', exist_ok=True)
+
+
+@torch.inference_mode()
+@spaces.GPU()
+def enhance_face(img, face_helper, has_aligned, only_center_face=False, paste_back=True, scale=2):
+    face_helper.clean_all()
+
+    if has_aligned:  # the inputs are already aligned
+        img = cv2.resize(img, (512, 512))
+        face_helper.cropped_faces = [img]
+    else:
+        face_helper.read_image(img)
+        face_helper.get_face_landmarks_5(only_center_face=only_center_face, eye_dist_threshold=5)
+        # eye_dist_threshold=5: skip faces whose eye distance is smaller than 5 pixels
+        # TODO: even with eye_dist_threshold, it will still introduce wrong detections and restorations.
+        # align and warp each face
+        face_helper.align_warp_face()
+
+    # face restoration
+    for cropped_face in face_helper.cropped_faces:
+        # prepare data
+        cropped_face_t = img2tensor(cropped_face / 255., bgr2rgb=True, float32=True)
+        cropped_face_t = cropped_face_t.unsqueeze(0).to(device)
+
+        try:
+            dummy_x = torch.zeros_like(cropped_face_t)
+            output = pmrf.generate_reconstructions(dummy_x, cropped_face_t, None, 25, device)
+            restored_face = tensor2img(output.squeeze(0), rgb2bgr=True, min_max=(0, 1))
+        except RuntimeError as error:
+            print(f'\tFailed inference for RestoreFormer: {error}.')
+            restored_face = cropped_face
+
+        restored_face = restored_face.astype('uint8')
+        face_helper.add_restored_face(restored_face)
+
+    if not has_aligned and paste_back:
+        # upsample the background
+        if upsampler is not None:
+            # Now only support RealESRGAN for upsampling background
+            bg_img = upsampler.enhance(img, outscale=scale)[0]
+        else:
+            bg_img = None
+
+        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
+    else:
+        return face_helper.cropped_faces, face_helper.restored_faces, None
+
+
+@torch.inference_mode()
+@spaces.GPU()
+def inference(img, aligned, scale):
+    if scale > 4:
+        scale = 4  # avoid too large scale value
+    try:
+
+        extension = os.path.splitext(os.path.basename(str(img)))[1]
+        img = cv2.imread(img, cv2.IMREAD_UNCHANGED)
+        if len(img.shape) == 3 and img.shape[2] == 4:
+            img_mode = 'RGBA'
+        elif len(img.shape) == 2:  # for gray inputs
+            img_mode = None
+            img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
+        else:
+            img_mode = None
+
+        h, w = img.shape[0:2]
+        if h > 3500 or w > 3500:
+            print('Image size too large.')
+            return None, None
+
+        if h < 300:
+            img = cv2.resize(img, (w * 2, h * 2), interpolation=cv2.INTER_LANCZOS4)
+
+        face_helper = FaceRestoreHelper(
+            scale,
+            face_size=512,
+            crop_ratio=(1, 1),
+            det_model='retinaface_resnet50',
+            save_ext='png',
+            use_parse=True,
+            device=device,
+            model_rootpath=None)
+
+        try:
+            has_aligned = True if aligned == 'aligned' else False
+            _, restored_aligned, restored_img = enhance_face(img, face_helper, has_aligned, only_center_face=False,
+                                                             paste_back=True)
+            if has_aligned:
+                output = restored_aligned[0]
+            else:
+                output = restored_img
+        except RuntimeError as error:
+            print('Error', error)
+
+        try:
+            if scale != 2:
+                interpolation = cv2.INTER_AREA if scale < 2 else cv2.INTER_LANCZOS4
+                h, w = img.shape[0:2]
+                output = cv2.resize(output, (int(w * scale / 2), int(h * scale / 2)), interpolation=interpolation)
+        except Exception as error:
+            print('wrong scale input.', error)
+        if img_mode == 'RGBA':  # RGBA images should be saved in png format
+            extension = 'png'
+        else:
+            extension = 'jpg'
+        save_path = f'output/out.{extension}'
+        cv2.imwrite(save_path, output)
+
+        output = cv2.cvtColor(output, cv2.COLOR_BGR2RGB)
+        return output, save_path
+    except Exception as error:
+        print('global exception', error)
+        return None, None
+
+
+css = r"""
+"""
+
+demo = gr.Interface(
+    inference, [
+        gr.Image(type="filepath", label="Input"),
+        gr.Radio(['aligned', 'unaligned'], type="value", value='unaligned', label='Image Alignment'),
+        gr.Number(label="Rescaling factor", value=2),
+    ], [
+        gr.Image(type="numpy", label="Output (The whole image)"),
+        gr.File(label="Download the output image")
+    ],
+)

arch/__init__.py

@@ -0,0 +1,2 @@
+from arch.hourglass.image_transformer_v2 import ImageTransformerDenoiserModelV2
+from arch.swinir.swinir import SwinIR

lightning_models/mmse_rectified_flow.py

@@ -0,0 +1,317 @@
+import os
+from contextlib import contextmanager, nullcontext
+
+import torch
+import wandb
+from pytorch_lightning import LightningModule
+from torch.nn.functional import mse_loss
+from torch.nn.functional import sigmoid
+from torch.optim import AdamW
+from torch_ema import ExponentialMovingAverage as EMA
+from torchmetrics.image import FrechetInceptionDistance, InceptionScore
+from torchvision.transforms.functional import to_pil_image
+from torchvision.utils import save_image
+
+from utils.create_arch import create_arch
+from utils.img_utils import create_grid
+from huggingface_hub import PyTorchModelHubMixin
+
+
+
+class MMSERectifiedFlow(LightningModule,
+                        PyTorchModelHubMixin,
+                        pipeline_tag="image-to-image",
+                        license="mit",
+                        ):
+    def __init__(self,
+                 stage,
+                 arch,
+                 conditional=False,
+                 mmse_model_ckpt_path=None,
+                 mmse_model_arch=None,
+                 lr=5e-4,
+                 weight_decay=1e-3,
+                 betas=(0.9, 0.95),
+                 mmse_noise_std=0.1,
+                 num_flow_steps=50,
+                 ema_decay=0.9999,
+                 eps=0.0,
+                 t_schedule='stratified_uniform',
+                 *args,
+                 **kwargs
+                 ):
+        super().__init__()
+        self.save_hyperparameters(logger=False)
+
+        if stage == 'flow':
+            if conditional:
+                condition_channels = 3
+            else:
+                condition_channels = 0
+            if mmse_model_arch is None and 'colorization' in kwargs and kwargs['colorization']:
+                condition_channels //= 3
+            self.model = create_arch(arch, condition_channels)
+            self.mmse_model = create_arch(mmse_model_arch, 0) if mmse_model_arch is not None else None
+            if mmse_model_ckpt_path is not None:
+                ckpt = torch.load(mmse_model_ckpt_path, map_location="cpu")
+                if mmse_model_arch is None:
+                    mmse_model_arch = ckpt['hyper_parameters']['arch']
+                self.mmse_model = create_arch(mmse_model_arch, 0)
+                if 'ema' in ckpt:
+                    # ema_decay doesn't affect anything here, because we are doing load_state_dict
+                    mmse_ema = EMA(self.mmse_model.parameters(), decay=ema_decay)
+                    mmse_ema.load_state_dict(ckpt['ema'])
+                    mmse_ema.copy_to()
+                elif 'params_ema' in ckpt:
+                    self.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()}
+                    self.mmse_model.load_state_dict(state_dict)
+                for param in self.mmse_model.parameters():
+                    param.requires_grad = False
+                self.mmse_model.eval()
+        else:
+            assert stage == 'mmse' or stage == 'naive_flow'
+            assert not conditional
+            self.model = create_arch(arch, 0)
+            self.mmse_model = None
+        if 'flow' in stage:
+            self.fid = FrechetInceptionDistance(reset_real_features=True, normalize=True)
+            self.inception_score = InceptionScore(normalize=True)
+
+        self.ema = EMA(self.model.parameters(), decay=ema_decay) if self.ema_wanted else None
+        self.test_results_path = None
+
+    @property
+    def ema_wanted(self):
+        return self.hparams.ema_decay != -1
+
+    def on_save_checkpoint(self, checkpoint: dict) -> None:
+        if self.ema_wanted:
+            checkpoint['ema'] = self.ema.state_dict()
+        return super().on_save_checkpoint(checkpoint)
+
+    def on_load_checkpoint(self, checkpoint: dict) -> None:
+        if self.ema_wanted:
+            self.ema.load_state_dict(checkpoint['ema'])
+        return super().on_load_checkpoint(checkpoint)
+
+    def on_before_zero_grad(self, optimizer) -> None:
+        if self.ema_wanted:
+            self.ema.update(self.model.parameters())
+        return super().on_before_zero_grad(optimizer)
+
+    def to(self, *args, **kwargs):
+        if self.ema_wanted:
+            self.ema.to(*args, **kwargs)
+        return super().to(*args, **kwargs)
+
+    # This will use the contextmanager of ema, to copy the EMA weights to the flow model during validation, and then restore them for training.
+    @contextmanager
+    def maybe_ema(self):
+        ema = self.ema
+        ctx = nullcontext if ema is None else ema.average_parameters
+        yield ctx
+
+    def forward_mmse(self, y):
+        return self.model(y).clip(0, 1)
+
+    def forward_flow(self, x_t, t, y=None):
+        if self.hparams.conditional:
+            if self.mmse_model is not None:
+                with torch.no_grad():
+                    self.mmse_model.eval()
+                    condition = self.mmse_model(y).clip(0, 1)
+            else:
+                condition = y
+            x_t = torch.cat((x_t, condition), dim=1)
+        return self.model(x_t, t)
+
+    def forward(self, x_t, t, y):
+        if 'flow' in self.hparams.stage:
+            return self.forward_flow(x_t, t, y)
+        else:
+            return self.forward_mmse(y)
+
+    @torch.no_grad()
+    def create_source_distribution_samples(self, x, y, non_noisy_z0):
+        with torch.no_grad():
+            if self.hparams.conditional:
+                source_dist_samples = torch.randn_like(x)
+            else:
+                if self.hparams.stage == 'flow':
+                    if non_noisy_z0 is None:
+                        self.mmse_model.eval()
+                        non_noisy_z0 = self.mmse_model(y).clip(0, 1)
+                    source_dist_samples = non_noisy_z0 + torch.randn_like(non_noisy_z0) * self.hparams.mmse_noise_std
+                else:
+                    assert self.hparams.stage == 'naive_flow'
+                    if non_noisy_z0 is not None:
+                        source_dist_samples = non_noisy_z0
+                    else:
+                        source_dist_samples = y
+                    if source_dist_samples.shape[1] != x.shape[1]:
+                        assert source_dist_samples.shape[1] == 1  # Colorization
+                        source_dist_samples = source_dist_samples.expand(-1, x.shape[1], -1, -1)
+                    if self.hparams.mmse_noise_std is not None:
+                        source_dist_samples = source_dist_samples + torch.randn_like(source_dist_samples) * self.hparams.mmse_noise_std
+        return source_dist_samples
+
+    @staticmethod
+    def stratified_uniform(bs, group=0, groups=1, dtype=None, device=None):
+        if groups <= 0:
+            raise ValueError(f"groups must be positive, got {groups}")
+        if group < 0 or group >= groups:
+            raise ValueError(f"group must be in [0, {groups})")
+        n = bs * groups
+        offsets = torch.arange(group, n, groups, dtype=dtype, device=device)
+        u = torch.rand(bs, dtype=dtype, device=device)
+        return ((offsets + u) / n).view(bs, 1, 1, 1)
+
+    def generate_random_t(self, bs, dtype=None):
+        if self.hparams.t_schedule == 'logit-normal':
+            return sigmoid(torch.randn(bs, 1, 1, 1, device=self.device)) * (1.0 - self.hparams.eps) + self.hparams.eps
+        elif self.hparams.t_schedule == 'uniform':
+            return torch.rand(bs, 1, 1, 1, device=self.device) * (1.0 - self.hparams.eps) + self.hparams.eps
+        elif self.hparams.t_schedule == 'stratified_uniform':
+            return self.stratified_uniform(bs, self.trainer.global_rank, self.trainer.world_size, dtype=dtype,
+                                           device=self.device) * (1.0 - self.hparams.eps) + self.hparams.eps
+        else:
+            raise NotImplementedError()
+
+    def training_step(self, batch, batch_idx):
+        x = batch['x']
+        y = batch['y']
+        non_noisy_z0 = batch['non_noisy_z0'] if 'non_noisy_z0' in batch else None
+        if 'flow' in self.hparams.stage:
+            with torch.no_grad():
+                t = self.generate_random_t(x.shape[0], dtype=x.dtype)
+                source_dist_samples = self.create_source_distribution_samples(x, y, non_noisy_z0)
+                x_t = t * x + (1.0 - t) * source_dist_samples
+            v_t = self(x_t, t.squeeze(), y)
+            loss = mse_loss(v_t, x - source_dist_samples)
+        else:
+            xhat = self(x_t=None, t=None, y=y)
+            loss = mse_loss(xhat, x)
+        self.log("train/loss", loss)
+        return loss
+
+    @torch.no_grad()
+    def generate_reconstructions(self, x, y, non_noisy_z0, num_flow_steps, result_device):
+        with self.maybe_ema():
+            if 'flow' in self.hparams.stage:
+                source_dist_samples = self.create_source_distribution_samples(x, y, non_noisy_z0)
+
+                dt = (1.0 / num_flow_steps) * (1.0 - self.hparams.eps)
+                x_t_next = source_dist_samples.clone()
+                x_t_seq = [x_t_next]
+                t_one = torch.ones(x.shape[0], device=self.device)
+                for i in range(num_flow_steps):
+                    num_t = (i / num_flow_steps) * (1.0 - self.hparams.eps) + self.hparams.eps
+                    v_t_next = self(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
+                    x_t_seq.append(x_t_next.to(result_device))
+
+                xhat = x_t_seq[-1].clip(0, 1).to(torch.float32)
+                source_dist_samples = source_dist_samples.to(result_device)
+            else:
+                xhat = self(x_t=None, t=None, y=y).to(torch.float32)
+                x_t_seq = None
+                source_dist_samples = None
+            return xhat.to(result_device), x_t_seq, source_dist_samples
+
+    def validation_step(self, batch, batch_idx):
+        x = batch['x']
+        y = batch['y']
+        non_noisy_z0 = batch['non_noisy_z0'] if 'non_noisy_z0' in batch else None
+        xhat, x_t_seq, source_dist_samples = self.generate_reconstructions(x, y, non_noisy_z0, self.hparams.num_flow_steps,
+                                                                           self.device)
+        x = x.to(torch.float32)
+        y = y.to(torch.float32)
+        self.log_dict({"val_metrics/mse": ((x - xhat) ** 2).mean()}, on_step=False, on_epoch=True, sync_dist=True,
+                      batch_size=x.shape[0])
+
+        if 'flow' in self.hparams.stage:
+            self.fid.update(x, real=True)
+            self.fid.update(xhat, real=False)
+            self.inception_score.update(xhat)
+
+        if batch_idx == 0:
+            wandb_logger = self.logger.experiment
+            wandb_logger.log({'val_images/x': [wandb.Image(to_pil_image(create_grid(x)))],
+                              'val_images/y': [wandb.Image(to_pil_image(create_grid(y.clip(0, 1))))],
+                              'val_images/xhat': [wandb.Image(to_pil_image(create_grid(xhat)))], })
+            if 'flow' in self.hparams.stage:
+                wandb_logger.log({'val_images/x_t_seq': [wandb.Image(to_pil_image(create_grid(
+                    torch.cat([elem[0].unsqueeze(0).to(torch.float32) for elem in x_t_seq], dim=0).clip(0, 1),
+                    num_images=len(x_t_seq))))], 'val_images/source_distribution_samples': [
+                    wandb.Image(to_pil_image(create_grid(source_dist_samples.clip(0, 1).to(torch.float32))))]})
+                if self.mmse_model is not None:
+                    xhat_mmse = self.mmse_model(y).clip(0, 1)
+                    wandb_logger.log({'val_images/xhat_mmse': [
+                        wandb.Image(to_pil_image(create_grid(xhat_mmse.to(torch.float32))))]})
+
+    def on_validation_epoch_end(self):
+        if 'flow' in self.hparams.stage:
+            inception_score_mean, inception_score_std = self.inception_score.compute()
+            self.log_dict(
+                {'val_metrics/fid': self.fid.compute(),
+                 'val_metrics/inception_score_mean': inception_score_mean,
+                 'val_metrics/inception_score_std': inception_score_std},
+                on_epoch=True, on_step=False, sync_dist=True,
+                batch_size=1)
+            self.fid.reset()
+            self.inception_score.reset()
+
+    def test_step(self, batch, batch_idx):
+        assert self.test_results_path is not None, "Please set test_results_path before testing."
+        assert os.path.isdir(self.test_results_path), 'Please make sure the test_result_path dir exists.'
+
+        def save_image_batch(images, folder, image_file_names):
+            os.makedirs(folder, exist_ok=True)
+            for i, img in enumerate(images):
+                save_image(images[i].clip(0, 1), os.path.join(folder, image_file_names[i]))
+
+        os.makedirs(self.test_results_path, exist_ok=True)
+        x = batch['x']
+        y = batch['y']
+        non_noisy_z0 = batch['non_noisy_z0'] if 'non_noisy_z0' in batch else None
+        y_path = os.path.join(self.test_results_path, 'y')
+        save_image_batch(y, y_path, batch['img_file_name'])
+
+        if 'flow' in self.hparams.stage:
+            source_dist_samples_to_save = None
+
+            for num_flow_steps in self.num_test_flow_steps:
+                xhat, x_t_seq, source_dist_samples = self.generate_reconstructions(x, y, non_noisy_z0, num_flow_steps,
+                                                                                   torch.device("cpu"))
+                xhat_path = os.path.join(self.test_results_path, f"num_flow_steps={num_flow_steps}", 'xhat')
+                save_image_batch(xhat, xhat_path, batch['img_file_name'])
+                if source_dist_samples_to_save is None:
+                    source_dist_samples_to_save = source_dist_samples
+
+            source_distribution_samples_path = os.path.join(self.test_results_path, 'source_distribution_samples')
+            save_image_batch(source_dist_samples_to_save, source_distribution_samples_path, batch['img_file_name'])
+            if self.mmse_model is not None:
+                mmse_estimates = self.mmse_model(y).clip(0, 1)
+                mmse_samples_path = os.path.join(self.test_results_path, 'mmse_samples')
+                save_image_batch(mmse_estimates, mmse_samples_path, batch['img_file_name'])
+
+
+        else:
+            xhat, _, _ = self.generate_reconstructions(x, y, non_noisy_z0, None, torch.device('cpu'))
+            xhat_path = os.path.join(self.test_results_path, 'xhat')
+            save_image_batch(xhat, xhat_path, batch['img_file_name'])
+
+    def configure_optimizers(self):
+        # Add here a learning rate scheduler if you wish to do so.
+        optimizer = AdamW(self.model.parameters(),
+                          betas=self.hparams.betas,
+                          eps=1e-8,
+                          lr=self.hparams.lr,
+                          weight_decay=self.hparams.weight_decay)
+        return optimizer