v1.1: feats a spec model

README.md

@@ -38,9 +38,13 @@ The selected directory must contain either `params/` (JAX checkpoint) or
 `model.safetensors` (PyTorch checkpoint), plus the training statistics at:
 
 ```text
-assets/ur_demo/norm_stats.json
+assets/**/norm_stats.json
 ```
 
+The Space also accepts a root-level `norm_stats.json`. The asset directory name
+is not required to be `ur_demo`; this supports checkpoints exported with the
+original dataset or robot name, such as `assets/F-Fer/ur-1/norm_stats.json`.
+
 Use a Space secret named `HF_TOKEN` when the model repository is private.
 
 ## Inputs and outputs
@@ -64,8 +68,12 @@ radians, matching the collected dataset.
 ## Deploy
 
 Create a Hugging Face Gradio Space with a CUDA GPU and push this repository.
-Model download and initialization happen lazily on the first prediction. Only
-one inference request runs at a time to protect GPU memory.
+When `PI05_MODEL_ID` and `PI05_CHECKPOINT_PATH` are configured as Space
+variables, the checkpoint is downloaded and validated during app startup,
+before the GPU-decorated prediction call. Model initialization remains lazy on
+the first prediction. If those variables are left empty, download falls back
+to the first prediction. Only one inference request runs at a time to protect
+GPU memory.
 
 For local use with all dependencies installed:
 
@@ -89,3 +97,5 @@ pytest tests/test_gpu_smoke.py -v
 
 Do not put a Hugging Face access token in this command; use `HF_TOKEN` as a
 local environment secret or a Hugging Face Space secret.
+
+哈基米

app.py

@@ -19,11 +19,21 @@ except ImportError:  # Local and dedicated-GPU environments omit this helper.
 
     spaces = _SpacesFallback()
 
-from artifacts import resolve_checkpoint_path, resolve_model_id
+from artifacts import download_checkpoint, resolve_checkpoint_path, resolve_model_id
 from inference import ACTION_LABELS, run_prediction
 from model_loader import DEFAULT_POLICY_CONFIG, MODEL_MANAGER, POLICY_CONFIGS
 
 
+def prefetch_configured_checkpoint() -> str:
+    """Download and validate configured weights during Space startup, before GPU use."""
+    model_id = resolve_model_id()
+    if not model_id:
+        return "No PI05_MODEL_ID configured; download will occur on first prediction."
+    checkpoint_path = resolve_checkpoint_path()
+    paths = download_checkpoint(model_id, checkpoint_path)
+    return f"Checkpoint ready: {model_id}/{checkpoint_path} ({paths.norm_stats.name})."
+
+
 def _gradio_integer(value, name: str) -> int:
     if isinstance(value, bool) or not isinstance(value, (int, float)):
         raise ValueError(f"{name} must be an integer")
@@ -89,6 +99,7 @@ def build_demo():
     with gr.Blocks(title="π₀.₅ UR Action Predictor") as demo:
         gr.Markdown(
             "# π₀.₅ UR Action Predictor\n"
+            "forked from ![XiangpengYang](https://huggingface.co/spaces/XiangpengYang/pi0.5), thx!"
             "Upload the fixed and wrist camera views, enter the current TCP/gripper "
             "state and a task instruction. This demo predicts actions only and does "
             "not directly control a robot."
@@ -127,15 +138,14 @@ def build_demo():
             tcp_y = gr.Number(value=0.0, label="TCP y")
             tcp_z = gr.Number(value=0.0, label="TCP z")
             tcp_roll = gr.Number(value=0.0, label="TCP roll")
+        gr.Markdown("哈基米")
         with gr.Row():
             tcp_pitch = gr.Number(value=0.0, label="TCP pitch")
             tcp_yaw = gr.Number(value=0.0, label="TCP yaw")
             gripper = gr.Number(value=0.0, label="Gripper")
             trial_index = gr.Number(value=0, precision=0, minimum=0, label="Trial index")
         predict_button = gr.Button("Predict actions", variant="primary")
-        status = gr.Markdown(
-            "The model loads on the first prediction; download and initialization may take several minutes."
-        )
+        status = gr.Markdown(STARTUP_STATUS)
         actions = gr.Dataframe(headers=list(ACTION_LABELS), interactive=False, label="Predicted actions")
         json_output = gr.File(label="Download JSON result")
         predict_button.click(
@@ -161,6 +171,13 @@ def build_demo():
     return demo
 
 
+# Hugging Face Spaces imports app.py during startup. Prefetching here moves the
+# multi-GB download out of the GPU-decorated prediction request when env vars are set.
+try:
+    STARTUP_STATUS = prefetch_configured_checkpoint()
+except Exception as exc:
+    STARTUP_STATUS = f"Checkpoint prefetch deferred: {exc}"
+
 demo = build_demo()
 
 

artifacts.py

@@ -13,6 +13,7 @@ DEFAULT_CHECKPOINT_PATH = "checkpoint"
 @dataclass(frozen=True)
 class ArtifactPaths:
     checkpoint: Path
+    norm_stats: Path
 
 
 def snapshot_download(**kwargs) -> str:
@@ -49,6 +50,22 @@ def resolve_checkpoint_path() -> str:
     return os.getenv("PI05_CHECKPOINT_PATH", DEFAULT_CHECKPOINT_PATH)
 
 
+def _find_norm_stats(checkpoint: Path) -> Path:
+    """Find checkpoint statistics without assuming the training repo name."""
+    preferred = checkpoint / "assets/ur_demo/norm_stats.json"
+    if preferred.is_file():
+        return preferred
+    candidates = sorted((checkpoint / "assets").rglob("norm_stats.json"))
+    if candidates:
+        return candidates[0]
+    root_stats = checkpoint / "norm_stats.json"
+    if root_stats.is_file():
+        return root_stats
+    raise FileNotFoundError(
+        f"UR normalization statistics not found under {checkpoint / 'assets'} or at {root_stats}"
+    )
+
+
 def download_checkpoint(model_id: str, checkpoint_path: str) -> ArtifactPaths:
     model_id = normalize_model_id(model_id)
     relative = normalize_checkpoint_path(checkpoint_path)
@@ -58,7 +75,5 @@ def download_checkpoint(model_id: str, checkpoint_path: str) -> ArtifactPaths:
         raise FileNotFoundError(
             f"checkpoint has neither params/ nor model.safetensors: {checkpoint}"
         )
-    statistics = checkpoint / "assets/ur_demo/norm_stats.json"
-    if not statistics.is_file():
-        raise FileNotFoundError(f"UR normalization statistics not found: {statistics}")
-    return ArtifactPaths(checkpoint=checkpoint)
+    statistics = _find_norm_stats(checkpoint)
+    return ArtifactPaths(checkpoint=checkpoint, norm_stats=statistics)

model_loader.py

@@ -98,6 +98,7 @@ class ModelManager:
         if runtime not in sys.path:
             sys.path.insert(0, runtime)
         from openpi.policies import policy_config
+        from openpi.shared import normalize
         from openpi.training import config as openpi_config
 
         paths = download_checkpoint(model_id, checkpoint_path)
@@ -105,6 +106,7 @@ class ModelManager:
         return policy_config.create_trained_policy(
             config,
             paths.checkpoint,
+            norm_stats=normalize.load(paths.norm_stats.parent),
             pytorch_device="cuda",
         )
 

tests/test_artifacts.py

@@ -24,6 +24,21 @@ class ArtifactTests(unittest.TestCase):
             with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
                 paths = artifacts.download_checkpoint("owner/model", "checkpoints/30000")
             self.assertEqual(paths.checkpoint, checkpoint)
+            self.assertEqual(paths.norm_stats, checkpoint / "assets/ur_demo/norm_stats.json")
+
+    def test_download_checkpoint_accepts_named_asset_directory(self):
+        import artifacts
+
+        with tempfile.TemporaryDirectory() as directory:
+            root = Path(directory)
+            checkpoint = root / "3000"
+            (checkpoint / "params/ocdbt.process_0/d").mkdir(parents=True)
+            stats = checkpoint / "assets/F-Fer/ur-1/norm_stats.json"
+            stats.parent.mkdir(parents=True)
+            stats.write_text("{}")
+            with mock.patch.object(artifacts, "snapshot_download", return_value=str(root)):
+                paths = artifacts.download_checkpoint("owner/model", "3000")
+            self.assertEqual(paths.norm_stats, stats)
 
     def test_download_checkpoint_reports_missing_weights(self):
         import artifacts