using HF hub to download model

app.py

@@ -3,8 +3,7 @@ from torch import nn
 from torchvision import models
 from torchvision.transforms import v2
 import os
-import requests
-import time
+from huggingface_hub import hf_hub_download
 
 
 labels = ['Pastel',
@@ -54,22 +53,10 @@ def predict(img, confidence):
     # If using GPU
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-    # Download model from GCS
-    model_path = os.getenv('MODEL_PATH')
-    print(model_path)
-    response = requests.get(model_path)
-    time.sleep(20)
+    hf_token = os.getenv('HF_token')
+    model_path = hf_hub_download(repo_id="fohy24/morphmarket_model", filename="model_v8_epoch9.pt", token=hf_token)
 
-    # # Define the cache directory path
-    # cache_dir = os.path.expanduser("~/.cache/torch/hub/checkpoints")
-    # model_file_name = os.path.basename(model_path)
-    # model_file_path = os.path.join(cache_dir, model_file_name)
-    # print(model_file_path)
-
-    with open('model.pt', 'wb') as f:
-        f.write(response.content)
-
-    checkpoint = torch.load('model.pt', map_location=device)
+    checkpoint = torch.load(model_path, map_location=device)
     densenet.load_state_dict(checkpoint['model_state_dict'])
 
     densenet.eval()