c

basicsr/models/archs/histoformer_arch.py

@@ -5,7 +5,7 @@ from pdb import set_trace as stx
 import numbers
 
 from einops import rearrange
-
+from huggingface_hub import PyTorchModelHubMixin
 #########################################################################
 
 Conv2d = nn.Conv2d
@@ -267,7 +267,7 @@ class Upsample(nn.Module):
         return self.body(x)
 
 ##########################################################################
-class Histoformer(nn.Module):
+class Histoformer(nn.Module, PyTorchModelHubMixin):
     def __init__(self, 
         inp_channels=3, 
         out_channels=3, 

demo.py

@@ -1,12 +1,14 @@
+import torch
 from basicsr.models.archs.histoformer_arch import Histoformer
-# from huggingface_hub import PyTorchModelHubMixin
+from huggingface_hub import PyTorchModelHubMixin
 from huggingface_hub import hf_hub_download
 # define model
 model = Histoformer()
 
 # equip with weights
 filepath = hf_hub_download(repo_id="sunsean/Histoformer", filename="Allweather/pretrained_models/net_g_real.pth")
-model.load_state_dict(filepath, map_location="gpu")
+print(filepath)
+model.load_state_dict(torch.load(filepath, map_location='cpu')['params'], )
 
 # push to hub
 model.push_to_hub("sunsean/Histoformer-real")