downlaod model before calling predict()

app.py

@@ -28,6 +28,13 @@ labels = ['Pastel',
             'Super Pastel']
 num_labels = len(labels)
 
+# If using GPU
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+hf_token = os.getenv('HF_token')
+model_path = hf_hub_download(repo_id="fohy24/morphmarket_model", filename="model_v9_epoch8.pt", token=hf_token)
+checkpoint = torch.load(model_path, map_location=device)
+
 def predict(img, confidence):
 
     new_layers = nn.Sequential(
@@ -46,17 +53,9 @@ def predict(img, confidence):
     v2.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
     ])
 
-
     efficientnet = models.efficientnet_v2_l(weights='EfficientNet_V2_L_Weights.DEFAULT')
     efficientnet.classifier = new_layers
 
-    # If using GPU
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
-    hf_token = os.getenv('HF_token')
-    model_path = hf_hub_download(repo_id="fohy24/morphmarket_model", filename="model_v9_epoch8.pt", token=hf_token)
-
-    checkpoint = torch.load(model_path, map_location=device)
     efficientnet.load_state_dict(checkpoint['model_state_dict'])
 
     efficientnet.eval()
@@ -64,7 +63,6 @@ def predict(img, confidence):
     input_img = transform(img)
     input_img = input_img.unsqueeze(0)
 
-
     with torch.no_grad():
         output = efficientnet(input_img)