From 396b5e1d92022ed776ed3d3174e26fd330ed0f74 Mon Sep 17 00:00:00 2001 From: sczhou Date: Wed, 5 Oct 2022 11:19:22 +0800 Subject: [PATCH] enable realesrgan in CPU (#40) --- README.md | 4 +-- app.py | 5 ++-- inference_codeformer.py | 54 +++++++++++++++++++++++++---------------- 3 files changed, 38 insertions(+), 25 deletions(-) diff --git a/README.md b/README.md index c5153d18..903cd31f 100644 --- a/README.md +++ b/README.md @@ -7,9 +7,9 @@ [Paper](https://arxiv.org/abs/2206.11253) | [Project Page](https://shangchenzhou.com/projects/CodeFormer/) | [Video](https://youtu.be/d3VDpkXlueI) -google colab logo [![Hugging Face](https://img.shields.io/badge/Demo-%F0%9F%A4%97%20Hugging%20Face-blue)](https://huggingface.co/spaces/sczhou/CodeFormer) [![Replicate](https://img.shields.io/badge/Demo-%F0%9F%9A%80%20Replicate-blue)](https://replicate.com/sczhou/codeformer) ![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer) +google colab logo [![Hugging Face](https://img.shields.io/badge/Demo-%F0%9F%A4%97%20Hugging%20Face-blue)](https://huggingface.co/spaces/sczhou/CodeFormer) [![Replicate](https://img.shields.io/badge/Demo-%F0%9F%9A%80%20Replicate-blue)](https://replicate.com/sczhou/codeformer) ![visitors](https://visitor-badge.laobi.icu/badge?page_id=sczhou/CodeFormer) - + [Shangchen Zhou](https://shangchenzhou.com/), [Kelvin C.K. Chan](https://ckkelvinchan.github.io/), [Chongyi Li](https://li-chongyi.github.io/), [Chen Change Loy](https://www.mmlab-ntu.com/person/ccloy/) diff --git a/app.py b/app.py index ce22e673..cb973afa 100644 --- a/app.py +++ b/app.py @@ -193,7 +193,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide imwrite(restored_img, str(save_path)) restored_img = cv2.cvtColor(restored_img, cv2.COLOR_BGR2RGB) - return restored_img + return restored_img, save_path @@ -230,7 +230,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide If you have any questions, please feel free to reach me out at shangchenzhou@gmail.com. -![visitors](https://visitor-badge.glitch.me/badge?page_id=sczhou/CodeFormer) +![visitors](https://visitor-badge.laobi.icu/badge?page_id=sczhou/CodeFormer) """ demo = gr.Interface( @@ -242,6 +242,7 @@ def inference(image, background_enhance, face_upsample, upscale, codeformer_fide gr.Slider(0, 1, value=0.5, step=0.01, label='Codeformer_Fidelity: 0 for better quality, 1 for better identity') ], [ gr.outputs.Image(type="numpy", label="Output"), + gr.outputs.File(label="Download the output") ], title=title, description=description, diff --git a/inference_codeformer.py b/inference_codeformer.py index 9525c995..9badb160 100644 --- a/inference_codeformer.py +++ b/inference_codeformer.py @@ -17,25 +17,36 @@ } def set_realesrgan(): - if not torch.cuda.is_available(): # CPU + from basicsr.archs.rrdbnet_arch import RRDBNet + from basicsr.utils.realesrgan_utils import RealESRGANer + + cuda_is_available = torch.cuda.is_available() + half = True if cuda_is_available else False + model = RRDBNet( + num_in_ch=3, + num_out_ch=3, + num_feat=64, + num_block=23, + num_grow_ch=32, + scale=2, + ) + upsampler = RealESRGANer( + scale=2, + model_path="https://github.com/sczhou/CodeFormer/releases/download/v0.1.0/RealESRGAN_x2plus.pth", + model=model, + tile=args.bg_tile, + tile_pad=40, + pre_pad=0, + half=half, # need to set False in CPU mode + ) + + if not cuda_is_available: # CPU import warnings - warnings.warn('The unoptimized RealESRGAN is slow on CPU. We do not use it. ' - 'If you really want to use it, please modify the corresponding codes.', + warnings.warn('Runing on CPU now... ' + 'The unoptimized RealESRGAN is slow on CPU. ' + 'If you want to disable it, please remove `--bg_upsampler` and `--face_upsample` in command.', category=RuntimeWarning) - bg_upsampler = None - else: - from basicsr.archs.rrdbnet_arch import RRDBNet - from basicsr.utils.realesrgan_utils import RealESRGANer - model = RRDBNet(num_in_ch=3, num_out_ch=3, num_feat=64, num_block=23, num_grow_ch=32, scale=2) - bg_upsampler = RealESRGANer( - scale=2, - model_path='https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.1/RealESRGAN_x2plus.pth', - model=model, - tile=args.bg_tile, - tile_pad=40, - pre_pad=0, - half=True) # need to set False in CPU mode - return bg_upsampler + return upsampler if __name__ == '__main__': device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') @@ -71,10 +82,11 @@ def set_realesrgan(): # ------------------ set up face upsampler ------------------ if args.face_upsample: - if bg_upsampler is not None: - face_upsampler = bg_upsampler - else: - face_upsampler = set_realesrgan() + face_upsampler = None + # if bg_upsampler is not None: + # face_upsampler = bg_upsampler + # else: + # face_upsampler = set_realesrgan() else: face_upsampler = None