app.py

"""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 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")
    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,
    config_name,
    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, config_name)
        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,
        )
        uses_discrete_state = config_name == "pi05_ur_demo_state"
        status = (
            f"{result.status} Config={config_name} "
            f"(discrete_state_input={uses_discrete_state})."
        )
        return result.actions, result.json_path, 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"
            "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."
        )
        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",
            )
            config_name = gr.Dropdown(
                choices=list(POLICY_CONFIGS),
                value=DEFAULT_POLICY_CONFIG,
                label="Policy config",
            )
        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\n"
            "These values are discrete state conditioning only when "
            "`pi05_ur_demo_state` is selected."
        )
        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")
        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(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(
            fn=predict_ui,
            inputs=[
                model_id,
                checkpoint_path,
                config_name,
                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


# 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()


if __name__ == "__main__":
    if demo is None:
        raise RuntimeError("Gradio is not installed; install requirements.txt first")
    demo.queue(default_concurrency_limit=1).launch()