Merge pull request #1641 from roboflow/feature/model-variant-in-roboflow-api

Add ability for inference-exp to parse modelVariant from RFAPI

Author: PawelPeczek-Roboflow

inference_experimental/inference_exp/models/auto_loaders/core.py

@@ -120,6 +120,7 @@ def describe_model(
             model_id=model_metadata.model_id,
             requested_model_id=model_id,
             model_architecture=model_metadata.model_architecture,
+            model_variant=model_metadata.model_variant,
             task_type=model_metadata.task_type,
             weights_provider=weights_provider,
             registered_packages=len(model_metadata.model_packages),

inference_experimental/inference_exp/models/auto_loaders/presentation_utils.py

@@ -20,6 +20,7 @@ def render_table_with_model_overview(
     model_id: str,
     requested_model_id: str,
     model_architecture: str,
+    model_variant: Optional[str],
     task_type: Optional[str],
     weights_provider: str,
     registered_packages: int,
@@ -32,6 +33,7 @@ def render_table_with_model_overview(
         model_id_str = f"{model_id_str} (alias: {requested_model_id})"
     table.add_row("Model ID:", model_id_str)
     table.add_row("Architecture:", model_architecture)
+    table.add_row("Variant:", model_variant or "N/A")
     table.add_row("Task:", task_type or "N/A")
     table.add_row("Weights provider:", weights_provider)
     table.add_row("Number of packages:", str(registered_packages))

inference_experimental/inference_exp/weights_providers/entities.py

@@ -156,3 +156,4 @@ class ModelMetadata:
     model_architecture: str
     model_packages: List[ModelPackageMetadata]
     task_type: Optional[str] = field(default=None)
+    model_variant: Optional[str] = field(default=None)

inference_experimental/inference_exp/weights_providers/roboflow.py

@@ -61,6 +61,7 @@ class RoboflowModelMetadata(BaseModel):
     type: Literal["external-model-metadata-v1"]
     model_id: str = Field(alias="modelId")
     model_architecture: str = Field(alias="modelArchitecture")
+    model_variant: Optional[str] = Field(alias="modelVariant", default=None)
     task_type: Optional[str] = Field(alias="taskType", default=None)
     model_packages: List[Union[RoboflowModelPackageV1, dict]] = Field(
         alias="modelPackages",
@@ -81,6 +82,7 @@ def get_roboflow_model(model_id: str, api_key: Optional[str] = None) -> ModelMet
         model_architecture=model_metadata.model_architecture,
         model_packages=parsed_model_packages,
         task_type=model_metadata.task_type,
+        model_variant=model_metadata.model_variant,
     )
 
 

inference_experimental/tests/unit_tests/weights_providers/test_roboflow.py

@@ -1317,6 +1317,7 @@ def test_get_roboflow_model(requests_mock: Mocker) -> None:
                         "type": "external-model-metadata-v1",
                         "modelId": "my-model",
                         "modelArchitecture": "yolov8",
+                        "modelVariant": "yolov8-n",
                         "taskType": "object-detection",
                         "modelPackages": [
                             {
@@ -1352,6 +1353,7 @@ def test_get_roboflow_model(requests_mock: Mocker) -> None:
                         "type": "external-model-metadata-v1",
                         "modelId": "my-model",
                         "modelArchitecture": "yolov8",
+                        "modelVariant": "yolov8-n",
                         "taskType": "object-detection",
                         "modelPackages": [
                             {
@@ -1401,5 +1403,6 @@ def test_get_roboflow_model(requests_mock: Mocker) -> None:
     # then
     assert result.model_id == "my-model"
     assert result.model_architecture == "yolov8"
+    assert result.model_variant == "yolov8-n"
     assert result.task_type == "object-detection"
     assert len(result.model_packages) == 2