Check hf_cache for mlx models

This commit is contained in:
jamesm131
2024-09-22 20:10:56 +10:00
parent e4eb4e29b7
commit e17a7dc443

View File

@@ -11,6 +11,7 @@ from pathlib import Path
from typing import Dict, List, Literal, NamedTuple, Optional, Sequence, Union from typing import Dict, List, Literal, NamedTuple, Optional, Sequence, Union
import mlx.core as mx import mlx.core as mx
from huggingface_hub import scan_cache_dir
from .utils import generate_step, load from .utils import generate_step, load
@@ -635,13 +636,21 @@ class APIHandler(BaseHTTPRequestHandler):
""" """
self._set_completion_headers(200) self._set_completion_headers(200)
self.end_headers() self.end_headers()
# Scan the cache directory for downloaded mlx models
hf_cache_info = scan_cache_dir()
downloaded_models = [
repo for repo in hf_cache_info.repos if "mlx" in repo.repo_id
]
# Create a list of available models # Create a list of available models
models = [ models = [
{ {
"id": self.model_provider.cli_args.model, "id": repo.repo_id,
"object": "model", "object": "model",
"created": self.created, "created": self.created,
} }
for repo in downloaded_models
] ]
response = {"object": "list", "data": models} response = {"object": "list", "data": models}