adding decorator

app.py

@@ -1,6 +1,7 @@
 import os
 import sys
 import uuid
+import spaces
 import torch
 import random
 import numpy as np
@@ -39,6 +40,7 @@ def extract_descriptions(doc, head):
             descriptions.append(chunk.text)
     return descriptions
 
+@spaces.GPU
 def caption_refiner(caption):
     doc = nlp(caption)
     subject, action_verb = find_subject(doc)
@@ -48,6 +50,7 @@ def caption_refiner(caption):
     else:
         return caption
 
+@spaces.GPU
 def sam2(image, input_boxes, model_id="facebook/sam-vit-base"):
     device = "cuda" if torch.cuda.is_available() else "cpu"
     model = SamModel.from_pretrained(model_id).to(device)
@@ -61,12 +64,14 @@ def sam2(image, input_boxes, model_id="facebook/sam-vit-base"):
     )
     return masks
 
+@spaces.GPU
 def load_florence2(model_id="microsoft/Florence-2-base-ft", device='cuda'):
     torch_dtype = torch.float16 if device == 'cuda' else torch.float32
     florence_model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch_dtype, trust_remote_code=True).to(device)
     florence_processor = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
     return florence_model, florence_processor
 
+@spaces.GPU
 def florence2(image, prompt="", task="<OD>"):
     device = florence_model.device
     torch_dtype = florence_model.dtype
@@ -83,7 +88,7 @@ def florence2(image, prompt="", task="<OD>"):
     return parsed_answer[task]
 
 
-# Load and preprocess an image.
+@spaces.GPU
 def depth_estimation(image_path):
     model.eval()
     image, _, f_px = depth_pro.load_rgb(image_path)
@@ -261,6 +266,7 @@ def filter_bboxes(data, iou_threshold=0.5):
 
     return {'bboxes': filtered_bboxes, 'labels': filtered_labels, 'caption': data['caption']}
 
+@spaces.GPU
 def process_image(image_path: str):
     depth, fx = depth_estimation(image_path)
 

requirements.txt

@@ -23,6 +23,7 @@ spacy==3.7.5
 onnxruntime-gpu
 pandas
 html5lib
+spaces
 
 git+https://github.com/apple/ml-depth-pro.git
 git+https://github.com/facebookresearch/sam2.git