Update app.py

app.py

@@ -1,6 +1,8 @@
 import torch
 import torch.nn.functional as F
 from torch import Tensor
+import spaces
+
 import numpy as np
 from PIL import Image
 import json, os, random
@@ -60,7 +62,8 @@ model_configs = {
     },
 }
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = "cuda"
 
 if truncation is None:  # regression, no truncation.
     bins, anchor_points = None, None
@@ -109,7 +112,6 @@ def load_model(model_choice: str):
     state_dict = load_file(weights_path)
     new_state_dict = {k.replace("model.", ""): v for k, v in state_dict.items()}
     model.load_state_dict(new_state_dict)
-    model.to(device)
     model.eval()
     
     loaded_model = model
@@ -139,6 +141,7 @@ def transform(image: Image.Image):
 # -----------------------------
 # Inference function
 # -----------------------------
+@spaces.GPU(duration=120)
 def predict(image: Image.Image, model_choice: str = "CLIP_EBC_ViT_B_16"):
     """
     Given an input image, preprocess it, run the model to obtain a density map,
@@ -149,6 +152,7 @@ def predict(image: Image.Image, model_choice: str = "CLIP_EBC_ViT_B_16"):
     if loaded_model is None or model_configs[model_choice]["model_name"] not in loaded_model.__class__.__name__:
         load_model(model_choice)
 
+    loaded_model.to(device)
     # Preprocess the image
     input_width, input_height = image.size
     input_tensor = transform(image).to(device)  # shape: (1, 3, H, W)