use hf_hub_download

app.py

@@ -1,5 +1,4 @@
 import os
-import json
 
 import torch
 import torch.nn.functional as F
@@ -8,14 +7,17 @@ from uniformer import uniformer_small
 from imagenet_class_index import imagenet_classnames
 
 import gradio as gr
+from huggingface_hub import hf_hub_download
 
 # Device on which to run the model
 # Set to cuda to load on GPU
 device = "cpu"
-os.system("wget https://cdn-lfs.huggingface.co/Andy1621/uniformer/fd192c31f8bd77670de8f171111bd51f56fd87e6aea45043ab2edc181e1fa775")
+# os.system("wget https://cdn-lfs.huggingface.co/Andy1621/uniformer/fd192c31f8bd77670de8f171111bd51f56fd87e6aea45043ab2edc181e1fa775")
+model_path = hf_hub_download(repo_id="https://huggingface.co/Andy1621/uniformer", filename="uniformer_small_in1k.pth")
 # Pick a pretrained model 
 model = uniformer_small()
-state_dict = torch.load('fd192c31f8bd77670de8f171111bd51f56fd87e6aea45043ab2edc181e1fa775', map_location='cpu')
+# state_dict = torch.load('fd192c31f8bd77670de8f171111bd51f56fd87e6aea45043ab2edc181e1fa775', map_location='cpu')
+state_dict = torch.load(model_path, map_location='cpu')
 model.load_state_dict(state_dict['model'])
 
 # Set to eval mode and move to desired device