Spaces:
Running on Zero
Running on Zero
File size: 3,675 Bytes
b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 380cc0c b24bc66 | 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 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 | """Thread-safe lazy lifecycle for the heavyweight π₀.₅ UR policy."""
from __future__ import annotations
import gc
from pathlib import Path
import sys
import threading
from collections.abc import Callable
from artifacts import (
download_checkpoint,
normalize_checkpoint_path,
normalize_model_id,
)
POLICY_CONFIGS = ("pi05_ur_demo_no_state", "pi05_ur_demo_state")
DEFAULT_POLICY_CONFIG = POLICY_CONFIGS[0]
def normalize_policy_config(value: str) -> str:
if value not in POLICY_CONFIGS:
choices = ", ".join(POLICY_CONFIGS)
raise ValueError(f"unsupported policy config {value!r}; choose one of: {choices}")
return value
class ModelUnavailableError(RuntimeError):
"""Raised when the requested policy cannot be initialized."""
def _release_gpu_memory() -> None:
gc.collect()
try:
import torch
if torch.cuda.is_available():
torch.cuda.empty_cache()
except ImportError:
pass
class ModelManager:
def __init__(self, loader: Callable[[str, str, str], object] | None = None):
self._loader = loader or self._load_default
self._lock = threading.Lock()
self._value = None
self._active_key: tuple[str, str, str] | None = None
self._error: str | None = None
@property
def active_key(self) -> tuple[str, str, str] | None:
return self._active_key
@property
def health_message(self) -> str:
if self._value is not None and self._active_key is not None:
model_id, checkpoint_path, config_name = self._active_key
return f"Model ready: {model_id}/{checkpoint_path} ({config_name})."
if self._error:
return f"Model unavailable: {self._error}"
return "Model has not been loaded yet."
def get(self, model_id: str, checkpoint_path: str, config_name: str):
key = (
normalize_model_id(model_id),
normalize_checkpoint_path(checkpoint_path),
normalize_policy_config(config_name),
)
if self._value is not None and self._active_key == key:
return self._value
with self._lock:
if self._value is not None and self._active_key == key:
return self._value
if self._value is not None:
self._value = None
self._active_key = None
_release_gpu_memory()
self._error = None
try:
value = self._loader(*key)
except Exception as exc:
_release_gpu_memory()
detail = str(exc) or exc.__class__.__name__
self._error = f"{key[0]}/{key[1]} ({key[2]}): {detail}"
raise ModelUnavailableError(self._error) from exc
self._value = value
self._active_key = key
return value
@staticmethod
def _load_default(model_id: str, checkpoint_path: str, config_name: str):
import torch
if not torch.cuda.is_available():
raise RuntimeError("CUDA GPU is required for π₀.₅ inference")
runtime = str(Path(__file__).resolve().parent / "openpi_runtime")
if runtime not in sys.path:
sys.path.insert(0, runtime)
from openpi.policies import policy_config
from openpi.training import config as openpi_config
paths = download_checkpoint(model_id, checkpoint_path)
config = openpi_config.get_config(config_name)
return policy_config.create_trained_policy(
config,
paths.checkpoint,
pytorch_device="cuda",
)
MODEL_MANAGER = ModelManager()
|