fix: different safetensors

app.py

@@ -5,11 +5,27 @@ from PIL import Image
 from net.CIDNet import CIDNet
 import torchvision.transforms as transforms
 import torch.nn.functional as F
-import os
+import safetensors.torch as sf
 import imquality.brisque as brisque
 from loss.niqe_utils import *
 import spaces
-import huggingface_hub
+from huggingface_hub import hf_hub_download
+import json
+
+def from_pretrained(cls, pretrained_model_name_or_path: str):
+    model_id = str(pretrained_model_name_or_path)
+
+    config_file = hf_hub_download(repo_id=model_id, filename="config.json", repo_type="model")
+    config = None
+    if config_file is not None:
+        with open(config_file, "r", encoding="utf-8") as f:
+            config = json.load(f)
+
+
+    model_file = hf_hub_download(repo_id=model_id, filename="model.safetensors", repo_type="model")
+    # instance = sf.load_model(cls, model_file, strict=False)
+    state_dict  = sf.load_file(model_file)
+    cls.load_state_dict(state_dict, strict=False) 
 
 eval_net = CIDNet().cuda()
 eval_net.trans.gated = True
@@ -20,7 +36,7 @@ def process_image(input_img,score,model_path,gamma=1.0,alpha_s=1.0,alpha_i=1.0):
     if model_path is None:
         return input_img,"Please choose a model weights."
     torch.set_grad_enabled(False)
-    eval_net.from_pretrained("Fediory/HVI-CIDNet-"+model_path)
+    from_pretrained(eval_net,"Fediory/HVI-CIDNet-"+model_path)
     # eval_net.load_state_dict(torch.load(os.path.join(directory,model_path), map_location=lambda storage, loc: storage))
     eval_net.eval()
     
@@ -81,4 +97,4 @@ interface = gr.Interface(
     allow_flagging="never"
 )
 
-interface.launch(share=True)
+interface.launch(share=False)