Update app.py

app.py

@@ -30,11 +30,12 @@ from UniVAD.models.grounded_sam import (
     load_model,
 )
 
+import spaces
+
+
 
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
 image_size = 336
-univad_model = UniVAD(image_size=image_size).to(device)
 
 
 transform = transforms.Compose(
@@ -44,22 +45,7 @@ transform = transforms.Compose(
     ]
 )
 
-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)
 
 
 def preprocess_image(img):
@@ -70,8 +56,27 @@ def preprocess_image(img):
 def update_image(image):
     if image is not None:
         return preprocess_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)