XiangpengYang/pi0.5

1"""Thread-safe lazy lifecycle for the heavyweight π₀.₅ UR policy."""2 3from __future__ import annotations4 5import gc6from pathlib import Path7import sys8import threading9from collections.abc import Callable10 11from artifacts import (12    download_checkpoint,13    normalize_checkpoint_path,14    normalize_model_id,15)16 17POLICY_CONFIGS = ("pi05_ur_demo_no_state", "pi05_ur_demo_state")18DEFAULT_POLICY_CONFIG = POLICY_CONFIGS[0]19 20 21def normalize_policy_config(value: str) -> str:22    if value not in POLICY_CONFIGS:23        choices = ", ".join(POLICY_CONFIGS)24        raise ValueError(f"unsupported policy config {value!r}; choose one of: {choices}")25    return value26 27 28class ModelUnavailableError(RuntimeError):29    """Raised when the requested policy cannot be initialized."""30 31 32def _release_gpu_memory() -> None:33    gc.collect()34    try:35        import torch36 37        if torch.cuda.is_available():38            torch.cuda.empty_cache()39    except ImportError:40        pass41 42 43class ModelManager:44    def __init__(self, loader: Callable[[str, str, str], object] | None = None):45        self._loader = loader or self._load_default46        self._lock = threading.Lock()47        self._value = None48        self._active_key: tuple[str, str, str] | None = None49        self._error: str | None = None50 51    @property52    def active_key(self) -> tuple[str, str, str] | None:53        return self._active_key54 55    @property56    def health_message(self) -> str:57        if self._value is not None and self._active_key is not None:58            model_id, checkpoint_path, config_name = self._active_key59            return f"Model ready: {model_id}/{checkpoint_path} ({config_name})."60        if self._error:61            return f"Model unavailable: {self._error}"62        return "Model has not been loaded yet."63 64    def get(self, model_id: str, checkpoint_path: str, config_name: str):65        key = (66            normalize_model_id(model_id),67            normalize_checkpoint_path(checkpoint_path),68            normalize_policy_config(config_name),69        )70        if self._value is not None and self._active_key == key:71            return self._value72        with self._lock:73            if self._value is not None and self._active_key == key:74                return self._value75            if self._value is not None:76                self._value = None77                self._active_key = None78                _release_gpu_memory()79            self._error = None80            try:81                value = self._loader(*key)82            except Exception as exc:83                _release_gpu_memory()84                detail = str(exc) or exc.__class__.__name__85                self._error = f"{key[0]}/{key[1]} ({key[2]}): {detail}"86                raise ModelUnavailableError(self._error) from exc87            self._value = value88            self._active_key = key89            return value90 91    @staticmethod92    def _load_default(model_id: str, checkpoint_path: str, config_name: str):93        import torch94 95        if not torch.cuda.is_available():96            raise RuntimeError("CUDA GPU is required for π₀.₅ inference")97        runtime = str(Path(__file__).resolve().parent / "openpi_runtime")98        if runtime not in sys.path:99            sys.path.insert(0, runtime)100        from openpi.policies import policy_config101        from openpi.training import config as openpi_config102 103        paths = download_checkpoint(model_id, checkpoint_path)104        config = openpi_config.get_config(config_name)105        return policy_config.create_trained_policy(106            config,107            paths.checkpoint,108            pytorch_device="cuda",109        )110 111 112MODEL_MANAGER = ModelManager()113