feat: deploy pi05 UR Gradio Space

.gitignore

@@ -0,0 +1,5 @@
+__pycache__/
+*.py[cod]
+.pytest_cache/
+.coverage
+htmlcov/

README.md

@@ -10,4 +10,57 @@ app_file: app.py
 pinned: false
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# π₀.₅ UR Action Predictor
+
+This Hugging Face Space deploys the state-conditioned `pi05_ur_demo_state`
+policy trained on the local UR LeRobot dataset. It predicts an action chunk for
+inspection or download; it never connects to or commands a robot.
+
+## Model repository
+
+Configure these Space variables (or enter both values in the UI):
+
+- `PI05_MODEL_ID`: Hugging Face repository containing the trained checkpoint.
+- `PI05_CHECKPOINT_PATH`: relative checkpoint directory, for example
+  `checkpoints/30000`.
+
+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
+```
+
+Use a Space secret named `HF_TOKEN` when the model repository is private.
+
+## Inputs and outputs
+
+The two image inputs correspond to training fields `video.image_0` (fixed
+camera) and `video.wrist` (wrist camera). State values must use this exact order:
+
+```text
+x, y, z, roll, pitch, yaw, gripper
+```
+
+The policy returns ten actions with columns:
+
+```text
+dx, dy, dz, droll, dpitch, dyaw, gripper
+```
+
+All state values must be finite. TCP translation uses metres and rotation uses
+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.
+
+For local use with all dependencies installed:
+
+```bash
+PYTHONPATH=openpi_runtime python app.py
+```
+
+The full OpenPI/JAX stack requires Python 3.11 and a compatible CUDA 12 GPU.

app.py

@@ -0,0 +1,154 @@
+"""Hugging Face Gradio Space for state-conditioned π₀.₅ UR inference."""
+
+from __future__ import annotations
+
+import gc
+
+try:
+    import gradio as gr
+except ImportError:  # Core inference tests can run without the UI dependency.
+    gr = None
+
+try:
+    import spaces
+except ImportError:  # Local and dedicated-GPU environments omit this helper.
+    class _SpacesFallback:
+        @staticmethod
+        def GPU(*args, **kwargs):
+            return lambda function: function
+
+    spaces = _SpacesFallback()
+
+from artifacts import resolve_checkpoint_path, resolve_model_id
+from inference import ACTION_LABELS, run_prediction
+from model_loader import MODEL_MANAGER
+
+
+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")
+    if not float(value).is_integer():
+        raise ValueError(f"{name} must be an integer")
+    result = int(value)
+    if result < 0:
+        raise ValueError(f"{name} must be non-negative")
+    return result
+
+
+@spaces.GPU(duration=120)
+def predict_ui(
+    model_id,
+    checkpoint_path,
+    fixed_image,
+    wrist_image,
+    instruction,
+    tcp_x,
+    tcp_y,
+    tcp_z,
+    tcp_roll,
+    tcp_pitch,
+    tcp_yaw,
+    gripper,
+    trial_index,
+):
+    try:
+        trial = _gradio_integer(trial_index, "trial index")
+        policy = MODEL_MANAGER.get(model_id, checkpoint_path)
+        result = run_prediction(
+            policy,
+            fixed_image,
+            wrist_image,
+            instruction,
+            [tcp_x, tcp_y, tcp_z, tcp_roll, tcp_pitch, tcp_yaw, gripper],
+            trial,
+            model_id,
+            checkpoint_path,
+        )
+        return result.actions, result.json_path, result.status
+    except Exception as exc:
+        gc.collect()
+        try:
+            import torch
+
+            if torch.cuda.is_available():
+                torch.cuda.empty_cache()
+        except ImportError:
+            pass
+        return None, None, f"Error: {exc}"
+
+
+def build_demo():
+    if gr is None:
+        return None
+    with gr.Blocks(title="π₀.₅ UR Action Predictor") as demo:
+        gr.Markdown(
+            "# π₀.₅ UR Action Predictor\n"
+            "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."
+        )
+        with gr.Row():
+            model_id = gr.Textbox(
+                value=resolve_model_id(),
+                label="Hugging Face model ID",
+                placeholder="owner/pi05-ur-checkpoint",
+            )
+            checkpoint_path = gr.Textbox(
+                value=resolve_checkpoint_path(),
+                label="Checkpoint path",
+                placeholder="checkpoints/30000",
+            )
+        with gr.Row():
+            fixed_image = gr.Image(type="pil", label="Fixed camera")
+            wrist_image = gr.Image(type="pil", label="Wrist camera")
+        instruction = gr.Textbox(
+            label="Task instruction",
+            placeholder="e.g. pick up the object and place it in the tray",
+            lines=2,
+        )
+        gr.Markdown("### Current state — metres/radians, followed by gripper state")
+        with gr.Row():
+            tcp_x = gr.Number(value=0.0, label="TCP x")
+            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")
+        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."
+        )
+        actions = gr.Dataframe(headers=list(ACTION_LABELS), interactive=False, label="Predicted actions")
+        json_output = gr.File(label="Download JSON result")
+        predict_button.click(
+            fn=predict_ui,
+            inputs=[
+                model_id,
+                checkpoint_path,
+                fixed_image,
+                wrist_image,
+                instruction,
+                tcp_x,
+                tcp_y,
+                tcp_z,
+                tcp_roll,
+                tcp_pitch,
+                tcp_yaw,
+                gripper,
+                trial_index,
+            ],
+            outputs=[actions, json_output, status],
+        )
+    return demo
+
+
+demo = build_demo()
+
+
+if __name__ == "__main__":
+    if demo is None:
+        raise RuntimeError("Gradio is not installed; install requirements.txt first")
+    demo.queue(default_concurrency_limit=1).launch()

tests/test_app.py

@@ -0,0 +1,46 @@
+from pathlib import Path
+import unittest
+from unittest import mock
+
+
+class AppTests(unittest.TestCase):
+    def test_predict_ui_returns_table_file_and_status(self):
+        import app
+
+        result = type(
+            "Result", (),
+            {"actions": "table", "json_path": "/tmp/result.json", "status": "done"},
+        )()
+        with mock.patch.object(app.MODEL_MANAGER, "get", return_value=object()), mock.patch.object(
+            app, "run_prediction", return_value=result
+        ):
+            actual = app.predict_ui(
+                "owner/model", "checkpoint", object(), object(), "task",
+                1, 2, 3, 4, 5, 6, 0, 0,
+            )
+        self.assertEqual(actual, ("table", "/tmp/result.json", "done"))
+
+    def test_predict_ui_turns_exceptions_into_status(self):
+        import app
+
+        with mock.patch.object(app.MODEL_MANAGER, "get", side_effect=RuntimeError("load failed")):
+            table, output_file, status = app.predict_ui(
+                "model", "checkpoint", object(), object(), "task",
+                0, 0, 0, 0, 0, 0, 0, 0,
+            )
+        self.assertIsNone(table)
+        self.assertIsNone(output_file)
+        self.assertEqual(status, "Error: load failed")
+
+    def test_source_exposes_required_prediction_controls(self):
+        source = Path("app.py").read_text()
+        for label in (
+            "Fixed camera", "Wrist camera", "Task instruction", "Predict actions",
+            "TCP x", "TCP y", "TCP z", "TCP roll", "TCP pitch", "TCP yaw", "Gripper",
+        ):
+            self.assertIn(label, source)
+        self.assertIn("default_concurrency_limit=1", source)
+
+
+if __name__ == "__main__":
+    unittest.main()