Release prep: ensemble-aware CLI/server, clean-install validation, compact Run locally section and local inference doc
5f15b97 verified | """Model-kind detection and loading for the CLI and HTTP server.""" | |
| from __future__ import annotations | |
| import json | |
| import os | |
| from .engine import PlutoEngine | |
| from .ensemble import EnsembleEngine | |
| from .task_specialist import TaskSpecialistEngine | |
| def resolve_model_path(model: str, token: str | None = None) -> str: | |
| """Return a local path, downloading a Hub repository if needed.""" | |
| if os.path.exists(model): | |
| return model | |
| from huggingface_hub import snapshot_download | |
| return snapshot_download(model, token=token) | |
| def _ensemble_config(directory: str) -> str | None: | |
| preferred = os.path.join(directory, "ensemble_config.json") | |
| if os.path.exists(preferred): | |
| return preferred | |
| legacy = os.path.join(directory, "config.json") | |
| if os.path.exists(legacy): | |
| try: | |
| data = json.load(open(legacy, encoding="utf-8")) | |
| except Exception: | |
| return None | |
| if isinstance(data, dict) and "models" in data and "weights" in data: | |
| return legacy | |
| return None | |
| def load_engine(model: str, device: str | None = None, token: str | None = None): | |
| """Load an ensemble, specialist, or general checkpoint by its layout.""" | |
| path = resolve_model_path(model, token=token) | |
| if os.path.isdir(path): | |
| ensemble = _ensemble_config(path) | |
| if ensemble: | |
| return EnsembleEngine(ensemble, device=device, token=token) | |
| if os.path.exists(os.path.join(path, "pluto_task_config.json")): | |
| return TaskSpecialistEngine(path, device=device) | |
| if os.path.exists(os.path.join(path, "pluto_config.json")): | |
| return PlutoEngine(path, device=device) | |
| raise ValueError( | |
| f"{path!r} is not a Pluto checkpoint: expected ensemble_config.json, " | |
| "pluto_task_config.json, or pluto_config.json" | |
| ) | |
| if os.path.isfile(path): | |
| base = os.path.dirname(path) | |
| name = os.path.basename(path) | |
| if name in ("ensemble_config.json", "config.json"): | |
| ensemble = _ensemble_config(base) | |
| if ensemble: | |
| return EnsembleEngine(ensemble, device=device, token=token) | |
| if name == "pluto_task_config.json": | |
| return TaskSpecialistEngine(base, device=device) | |
| if name == "pluto_config.json": | |
| return PlutoEngine(base, device=device) | |
| raise ValueError(f"unsupported model path: {path!r}") | |