Use mixin

app.py

@@ -9,7 +9,6 @@ from PIL import Image
 from copy import deepcopy
 from torch.nn.parallel import DataParallel, DistributedDataParallel
 
-from huggingface_hub import hf_hub_download
 from gradio_imageslider import ImageSlider
 
 ## local code
@@ -71,7 +70,6 @@ def load_network(net, load_path, strict=True, param_key='params'):
     net.load_state_dict(load_net, strict=strict)
 
 CONFIG = "configs/eval_seemore_t_x4.yml"
-hf_hub_download(repo_id="eduardzamfir/SeemoRe-T", filename="SeemoRe_T_X4.pth", local_dir="./")
 MODEL_NAME = "SeemoRe_T_X4.pth"
 
 # parse config file
@@ -81,10 +79,7 @@ with open(os.path.join(CONFIG), "r") as f:
 cfg = dict2namespace(config)
 
 device = torch.device("cpu")
-model = seemore.SeemoRe(scale=cfg.model.scale, in_chans=cfg.model.in_chans,
-                        num_experts=cfg.model.num_experts, num_layers=cfg.model.num_layers, embedding_dim=cfg.model.embedding_dim, 
-                        img_range=cfg.model.img_range, use_shuffle=cfg.model.use_shuffle, global_kernel_size=cfg.model.global_kernel_size, 
-                        recursive=cfg.model.recursive, lr_space=cfg.model.lr_space, topk=cfg.model.topk)
+model = seemore.SeemoRe.from_pretrained("eduardzamfir/SeemoRe-T")
 
 model = model.to(device)
 print ("IMAGE MODEL CKPT:", MODEL_NAME)

models/seemore.py

@@ -6,11 +6,13 @@ import torch.nn as nn
 import torch.nn.functional as F
 from einops.layers.torch import Rearrange
 
+from huggingface_hub import PyTorchModelHubMixin
+
 
 ######################
 # Meta Architecture
 ######################
-class SeemoRe(nn.Module):
+class SeemoRe(nn.Module, PyTorchModelHubMixin):
     def __init__(self,
                  scale: int = 4,
                  in_chans: int = 3,