46 lines
1.3 KiB
Python
46 lines
1.3 KiB
Python
import time
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Dict, List, Optional
|
|
|
|
|
|
@dataclass
|
|
class ModelInfo:
|
|
model_id: str
|
|
created: int
|
|
size: int
|
|
path: Path
|
|
|
|
|
|
def scan_models(model_dir: str) -> List[ModelInfo]:
|
|
base = Path(model_dir)
|
|
if not base.exists():
|
|
return []
|
|
models: List[ModelInfo] = []
|
|
now = int(time.time())
|
|
for entry in base.iterdir():
|
|
if entry.name.endswith(".partial"):
|
|
continue
|
|
if entry.is_file():
|
|
size = entry.stat().st_size
|
|
models.append(ModelInfo(model_id=entry.name, created=now, size=size, path=entry))
|
|
elif entry.is_dir():
|
|
models.append(ModelInfo(model_id=entry.name, created=now, size=0, path=entry))
|
|
models.sort(key=lambda m: m.model_id.lower())
|
|
return models
|
|
|
|
|
|
def find_model(model_dir: str, model_id: str) -> Optional[ModelInfo]:
|
|
for model in scan_models(model_dir):
|
|
if model.model_id == model_id:
|
|
return model
|
|
return None
|
|
|
|
|
|
def resolve_model(model_dir: str, requested: str, aliases: Dict[str, str]) -> Optional[ModelInfo]:
|
|
if not requested:
|
|
return None
|
|
if requested in aliases:
|
|
requested = aliases[requested]
|
|
return find_model(model_dir, requested)
|