Update app.py (#4)

- Update app.py (ef28fa7ef601ab860d6c4a2f0e5f86b00e2cd1d7)


Co-authored-by: YuqianYuan <[email protected]>

app.py

@@ -13,6 +13,7 @@ sys.path.append('./')
 from videollama3 import disable_torch_init, model_init, mm_infer, get_model_output
 from videollama3.mm_utils import load_images
 from videollama3.mm_utils import load_video
+import spaces
 
 
 color_rgb = (1.0, 1.0, 1.0)
@@ -54,6 +55,7 @@ def add_contour(img, mask, color=(1., 1., 1.)):
 
     return img
 
+@spaces.GPU(duration=120)
 def generate_masks(image):
     global mask_list
     global mask_raw_list
@@ -82,6 +84,7 @@ def generate_masks(image):
     return mask_list, image
 
 
+@spaces.GPU(duration=120)
 def generate_masks_video(image):
     global mask_list_video
     global mask_raw_list_video
@@ -110,7 +113,7 @@ def generate_masks_video(image):
     return mask_list_video, image
 
 
-
+@spaces.GPU(duration=120)
 def describe(image, mode, query, masks):
     # Create an image object from the uploaded image
     # print(image.keys())
@@ -211,6 +214,7 @@ def load_first_frame(video_path):
     image = Image.fromarray(frame)  
     return image
 
+@spaces.GPU(duration=120)
 def describe_video(video_path, mode, query, annotated_frame, masks):
     global mask_list_video
     # Create a temporary directory to save extracted video frames
@@ -319,7 +323,7 @@ def describe_video(video_path, mode, query, annotated_frame, masks):
         yield gr.update(), text, gr.update()
 
 
-
+@spaces.GPU(duration=120)
 def apply_sam(image, input_points):
     inputs = sam_processor(image, input_points=input_points, return_tensors="pt").to(device)