Enhance model handling and CLI functionality in pipeline generator

- Added support for listing all models regardless of file extension with the `--all` flag in the CLI.
- Updated the model listing logic to differentiate between verified and non-verified models.
- Enhanced the model information structure to include detected file extensions and primary extension for better clarity.
- Improved documentation to reflect the new model handling capabilities and supported file formats.

Signed-off-by: Victor Chang <[email protected]>

Author: mocsharp

tools/pipeline-generator/docs/design.md

@@ -30,6 +30,7 @@ As of MONAI Core, bundles can also be exported in a **Hugging Face-compatible fo
 
 - The tool does not convert input formats given that each model may expect a different type of input
 - The tool does not convert output formats given that each model may output a different type of result
+- The tool supports only torchscript (ts) models
 
 ## **Scope**
 

tools/pipeline-generator/pipeline_generator/cli/main.py

@@ -64,61 +64,74 @@ def cli(ctx: click.Context, config: Optional[str]) -> None:
 )
 @click.option("--bundles-only", "-b", is_flag=True, help="Show only MONAI Bundles")
 @click.option("--tested-only", "-t", is_flag=True, help="Show only tested models")
+@click.option("--all", is_flag=True, help="Show all models regardless of file extension")
 @click.pass_context
-def list(ctx: click.Context, format: str, bundles_only: bool, tested_only: bool) -> None:
+def list(ctx: click.Context, format: str, bundles_only: bool, tested_only: bool, all: bool) -> None:
     """List available models from configured endpoints.
 
+    By default, only shows models with TorchScript (.ts) files.
+    Use --all to show models with any supported extension.
+
     Args:
         ctx: Click context containing configuration
         format: Output format (table, simple, or json)
         bundles_only: If True, show only MONAI Bundles
         tested_only: If True, show only tested models
+        all: If True, show all models regardless of file extension
 
     Example:
         pg list --format table --bundles-only
+        pg list --all  # Show all models
     """
 
     # Load configuration
     config_path = ctx.obj.get("config_path")
     settings = load_config(config_path)
 
-    # Get set of tested model IDs from configuration
-    tested_models = set()
+    # Get set of verified model IDs from configuration
+    verified_models = set()
     for endpoint in settings.endpoints:
         for model in endpoint.models:
-            tested_models.add(model.model_id)
+            verified_models.add(model.model_id)
 
-    # Create HuggingFace client
-    client = HuggingFaceClient()
+    # Create HuggingFace client with settings
+    client = HuggingFaceClient(settings=settings)
 
     # Fetch models from all endpoints
     console.print("[blue]Fetching models from HuggingFace...[/blue]")
-    models = client.list_models_from_endpoints(settings.get_all_endpoints())
+    
+    if all:
+        # Show all models, but fetch detailed info for MONAI Bundles to get accurate extension data
+        console.print("[yellow]Note: Fetching detailed info for MONAI Bundles to show accurate extension data[/yellow]")
+        models = client.list_models_from_endpoints(settings.get_all_endpoints(), fetch_details_for_bundles=True)
+    else:
+        # Show only models with TorchScript (.ts) files by default
+        models = client.list_torchscript_models(settings.get_all_endpoints())
 
     # Filter for bundles if requested
     if bundles_only:
         models = [m for m in models if m.is_monai_bundle]
 
-    # Filter for tested models if requested
+    # Filter for verified models if requested
     if tested_only:
-        models = [m for m in models if m.model_id in tested_models]
+        models = [m for m in models if m.model_id in verified_models]
 
     # Sort models by name
     models.sort(key=lambda m: m.display_name)
 
     # Display results based on format
     if format == "table":
-        _display_table(models, tested_models)
+        _display_table(models, verified_models)
     elif format == "simple":
-        _display_simple(models, tested_models)
+        _display_simple(models, verified_models)
     elif format == "json":
-        _display_json(models, tested_models)
+        _display_json(models, verified_models)
 
     # Summary
     bundle_count = sum(1 for m in models if m.is_monai_bundle)
-    tested_count = sum(1 for m in models if m.model_id in tested_models)
+    verified_count = sum(1 for m in models if m.model_id in verified_models)
     console.print(
-        f"\n[green]Total models: {len(models)} (MONAI Bundles: {bundle_count}, Verified: {tested_count})[/green]"
+        f"\n[green]Total models: {len(models)} (MONAI Bundles: {bundle_count}, Verified: {verified_count})[/green]"
     )
 
 
@@ -225,28 +238,37 @@ def gen(
         raise click.Abort() from e
 
 
-def _display_table(models: List[ModelInfo], tested_models: Set[str]) -> None:
+def _display_table(models: List[ModelInfo], verified_models: Set[str]) -> None:
     """Display models in a rich table format.
 
     Args:
         models: List of ModelInfo objects to display
-        tested_models: Set of tested model IDs
+        verified_models: Set of verified model IDs
     """
     table = Table(title="Available Models", show_header=True, header_style="bold magenta")
     table.add_column("Model ID", style="cyan", width=40)
     table.add_column("Name", style="white")
-    table.add_column("Type", style="green")
+    table.add_column("MONAI Bundle", style="green")
     table.add_column("Status", style="blue", width=10)
     table.add_column("Downloads", justify="right", style="yellow")
     table.add_column("Likes", justify="right", style="red")
 
     for model in models:
-        model_type = "[green]MONAI Bundle[/green]" if model.is_monai_bundle else "Model"
-        status = "[bold green]✓ Verified[/bold green]" if model.model_id in tested_models else ""
+        # MONAI Bundle column logic: "Yes" if is_monai_bundle (has .ts), "No (extension)" otherwise
+        if model.is_monai_bundle:
+            bundle_status = "[green]✓ Yes[/green]"
+        else:
+            primary_ext = model.primary_extension
+            if primary_ext:
+                bundle_status = f"[dim]✗ No ({primary_ext})[/dim]"
+            else:
+                bundle_status = "[dim]✗ No[/dim]"
+        
+        status = "[bold green]✓ Verified[/bold green]" if model.model_id in verified_models else ""
         table.add_row(
             model.model_id,
             model.display_name,
-            model_type,
+            bundle_status,
             status,
             str(model.downloads or "N/A"),
             str(model.likes or "N/A"),
@@ -255,31 +277,31 @@ def _display_table(models: List[ModelInfo], tested_models: Set[str]) -> None:
     console.print(table)
 
 
-def _display_simple(models: List[ModelInfo], tested_models: Set[str]) -> None:
+def _display_simple(models: List[ModelInfo], verified_models: Set[str]) -> None:
     """Display models in a simple list format.
 
     Shows each model with emoji indicators:
     - 📦 for MONAI Bundle, 📄 for regular model
-    - ✓ for tested models
+    - ✓ for verified models
 
     Args:
         models: List of ModelInfo objects to display
-        tested_models: Set of tested model IDs
+        verified_models: Set of verified model IDs
     """
     for model in models:
         bundle_marker = "📦" if model.is_monai_bundle else "📄"
-        tested_marker = " ✓" if model.model_id in tested_models else ""
-        console.print(f"{bundle_marker} {model.model_id} - {model.display_name}{tested_marker}")
+        verified_marker = " ✓" if model.model_id in verified_models else ""
+        console.print(f"{bundle_marker} {model.model_id} - {model.display_name}{verified_marker}")
 
 
-def _display_json(models: List[ModelInfo], tested_models: Set[str]) -> None:
+def _display_json(models: List[ModelInfo], verified_models: Set[str]) -> None:
     """Display models in JSON format.
 
     Outputs a JSON array of model information suitable for programmatic consumption.
 
     Args:
         models: List of ModelInfo objects to display
-        tested_models: Set of tested model IDs
+        verified_models: Set of verified model IDs
     """
     import json
 
@@ -288,7 +310,10 @@ def _display_json(models: List[ModelInfo], tested_models: Set[str]) -> None:
             "model_id": m.model_id,
             "name": m.display_name,
             "is_monai_bundle": m.is_monai_bundle,
-            "is_tested": m.model_id in tested_models,
+            "has_torchscript": m.has_torchscript,
+            "model_extensions": m.model_extensions,
+            "primary_extension": m.primary_extension,
+            "is_verified": m.model_id in verified_models,
             "downloads": m.downloads,
             "likes": m.likes,
             "tags": m.tags,

tools/pipeline-generator/pipeline_generator/config/config.yaml

@@ -11,6 +11,13 @@
 
 # Pipeline Generator Configuration
 
+# Supported model file extensions
+supported_models:
+  - ".ts"      # TorchScript (MONAI Bundle)
+  - ".pt"      # PyTorch state dict
+  - ".pth"     # PyTorch model
+  - ".safetensors"  # SafeTensor format
+
 # HuggingFace endpoints to scan for MONAI models
 endpoints:
   - organization: "MONAI"

tools/pipeline-generator/pipeline_generator/config/settings.py

@@ -54,6 +54,10 @@ class Endpoint(BaseModel):
 class Settings(BaseModel):
     """Application settings."""
 
+    supported_models: List[str] = Field(
+        default_factory=lambda: [".ts", ".pt", ".pth", ".safetensors"],
+        description="Supported model file extensions"
+    )
     endpoints: List[Endpoint] = Field(default_factory=list)
     additional_models: List[Endpoint] = Field(default_factory=list)
 

tools/pipeline-generator/pipeline_generator/core/hub_client.py

@@ -17,7 +17,7 @@
 from huggingface_hub import HfApi, list_models, model_info
 from huggingface_hub.utils import HfHubHTTPError
 
-from ..config import Endpoint
+from ..config import Endpoint, Settings
 from .models import ModelInfo
 
 logger = logging.getLogger(__name__)
@@ -26,9 +26,17 @@
 class HuggingFaceClient:
     """Client for interacting with HuggingFace Hub."""
 
-    def __init__(self) -> None:
-        """Initialize the HuggingFace Hub client."""
+    def __init__(self, settings: Optional[Settings] = None) -> None:
+        """Initialize the HuggingFace Hub client.
+        
+        Args:
+            settings: Pipeline generator settings containing supported extensions
+        """
         self.api = HfApi()
+        self.settings = settings
+        self.supported_extensions = (
+            settings.supported_models if settings else [".ts", ".pt", ".pth", ".safetensors"]
+        )
 
     def list_models_from_organization(self, organization: str) -> List[ModelInfo]:
         """List all models from a HuggingFace organization.
@@ -71,11 +79,12 @@ def get_model_info(self, model_id: str) -> Optional[ModelInfo]:
             logger.error(f"Error fetching model {model_id}: {e}")
             return None
 
-    def list_models_from_endpoints(self, endpoints: List[Endpoint]) -> List[ModelInfo]:
+    def list_models_from_endpoints(self, endpoints: List[Endpoint], fetch_details_for_bundles: bool = False) -> List[ModelInfo]:
         """List models from all configured endpoints.
 
         Args:
             endpoints: List of endpoint configurations
+            fetch_details_for_bundles: If True, fetch detailed info for potential MONAI Bundles to get accurate extension info
 
         Returns:
             List of ModelInfo objects from all endpoints
@@ -87,6 +96,29 @@ def list_models_from_endpoints(self, endpoints: List[Endpoint]) -> List[ModelInf
                 # List all models from organization
                 logger.info(f"Fetching models from organization: {endpoint.organization}")
                 models = self.list_models_from_organization(endpoint.organization)
+                
+                if fetch_details_for_bundles:
+                    # For MONAI organization, fetch detailed info for all models to get accurate extension data
+                    # This is needed because bulk API doesn't provide file information (siblings=None)
+                    if endpoint.organization == "MONAI":
+                        enhanced_models = []
+                        for model in models:
+                            # Fetch detailed model info to get file extensions and accurate MONAI Bundle detection
+                            detailed_model = self.get_model_info(model.model_id)
+                            enhanced_models.append(detailed_model if detailed_model else model)
+                        models = enhanced_models
+                    else:
+                        # For non-MONAI organizations, only fetch details for models that might be bundles
+                        enhanced_models = []
+                        for model in models:
+                            if any("monai" in tag.lower() for tag in model.tags):
+                                # Fetch detailed model info to get file extensions
+                                detailed_model = self.get_model_info(model.model_id)
+                                enhanced_models.append(detailed_model if detailed_model else model)
+                            else:
+                                enhanced_models.append(model)
+                        models = enhanced_models
+                
                 all_models.extend(models)
 
             elif endpoint.model_id:
@@ -98,6 +130,66 @@ def list_models_from_endpoints(self, endpoints: List[Endpoint]) -> List[ModelInf
 
         return all_models
 
+    def _detect_model_extensions(self, model_data: Any) -> List[str]:
+        """Detect model file extensions in a HuggingFace repository.
+
+        Args:
+            model_data: Model data from HuggingFace API
+
+        Returns:
+            List of detected model file extensions
+        """
+        extensions = []
+        
+        try:
+            if hasattr(model_data, "siblings") and model_data.siblings is not None:
+                file_names = [f.rfilename for f in model_data.siblings]
+                for filename in file_names:
+                    for ext in self.supported_extensions:
+                        if filename.endswith(ext):
+                            if ext not in extensions:
+                                extensions.append(ext)
+        except Exception as e:
+            logger.debug(f"Could not detect extensions for {getattr(model_data, 'modelId', 'unknown')}: {e}")
+        
+        return extensions
+
+    def list_torchscript_models(self, endpoints: List[Endpoint]) -> List[ModelInfo]:
+        """List models that have TorchScript (.ts) files.
+
+        This method fetches detailed information for each model individually to 
+        check for TorchScript files, which is slower than bulk listing but accurate.
+
+        Args:
+            endpoints: List of endpoint configurations
+
+        Returns:
+            List of ModelInfo objects that contain .ts files
+        """
+        torchscript_models = []
+        
+        for endpoint in endpoints:
+            if endpoint.organization:
+                # List all models from organization first (bulk)
+                logger.info(f"Checking TorchScript models from organization: {endpoint.organization}")
+                try:
+                    for model in list_models(author=endpoint.organization):
+                        # Fetch detailed model info to get file information
+                        detailed_model = self.get_model_info(model.modelId)
+                        if detailed_model and detailed_model.has_torchscript:
+                            torchscript_models.append(detailed_model)
+                except Exception as e:
+                    logger.error(f"Error checking TorchScript models from {endpoint.organization}: {e}")
+
+            elif endpoint.model_id:
+                # Get specific model
+                logger.info(f"Checking TorchScript model: {endpoint.model_id}")
+                model = self.get_model_info(endpoint.model_id)
+                if model and model.has_torchscript:
+                    torchscript_models.append(model)
+
+        return torchscript_models
+
     def _extract_model_info(self, model_data: Any) -> ModelInfo:
         """Extract ModelInfo from HuggingFace model data.
 
@@ -107,23 +199,14 @@ def _extract_model_info(self, model_data: Any) -> ModelInfo:
         Returns:
             ModelInfo object
         """
-        # Check if this is a MONAI Bundle
-        is_monai_bundle = False
+        # Detect model extensions
+        model_extensions = self._detect_model_extensions(model_data)
+
+        # Check if this is a MONAI Bundle - defined as having TorchScript (.ts) files
+        is_monai_bundle = ".ts" in model_extensions
         bundle_metadata = None
 
-        # Check tags for MONAI-related tags
         tags = getattr(model_data, "tags", [])
-        if any("monai" in tag.lower() for tag in tags):
-            is_monai_bundle = True
-
-        # Check if metadata.json exists in the model files
-        try:
-            if hasattr(model_data, "siblings"):
-                file_names = [f.rfilename for f in model_data.siblings]
-                if any("metadata.json" in f for f in file_names):
-                    is_monai_bundle = True
-        except Exception:
-            pass
 
         # Extract description from cardData if available
         description = None
@@ -145,4 +228,5 @@ def _extract_model_info(self, model_data: Any) -> ModelInfo:
             tags=tags,
             is_monai_bundle=is_monai_bundle,
             bundle_metadata=bundle_metadata,
+            model_extensions=model_extensions,
         )