Merge pull request #466 from AlexWorland/feature/hf-cache-discovery

feat: resolve HF Hub cache layout in model directory scanning

Author: jundot

.gitignore

@@ -110,3 +110,6 @@ omlx/_build_info.py
 
 # Tailwind CSS standalone CLI binary (download on demand via build_css.py)
 omlx/admin/tailwindcss-*
+
+# Git worktrees
+.worktrees/

omlx/admin/routes.py

@@ -2970,6 +2970,24 @@ async def list_hf_models(is_admin: bool = Depends(require_admin)):
 
     model_dirs = global_settings.model.get_model_dirs(global_settings.base_path)
 
+    from ..model_discovery import _resolve_hf_cache_entry
+
+    def _add_model(model_path: Path, model_name: str) -> None:
+        if model_name in seen_names:
+            return
+        seen_names.add(model_name)
+        total_size = sum(
+            f.stat().st_size for f in model_path.rglob("*") if f.is_file()
+        )
+        models.append(
+            {
+                "name": model_name,
+                "path": str(model_path),
+                "size": total_size,
+                "size_formatted": format_size(total_size),
+            }
+        )
+
     models = []
     seen_names: set[str] = set()
     for model_dir in model_dirs:
@@ -2981,44 +2999,22 @@ async def list_hf_models(is_admin: bool = Depends(require_admin)):
 
             if (subdir / "config.json").exists():
                 # Level 1: direct model folder
-                if subdir.name in seen_names:
-                    continue
-                seen_names.add(subdir.name)
-                total_size = sum(
-                    f.stat().st_size for f in subdir.rglob("*") if f.is_file()
-                )
-                models.append(
-                    {
-                        "name": subdir.name,
-                        "path": str(subdir),
-                        "size": total_size,
-                        "size_formatted": format_size(total_size),
-                    }
-                )
+                _add_model(subdir, subdir.name)
             else:
+                # HF Hub cache entry: models--Org--Name/snapshots/<hash>/
+                hf_resolved = _resolve_hf_cache_entry(subdir)
+                if hf_resolved is not None:
+                    snapshot_path, model_name = hf_resolved
+                    if (snapshot_path / "config.json").exists():
+                        _add_model(snapshot_path, model_name)
+                    continue
+
                 # Level 2: organization folder — scan children
                 for child in sorted(subdir.iterdir()):
                     if not child.is_dir() or child.name.startswith("."):
                         continue
-                    if not (child / "config.json").exists():
-                        continue
-                    if child.name in seen_names:
-                        continue
-                    seen_names.add(child.name)
-
-                    total_size = sum(
-                        f.stat().st_size
-                        for f in child.rglob("*")
-                        if f.is_file()
-                    )
-                    models.append(
-                        {
-                            "name": child.name,
-                            "path": str(child),
-                            "size": total_size,
-                            "size_formatted": format_size(total_size),
-                        }
-                    )
+                    if (child / "config.json").exists():
+                        _add_model(child, child.name)
 
     return {"models": models}
 

omlx/model_discovery.py

@@ -501,6 +501,30 @@ def _is_model_dir(path: Path) -> bool:
     return (path / "config.json").exists() and not _is_adapter_dir(path)
 
 
+def _resolve_hf_cache_entry(path: Path) -> tuple[Path, str] | None:
+    """Resolve an HF Hub cache entry (models--Org--Name/) to its active snapshot.
+
+    Returns (snapshot_path, model_name) or None if not a valid HF cache entry.
+    """
+    name = path.name
+    if not name.startswith("models--") or name.count("--") < 2:
+        return None
+
+    # "models--Org--Name" → "Name"
+    model_name = name.split("--", 2)[2]
+
+    try:
+        commit_hash = (path / "refs" / "main").read_text().strip()
+    except OSError:
+        return None
+
+    snapshot = path / "snapshots" / commit_hash
+    if not snapshot.is_dir():
+        return None
+
+    return snapshot, model_name
+
+
 def _register_model(
     models: dict[str, DiscoveredModel],
     model_dir: Path,
@@ -607,6 +631,14 @@ def discover_models(model_dir: Path) -> dict[str, DiscoveredModel]:
             # Level 1: direct model folder
             _register_model(models, subdir, subdir.name)
         else:
+            # HF Hub cache entry: models--Org--Name/snapshots/<hash>/
+            hf_resolved = _resolve_hf_cache_entry(subdir)
+            if hf_resolved is not None:
+                snapshot_path, model_name = hf_resolved
+                if _is_model_dir(snapshot_path):
+                    _register_model(models, snapshot_path, model_name)
+                continue
+
             # Level 2: organization folder — scan children
             has_children = False
             for child in sorted(subdir.iterdir()):

tests/test_model_discovery.py

@@ -11,6 +11,7 @@
     DiscoveredModel,
     _is_adapter_dir,
     _is_unsupported_model,
+    _resolve_hf_cache_entry,
     detect_model_type,
     discover_models,
     discover_models_from_dirs,
@@ -869,3 +870,139 @@ def test_audio_models_included_in_discovery(self, tmp_path):
         assert models["whisper-large-v3"].model_type == "audio_stt"
         assert "Qwen3-TTS" in models
         assert models["Qwen3-TTS"].model_type == "audio_tts"
+
+
+class TestHfCacheDiscovery:
+    """Tests for HF Hub cache entry resolution and discovery."""
+
+    FAKE_COMMIT = "abc123def456"
+
+    def _make_model(self, path: Path, model_type: str = "llama"):
+        """Helper to create a valid model directory."""
+        path.mkdir(parents=True, exist_ok=True)
+        (path / "config.json").write_text(json.dumps({"model_type": model_type}))
+        (path / "model.safetensors").write_bytes(b"0" * 1000)
+
+    def _make_hf_cache_entry(self, parent: Path, org: str, name: str):
+        """Helper to create a bare HF Hub cache directory layout (no model files)."""
+        entry = parent / f"models--{org}--{name}"
+        refs = entry / "refs"
+        refs.mkdir(parents=True)
+        (refs / "main").write_text(self.FAKE_COMMIT)
+        snapshot = entry / "snapshots" / self.FAKE_COMMIT
+        snapshot.mkdir(parents=True)
+        return entry, snapshot
+
+    def _make_hf_cache_model(self, parent: Path, org: str, name: str, model_type: str = "llama"):
+        """Helper to create an HF cache entry with a valid model in the snapshot."""
+        _, snapshot = self._make_hf_cache_entry(parent, org, name)
+        (snapshot / "config.json").write_text(json.dumps({"model_type": model_type}))
+        (snapshot / "model.safetensors").write_bytes(b"0" * 1000)
+
+    def test_resolve_valid_entry(self, tmp_path):
+        """Valid HF cache entry resolves to snapshot path and model name."""
+        entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+
+        result = _resolve_hf_cache_entry(entry)
+        assert result is not None
+        assert result[0] == snapshot
+        assert result[1] == "Qwen3-8B-4bit"
+
+    def test_resolve_regular_dir_returns_none(self, tmp_path):
+        """Regular directory without models-- prefix returns None."""
+        regular = tmp_path / "mlx-community"
+        regular.mkdir()
+        assert _resolve_hf_cache_entry(regular) is None
+
+    def test_resolve_single_separator_returns_none(self, tmp_path):
+        """models--Name (no org separator) returns None."""
+        entry = tmp_path / "models--NoOrg"
+        entry.mkdir()
+        assert _resolve_hf_cache_entry(entry) is None
+
+    def test_resolve_missing_refs_main_returns_none(self, tmp_path):
+        """Missing refs/main returns None."""
+        entry = tmp_path / "models--mlx-community--Qwen3-8B"
+        entry.mkdir(parents=True)
+        assert _resolve_hf_cache_entry(entry) is None
+
+    def test_resolve_missing_snapshot_returns_none(self, tmp_path):
+        """Valid refs/main but missing snapshot directory returns None."""
+        entry = tmp_path / "models--mlx-community--Qwen3-8B"
+        refs = entry / "refs"
+        refs.mkdir(parents=True)
+        (refs / "main").write_text("deadbeef")
+        assert _resolve_hf_cache_entry(entry) is None
+
+    def test_resolve_strips_whitespace_from_refs(self, tmp_path):
+        """Trailing newline in refs/main is stripped (matches real HF cache)."""
+        entry, snapshot = self._make_hf_cache_entry(tmp_path, "mlx-community", "Qwen3-8B")
+        # Overwrite with trailing newline (like real HF cache)
+        (entry / "refs" / "main").write_text(self.FAKE_COMMIT + "\n")
+
+        result = _resolve_hf_cache_entry(entry)
+        assert result is not None
+        assert result[0] == snapshot
+
+    def test_discover_hf_cache_model(self, tmp_path):
+        """HF cache entries are discovered as models."""
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 1
+        assert "Qwen3-8B-4bit" in models
+        assert models["Qwen3-8B-4bit"].model_type == "llm"
+
+    def test_discover_multiple_hf_cache_models(self, tmp_path):
+        """Multiple HF cache entries are all discovered."""
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B-v0.3")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 2
+        assert "Qwen3-8B-4bit" in models
+        assert "Mistral-7B-v0.3" in models
+
+    def test_hf_cache_model_path_points_to_snapshot(self, tmp_path):
+        """model_path points to the snapshot dir, not the cache entry."""
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+
+        models = discover_models(tmp_path)
+        assert models["Qwen3-8B-4bit"].model_path == str(
+            tmp_path / "models--mlx-community--Qwen3-8B-4bit" / "snapshots" / self.FAKE_COMMIT
+        )
+
+    def test_hf_cache_without_config_json_skipped(self, tmp_path):
+        """HF cache entries without config.json in snapshot are skipped."""
+        self._make_hf_cache_entry(tmp_path, "mlx-community", "NoConfig")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 0
+
+    def test_mixed_flat_and_hf_cache(self, tmp_path):
+        """Mix of flat models and HF cache entries."""
+        self._make_model(tmp_path / "mistral-7b")
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 2
+        assert "mistral-7b" in models
+        assert "Qwen3-8B-4bit" in models
+
+    def test_mixed_org_and_hf_cache(self, tmp_path):
+        """Mix of org folders and HF cache entries."""
+        self._make_model(tmp_path / "Qwen" / "Qwen3-8B", model_type="qwen2")
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Mistral-7B")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 2
+        assert "Qwen3-8B" in models
+        assert "Mistral-7B" in models
+
+    def test_hf_cache_does_not_fall_through_to_org_scan(self, tmp_path):
+        """HF cache entries don't get scanned as org folders."""
+        self._make_hf_cache_model(tmp_path, "mlx-community", "Qwen3-8B-4bit")
+
+        models = discover_models(tmp_path)
+        assert len(models) == 1
+        assert "Qwen3-8B-4bit" in models