🤗 Add support for downloading Huggingface weights.

Trainer_finetune.py

@@ -5,7 +5,13 @@
 
 from config import *
 
-
+def convert(param):
+    return {
+        k.replace("module.", ""): v
+        for k, v in param.items()
+        if "module." in k and 'attn_mask' not in k and 'HW' not in k
+    }
+
 class Model:
     def __init__(self, local_rank):
         backbonetype, multiscaletype = MODEL_CONFIG['MODEL_TYPE']
@@ -30,18 +36,28 @@ def device(self):
         self.net.to(torch.device("cuda"))
 
     def load_model(self, name=None, rank=0, real=False):
-        def convert(param):
-            return {
-            k.replace("module.", ""): v
-                for k, v in param.items()
-                if "module." in k and 'attn_mask' not in k and 'HW' not in k
-            }
         if rank <= 0 :
             if name is None:
                 name = self.name
             print(f"loading {name} ckpt")
             self.net.load_state_dict(convert(torch.load(f'ckpt/{name}.pkl')), strict=True)
 
+    def from_pretrained(self, model_name):
+        try:
+            from huggingface_hub import hf_hub_download
+
+            ckpt_path = hf_hub_download(
+                repo_id="MCG-NJU/VFIMamba", filename="ckpt/" + model_name + ".pkl"
+            )
+            checkpoint = torch.load(ckpt_path)
+        except:
+            # In case the model is not hosted on huggingface
+            # or the user cannot import huggingface_hub correctly, model_name option: VFIMamba, VFIMamba_S
+            _VFIMAMBA_URL = f"https://huggingface.co/MCG-NJU/VFIMamba/resolve/main/ckpt/{model_name}.pkl"
+            checkpoint = torch.hub.load_state_dict_from_url(_VFIMAMBA_URL)
+
+        self.net.load_state_dict(convert(checkpoint), strict=True)
+
     @torch.no_grad()
     def hr_inference(self, img0, img1, local, TTA = False, down_scale = 1.0, timestep = 0.5, fast_TTA = False):
         '''

hf_demo_2x.py

@@ -0,0 +1,54 @@
+import cv2
+import math
+import sys
+import torch
+import numpy as np
+import argparse
+from imageio import mimsave
+
+'''==========import from our code=========='''
+sys.path.append('.')
+import config as cfg
+from Trainer_finetune import Model
+from benchmark.utils.padder import InputPadder
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--model', default='VFIMamba_S', type=str)
+parser.add_argument('--scale', default=0, type=float)
+
+args = parser.parse_args()
+assert args.model in ['VFIMamba_S', 'VFIMamba'], 'Model not exists!'
+
+
+'''==========Model setting=========='''
+TTA = False
+if args.model == 'VFIMamba':
+    TTA = True
+    cfg.MODEL_CONFIG['LOGNAME'] = 'VFIMamba'
+    cfg.MODEL_CONFIG['MODEL_ARCH'] = cfg.init_model_config(
+        F = 32,
+        depth = [2, 2, 2, 3, 3]
+    )
+model = Model(-1)
+model.from_pretrained(args.model)
+model.eval()
+model.device()
+
+
+print(f'=========================Start Generating=========================')
+
+I0 = cv2.imread('example/im1.png')
+I2 = cv2.imread('example/im2.png')
+
+I0_ = (torch.tensor(I0.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0)
+I2_ = (torch.tensor(I2.transpose(2, 0, 1)).cuda() / 255.).unsqueeze(0)
+
+padder = InputPadder(I0_.shape, divisor=32)
+I0_, I2_ = padder.pad(I0_, I2_)
+
+mid = (padder.unpad(model.inference(I0_, I2_, True, TTA=TTA, fast_TTA=TTA, scale=args.scale))[0].detach().cpu().numpy().transpose(1, 2, 0) * 255.0).astype(np.uint8)
+images = [I0[:, :, ::-1], mid[:, :, ::-1], I2[:, :, ::-1]]
+mimsave('example/out_2x_hf.gif', images, fps=3)
+
+
+print(f'=========================Done=========================')