File size: 2,417 Bytes
5f15b97 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 | """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}")
|