run with concurrency

app.py

@@ -1,12 +1,31 @@
-import os
 import torch
 import gradio as gr
 
+
+def get_free_memory_gb():
+    gpu_index = torch.cuda.current_device()
+    # Get the GPU's properties
+    gpu_properties = torch.cuda.get_device_properties(gpu_index)
+
+    # Get the total and allocated memory
+    total_memory = gpu_properties.total_memory
+    allocated_memory = torch.cuda.memory_allocated(gpu_index)
+
+    # Calculate the free memory
+    free_memory = total_memory - allocated_memory
+    return free_memory / 1024**3
+
+
 model = torch.hub.load("PeterL1n/RobustVideoMatting", "mobilenetv3")
 
 if torch.cuda.is_available():
-    print("Using GPU")
+    free_memory = get_free_memory_gb()
+    concurrency_count = int(free_memory // 7.4)
     model = model.cuda()
+    print(f"Using GPU with concurrency: {concurrency_count}")
+else:
+    print("Using CPU")
+    concurrency_count = 1
 
 convert_video = torch.hub.load("PeterL1n/RobustVideoMatting", "converter")
 
@@ -47,4 +66,4 @@ with gr.Blocks(title="Robust Video Matting") as block:
         "<p style='text-align: center'><a href='https://arxiv.org/abs/2108.11515'>Robust High-Resolution Video Matting with Temporal Guidance</a> | <a href='https://github.com/PeterL1n/RobustVideoMatting'>Github Repo</a></p>"
     )
 
-block.queue(api_open=False, max_size=5).launch()
+block.queue(api_open=False, max_size=5, concurrency_count=concurrency_count).launch()