changes for zerogpu

app.py

@@ -11,6 +11,7 @@ import traceback
 import matplotlib.pyplot as plt
 import ffmpeg
 from utils import load_model_without_flash_attn
+import spaces
 
 
 # CUDA optimizations
@@ -63,6 +64,7 @@ def apply_color_mask(frame, mask, obj_id):
  colored_mask = mask * color
  return frame * (1 - mask) + colored_mask * 255
 
+@spaces.GPU
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def run_florence(image, text_input):
@@ -99,6 +101,7 @@ def remove_directory_contents(directory):
  for name in dirs:
  os.rmdir(os.path.join(root, name))
 
+@spaces.GPU
 @torch.inference_mode()
 @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process_video(video_path, prompt):
@@ -216,6 +219,7 @@ def process_video(video_path, prompt):
  print(traceback.format_exc()) # This will print the full stack trace
  return None
  
+@spaces.GPU(duration=1200)
 def segment_video(video_file, prompt):
  if video_file is None:
  return None

requirements.txt

@@ -9,4 +9,5 @@ matplotlib
 einops
 timm
 pytest
-ffmpeg-python
+ffmpeg-python
+spaces