update app with model selection

app.py

@@ -23,7 +23,7 @@ matplotlib.use("Agg")  # Use non-interactive backend
 import matplotlib.pyplot as plt
 import numpy as np
 import requests
-from typing import Any, Optional, Tuple
+from typing import Any, List, Optional, Tuple
 
 from rfm.data.dataset_types import Trajectory, ProgressSample, PreferenceSample, SimilaritySample
 from rfm.evals.eval_utils import build_payload, post_batch_npy
@@ -62,14 +62,72 @@ PREDEFINED_DATASETS = [
     "aliangdw/usc_franka_policy_ranking",
     "aliangdw/utd_so101_policy_ranking",
     "aliangdw/utd_so101_human",
+    "jesbu1/utd_so101_clean_policy_ranking_top",
+    "jesbu1/utd_so101_clean_policy_ranking_wrist",
+    "jesbu1/mit_franka_p-rank_rfm",
+    "jesbu1/usc_koch_p_ranking_rfm",
 ]
 
 # Global server state
 _server_state = {
     "server_url": None,
+    "base_url": "http://40.119.56.66",  # Default base URL
 }
 
 
+def discover_available_models(base_url: str = "http://40.119.56.66", port_range: tuple = (8000, 8010)) -> List[Tuple[str, str]]:
+    """Discover available models by pinging ports in the specified range.
+    
+    Returns:
+        List of tuples: [(server_url, model_name), ...]
+    """
+    available_models = []
+    start_port, end_port = port_range
+    
+    for port in range(start_port, end_port + 1):
+        server_url = f"{base_url.rstrip('/')}:{port}"
+        try:
+            # Check health endpoint
+            health_url = f"{server_url}/health"
+            health_response = requests.get(health_url, timeout=2.0)
+            if health_response.status_code == 200:
+                # Try to get model info for model name
+                try:
+                    model_info_url = f"{server_url}/model_info"
+                    model_info_response = requests.get(model_info_url, timeout=2.0)
+                    if model_info_response.status_code == 200:
+                        model_info_data = model_info_response.json()
+                        model_name = model_info_data.get("model_path", f"Model on port {port}")
+                        available_models.append((server_url, model_name))
+                    else:
+                        # Health check passed but no model info, use port as name
+                        available_models.append((server_url, f"Model on port {port}"))
+                except:
+                    # Health check passed but couldn't get model info
+                    available_models.append((server_url, f"Model on port {port}"))
+        except requests.exceptions.RequestException:
+            # Port not available, continue
+            continue
+    
+    return available_models
+
+
+def get_model_info_for_url(server_url: str) -> Optional[str]:
+    """Get formatted model info for a given server URL."""
+    if not server_url:
+        return None
+    
+    try:
+        model_info_url = server_url.rstrip("/") + "/model_info"
+        model_info_response = requests.get(model_info_url, timeout=5.0)
+        if model_info_response.status_code == 200:
+            model_info_data = model_info_response.json()
+            return format_model_info(model_info_data)
+    except Exception as e:
+        logger.warning(f"Could not fetch model info: {e}")
+    return None
+
+
 def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[str]]:
     """Check server health and get model info."""
     if not server_url:
@@ -92,15 +150,7 @@ def check_server_health(server_url: str) -> Tuple[str, Optional[dict], Optional[
             pass
 
         # Try to get model info
-        model_info_text = None
-        try:
-            model_info_url = server_url.rstrip("/") + "/model_info"
-            model_info_response = requests.get(model_info_url, timeout=5.0)
-            if model_info_response.status_code == 200:
-                model_info_data = model_info_response.json()
-                model_info_text = format_model_info(model_info_data)
-        except Exception as e:
-            logger.warning(f"Could not fetch model info: {e}")
+        model_info_text = get_model_info_for_url(server_url)
 
         _server_state["server_url"] = server_url
         return (
@@ -271,11 +321,12 @@ def process_single_video(
     fps: float = 1.0,
 ) -> Tuple[Optional[str], Optional[str]]:
     """Process single video for progress and success predictions using eval server."""
+    # Get server URL from state if not provided
     if not server_url:
-        return None, "Please provide a server URL and check connection first."
-
-    if not _server_state.get("server_url"):
-        return None, "Server not connected. Please check server connection first."
+        server_url = _server_state.get("server_url")
+    
+    if not server_url:
+        return None, "Please select a model from the dropdown above and ensure it's connected."
 
     if video_path is None:
         return None, "Please provide a video."
@@ -377,11 +428,12 @@ def process_two_videos(
     fps: float = 1.0,
 ) -> Tuple[Optional[str], Optional[str], Optional[str]]:
     """Process two videos for preference, similarity, or progress prediction using eval server."""
+    # Get server URL from state if not provided
     if not server_url:
-        return "Please provide a server URL and check connection first.", None, None
-
-    if not _server_state.get("server_url"):
-        return "Server not connected. Please check server connection first.", None, None
+        server_url = _server_state.get("server_url")
+    
+    if not server_url:
+        return "Please select a model from the dropdown above and ensure it's connected.", None, None
 
     if video_a_path is None or video_b_path is None:
         return "Please provide both videos.", None, None
@@ -558,41 +610,130 @@ with demo:
         """
         # RFM (Reward Foundation Model) Evaluation Server
         
-        **Note:** This app connects to an eval server. Please provide the server URL and check connection before use.
+        Select a model from the dropdown below. The app will automatically discover available models.
         """
     )
 
-    with gr.Tab("Server Setup"):
-        gr.Markdown("### Connect to Eval Server")
-        gr.Markdown("Enter the eval server URL and check connection.")
-
-        with gr.Row():
-            with gr.Column(scale=3):
-                server_url_input = gr.Textbox(
-                    label="Server URL",
-                    placeholder="http://40.119.56.66:8000",
-                    value="http://40.119.56.66:8000",
-                    interactive=True,
-                )
-            with gr.Column(scale=1):
-                check_connection_btn = gr.Button("Check Connection", variant="primary", size="sm")
+    # Model selector at the top
+    with gr.Row():
+        with gr.Column(scale=4):
+            base_url_input = gr.Textbox(
+                label="Base Server URL",
+                placeholder="http://40.119.56.66",
+                value="http://40.119.56.66",
+                interactive=True,
+            )
+            model_dropdown = gr.Dropdown(
+                label="Select Model",
+                choices=[],
+                value=None,
+                interactive=True,
+                info="Click 'Discover Models' to find available models on ports 8000-8010",
+            )
+        with gr.Column(scale=1):
+            discover_btn = gr.Button("🔍 Discover Models", variant="primary", size="lg")
+    
+    with gr.Row():
+        server_status = gr.Markdown("Click 'Discover Models' to find available models", visible=True)
+    
+    with gr.Accordion("📋 Model Information", open=False) as model_info_accordion:
+        model_info_display = gr.Markdown("", visible=True)
+
+    # Hidden state to store server URL and model mapping
+    server_url_state = gr.State(value=None)
+    model_url_mapping_state = gr.State(value={})  # Maps model_name -> server_url
+
+    def discover_and_select_models(base_url: str):
+        """Discover models and update dropdown."""
+        if not base_url:
+            return (
+                gr.update(choices=[], value=None),
+                gr.update(value="Please provide a base URL", visible=True),
+                gr.update(value="", visible=True),
+                None,
+                {},  # Empty mapping
+            )
+        
+        _server_state["base_url"] = base_url
+        models = discover_available_models(base_url, port_range=(8000, 8010))
+        
+        if not models:
+            return (
+                gr.update(choices=[], value=None),
+                gr.update(value="❌ No models found on ports 8000-8010. Make sure servers are running.", visible=True),
+                gr.update(value="", visible=True),
+                None,
+                {},  # Empty mapping
+            )
+        
+        # Format choices: show model_name in dropdown
+        # Store mapping of model_name to URL in state
+        choices = []
+        url_map = {}
+        for url, name in models:
+            choices.append(name)
+            url_map[name] = url
+        
+        # Auto-select first model
+        selected_choice = choices[0] if choices else None
+        selected_url = url_map.get(selected_choice) if selected_choice else None
+        
+        # Get model info for selected model
+        model_info_text = get_model_info_for_url(selected_url) if selected_url else ""
+        status_text = f"✅ Found {len(models)} model(s). Auto-selected first model."
+        
+        _server_state["server_url"] = selected_url
+        
+        return (
+            gr.update(choices=choices, value=selected_choice),
+            gr.update(value=status_text, visible=True),
+            gr.update(value=model_info_text, visible=True),
+            selected_url,
+            url_map,  # Return mapping for state
+        )
 
-        server_status = gr.Markdown("Enter server URL and click 'Check Connection'")
-        model_info_display = gr.Markdown("", visible=False)
+    def on_model_selected(model_choice: str, url_mapping: dict):
+        """Handle model selection change."""
+        if not model_choice:
+            return (
+                gr.update(value="No model selected", visible=True),
+                gr.update(value="", visible=True),
+                None,
+            )
+        
+        # Get URL from mapping
+        server_url = url_mapping.get(model_choice) if url_mapping else None
+        
+        if not server_url:
+            return (
+                gr.update(value="Could not find server URL for selected model. Please rediscover models.", visible=True),
+                gr.update(value="", visible=True),
+                None,
+            )
+        
+        # Get model info
+        model_info_text = get_model_info_for_url(server_url) or ""
+        status, health_data, _ = check_server_health(server_url)
+        
+        _server_state["server_url"] = server_url
+        
+        return (
+            gr.update(value=status, visible=True),
+            gr.update(value=model_info_text, visible=True),
+            server_url,
+        )
 
-        def on_check_connection(server_url: str):
-            """Handle server connection check."""
-            status, health_data, model_info_text = check_server_health(server_url)
-            if model_info_text:
-                return status, gr.update(value=model_info_text, visible=True)
-            else:
-                return status, gr.update(visible=False)
+    discover_btn.click(
+        fn=discover_and_select_models,
+        inputs=[base_url_input],
+        outputs=[model_dropdown, server_status, model_info_display, server_url_state, model_url_mapping_state],
+    )
 
-        check_connection_btn.click(
-            fn=on_check_connection,
-            inputs=[server_url_input],
-            outputs=[server_status, model_info_display],
-        )
+    model_dropdown.change(
+        fn=on_model_selected,
+        inputs=[model_dropdown, model_url_mapping_state],
+        outputs=[server_status, model_info_display, server_url_state],
+    )
 
     with gr.Tab("Progress Prediction"):
         gr.Markdown("### Progress & Success Prediction")
@@ -851,7 +992,7 @@ with demo:
 
         analyze_single_btn.click(
             fn=process_single_video,
-            inputs=[single_video_input, task_text_input, server_url_input, fps_input_single],
+            inputs=[single_video_input, task_text_input, server_url_state, fps_input_single],
             outputs=[progress_plot, info_output],
             api_name="process_single_video",
         )
@@ -1329,7 +1470,7 @@ with demo:
 
         analyze_dual_btn.click(
             fn=process_two_videos,
-            inputs=[video_a_input, video_b_input, task_text_dual, prediction_type, server_url_input, fps_input_dual],
+            inputs=[video_a_input, video_b_input, task_text_dual, prediction_type, server_url_state, fps_input_dual],
             outputs=[result_text, video_a_display, video_b_display],
             api_name="process_two_videos",
         )