Update app.py

app.py

@@ -37,6 +37,25 @@ import spaces
 
 image_size = 336
 
+device = "cuda" if torch.cuda.is_available() else "cpu"
+univad_model = UniVAD(image_size=image_size).to(device)
+ram_model = ram_plus(
+    pretrained="./ram_plus_swin_large_14m.pth",
+    image_size=384,
+    vit="swin_l",
+)
+ram_model.eval()
+ram_model = ram_model.to(device)
+
+
+grounding_model = load_model(
+    "./UniVAD/models/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py",
+    "./groundingdino_swint_ogc.pth",
+    "cuda" if torch.cuda.is_available() else "cpu"
+)
+sam = sam_hq_model_registry["vit_h"]("./sam_hq_vit_h.pth").to(device)
+sam_predictor = SamPredictor(sam)
+
 
 transform = transforms.Compose(
     [
@@ -59,24 +78,6 @@ def update_image(image):
 
 @spaces.GPU
 def ad(image_pil, normal_image, box_threshold, text_threshold, text_prompt, background_prompt, cluster_num):
-    device = "cuda" if torch.cuda.is_available() else "cpu"
-    univad_model = UniVAD(image_size=image_size).to(device)
-    ram_model = ram_plus(
-        pretrained="./ram_plus_swin_large_14m.pth",
-        image_size=384,
-        vit="swin_l",
-    )
-    ram_model.eval()
-    ram_model = ram_model.to(device)
-    
-    
-    grounding_model = load_model(
-        "./UniVAD/models/GroundingDINO/groundingdino/config/GroundingDINO_SwinT_OGC.py",
-        "./groundingdino_swint_ogc.pth",
-        "cuda" if torch.cuda.is_available() else "cpu"
-    )
-    sam = sam_hq_model_registry["vit_h"]("./sam_hq_vit_h.pth").to(device)
-    sam_predictor = SamPredictor(sam)
     return process_image(image_pil, normal_image, box_threshold, text_threshold, sam_predictor, grounding_model, univad_model, ram_model, text_prompt, background_prompt, cluster_num, image_size)