feat: Enable MCP

Hello! This is an automated PR adding MCP compatibility to your AI App 🤖.

This PR introduces two improvements:
1. Adds docstrings to the functions in the app file that are directly connected to the Gradio UI, for the downstream LLM to use.
2. Enables the Model-Compute-Platform by adding `mcp_server=True` to the `.launch()` call.

No other logic has been changed. Please review and merge if it looks good!Learn more about MCP compatibility in Spaces here: https://huggingface.co/changelog/add-compatible-spaces-to-your-mcp-tools

app.py

@@ -51,6 +51,17 @@ def add_contour(img, mask, color=(1., 1., 1.)):
 
 @spaces.GPU(duration=120)
 def generate_masks(image, mask_list, mask_raw_list):
+    """
+    Generate masks from user-drawn annotations on an image.
+    
+    Args:
+        image: Dictionary containing the image editor state with background and layers
+        mask_list: List of generated mask images with labels
+        mask_raw_list: List of raw numpy arrays of masks
+        
+    Returns:
+        Tuple containing updated mask_list, image editor state, mask_list, and mask_raw_list
+    """
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
     assert len(image['layers']) == 1, f"Expected 1 layer, got {len(image['layers'])}"
@@ -77,6 +88,17 @@ def generate_masks(image, mask_list, mask_raw_list):
 
 @spaces.GPU(duration=120)
 def generate_masks_video(image, mask_list_video, mask_raw_list_video):
+    """
+    Generate masks from user-drawn annotations on a video frame.
+    
+    Args:
+        image: Dictionary containing the image editor state with background and layers
+        mask_list_video: List of generated mask images with labels for video
+        mask_raw_list_video: List of raw numpy arrays of masks for video
+        
+    Returns:
+        Tuple containing updated mask_list_video, image editor state, mask_list_video, and mask_raw_list_video
+    """
     image['image'] = image['background'].convert('RGB')
     # del image['background'], image['composite']
     assert len(image['layers']) == 1, f"Expected 1 layer, got {len(image['layers'])}"
@@ -104,6 +126,18 @@ def generate_masks_video(image, mask_list_video, mask_raw_list_video):
 
 @spaces.GPU(duration=120)
 def describe(image, mode, query, masks):
+    """
+    Generate descriptions or answer questions about regions in an image.
+    
+    Args:
+        image: Dictionary containing the image editor state
+        mode: Either "Caption" or "QA" mode
+        query: Question to ask about the image (used in QA mode)
+        masks: List of mask arrays for the regions
+        
+    Returns:
+        Generator yielding image with contours, generated text, and updated image state
+    """
     # Create an image object from the uploaded image
     # print(image.keys())
 
@@ -194,6 +228,15 @@ def describe(image, mode, query, masks):
 
   
 def load_first_frame(video_path):
+    """
+    Load and return the first frame of a video.
+    
+    Args:
+        video_path: Path to the video file
+        
+    Returns:
+        PIL Image of the first frame
+    """
     cap = cv2.VideoCapture(video_path)
     ret, frame = cap.read()
     cap.release()
@@ -205,6 +248,20 @@ def load_first_frame(video_path):
 
 @spaces.GPU(duration=120)
 def describe_video(video_path, mode, query, annotated_frame, masks, mask_list_video):
+    """
+    Generate descriptions or answer questions about regions in a video.
+    
+    Args:
+        video_path: Path to the video file
+        mode: Either "Caption" or "QA" mode
+        query: Question to ask about the video (used in QA mode)
+        annotated_frame: Dictionary containing the annotated first frame
+        masks: List of mask arrays for the regions
+        mask_list_video: List of mask images with labels
+        
+    Returns:
+        Generator yielding frame image, generated text, and updated mask lists
+    """
     # Create a temporary directory to save extracted video frames
     cap = cv2.VideoCapture(video_path)
 
@@ -294,247 +351,4 @@ def describe_video(video_path, mode, query, annotated_frame, masks, mask_list_vi
         mask_image = Image.fromarray((mask_np[:,:,np.newaxis] * np.array(annotated_frame['image'])).astype(np.uint8))
         mask_list_video.append((mask_image, f"<object{len(mask_list_video)}>"))
     text = ""
-    yield frame_img, text, mask_list_video, mask_list_video
-
-    for token in get_model_output(
-        video_tensor,
-        query,
-        model=model,
-        tokenizer=tokenizer,
-        masks=masks,
-        mask_ids=mask_ids,
-        modal='video',
-        streaming=True,
-    ):
-        text += token
-        yield gr.update(), text, gr.update(), 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)
-
-    with torch.no_grad():
-        outputs = sam_model(**inputs)
-
-    masks = sam_processor.image_processor.post_process_masks(outputs.pred_masks.cpu(), inputs["original_sizes"].cpu(), inputs["reshaped_input_sizes"].cpu())[0][0]
-    scores = outputs.iou_scores[0, 0]
-
-    mask_selection_index = scores.argmax()
-
-    mask_np = masks[mask_selection_index].numpy()
-
-    return mask_np
-
-
-def clear_masks():
-    return [], [], []
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="VideoRefer gradio demo")
-    parser.add_argument("--model-path", type=str, default="DAMO-NLP-SG/VideoRefer-VideoLLaMA3-7B", help="Path to the model checkpoint")
-    parser.add_argument("--prompt-mode", type=str, default="focal_prompt", help="Prompt mode")
-    parser.add_argument("--conv-mode", type=str, default="v1", help="Conversation mode")
-    parser.add_argument("--temperature", type=float, default=0.2, help="Sampling temperature")
-    parser.add_argument("--top_p", type=float, default=0.5, help="Top-p for sampling")
-
-    args_cli = parser.parse_args()
-
-    with gr.Blocks(theme=gr.themes.Soft(primary_hue="amber")) as demo:
-
-        mask_list = gr.State([])  
-        mask_raw_list = gr.State([])  
-        mask_list_video = gr.State([])  
-        mask_raw_list_video = gr.State([])  
-
-
-        HEADER = ("""
-            <div>
-                <h1>VideoRefer X VideoLLaMA3 Demo</h1>
-                <h5 style="margin: 0;">Feel free to click on anything that grabs your interest!</h5>
-                <h5 style="margin: 0;">If this demo please you, please give us a star ⭐ on Github or 💖 on this space.</h5>
-            </div>
-            </div>
-            <div style="display: flex; justify-content: left; margin-top: 10px;">
-            <a href="https://arxiv.org/pdf/2501.00599"><img src="https://img.shields.io/badge/Arxiv-2501.00599-ECA8A7" style="margin-right: 5px;"></a>
-            <a href="https://github.com/DAMO-NLP-SG/VideoRefer"><img src='https://img.shields.io/badge/Github-VideoRefer-F7C97E' style="margin-right: 5px;"></a>
-            <a href="https://github.com/DAMO-NLP-SG/VideoLLaMA3"><img src='https://img.shields.io/badge/Github-VideoLLaMA3-9DC3E6' style="margin-right: 5px;"></a>
-            </div>
-            """)
-
-        with gr.Row():
-            with gr.Column():
-                gr.HTML(HEADER)
-    
-
-        image_tips = """
-                ### 💡 Tips:
-
-                🧸 Upload an image, and you can use the drawing tool✍️ to highlight the areas you're interested in.
-            
-                🔖 For single-object caption mode, simply select the area and click the 'Generate Caption' button to receive a caption for the object.
-                
-                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
-                
-                📌 Click the button 'Clear Masks' to clear the current generated masks.
-                
-                """
-        
-        video_tips = """
-                ### 💡 Tips:
-                ⚠️ For video mode, we only support masking on the first frame in this demo.
-
-                🧸 Upload an video, and you can use the drawing tool✍️ to highlight the areas you're interested in the first frame.
-            
-                🔖 For single-object caption mode, simply select the area and click the 'Generate Caption' button to receive a caption for the object.
-                
-                🔔 In QA mode, you can generate multiple masks by clicking the 'Generate Mask' button multiple times. Afterward, use the corresponding object id to ask questions.
-                
-                📌 Click the button 'Clear Masks' to clear the current generated masks.
-                
-                """
-  
-
-        with gr.TabItem("Image"):
-            with gr.Row():
-                with gr.Column():
-                    image_input = gr.ImageEditor(
-                        label="Image",
-                        type="pil", 
-                        sources=['upload'], 
-                        brush=gr.Brush(colors=["#ED7D31"], color_mode="fixed", default_size=10),
-                        eraser=True,
-                        layers=False,
-                        transforms=[],
-                        height=300,
-                    )
-                    generate_mask_btn = gr.Button("1️⃣ Generate Mask", visible=False, variant="primary")
-                    mode = gr.Radio(label="Mode", choices=["Caption", "QA"], value="Caption")
-                    query = gr.Textbox(label="Question", value="What is the relationship between <region0> and <region1>?", interactive=True, visible=False)
-                    
-                    submit_btn = gr.Button("Generate Caption", variant="primary")
-                    submit_btn1 = gr.Button("2️⃣ Generate Answer", variant="primary", visible=False)
-                    gr.Examples([f"./demo/images/{i+1}.jpg" for i in range(8)], inputs=image_input, label="Examples")
-    
-                with gr.Column():
-                    mask_output = gr.Gallery(label="Referred Masks", object_fit='scale-down', visible=False)
-                    output_image = gr.Image(label="Image with Mask", visible=True, height=400)
-                    description = gr.Textbox(label="Output", visible=True)
-                    
-                    clear_masks_btn = gr.Button("Clear Masks", variant="secondary", visible=False)
-            gr.Markdown(image_tips)
-
-        with gr.TabItem("Video"):
-            with gr.Row():
-                with gr.Column():
-                    video_input = gr.Video(label="Video")
-                    # load_btn = gr.Button("🖼️ Load First Frame", variant="secondary")
-                    first_frame = gr.ImageEditor(
-                        label="Annotate First Frame",
-                        type="pil", 
-                        sources=['upload'], 
-                        brush=gr.Brush(colors=["#ED7D31"], color_mode="fixed", default_size=10),
-                        eraser=True,
-                        layers=False,
-                        transforms=[],
-                        height=300,
-                    )
-                    generate_mask_btn_video = gr.Button("1️⃣ Generate Mask", visible=False, variant="primary")
-                    gr.Examples([f"./demo/videos/{i+1}.mp4" for i in range(4)], inputs=video_input, label="Examples")
-
-                with gr.Column():
-                    mode_video = gr.Radio(label="Mode", choices=["Caption", "QA"], value="Caption")
-                    mask_output_video = gr.Gallery(label="Referred Masks", object_fit='scale-down')
-
-                    query_video = gr.Textbox(label="Question", value="What is the relationship between <object0> and <object1>?", interactive=True, visible=False)
-
-                    submit_btn_video = gr.Button("Generate Caption", variant="primary")
-                    submit_btn_video1 = gr.Button("2️⃣ Generate Answer", variant="primary", visible=False)
-                    description_video = gr.Textbox(label="Output", visible=True)
-                    
-                    clear_masks_btn_video = gr.Button("Clear Masks", variant="secondary")
-
-            gr.Markdown(video_tips)
-
-        
-        def toggle_query_and_generate_button(mode):
-            query_visible = mode == "QA"
-            caption_visible = mode == "Caption"
-            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), gr.update(visible=caption_visible), [], "", [], [],[],[]
-
-        video_input.change(load_first_frame, inputs=video_input, outputs=first_frame)
-
-        mode.change(toggle_query_and_generate_button, inputs=mode, outputs=[query, generate_mask_btn, clear_masks_btn, submit_btn1, mask_output, output_image, submit_btn, mask_output, description, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
-        
-        def toggle_query_and_generate_button_video(mode):
-            query_visible = mode == "QA"
-            caption_visible = mode == "Caption"
-            return gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=query_visible), gr.update(visible=caption_visible), [], [], [], [], []
-
-
-        mode_video.change(toggle_query_and_generate_button_video, inputs=mode_video, outputs=[query_video, generate_mask_btn_video, submit_btn_video1, submit_btn_video, mask_output_video, mask_list, mask_raw_list, mask_list_video, mask_raw_list_video])
-
-        submit_btn.click(
-            fn=describe,
-            inputs=[image_input, mode, query, mask_raw_list],
-            outputs=[output_image, description, image_input],
-            api_name="describe"
-        )
-
-        submit_btn1.click(
-            fn=describe,
-            inputs=[image_input, mode, query, mask_raw_list],
-            outputs=[output_image, description, image_input],
-            api_name="describe"
-        )
-
-        generate_mask_btn.click(
-            fn=generate_masks,
-            inputs=[image_input, mask_list, mask_raw_list],
-            outputs=[mask_output, image_input, mask_list, mask_raw_list]
-        )
-
-        generate_mask_btn_video.click(
-            fn=generate_masks_video,
-            inputs=[first_frame, mask_list_video, mask_raw_list_video],
-            outputs=[mask_output_video, first_frame, mask_list_video, mask_raw_list_video]
-        )
-
-        clear_masks_btn.click(
-            fn=clear_masks,
-            outputs=[mask_output, mask_list, mask_raw_list]
-        )
-
-        clear_masks_btn_video.click(
-            fn=clear_masks,
-            outputs=[mask_output_video, mask_list_video, mask_raw_list_video]
-        )
-
-        submit_btn_video.click(
-            fn=describe_video,
-            inputs=[video_input, mode_video, query_video, first_frame, mask_raw_list_video, mask_list_video],
-            outputs=[first_frame, description_video, mask_output_video, mask_list_video],
-            api_name="describe_video"
-        )
-
-        submit_btn_video1.click(
-            fn=describe_video,
-            inputs=[video_input, mode_video, query_video, first_frame, mask_raw_list_video, mask_list_video],
-            outputs=[first_frame, description_video, mask_output_video, mask_list_video],
-            api_name="describe_video"
-        )
-
-
-
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    sam_model = SamModel.from_pretrained("facebook/sam-vit-huge").to(device)
-    sam_processor = SamProcessor.from_pretrained("facebook/sam-vit-huge")
-
-    disable_torch_init()
-
-
-    model, processor, tokenizer = model_init(args_cli.model_path)
-    
-
-    demo.launch()
+    yield frame_img, text, mask_list_video, mask_list_video