add zero gpu

app.py

@@ -2,6 +2,8 @@ import os
 import gradio as gr
 import numpy as np
 from PIL import Image
+import cv2
+import spaces
 
 from inference.seg import process_image_or_video
 from config import SAPIENS_LITE_MODELS_PATH
@@ -10,7 +12,8 @@ def update_model_choices(task):
     model_choices = list(SAPIENS_LITE_MODELS_PATH[task.lower()].keys())
     return gr.Dropdown(choices=model_choices, value=model_choices[0] if model_choices else None)
 
-def gradio_wrapper(input_image, task, version):
+@spaces.GPU(duration=120)
+def process_image(input_image, task, version):
     if isinstance(input_image, np.ndarray):
         input_image = Image.fromarray(input_image)
     
@@ -18,38 +21,86 @@ def gradio_wrapper(input_image, task, version):
     
     return result
 
+def process_video(input_video, task, version):
+    cap = cv2.VideoCapture(input_video)
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    output_video = cv2.VideoWriter('output_video.mp4', cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
+    
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        
+        frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+        processed_frame = process_image_or_video(frame_rgb, task=task.lower(), version=version)
+        
+        if processed_frame is not None:
+            processed_frame_bgr = cv2.cvtColor(np.array(processed_frame), cv2.COLOR_RGB2BGR)
+            output_video.write(processed_frame_bgr)
+    
+    cap.release()
+    output_video.release()
+    
+    return 'output_video.mp4'
+
 with gr.Blocks() as demo:
-    gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode- Not yet available")
+    gr.Markdown("# Sapiens Arena 🤸🏽‍♂️ - WIP devmode")
     with gr.Tabs():
         with gr.TabItem('Image'):
             with gr.Row():
                 with gr.Column():
                     input_image = gr.Image(label="Input Image", type="pil")
-                    select_task = gr.Radio(
+                    select_task_image = gr.Radio(
                         ["seg", "pose", "depth", "normal"], 
                         label="Task", 
                         info="Choose the task to perform",
                         value="seg"
                     )
-                    model_name = gr.Dropdown(
+                    model_name_image = gr.Dropdown(
                         label="Model Version",
                         choices=list(SAPIENS_LITE_MODELS_PATH["seg"].keys()),
                         value="sapiens_0.3b",
                     )
                 with gr.Column():
                     result_image = gr.Image(label="Result")
-                    run_button = gr.Button("Run")
+                    run_button_image = gr.Button("Run")
         
         with gr.TabItem('Video'):
-            gr.Markdown("In construction")
+            with gr.Row():
+                with gr.Column():
+                    input_video = gr.Video(label="Input Video")
+                    select_task_video = gr.Radio(
+                        ["seg", "pose", "depth", "normal"], 
+                        label="Task", 
+                        info="Choose the task to perform",
+                        value="seg"
+                    )
+                    model_name_video = gr.Dropdown(
+                        label="Model Version",
+                        choices=list(SAPIENS_LITE_MODELS_PATH["seg"].keys()),
+                        value="sapiens_0.3b",
+                    )
+                with gr.Column():
+                    result_video = gr.Video(label="Result")
+                    run_button_video = gr.Button("Run")
 
-    select_task.change(fn=update_model_choices, inputs=select_task, outputs=model_name)
+    select_task_image.change(fn=update_model_choices, inputs=select_task_image, outputs=model_name_image)
+    select_task_video.change(fn=update_model_choices, inputs=select_task_video, outputs=model_name_video)
 
-    run_button.click(
-        fn=gradio_wrapper,
-        inputs=[input_image, select_task, model_name],
+    run_button_image.click(
+        fn=process_image,
+        inputs=[input_image, select_task_image, model_name_image],
         outputs=[result_image],
     )
 
+    run_button_video.click(
+        fn=process_video,
+        inputs=[input_video, select_task_video, model_name_video],
+        outputs=[result_video],
+    )
+
 if __name__ == "__main__":
-    demo.launch(share=True)
+    demo.launch(share=False)