kinfer 0.3.2__tar.gz → 0.4.1__tar.gz

{kinfer-0.3.2 → kinfer-0.4.1}/.cargo/config.toml

@@ -1,5 +1,7 @@
-[target.aarch64-apple-darwin]
+[target.x86_64-unknown-linux-gnu]
+rustflags = [ "-Clink-args=-Wl,-rpath,\\$ORIGIN" ]
 
+[target.aarch64-apple-darwin]
 rustflags = [
   "-C", "link-arg=-undefined",
   "-C", "link-arg=dynamic_lookup",

{kinfer-0.3.2 → kinfer-0.4.1}/Cargo.toml

@@ -8,7 +8,7 @@ resolver = "2"
 
 [workspace.package]
 
-version = "0.3.2"
+version = "0.4.1"
 authors = ["Wesley Maa <wesley@kscale.dev>", "Benjamin Bolte <ben@kscale.dev>"]
 edition = "2021"
 description = "K-Scale Inference Library"

kinfer-0.4.1/PKG-INFO

@@ -0,0 +1,55 @@
+Metadata-Version: 2.4
+Name: kinfer
+Version: 0.4.1
+Summary: Tool to make it easier to run a model on a real robot
+Home-page: https://github.com/kscalelabs/kinfer.git
+Author: K-Scale Labs
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: onnx
+Requires-Dist: onnxruntime==1.20.0
+Requires-Dist: pydantic
+Provides-Extra: dev
+Requires-Dist: black; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: ruff; extra == "dev"
+Requires-Dist: types-tensorflow; extra == "dev"
+Provides-Extra: pytorch
+Requires-Dist: torch; extra == "pytorch"
+Provides-Extra: jax
+Requires-Dist: tensorflow; extra == "jax"
+Requires-Dist: tf2onnx; extra == "jax"
+Requires-Dist: jax; extra == "jax"
+Requires-Dist: equinox; extra == "jax"
+Requires-Dist: numpy<2; extra == "jax"
+Provides-Extra: all
+Requires-Dist: black; extra == "all"
+Requires-Dist: darglint; extra == "all"
+Requires-Dist: mypy; extra == "all"
+Requires-Dist: pytest; extra == "all"
+Requires-Dist: ruff; extra == "all"
+Requires-Dist: types-tensorflow; extra == "all"
+Requires-Dist: torch; extra == "all"
+Requires-Dist: tensorflow; extra == "all"
+Requires-Dist: tf2onnx; extra == "all"
+Requires-Dist: jax; extra == "all"
+Requires-Dist: equinox; extra == "all"
+Requires-Dist: numpy<2; extra == "all"
+Dynamic: author
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# kinfer
+
+This package is designed to support running real-time robotics models.
+
+For more information, see the documentation [here](https://docs.kscale.dev/docs/k-infer).

kinfer-0.4.1/README.md

@@ -0,0 +1,5 @@
+# kinfer
+
+This package is designed to support running real-time robotics models.
+
+For more information, see the documentation [here](https://docs.kscale.dev/docs/k-infer).

{kinfer-0.3.2 → kinfer-0.4.1}/kinfer/__init__.py

@@ -1,6 +1,5 @@
 """Defines the kinfer API."""
 
-from . import export, inference
 from .rust_bindings import get_version
 
 __version__ = get_version()

kinfer-0.4.1/kinfer/common/types.py

@@ -0,0 +1,12 @@
+"""Defines common types."""
+
+__all__ = [
+    "Metadata",
+]
+
+from pydantic import BaseModel
+
+
+class Metadata(BaseModel):
+    joint_names: list[str]
+    num_commands: int | None

kinfer-0.4.1/kinfer/export/common.py

@@ -0,0 +1,41 @@
+"""Defines common utilities for exporting models."""
+
+
+def get_shape(
+    name: str,
+    num_joints: int | None = None,
+    num_commands: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+) -> tuple[int, ...]:
+    match name:
+        case "joint_angles":
+            if num_joints is None:
+                raise ValueError("`num_joints` must be provided when using `joint_angles`")
+            return (num_joints,)
+
+        case "joint_angular_velocities":
+            if num_joints is None:
+                raise ValueError("`num_joints` must be provided when using `joint_angular_velocities`")
+            return (num_joints,)
+
+        case "projected_gravity":
+            return (3,)
+
+        case "accelerometer":
+            return (3,)
+
+        case "gyroscope":
+            return (3,)
+
+        case "command":
+            if num_commands is None:
+                raise ValueError("`num_commands` must be provided when using `command`")
+            return (num_commands,)
+
+        case "carry":
+            if carry_shape is None:
+                raise ValueError("`carry_shape` must be provided for `carry`")
+            return carry_shape
+
+        case _:
+            raise ValueError(f"Unknown tensor name: {name}")

kinfer-0.4.1/kinfer/export/jax.py

@@ -0,0 +1,53 @@
+"""Jax model export utilities."""
+
+import inspect
+import logging
+
+import tensorflow as tf
+import tf2onnx
+from equinox.internal._finalise_jaxpr import finalise_fn
+from jax._src.stages import Wrapped
+from jax.experimental import jax2tf
+from onnx.onnx_pb import ModelProto
+
+from kinfer.export.common import get_shape
+
+logger = logging.getLogger(__name__)
+
+
+def export_fn(
+    model: Wrapped,
+    *,
+    num_joints: int | None = None,
+    num_commands: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+    opset: int = 13,
+) -> ModelProto:
+    """Export a JAX function to ONNX."""
+    if not isinstance(model, Wrapped):
+        raise ValueError("Model must be a Wrapped function")
+
+    params = inspect.signature(model).parameters
+    input_names = list(params.keys())
+
+    # Gets the dummy input tensors for exporting the model.
+    tf_args = []
+    for name in input_names:
+        shape = get_shape(
+            name,
+            num_joints=num_joints,
+            num_commands=num_commands,
+            carry_shape=carry_shape,
+        )
+        tf_args.append(tf.TensorSpec(shape, tf.float32, name=name))
+
+    finalised_fn = finalise_fn(model)
+    tf_fn = tf.function(jax2tf.convert(finalised_fn, enable_xla=False))
+
+    model_proto, _ = tf2onnx.convert.from_function(
+        tf_fn,
+        input_signature=tf_args,
+        opset=opset,
+        large_model=False,
+    )
+    return model_proto

kinfer-0.4.1/kinfer/export/pytorch.py

@@ -0,0 +1,63 @@
+"""PyTorch model export utilities."""
+
+__all__ = [
+    "export_fn",
+]
+
+import io
+from typing import cast
+
+import onnx
+import torch
+from onnx.onnx_pb import ModelProto
+from torch._C import FunctionSchema
+
+from kinfer.export.common import get_shape
+
+
+def export_fn(
+    model: torch.jit.ScriptFunction,
+    *,
+    num_joints: int | None = None,
+    num_commands: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+) -> ModelProto:
+    """Exports a PyTorch function to ONNX.
+
+    Args:
+        model: The model to export.
+        num_joints: The number of joints in the model.
+        num_commands: The number of commands in the model.
+        carry_shape: The shape of the carry tensor.
+
+    Returns:
+        The ONNX model as a `ModelProto`.
+    """
+    if not isinstance(model, torch.jit.ScriptFunction):
+        raise ValueError("Model must be a torch.jit.ScriptFunction")
+
+    schema = cast(FunctionSchema, model.schema)
+    input_names = [arg.name for arg in schema.arguments]
+
+    # Gets the dummy input tensors for exporting the model.
+    args = []
+    for name in input_names:
+        shape = get_shape(
+            name,
+            num_joints=num_joints,
+            num_commands=num_commands,
+            carry_shape=carry_shape,
+        )
+        args.append(torch.zeros(shape))
+
+    buffer = io.BytesIO()
+    torch.onnx.export(
+        model=model,
+        f=buffer,  # type: ignore[arg-type]
+        args=tuple(args),
+        input_names=input_names,
+        external_data=False,
+    )
+    buffer.seek(0)
+    model_bytes = buffer.read()
+    return onnx.load_from_string(model_bytes)

kinfer-0.4.1/kinfer/export/serialize.py

@@ -0,0 +1,93 @@
+"""Functions for serializing and deserializing models."""
+
+__all__ = [
+    "pack",
+]
+
+
+import io
+import tarfile
+
+from onnx.onnx_pb import ModelProto
+
+from kinfer.common.types import Metadata
+from kinfer.export.common import get_shape
+
+
+def pack(
+    init_fn: ModelProto,
+    step_fn: ModelProto,
+    joint_names: list[str],
+    num_commands: int | None = None,
+    carry_shape: tuple[int, ...] | None = None,
+) -> bytes:
+    """Packs the initialization function and step function into a directory.
+
+    Args:
+        init_fn: The initialization function.
+        step_fn: The step function.
+        joint_names: The list of joint names, in the order that the model
+            expects them to be provided.
+        num_commands: The number of commands in the model.
+        carry_shape: The shape of the carry tensor.
+    """
+    num_joints = len(joint_names)
+
+    # Checks the `init` function.
+    if len(init_fn.graph.input) > 0:
+        raise ValueError(f"`init` function should not have any inputs! Got {len(init_fn.graph.input)}")
+    if len(init_fn.graph.output) != 1:
+        raise ValueError(f"`init` function should have exactly 1 output! Got {len(init_fn.graph.output)}")
+    init_carry = init_fn.graph.output[0]
+    init_carry_shape = tuple(dim.dim_value for dim in init_carry.type.tensor_type.shape.dim)
+    if carry_shape is not None and init_carry_shape != carry_shape:
+        raise ValueError(f"Expected carry shape {carry_shape} for output `{init_carry.name}`, got {init_carry_shape}")
+
+    # Checks the `step` function.
+    for step_input in step_fn.graph.input:
+        step_input_type = step_input.type.tensor_type
+        shape = tuple(dim.dim_value for dim in step_input_type.shape.dim)
+        expected_shape = get_shape(
+            step_input.name,
+            num_joints=num_joints,
+            num_commands=num_commands,
+            carry_shape=carry_shape,
+        )
+        if shape != expected_shape:
+            raise ValueError(f"Expected shape {expected_shape} for input `{step_input.name}`, got {shape}")
+
+    if len(step_fn.graph.output) != 2:
+        raise ValueError(f"Step function must have exactly 2 outputs, got {len(step_fn.graph.output)}")
+
+    output_actions = step_fn.graph.output[0]
+    actions_shape = tuple(dim.dim_value for dim in output_actions.type.tensor_type.shape.dim)
+    if actions_shape != (num_joints,):
+        raise ValueError(f"Expected output shape {num_joints} for output `{output_actions.name}`, got {actions_shape}")
+
+    output_carry = step_fn.graph.output[1]
+    output_carry_shape = tuple(dim.dim_value for dim in output_carry.type.tensor_type.shape.dim)
+    if output_carry_shape != init_carry_shape:
+        raise ValueError(f"Expected carry shape {init_carry_shape} for output carry, got {output_carry_shape}")
+
+    # Builds the metadata object.
+    metadata = Metadata(
+        joint_names=joint_names,
+        num_commands=num_commands,
+    )
+
+    buffer = io.BytesIO()
+
+    with tarfile.open(fileobj=buffer, mode="w:gz") as tar:
+
+        def add_file_bytes(name: str, data: bytes) -> None:  # noqa: ANN401
+            info = tarfile.TarInfo(name=name)
+            info.size = len(data)
+            tar.addfile(info, io.BytesIO(data))
+
+        add_file_bytes("init_fn.onnx", init_fn.SerializeToString())
+        add_file_bytes("step_fn.onnx", step_fn.SerializeToString())
+        add_file_bytes("metadata.json", metadata.model_dump_json().encode("utf-8"))
+
+    buffer.seek(0)
+
+    return buffer.read()

kinfer-0.4.1/kinfer/requirements.txt

@@ -0,0 +1,8 @@
+# requirements.txt
+
+# Machine Learning
+onnx
+onnxruntime==1.20.0
+
+# Serialization
+pydantic

kinfer-0.4.1/kinfer/rust/Cargo.toml

@@ -0,0 +1,32 @@
+[package]
+
+name = "kinfer"
+version.workspace = true
+edition.workspace = true
+description.workspace = true
+authors.workspace = true
+repository.workspace = true
+license.workspace = true
+readme.workspace = true
+
+[lib]
+
+name = "kinfer"
+crate-type = ["cdylib", "rlib"]
+
+[dependencies]
+
+async-trait = "0.1"
+flate2 = "1.0"
+futures-util = "0.3.30"
+ndarray = "0.16.1"
+ort = { version = "2.0.0-rc.9", features = [ "load-dynamic" ] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+tar = "0.4"
+thiserror = "1.0"
+tokio = { version = "1.0", features = ["full"] }
+
+[dev-dependencies]
+
+rand = "0.8"

kinfer-0.4.1/kinfer/rust/src/lib.rs

@@ -0,0 +1,5 @@
+pub mod model;
+pub mod runtime;
+
+pub use model::*;
+pub use runtime::*;