enable realesrgan in CPU (#40)

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)
 
 
-<a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a> [![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)
+<a href="https://colab.research.google.com/drive/1m52PNveE4PBhYrecj34cnpEeiHcC5LTb?usp=sharing"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="google colab logo"></a> [![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)
 
-<!-- ![visitors](https://visitor-badge.laobi.icu/badge?page_id=sczhou/CodeFormer) -->
+<!-- ![visitors](https://visitor-badge.glitch.me/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/) 

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 @@ Redistribution and use for non-commercial purposes should follow this license.
 
 If you have any questions, please feel free to reach me out at <b>shangchenzhou@gmail.com</b>.
 
-![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 @@ demo = gr.Interface(
         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,

inference_codeformer.py

@@ -17,25 +17,36 @@ pretrain_model_url = {
 }
 
 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 @@ if __name__ == '__main__':
 
     # ------------------ 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