first upload

app.py

@@ -0,0 +1,139 @@
+#!/usr/bin/env python
+# -*- coding:utf-8 -*-
+# Power by Zongsheng Yue 2022-12-16 16:17:14
+
+import os
+import torch
+import argparse
+import numpy as np
+import grdio as gr
+from pathlib import Path
+from einops import rearrange
+from omegaconf import OmegaConf
+from skimage import img_as_ubyte
+
+from utils import util_opts
+from utils import util_image
+from utils import util_common
+
+from sampler import DifIRSampler
+from ResizeRight.resize_right import resize
+from basicsr.utils.download_util import load_file_from_url
+
+def predict(im_path, background_enhance, face_upsample, upscale, started_timesteps):
+    cfg_path = 'configs/sample/iddpm_ffhq512_swinir.yaml'
+
+    # setting configurations
+    configs = OmegaConf.load(cfg_path)
+    configs.aligned = False
+    configs.background_enhance = background_enhance
+    configs.face_upsample = face_upsample
+
+    started_timesteps = int(started_timesteps)
+    assert started_timesteps < int(configs.diffusion.params.timestep_respacing)
+
+    # prepare the checkpoint
+    if not Path(configs.model.ckpt_path).exists():
+        load_file_from_url(
+            url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/iddpm_ffhq512_ema500000.pth",
+            model_dir=str(Path(configs.model.ckpt_path).parent),
+            progress=True,
+            file_name=Path(configs.model.ckpt_path).name,
+            )
+    if not Path(configs.model_ir.ckpt_path).exists():
+        load_file_from_url(
+            url="https://github.com/zsyOAOA/DifFace/releases/download/V1.0/General_Face_ffhq512.pth",
+            model_dir=str(Path(configs.model_ir.ckpt_path).parent),
+            progress=True,
+            file_name=Path(configs.model_ir.ckpt_path).name,
+            )
+
+    # Load image
+    im_lq = util_image.imread(im_path, chn='bgr', dtype='uint8')
+    if upscale > 4:
+        upscale = 4  # avoid momory exceeded due to too large upscale
+    if upscale > 2 and min(im_lq.shape[:2])>1280:
+        upscale = 2  # avoid momory exceeded due to too large img resolution
+    configs.detection.upscale = int(upscale)
+
+    # build the sampler for diffusion
+    sampler_dist = DifIRSampler(configs)
+
+    image_restored, face_restored, face_cropped = sampler_dist.sample_func_bfr_unaligned(
+            y0=im_lq,
+            start_timesteps=started_timesteps,
+            need_restoration=True,
+            draw_box=False,
+            )
+
+    restored_image_dir = Path('restored_output')
+    if not restored_image_dir.exists():
+        restored_image_dir.mkdir()
+    # save the whole image
+    save_path = restored_image_dir / Path(im_path).name
+    util_image.imwrite(image_restored, save_path, chn='bgr', dtype_in='uint8')
+
+    return image_restored, str(save_path)
+
+# im_path = './testdata/whole_imgs/00.jpg'
+# predict(im_path, True, True, 3, 100)
+
+title = "DifFace: Blind Face Restoration with Diffused Error Contraction"
+description = r"""<center><img src='./assets/DifFace_Framework.png' alt='DifFace logo'></center>
+<b>Official Gradio demo</b> for <a href='https://github.com/zsyOAOA/DifFace' target='_blank'><b>DifFace: Blind Face Restoration with Diffused Error Contraction</b></a>.<br>
+🔥 DifFace is a robust face restoration algorithm for old or corrupted photos.<br>
+"""
+article = r"""
+If DifFace is helpful for your work, please help to ⭐ the <a href='https://github.com/zsyOAOA/DifFace' target='_blank'>Github Repo</a>. Thanks!
+[![GitHub Stars](https://img.shields.io/github/stars/zsyOAOA/DifFace?affiliations=OWNER&color=green&style=social)](https://github.com/zsyOAOA/DifFace)
+
+---
+
+📝 **Citation**
+
+If our work is useful for your research, please consider citing:
+```bibtex
+@article{yue2022difface,
+  title={DifFace: Blind Face Restoration with Diffused Error Contraction},
+  author={Yue, Zongsheng and Loy, Chen Change},
+  journal={arXiv preprint arXiv:2212.06512},
+  year={2022}
+}
+```
+
+📋 **License**
+
+This project is licensed under <a rel="license" href="https://github.com/zsyOAOA/DifFace/blob/master/LICENSE">S-Lab License 1.0</a>.
+Redistribution and use for non-commercial purposes should follow this license.
+
+📧 **Contact**
+If you have any questions, please feel free to contact me via <b>[email protected]</b>.
+![visitors](https://visitor-badge.laobi.icu/badge?page_id=zsyOAOA/DifFace)
+"""
+
+demo = gr.Interface(
+    inference,
+    inputs=[
+        gr.inputs.Image(type="filepath", label="Input"),
+        gr.inputs.Checkbox(default=True, label="Background_Enhance"),
+        gr.inputs.Checkbox(default=True, label="Face_Upsample"),
+        gr.inputs.Number(default=2, label="Rescaling_Factor (up to 4)"),
+        gr.Slider(1, 200, value=100, step=10, label='Realism-Fidelity Trade-off')
+    ],
+    outputs=[
+        gr.outputs.Image(type="numpy", label="Output"),
+        gr.outputs.File(label="Download the output")
+    ],
+    title=title,
+    description=description,
+    article=article,
+    examples=[
+        ['./testdata/whole_imgs/00.jpg', True, True, 2, 100],
+        ['./testdata/whole_imgs/01.jpg', True, True, 2, 100],
+        ['./testdata/whole_imgs/04.jpg', True, True, 2, 100],
+        ['./testdata/whole_imgs/Solvay_conference_1927.png', True, True, 2, 100],
+      ]
+    )
+
+demo.queue(concurrency_count=4)
+demo.launch()

configs/sample/iddpm_ffhq512.yaml

@@ -0,0 +1,39 @@
+gpu_id: ""
+seed: 10000
+display: True
+im_size: 512
+
+diffusion:
+  target: models.script_util.create_gaussian_diffusion
+  params:
+    steps: 1000
+    learn_sigma: True
+    sigma_small: False
+    noise_schedule: linear
+    use_kl: False
+    predict_xstart: False
+    rescale_timesteps: False
+    rescale_learned_sigmas: True
+    timestep_respacing: "1000"
+
+model:
+  target: models.unet.UNetModel
+  ckpt_path: pretrained_zoo/iddpm_ffhq512/ema0999_model_500000.pth
+  params:
+    image_size: 512
+    in_channels: 3
+    model_channels: 32
+    out_channels: 6
+    attention_resolutions: [32, 16, 8]
+    dropout: 0
+    channel_mult: [1, 2, 4, 8, 8, 16, 16]
+    num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
+    conv_resample: True
+    dims: 2
+    use_fp16: False
+    num_head_channels: 64
+    use_scale_shift_norm: True
+    resblock_updown: False
+    use_new_attention_order: False
+
+model_ir: ~

configs/sample/iddpm_ffhq512_swinir.yaml

@@ -0,0 +1,64 @@
+gpu_id: ""
+seed: 10000
+display: True
+im_size: 512
+aligned: True
+background_enhance: True
+face_upsample: True
+
+diffusion:
+  target: models.script_util.create_gaussian_diffusion
+  params:
+    steps: 1000
+    learn_sigma: True
+    sigma_small: False
+    noise_schedule: linear
+    use_kl: False
+    predict_xstart: False
+    rescale_timesteps: False
+    rescale_learned_sigmas: True
+    timestep_respacing: "250"
+
+model:
+  target: models.unet.UNetModel
+  ckpt_path: ./weights/diffusion/iddpm_ffhq512_ema500000.pth
+  params:
+    image_size: 512
+    in_channels: 3
+    model_channels: 32
+    out_channels: 6
+    attention_resolutions: [32, 16, 8]
+    dropout: 0
+    channel_mult: [1, 2, 4, 8, 8, 16, 16]
+    num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
+    conv_resample: True
+    dims: 2
+    use_fp16: False
+    num_head_channels: 64
+    use_scale_shift_norm: True
+    resblock_updown: False
+    use_new_attention_order: False
+
+model_ir:
+  target: models.swinir.SwinIR
+  ckpt_path: ./weights/SwinIR/General_Face_ffhq512.pth
+  params:
+    img_size: 64
+    patch_size: 1
+    in_chans: 3
+    embed_dim: 180
+    depths: [6, 6, 6, 6, 6, 6, 6, 6]
+    num_heads: [6, 6, 6, 6, 6, 6, 6, 6]
+    window_size: 8
+    mlp_ratio: 2
+    sf: 8
+    img_range: 1.0
+    upsampler: "nearest+conv"
+    resi_connection: "1conv"
+    unshuffle: True
+    unshuffle_scale: 8
+
+# face detection model for unaligned face
+detection:
+  det_model: "YOLOv5l"  # large model: 'YOLOv5l', 'retinaface_resnet50'; small model: 'YOLOv5n', 'retinaface_mobile0.25'
+  upscale: 2            # The final upscaling factor for the whole image

configs/training/diffusion_ffhq512.yaml

@@ -0,0 +1,55 @@
+model:
+  target: models.unet.UNetModel
+  params:
+    image_size: 512
+    in_channels: 3
+    model_channels: 32
+    out_channels: 6
+    attention_resolutions: [32, 16, 8]
+    dropout: 0
+    channel_mult: [1, 2, 4, 8, 8, 16, 16]
+    num_res_blocks: [1, 2, 2, 2, 2, 3, 4]
+    conv_resample: True
+    dims: 2
+    use_fp16: False
+    num_head_channels: 64
+    use_scale_shift_norm: True
+    resblock_updown: False
+    use_new_attention_order: False
+
+diffusion:
+  target: models.script_util.create_gaussian_diffusion
+  params:
+    steps: 1000
+    learn_sigma: True
+    sigma_small: False
+    noise_schedule: linear
+    use_kl: False
+    predict_xstart: False
+    rescale_timesteps: False
+    rescale_learned_sigmas: True
+    timestep_respacing: ""
+
+train:
+  lr: 1e-4
+  batch: [32, 4]   # batchsize for training and validation
+  microbatch: 8
+  use_fp16: False
+  num_workers: 16
+  prefetch_factor: 2
+  iterations: 800000
+  weight_decay: 0
+  scheduler: step   # step or cosin
+  milestones: [10000, 800000]
+  ema_rates: [0.999]
+  save_freq: 10000
+  val_freq: 5000
+  log_freq: [1000, 2000]
+
+data:
+  train:
+    type: face
+    params:
+      ffhq_txt: ./datapipe/files_txt/ffhq512.txt
+      out_size: 512
+      transform_type: face

configs/training/swinir_ffhq512.yaml

@@ -0,0 +1,67 @@
+model:
+  target: models.swinir.SwinIR
+  params:
+    img_size: 64
+    patch_size: 1
+    in_chans: 3
+    embed_dim: 180
+    depths: [6, 6, 6, 6, 6, 6, 6, 6]
+    num_heads: [6, 6, 6, 6, 6, 6, 6, 6]
+    window_size: 8
+    mlp_ratio: 2
+    sf: 8
+    img_range: 1.0
+    upsampler: "nearest+conv"
+    resi_connection: "1conv"
+    unshuffle: True
+    unshuffle_scale: 8
+
+train:
+  lr: 1e-4
+  lr_min: 5e-6
+  batch: [16, 4]   # batchsize for training and validation
+  microbatch: 2
+  num_workers: 8
+  prefetch_factor: 2
+  iterations: 800000
+  weight_decay: 0
+  save_freq: 20000
+  val_freq: 20000
+  log_freq: [100, 2000, 50]
+
+data:
+  train:
+    type: gfpgan 
+    params:
+      files_txt: ./datapipe/files_txt/ffhq512_train.txt
+      io_backend:
+        type: disk
+   
+      use_hflip: true
+      mean: [0.0, 0.0, 0.0]
+      std: [1.0, 1.0, 1.0]
+      out_size: 512
+   
+      blur_kernel_size: 41
+      kernel_list: ['iso', 'aniso']
+      kernel_prob: [0.5, 0.5]
+      blur_sigma: [0.1, 15]
+      downsample_range: [0.8, 32]
+      noise_range: [0, 20]
+      jpeg_range: [30, 100]
+   
+      color_jitter_prob: ~
+      color_jitter_pt_prob: ~
+      gray_prob: 0.01
+      gt_gray: True
+
+      need_gt_path: False
+  val:
+    type: folder 
+    params:
+      dir_path: /mnt/vdb/IRDiff/Face/testing_data/celeba512_lq
+      dir_path_gt: /mnt/vdb/IRDiff/Face/testing_data/celeba512_hq
+      ext: png
+      need_gt_path: False
+      length: ~
+