Fix dlc parsing on Benchmark __init__

luxonis · Jan 14, 2025 · e8bc974 · e8bc974
1 parent 4d3bc5b
commit e8bc974
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 6 deletions.
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -1,3 +1,5 @@
+default_language_version:
+    python: python3
 repos:
   - repo: https://github.com/astral-sh/ruff-pre-commit
     rev: v0.1.2

diff --git a/modelconverter/packages/base_benchmark.py b/modelconverter/packages/base_benchmark.py
@@ -24,7 +24,7 @@
 
 
 class Benchmark(ABC):
-    VALID_EXTENSIONS = (".tar.xz", ".blob")
+    VALID_EXTENSIONS = (".tar.xz", ".blob", ".dlc")
     HUB_MODEL_PATTERN = re.compile(r"^(?:([^/]+)/)?([^:]+):(.+)$")
 
     def __init__(
@@ -40,7 +40,7 @@ def __init__(
             if not hub_match:
                 raise ValueError(
                     "Invalid 'model-path' format. Expected either:\n"
-                    "- Model file path: path/to/model.blob or path/to/model.tar.xz\n"
+                    "- Model file path: path/to/model.blob, path/to/model.dlc or path/to/model.tar.xz\n"
                     "- HubAI model slug: [team_name/]model_name:variant"
                 )
             team_name, model_name, model_variant = hub_match.groups()

diff --git a/modelconverter/utils/hubai_utils.py b/modelconverter/utils/hubai_utils.py
@@ -10,10 +10,13 @@ def is_hubai_available(model_name: str, model_variant: str) -> bool:
 
     model_variants = []
     for is_public in [True, False]:
-        model_variants += Request.get(
-            "modelVersions/",
-            params={"model_id": model_id, "is_public": is_public},
-        )
+        try:
+            model_variants += Request.get(
+                "modelVersions/",
+                params={"model_id": model_id, "is_public": is_public},
+            )
+        except Exception:
+            pass
 
     for version in model_variants:
         if f"{model_name}:{version['variant_slug']}" == model_slug: