uranus-sdk 0.1.0__tar.gz

uranus_sdk-0.1.0/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.3
+Name: uranus-sdk
+Version: 0.1.0
+Summary: Python SDK for the Uranus video generation serving API
+Author: D-Robotics Large Model Team
+Author-email: D-Robotics Large Model Team <vincent.qin@d-robotics.cc>
+Requires-Dist: requests>=2.31
+Requires-Dist: numpy>=1.24
+Requires-Dist: opencv-python>=4.8
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+
+# Uranus SDK
+
+Python SDK for the Uranus video-generation serving API.  Provides a
+high-level `UranusSimulationSession` class that wraps HTTP communication,
+multipart parsing, session management, and authentication.
+
+## Installation
+
+```bash
+pip install uranus-sdk
+```
+
+## Quick Start (simplest)
+
+Load a sample from HuggingFace, run inference, save to mp4:
+
+```python
+import os
+import cv2
+import numpy as np
+from uranus_sdk import load_sample, UranusSimulationSession
+
+os.environ["URANUS_BASE_URL"] = "http://localhost:8000"
+os.environ["URANUS_API_KEY"] = "sk-your-token"  # optional
+
+# 1. Load sample from HuggingFace Hub (set HF_ENDPOINT to use a mirror)
+sample = load_sample("D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000",
+                     num_chunks=8)
+
+# 2. Create session + run step loop
+with UranusSimulationSession(**sample["create"]) as session:
+    all_frames = {}
+    for step in sample["steps"]:
+        frames = session.step(**step)
+        for cam, imgs in frames.items():
+            all_frames.setdefault(cam, []).extend(imgs)
+
+# 3. Save to mp4
+for cam, imgs in all_frames.items():
+    h, w = imgs[0].shape[:2]
+    writer = cv2.VideoWriter(f"{cam}.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 10, (w, h))
+    for img in imgs:
+        writer.write(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+    writer.release()
+```
+
+## Configuration
+
+| Variable | Required | Default | Description |
+|---|---|---|---|
+| `URANUS_BASE_URL` | No | `http://localhost:8000` | Server address |
+| `URANUS_API_KEY` | No | -- | Bearer token |
+| `HF_ENDPOINT` | No | `https://huggingface.co` | HuggingFace endpoint (use `https://hf-mirror.com` for mirror) |
+
+## Configuration
+
+Server URL and authentication are read from environment variables:
+
+| Variable | Required | Default | Description |
+|---|---|---|---|
+| `URANUS_BASE_URL` | No | `http://localhost:8000` | Server address |
+| `URANUS_API_KEY` | No | — | Bearer token; when set, every request carries `Authorization: Bearer <token>` |
+
+```bash
+export URANUS_BASE_URL=http://uranus-server:8000
+export URANUS_API_KEY=sk-xxxx
+```
+
+## API
+
+### `UranusSimulationSession(...)`
+
+Construction triggers `POST /create` — the session is ready immediately.
+
+**Parameters:**
+
+| Parameter | Type | Required | Description |
+|---|---|---|---|
+| `prompt` | `str` | Yes | Natural-language task description |
+| `mjcf_path` | `str \| Path` | Yes | MJCF XML file path (read and sent inline) |
+| `cameras` | `list[CameraSpec]` | Yes | Ordered camera specs; order fixes output frame order |
+| `ref_cam_images` | `list[str \| bytes \| ndarray]` | Yes | One reference image per camera (path, bytes, or RGB array) |
+| `ref_qpos` | `list[float] \| ndarray` | Yes | Complete MuJoCo qpos vector for the reference frame |
+| `robot2world_transform` | `ndarray \| None` | No | 4x4 robot-to-world matrix; omit for fixed-base |
+| `end_effectors` | `list[EESpec] \| None` | No | End-effector specs (aligned with OSS) |
+| `skeleton` | `SkeletonSpec \| None` | No | Skeleton topology spec (aligned with OSS) |
+| `target_size` | `tuple[int, int]` | No | Generation resolution (height, width); default (384, 640) |
+| `seed` | `int` | No | RNG seed; default 1 |
+| `session_id` | `str \| None` | No | Custom session ID; auto-generated if omitted |
+| `timeout` | `float` | No | HTTP timeout in seconds; default 600 |
+
+### `session.step(qpos, *, robot2world_transforms=None)`
+
+Generate video frames for a chunk of motion.
+
+**Parameters:**
+
+| Parameter | Type | Required | Description |
+|---|---|---|---|
+| `qpos` | `ndarray` or `list` | Yes | Motion frames `(N, nq)`. Server rounds N up to a multiple of 4 |
+| `robot2world_transforms` | `list[ndarray \| None] \| None` | No | Per-frame 4x4 transforms; None = identity |
+
+**Returns:** `dict[str, list[np.ndarray]]` -- `{camera_name: [RGB uint8 (H,W,3), ...]}`
+
+### `session.finish()`
+
+Release the session. Idempotent. Called automatically by `__exit__`.
+
+## Spec Classes
+
+The SDK reuses Spec classes from `uranus.skeleton` when the full Uranus
+package is installed.  Lightweight fallbacks are provided when Uranus is
+not available -- they serialize to the same JSON.
+
+### `CameraSpec(name: str)`
+
+### `EESpec(object_type, object_name, radius_mode, pad_bodies?, radius?, sh_correction?)`
+
+### `SkeletonSpec(mode, chains, skip_bodies, gripper_keypoint_overrides?)`
+
+## Error Handling
+
+```python
+from uranus_sdk import UranusSDKError
+
+try:
+    session = UranusSimulationSession(...)
+except UranusSDKError as e:
+    print(f"Failed: {e} (code={e.code}, status={e.status_code})")
+```
+
+| Scenario | HTTP Status | Code |
+|---|---|---|
+| Invalid parameters | 400 | `invalid_payload` |
+| Session not found | 404 | `unknown_session` |
+| Session already exists | 409 | `session_exists` |
+| Server busy | 503 | `server_busy` |
+| Auth failure | 401 | -- |
+
+## License
+
+MIT

uranus_sdk-0.1.0/README.md

@@ -0,0 +1,142 @@
+# Uranus SDK
+
+Python SDK for the Uranus video-generation serving API.  Provides a
+high-level `UranusSimulationSession` class that wraps HTTP communication,
+multipart parsing, session management, and authentication.
+
+## Installation
+
+```bash
+pip install uranus-sdk
+```
+
+## Quick Start (simplest)
+
+Load a sample from HuggingFace, run inference, save to mp4:
+
+```python
+import os
+import cv2
+import numpy as np
+from uranus_sdk import load_sample, UranusSimulationSession
+
+os.environ["URANUS_BASE_URL"] = "http://localhost:8000"
+os.environ["URANUS_API_KEY"] = "sk-your-token"  # optional
+
+# 1. Load sample from HuggingFace Hub (set HF_ENDPOINT to use a mirror)
+sample = load_sample("D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000",
+                     num_chunks=8)
+
+# 2. Create session + run step loop
+with UranusSimulationSession(**sample["create"]) as session:
+    all_frames = {}
+    for step in sample["steps"]:
+        frames = session.step(**step)
+        for cam, imgs in frames.items():
+            all_frames.setdefault(cam, []).extend(imgs)
+
+# 3. Save to mp4
+for cam, imgs in all_frames.items():
+    h, w = imgs[0].shape[:2]
+    writer = cv2.VideoWriter(f"{cam}.mp4", cv2.VideoWriter_fourcc(*"mp4v"), 10, (w, h))
+    for img in imgs:
+        writer.write(cv2.cvtColor(img, cv2.COLOR_RGB2BGR))
+    writer.release()
+```
+
+## Configuration
+
+| Variable | Required | Default | Description |
+|---|---|---|---|
+| `URANUS_BASE_URL` | No | `http://localhost:8000` | Server address |
+| `URANUS_API_KEY` | No | -- | Bearer token |
+| `HF_ENDPOINT` | No | `https://huggingface.co` | HuggingFace endpoint (use `https://hf-mirror.com` for mirror) |
+
+## Configuration
+
+Server URL and authentication are read from environment variables:
+
+| Variable | Required | Default | Description |
+|---|---|---|---|
+| `URANUS_BASE_URL` | No | `http://localhost:8000` | Server address |
+| `URANUS_API_KEY` | No | — | Bearer token; when set, every request carries `Authorization: Bearer <token>` |
+
+```bash
+export URANUS_BASE_URL=http://uranus-server:8000
+export URANUS_API_KEY=sk-xxxx
+```
+
+## API
+
+### `UranusSimulationSession(...)`
+
+Construction triggers `POST /create` — the session is ready immediately.
+
+**Parameters:**
+
+| Parameter | Type | Required | Description |
+|---|---|---|---|
+| `prompt` | `str` | Yes | Natural-language task description |
+| `mjcf_path` | `str \| Path` | Yes | MJCF XML file path (read and sent inline) |
+| `cameras` | `list[CameraSpec]` | Yes | Ordered camera specs; order fixes output frame order |
+| `ref_cam_images` | `list[str \| bytes \| ndarray]` | Yes | One reference image per camera (path, bytes, or RGB array) |
+| `ref_qpos` | `list[float] \| ndarray` | Yes | Complete MuJoCo qpos vector for the reference frame |
+| `robot2world_transform` | `ndarray \| None` | No | 4x4 robot-to-world matrix; omit for fixed-base |
+| `end_effectors` | `list[EESpec] \| None` | No | End-effector specs (aligned with OSS) |
+| `skeleton` | `SkeletonSpec \| None` | No | Skeleton topology spec (aligned with OSS) |
+| `target_size` | `tuple[int, int]` | No | Generation resolution (height, width); default (384, 640) |
+| `seed` | `int` | No | RNG seed; default 1 |
+| `session_id` | `str \| None` | No | Custom session ID; auto-generated if omitted |
+| `timeout` | `float` | No | HTTP timeout in seconds; default 600 |
+
+### `session.step(qpos, *, robot2world_transforms=None)`
+
+Generate video frames for a chunk of motion.
+
+**Parameters:**
+
+| Parameter | Type | Required | Description |
+|---|---|---|---|
+| `qpos` | `ndarray` or `list` | Yes | Motion frames `(N, nq)`. Server rounds N up to a multiple of 4 |
+| `robot2world_transforms` | `list[ndarray \| None] \| None` | No | Per-frame 4x4 transforms; None = identity |
+
+**Returns:** `dict[str, list[np.ndarray]]` -- `{camera_name: [RGB uint8 (H,W,3), ...]}`
+
+### `session.finish()`
+
+Release the session. Idempotent. Called automatically by `__exit__`.
+
+## Spec Classes
+
+The SDK reuses Spec classes from `uranus.skeleton` when the full Uranus
+package is installed.  Lightweight fallbacks are provided when Uranus is
+not available -- they serialize to the same JSON.
+
+### `CameraSpec(name: str)`
+
+### `EESpec(object_type, object_name, radius_mode, pad_bodies?, radius?, sh_correction?)`
+
+### `SkeletonSpec(mode, chains, skip_bodies, gripper_keypoint_overrides?)`
+
+## Error Handling
+
+```python
+from uranus_sdk import UranusSDKError
+
+try:
+    session = UranusSimulationSession(...)
+except UranusSDKError as e:
+    print(f"Failed: {e} (code={e.code}, status={e.status_code})")
+```
+
+| Scenario | HTTP Status | Code |
+|---|---|---|
+| Invalid parameters | 400 | `invalid_payload` |
+| Session not found | 404 | `unknown_session` |
+| Session already exists | 409 | `session_exists` |
+| Server busy | 503 | `server_busy` |
+| Auth failure | 401 | -- |
+
+## License
+
+MIT

uranus_sdk-0.1.0/pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "uranus-sdk"
+version = "0.1.0"
+description = "Python SDK for the Uranus video generation serving API"
+readme = "README.md"
+authors = [
+    { name = "D-Robotics Large Model Team", email = "vincent.qin@d-robotics.cc" }
+]
+requires-python = ">=3.11"
+dependencies = [
+    "requests>=2.31",
+    "numpy>=1.24",
+    "opencv-python>=4.8",
+]
+
+[build-system]
+requires = ["uv_build>=0.11.6,<0.12.0"]
+build-backend = "uv_build"
+
+[tool.uv]
+dev-dependencies = [
+    "pytest>=8.0",
+]

uranus_sdk-0.1.0/src/uranus_sdk/__init__.py

@@ -0,0 +1,21 @@
+"""Uranus SDK — Python client for the Uranus video generation serving API."""
+
+from uranus_sdk.session import UranusSimulationSession, UranusSDKError
+from uranus_sdk.specs import (
+    CameraSpec,
+    EESpec,
+    GripperKeypointOverride,
+    SkeletonSpec,
+)
+from uranus_sdk.dataset import load_sample
+
+__all__ = [
+    "UranusSimulationSession",
+    "UranusSDKError",
+    "CameraSpec",
+    "EESpec",
+    "GripperKeypointOverride",
+    "SkeletonSpec",
+    "load_sample",
+]
+__version__ = "0.1.0"

uranus_sdk-0.1.0/src/uranus_sdk/dataset.py

@@ -0,0 +1,182 @@
+"""Load sample data from HuggingFace Hub via direct HTTP (no huggingface_hub).
+
+Usage:
+    sample = load_sample("D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000")
+    # sample["create"] → kwargs for UranusSimulationSession
+    # sample["steps"]  → list of {qpos, robot2world_transforms} dicts
+"""
+
+from __future__ import annotations
+
+import json
+import os
+from typing import Any
+
+import numpy as np
+import requests
+
+from uranus_sdk.specs import (
+    CameraSpec,
+    EESpec,
+    GripperKeypointOverride,
+    SkeletonSpec,
+)
+
+_HF_BASE = os.environ.get("HF_ENDPOINT", "https://huggingface.co").rstrip("/") + "/datasets"
+
+
+def _hf_url(hf_path: str, filename: str) -> str:
+    """Build a HuggingFace direct-download URL.
+
+    ``hf_path`` is everything after ``huggingface.co/datasets/``, e.g.
+    ``"D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000"``.
+    """
+    return f"{_HF_BASE}/{hf_path}/{filename}"
+
+
+def _download_text(hf_path: str, filename: str) -> str:
+    """Download a text file from HF Hub."""
+    resp = requests.get(_hf_url(hf_path, filename), timeout=60)
+    resp.raise_for_status()
+    return resp.text
+
+
+def _download_bytes(hf_path: str, filename: str) -> bytes:
+    """Download a binary file from HF Hub."""
+    resp = requests.get(_hf_url(hf_path, filename), timeout=120)
+    resp.raise_for_status()
+    return resp.content
+
+
+def load_sample(
+    hf_path: str,
+    *,
+    num_chunks: int | None = None,
+    step_length: int = 4,
+) -> dict[str, Any]:
+    """Load a sample from HuggingFace Hub and build SDK-ready data.
+
+    Args:
+        hf_path:    The path after ``huggingface.co/datasets/``, e.g.
+                    ``"D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000"``.
+        num_chunks: Number of step chunks to generate. ``None`` → all available.
+        step_length: Frames per step (default 4, = temporal_interval).
+
+    Returns:
+        A dict with two keys:
+
+        ``"create"``:  kwargs dict for ``UranusSimulationSession(**sample["create"])``
+        ``"steps"``:  list of ``{"qpos": ndarray, "robot2world_transforms": list}`` dicts
+
+    Example::
+
+        sample = load_sample("D-Robotics/Uranus-Demo-Data/resolve/main/rc-table30-v1__rc-aloha/000000")
+        with UranusSimulationSession(**sample["create"]) as session:
+            for step in sample["steps"]:
+                frames = session.step(**step)
+    """
+    # --- download meta + temporal ---
+    meta = json.loads(_download_text(hf_path, "meta.json"))
+    temporal = json.loads(_download_text(hf_path, "temporal.json"))
+
+    camera_names = list(meta["cameras"])
+    cameras = [CameraSpec(name=name) for name in camera_names]
+
+    # --- download MJCF ---
+    mjcf_filename = meta["mjcf_path"]
+    # mjcf_path may be relative ("mjcf/robot.xml") — strip any dir prefix
+    if "/" in mjcf_filename:
+        mjcf_filename = mjcf_filename.split("/")[-1]
+    # try with the original path first, then fallback to basename
+    try:
+        mjcf_xml = _download_text(hf_path, meta["mjcf_path"])
+    except requests.HTTPError:
+        mjcf_xml = _download_text(hf_path, f"mjcf/{mjcf_filename}")
+
+    # --- download reference images ---
+    ref_cam_images = []
+    for name in camera_names:
+        img_path = meta["ref_cam_images"][name]
+        try:
+            ref_cam_images.append(_download_bytes(hf_path, img_path))
+        except requests.HTTPError:
+            # fallback: strip directory prefix
+            basename = img_path.split("/")[-1]
+            ref_cam_images.append(_download_bytes(hf_path, f"ref_images/{basename}"))
+
+    # --- build end_effectors ---
+    end_effectors = []
+    for ee in meta.get("end_effectors", []):
+        end_effectors.append(EESpec(
+            object_type=ee["object_type"],
+            object_name=ee["object_name"],
+            radius_mode=ee["radius_mode"],
+            pad_bodies=tuple(ee.get("pad_bodies", [])) if ee.get("pad_bodies") else None,
+            radius=float(ee["radius"]) if ee.get("radius") is not None else None,
+            sh_correction=np.array(ee["sh_correction"]) if ee.get("sh_correction") else None,
+        ))
+
+    # --- build skeleton ---
+    skel = meta.get("skeleton", {})
+    overrides = tuple(
+        GripperKeypointOverride(
+            ee_object_type=str(ov["ee_object_type"]),
+            ee_object_name=str(ov["ee_object_name"]),
+            finger_bodies=tuple(ov["finger_bodies"]),
+            closing_axis=int(ov.get("closing_axis", 1)),
+        )
+        for ov in skel.get("gripper_keypoint_overrides", [])
+    )
+    skeleton = SkeletonSpec(
+        mode=skel.get("mode", "full_tree"),
+        chains=tuple(tuple(c) for c in skel.get("chains", [])),
+        skip_bodies=tuple(skel.get("skip_bodies", [])),
+        gripper_keypoint_overrides=overrides,
+    )
+
+    # --- reference frame ---
+    states = temporal["step_qpos"]
+    ref_frame = states[0]
+    ref_qpos = ref_frame["state"]
+    r2w = (
+        np.array(ref_frame["robot2world_transform"])
+        if "robot2world_transform" in ref_frame
+        else None
+    )
+
+    # --- build create kwargs ---
+    create_kwargs = {
+        "prompt": meta.get("prompt", ""),
+        "mjcf_xml": mjcf_xml,
+        "cameras": cameras,
+        "ref_cam_images": ref_cam_images,
+        "ref_qpos": ref_qpos,
+        "robot2world_transform": r2w,
+        "end_effectors": end_effectors,
+        "skeleton": skeleton,
+        "target_size": (int(meta.get("height", 384)), int(meta.get("width", 640))),
+        "seed": 1,
+    }
+
+    # --- build steps ---
+    available = len(states) - 1  # frame 0 is reference
+    if num_chunks is None:
+        num_chunks = available // step_length
+
+    steps = []
+    for i in range(num_chunks):
+        start = i * step_length + 1
+        chunk = states[start : start + step_length]
+        if len(chunk) < step_length:
+            chunk = chunk + [chunk[-1]] * (step_length - len(chunk))
+
+        qpos = np.array([f["state"] for f in chunk])
+        r2w_list = [
+            np.array(f["robot2world_transform"])
+            if "robot2world_transform" in f
+            else None
+            for f in chunk
+        ]
+        steps.append({"qpos": qpos, "robot2world_transforms": r2w_list})
+
+    return {"create": create_kwargs, "steps": steps}

uranus_sdk-0.1.0/src/uranus_sdk/session.py

@@ -0,0 +1,350 @@
+"""UranusSimulationSession — high-level SDK for the Uranus serving API.
+
+Construction triggers ``create``; ``step`` returns RGB frames; ``finish``
+releases the session.  Server URL and auth are read from environment
+variables ``URANUS_BASE_URL`` and ``URANUS_API_KEY``.
+
+Spec classes (``CameraSpec``, ``EESpec``, ``SkeletonSpec``) are defined
+in :mod:`uranus_sdk.specs` and are structurally compatible with the
+``uranus.skeleton`` Spec classes used by the server.
+"""
+
+from __future__ import annotations
+
+import base64
+import json
+import os
+from dataclasses import asdict
+from pathlib import Path
+from typing import Any, Sequence
+
+import cv2
+import numpy as np
+import requests
+
+from uranus_sdk.specs import (
+    CameraSpec,
+    EESpec,
+    GripperKeypointOverride,
+    SkeletonSpec,
+)
+
+
+# ---------------------------------------------------------------------------
+# Error
+# ---------------------------------------------------------------------------
+
+class UranusSDKError(RuntimeError):
+    """Raised when the server returns an error or the network fails."""
+
+    def __init__(
+        self,
+        message: str,
+        *,
+        code: str | None = None,
+        status_code: int | None = None,
+    ) -> None:
+        super().__init__(message)
+        self.code = code
+        self.status_code = status_code
+
+
+# ---------------------------------------------------------------------------
+# Helpers
+# ---------------------------------------------------------------------------
+
+def _spec_to_dict(obj: Any) -> Any:
+    """Recursively convert a Spec dataclass (or nested structure) to a JSON-
+    safe dict, converting numpy arrays to nested lists."""
+    if hasattr(obj, "__dataclass_fields__"):
+        return {k: _spec_to_dict(v) for k, v in asdict(obj).items()}  # type: ignore[arg-type]
+    if isinstance(obj, np.ndarray):
+        return obj.tolist()
+    if isinstance(obj, (list, tuple)):
+        return [_spec_to_dict(v) for v in obj]
+    return obj
+
+
+def _load_image(img: str | Path | bytes | np.ndarray) -> bytes:
+    """Accept a file path, raw bytes, or RGB ndarray and return PNG bytes."""
+    if isinstance(img, (str, Path)):
+        return Path(img).read_bytes()
+    if isinstance(img, np.ndarray):
+        return cv2.imencode(".png", cv2.cvtColor(img, cv2.COLOR_RGB2BGR))[1].tobytes()
+    return img  # assume raw bytes
+
+
+# ---------------------------------------------------------------------------
+# Multipart decode
+# ---------------------------------------------------------------------------
+
+def _decode_multipart(
+    resp: requests.Response,
+) -> tuple[dict[str, Any], dict[str, list[np.ndarray]]]:
+    """Parse a ``multipart/mixed`` step response.
+
+    Returns ``(info, frames)`` where *info* is the JSON metadata dict and
+    *frames* maps ``{camera_name: [BGR ndarray, ...]}``.
+    """
+    raw = resp.content
+    boundary = raw[: raw.index(b"\r\n")]
+    sections = raw.split(b"\r\n" + boundary)[:-1]
+
+    parts: list[tuple[dict[bytes, bytes], bytes]] = []
+    for section in sections:
+        header_bytes, _, payload = section.partition(b"\r\n\r\n")
+        headers = {}
+        for line in header_bytes.split(b"\r\n"):
+            if b": " in line:
+                k, v = line.split(b": ", 1)
+                headers[k] = v
+        parts.append((headers, payload))
+
+    info = json.loads(parts[0][1])
+    frames: dict[str, list[np.ndarray]] = {}
+    for headers, jpeg in parts[1:]:
+        cam = headers[b"X-Camera"].decode()
+        img = cv2.imdecode(np.frombuffer(jpeg, dtype=np.uint8), cv2.IMREAD_COLOR)
+        frames.setdefault(cam, []).append(img)
+    return info, frames
+
+
+# ---------------------------------------------------------------------------
+# UranusSimulationSession
+# ---------------------------------------------------------------------------
+
+class UranusSimulationSession:
+    """High-level SDK for the Uranus video-generation serving API.
+
+    Construction triggers ``POST /create`` — the session is ready to use as
+    soon as ``__init__`` returns.  Call :meth:`step` to generate video frames
+    and :meth:`finish` (or use ``with``) to release the session.
+
+    Environment variables:
+        ``URANUS_BASE_URL`` — server address (default ``http://localhost:8000``)
+        ``URANUS_API_KEY``  — Bearer token; when set, every request carries
+        ``Authorization: Bearer <token>``.
+    """
+
+    def __init__(
+        self,
+        *,
+        prompt: str,
+        cameras: Sequence[CameraSpec],
+        ref_cam_images: Sequence[str | Path | bytes | np.ndarray],
+        ref_qpos: list[float] | np.ndarray,
+        mjcf_path: str | Path | None = None,
+        mjcf_xml: str | None = None,
+        robot2world_transform: np.ndarray | None = None,
+        end_effectors: Sequence[EESpec] | None = None,
+        skeleton: SkeletonSpec | None = None,
+        target_size: tuple[int, int] = (384, 640),
+        seed: int = 1,
+        session_id: str | None = None,
+        timeout: float = 600.0,
+    ) -> None:
+        """Create a new generation session.
+
+        Args:
+            prompt:  Natural-language task description.
+            cameras:  Ordered :class:`CameraSpec` list; order fixes output
+                frame ordering for the whole session.
+            ref_cam_images:  One reference image per camera.  Each entry may
+                be a file path, raw bytes, or an RGB ``ndarray``.
+            ref_qpos:  Complete MuJoCo ``qpos`` vector for the reference
+                frame (length = ``model.nq``).
+            mjcf_path:  Path to the MJCF XML file (read and sent inline).
+                Mutually exclusive with *mjcf_xml*.
+            mjcf_xml:  Inline MJCF XML string (use when loading from a URL
+                via :func:`load_sample`).  Mutually exclusive with *mjcf_path*.
+            cameras:  Ordered :class:`CameraSpec` list; order fixes output
+                frame ordering for the whole session.
+            ref_cam_images:  One reference image per camera.  Each entry may
+                be a file path, raw bytes, or an RGB ``ndarray``.
+            ref_qpos:  Complete MuJoCo ``qpos`` vector for the reference
+                frame (length = ``model.nq``).
+            robot2world_transform:  Optional 4x4 robot-to-world matrix.
+                Omit for fixed-base robots (defaults to identity).
+            end_effectors:  :class:`EESpec` list (aligned with OSS).
+            skeleton:  :class:`SkeletonSpec` (aligned with OSS).
+            target_size:  Generation resolution ``(height, width)``.
+            seed:  RNG seed for reproducible generation.
+            session_id:  Optional custom session ID; auto-generated if
+                omitted.
+            timeout:  Per-request HTTP timeout in seconds.
+
+        Raises:
+            UranusSDKError: If the create request fails.
+        """
+        # --- server config from env ---
+        base_url = os.environ.get("URANUS_BASE_URL", "http://localhost:8000")
+        api_key = os.environ.get("URANUS_API_KEY")
+        self._base_url = base_url.rstrip("/")
+        self._timeout = timeout
+        self._headers: dict[str, str] = {"Content-Type": "application/json"}
+        if api_key:
+            self._headers["Authorization"] = f"Bearer {api_key}"
+        self._session_id: str | None = None
+
+        # --- read MJCF ---
+        if mjcf_xml is not None:
+            pass  # use inline string directly
+        elif mjcf_path is not None:
+            mjcf_xml = Path(mjcf_path).read_text(encoding="utf-8")
+        else:
+            raise ValueError("Either mjcf_path or mjcf_xml must be provided")
+        camera_names = [c.name for c in cameras]
+
+        # --- ref_qpos -> frame dict ---
+        ref_frame: dict[str, Any] = {
+            "state": np.asarray(ref_qpos, dtype=np.float64).tolist(),
+        }
+        if robot2world_transform is not None:
+            ref_frame["robot2world_transform"] = np.asarray(
+                robot2world_transform, dtype=np.float64
+            ).tolist()
+
+        # --- ref_cam_images -> bytes ---
+        ref_images_bytes = [_load_image(img) for img in ref_cam_images]
+
+        # --- Spec -> dict ---
+        ee_dicts = [_spec_to_dict(ee) for ee in (end_effectors or [])]
+        skel_dict = _spec_to_dict(skeleton) if skeleton is not None else {}
+
+        # --- payload ---
+        payload = {
+            "camera_names": camera_names,
+            "prompt": prompt,
+            "seed": seed,
+            "target_size": list(target_size),
+            "mjcf_xml": mjcf_xml,
+            "end_effectors": ee_dicts,
+            "skeleton": skel_dict,
+            "ref_qpos": ref_frame,
+        }
+        body: dict[str, Any] = {
+            "payload": payload,
+            "attachments": {
+                "ref_cam_images": [
+                    base64.b64encode(img).decode() for img in ref_images_bytes
+                ]
+            },
+        }
+        if session_id is not None:
+            body["session_id"] = session_id
+
+        # --- POST /create ---
+        resp = requests.post(
+            f"{self._base_url}/create",
+            json=body,
+            headers=self._headers,
+            timeout=self._timeout,
+        )
+        self._raise_for_error(resp)
+        self._session_id = resp.json()["session_id"]
+
+    # ------------------------------------------------------------------
+    # step
+    # ------------------------------------------------------------------
+
+    def step(
+        self,
+        qpos: list[list[float]] | np.ndarray,
+        *,
+        robot2world_transforms: list[np.ndarray | None] | None = None,
+    ) -> dict[str, list[np.ndarray]]:
+        """Generate video frames for a chunk of motion.
+
+        Args:
+            qpos:  Motion frame sequence ``(N, nq)``.  Each row is a complete
+                MuJoCo ``qpos`` vector.  The server rounds ``N`` up to a
+                multiple of ``temporal_interval`` (=4).
+            robot2world_transforms:  Optional per-frame 4x4 transforms.
+                ``None`` or a frame entry ``None`` → identity matrix.
+
+        Returns:
+            ``{camera_name: [RGB uint8 ndarray (H, W, 3), ...]}``
+
+        Raises:
+            UranusSDKError: On server error or if the session was finished.
+        """
+        if self._session_id is None:
+            raise UranusSDKError("Session not created or already finished")
+
+        qpos_arr = np.asarray(qpos, dtype=np.float64)
+        if qpos_arr.ndim == 1:
+            qpos_arr = qpos_arr[np.newaxis, :]
+
+        frames: list[dict[str, Any]] = []
+        for i in range(len(qpos_arr)):
+            frame: dict[str, Any] = {"state": qpos_arr[i].tolist()}
+            if robot2world_transforms is not None:
+                t = robot2world_transforms[i]
+                if t is not None:
+                    frame["robot2world_transform"] = np.asarray(
+                        t, dtype=np.float64
+                    ).tolist()
+            frames.append(frame)
+
+        payload = {"qpos": frames, "num_step": len(qpos_arr)}
+
+        resp = requests.post(
+            f"{self._base_url}/step",
+            json={"session_id": self._session_id, "payload": payload},
+            headers=self._headers,
+            timeout=self._timeout,
+        )
+        self._raise_for_error(resp)
+        _info, bgr_frames = _decode_multipart(resp)
+
+        return {
+            cam: [cv2.cvtColor(img, cv2.COLOR_BGR2RGB) for img in imgs]
+            for cam, imgs in bgr_frames.items()
+        }
+
+    # ------------------------------------------------------------------
+    # finish
+    # ------------------------------------------------------------------
+
+    def finish(self) -> None:
+        """Release the session (GPU/CPU state).  Idempotent."""
+        if self._session_id is not None:
+            try:
+                requests.delete(
+                    f"{self._base_url}/sessions/{self._session_id}",
+                    headers=self._headers,
+                    timeout=self._timeout,
+                )
+            except Exception:
+                pass
+            self._session_id = None
+
+    # ------------------------------------------------------------------
+    # context manager
+    # ------------------------------------------------------------------
+
+    def __enter__(self) -> UranusSimulationSession:
+        return self
+
+    def __exit__(self, *args: object) -> None:
+        self.finish()
+
+    # ------------------------------------------------------------------
+    # error handling
+    # ------------------------------------------------------------------
+
+    @staticmethod
+    def _raise_for_error(resp: requests.Response) -> None:
+        if resp.status_code >= 400:
+            try:
+                data = resp.json()
+                err = data.get("error", data)
+                msg = err.get("message", resp.text)
+                code = err.get("code")
+            except Exception:
+                msg, code = resp.text, None
+            raise UranusSDKError(
+                f"HTTP {resp.status_code}: {msg}",
+                code=code,
+                status_code=resp.status_code,
+            )

uranus_sdk-0.1.0/src/uranus_sdk/specs.py

@@ -0,0 +1,98 @@
+"""Spec classes for rig configuration — structurally compatible with
+``uranus.skeleton.specs``.
+
+These dataclasses serialize (via :func:`dataclasses.asdict`) to the same
+JSON shape that the server's ``compat.from_create_payload`` expects.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+
+
+def _require(condition: bool, message: str) -> None:
+    if not condition:
+        raise ValueError(message)
+
+
+@dataclass(frozen=True)
+class CameraSpec:
+    """One XML-defined camera of the rig."""
+
+    name: str
+
+    def validate(self) -> None:
+        _require(bool(self.name), "camera name must be non-empty")
+
+
+@dataclass(frozen=True)
+class EESpec:
+    """One end-effector SH sphere: anchor object + radius source.
+
+    Fields mirror ``uranus.skeleton.specs.EESpec``.
+    """
+
+    object_type: str  # "body" | "site"
+    object_name: str
+    radius_mode: str  # "pad_pair" | "fixed"
+    pad_bodies: tuple[str, str] | None = None
+    radius: float | None = None  # fixed mode, meters
+    sh_correction: np.ndarray | None = None  # 3x3, default identity
+
+    def validate(self) -> None:
+        _require(
+            self.object_type in ("body", "site"),
+            f"end_effectors[{self.object_name!r}].object_type must be 'body' or 'site'",
+        )
+        _require(bool(self.object_name), "end_effector object_name must be non-empty")
+        if self.radius_mode == "pad_pair":
+            _require(
+                self.pad_bodies is not None and len(self.pad_bodies) == 2,
+                f"end_effectors[{self.object_name!r}] pad_pair requires two pad bodies",
+            )
+        elif self.radius_mode == "fixed":
+            _require(
+                self.radius is not None and self.radius > 0.0,
+                f"end_effectors[{self.object_name!r}] fixed radius must be > 0",
+            )
+        else:
+            raise ValueError(
+                f"end_effectors[{self.object_name!r}].radius_mode must be 'pad_pair' or 'fixed'"
+            )
+
+
+@dataclass(frozen=True)
+class GripperKeypointOverride:
+    """Legacy ``render_gripper_keypoints_from_width`` semantics, expressed via
+    the rig: the two finger-body keypoints are repositioned to
+    ``EE_pos +/- EE_rot[:, closing_axis] * width/2``."""
+
+    ee_object_type: str  # "body" | "site"
+    ee_object_name: str
+    finger_bodies: tuple[str, str]
+    closing_axis: int = 1  # column index into the EE rotation
+
+    def validate(self) -> None:
+        _require(self.ee_object_type in ("body", "site"), "gripper override ee_object_type")
+        _require(len(self.finger_bodies) == 2, "gripper override needs two finger bodies")
+        _require(0 <= self.closing_axis < 3, "closing_axis must be 0/1/2")
+
+
+@dataclass(frozen=True)
+class SkeletonSpec:
+    """Keypoint topology for rendering."""
+
+    mode: str = "full_tree"  # "chains" | "full_tree"
+    chains: tuple[tuple[str, ...], ...] = ()
+    skip_bodies: tuple[str, ...] = ()
+    gripper_keypoint_overrides: tuple[GripperKeypointOverride, ...] = ()
+
+    def validate(self) -> None:
+        if self.mode not in ("chains", "full_tree"):
+            raise ValueError(f"skeleton.mode must be 'chains' or 'full_tree', got {self.mode!r}")
+        if self.mode == "chains":
+            _require(len(self.chains) > 0, "skeleton chains must be non-empty in 'chains' mode")
+        for ov in self.gripper_keypoint_overrides:
+            ov.validate()