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
import mlx.core as mx
from huggingface_hub import scan_cache_dir
from .utils import generate_step, load
@@ -635,13 +636,21 @@ class APIHandler(BaseHTTPRequestHandler):
"""
self._set_completion_headers(200)
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
models = [
{
"id": self.model_provider.cli_args.model,
"id": repo.repo_id,
"object": "model",
"created": self.created,
}
for repo in downloaded_models
]
response = {"object": "list", "data": models}