Add huggingface support

RealESRGAN/model.py

+import os
 import torch
 from torch.nn import functional as F
 from PIL import Image
 import numpy as np
 import cv2
+from huggingface_hub import hf_hub_url, cached_download
 
 from .rrdbnet_arch import RRDBNet
 from .utils import pad_reflect, split_image_into_overlapping_patches, stich_together, \
                    unpad_image
 
 
+HF_MODELS = {
+    2: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x2.pth',
+    ),
+    4: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x4.pth',
+    ),
+    8: dict(
+        repo_id='sberbank-ai/Real-ESRGAN',
+        filename='RealESRGAN_x8.pth',
+    ),
+}
+
+
 class RealESRGAN:
     def __init__(self, device, scale=4):
         self.device = device
@@ -18,7 +36,16 @@ class RealESRGAN:
             num_block=23, num_grow_ch=32, scale=scale
         )
         
-    def load_weights(self, model_path):
+    def load_weights(self, model_path, download=True):
+        if not os.path.exists(model_path) and download:
+            assert self.scale in [2,4,8], 'You can download models only with scales: 2, 4, 8'
+            config = HF_MODELS[self.scale]
+            cache_dir = os.path.dirname(model_path)
+            local_filename = os.path.basename(model_path)
+            config_file_url = hf_hub_url(repo_id=config['repo_id'], filename=config['filename'])
+            cached_download(config_file_url, cache_dir=cache_dir, force_filename=local_filename)
+            print('Weights downloaded to:', os.path.join(cache_dir, local_filename))
+        
         loadnet = torch.load(model_path)
         if 'params' in loadnet:
             self.model.load_state_dict(loadnet['params'], strict=True)

main.py

@@ -8,7 +8,7 @@ from RealESRGAN import RealESRGAN
 def main() -> int:
     device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
     model = RealESRGAN(device, scale=4)
-    model.load_weights('weights/RealESRGAN_x4.pth')
+    model.load_weights('weights/RealESRGAN_x4.pth', download=True)
     for i, image in enumerate(os.listdir("inputs")):
         image = Image.open(f"inputs/{image}").convert('RGB')
         sr_image = model.predict(image)

requirements.txt

@@ -4,3 +4,4 @@ Pillow
 torch>=1.7
 torchvision>=0.8.0
 tqdm
+huggingface-hub
\ No newline at end of file