kinfer 0.4.2__cp311-cp311-macosx_11_0_arm64.whl → 0.4.3__cp311-cp311-macosx_11_0_arm64.whl

kinfer/export/common.py

@@ -37,5 +37,8 @@ def get_shape(
                 raise ValueError("`carry_shape` must be provided for `carry`")
             return carry_shape
 
+        case "time":
+            return (1,)
+
         case _:
             raise ValueError(f"Unknown tensor name: {name}")

kinfer/rust/src/model.rs

@@ -47,6 +47,7 @@ pub trait ModelProvider: Send + Sync {
     async fn get_accelerometer(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_gyroscope(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_command(&self) -> Result<Array<f32, IxDyn>, ModelError>;
+    async fn get_time(&self) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError>;
     async fn take_action(
         &self,
@@ -188,6 +189,15 @@ impl ModelRunner {
                         .into());
                     }
                 }
+                "time" => {
+                    if *dims != vec![1] {
+                        return Err(format!(
+                            "Expected shape [1] for input `{}`, got {:?}",
+                            input.name, dims
+                        )
+                        .into());
+                    }
+                }
                 "carry" => {
                     if dims != carry_shape {
                         return Err(format!(
@@ -269,6 +279,7 @@ impl ModelRunner {
                 "gyroscope" => futures.push(self.provider.get_gyroscope()),
                 "command" => futures.push(self.provider.get_command()),
                 "carry" => futures.push(self.provider.get_carry(carry.clone())),
+                "time" => futures.push(self.provider.get_time()),
                 _ => return Err(format!("Unknown input name: {}", name).into()),
             }
         }

kinfer/rust/src/runtime.rs

@@ -3,8 +3,8 @@ use std::sync::Arc;
 use tokio::runtime::Runtime;
 
 use crate::model::{ModelError, ModelRunner};
-use std::time::{Duration, Instant};
-use tokio::time::sleep;
+use std::time::Duration;
+use tokio::time::interval;
 
 pub struct ModelRuntime {
     model_runner: Arc<ModelRunner>,
@@ -61,7 +61,10 @@ impl ModelRuntime {
                 .get_joint_angles()
                 .await
                 .map_err(|e| ModelError::Provider(e.to_string()))?;
-            let mut last_time = Instant::now();
+
+            // Wait for the first tick, since it happens immediately.
+            let mut interval = interval(dt);
+            interval.tick().await;
 
             while running.load(Ordering::Relaxed) {
                 let (output, next_carry) = model_runner
@@ -80,10 +83,7 @@ impl ModelRuntime {
                         .take_action(interp_joint_positions * magnitude_factor)
                         .await
                         .map_err(|e| ModelError::Provider(e.to_string()))?;
-                    last_time = last_time + dt;
-                    if let Some(sleep_duration) = last_time.checked_duration_since(Instant::now()) {
-                        sleep(sleep_duration).await;
-                    }
+                    interval.tick().await;
                 }
 
                 joint_positions = output;

kinfer/rust_bindings/rust_bindings.pyi

@@ -14,6 +14,7 @@ class ModelProviderABC:
     def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def get_time(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
 
 class PyModelProvider:

kinfer/rust_bindings/src/lib.rs

@@ -10,7 +10,7 @@ use pyo3_stub_gen::define_stub_info_gatherer;
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pyfunction, gen_stub_pymethods};
 use std::sync::Arc;
 use std::sync::Mutex;
-
+use std::time::Instant;
 #[pyfunction]
 #[gen_stub_pyfunction]
 fn get_version() -> String {
@@ -73,6 +73,10 @@ impl ModelProviderABC {
         Err(PyNotImplementedError::new_err("Must override get_command"))
     }
 
+    fn get_time<'py>(&self) -> PyResult<Bound<'py, PyArray1<f32>>> {
+        Err(PyNotImplementedError::new_err("Must override get_time"))
+    }
+
     fn take_action<'py>(
         &self,
         joint_names: Vec<String>,
@@ -92,13 +96,17 @@ impl ModelProviderABC {
 #[derive(Clone)]
 struct PyModelProvider {
     obj: Arc<Py<ModelProviderABC>>,
+    start_time: Instant,
 }
 
 #[pymethods]
 impl PyModelProvider {
     #[new]
     fn new(obj: Py<ModelProviderABC>) -> Self {
-        Self { obj: Arc::new(obj) }
+        Self {
+            obj: Arc::new(obj),
+            start_time: Instant::now(),
+        }
     }
 }
 
@@ -182,6 +190,18 @@ impl ModelProvider for PyModelProvider {
         Ok(args)
     }
 
+    async fn get_time(&self) -> Result<Array<f32, IxDyn>, ModelError> {
+        let args = Python::with_gil(|py| -> PyResult<Array<f32, IxDyn>> {
+            let obj = self.obj.clone();
+            let args = ();
+            let result = obj.call_method(py, "get_time", args, None)?;
+            let array = result.extract::<Vec<f32>>(py)?;
+            Ok(Array::from_vec(array).into_dyn())
+        })
+        .map_err(|e| ModelError::Provider(e.to_string()))?;
+        Ok(args)
+    }
+
     async fn get_carry(&self, carry: Array<f32, IxDyn>) -> Result<Array<f32, IxDyn>, ModelError> {
         Ok(carry)
     }
@@ -217,9 +237,7 @@ struct PyModelRunner {
 impl PyModelRunner {
     #[new]
     fn new(model_path: String, provider: Py<ModelProviderABC>) -> PyResult<Self> {
-        let input_provider = Arc::new(PyModelProvider {
-            obj: Arc::new(provider),
-        });
+        let input_provider = Arc::new(PyModelProvider::new(provider));
 
         let runner = tokio::runtime::Runtime::new().unwrap().block_on(async {
             ModelRunner::new(model_path, input_provider)

kinfer/rust_bindings.pyi

@@ -14,6 +14,7 @@ class ModelProviderABC:
     def get_accelerometer(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_gyroscope(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def get_command(self) -> numpy.typing.NDArray[numpy.float32]: ...
+    def get_time(self) -> numpy.typing.NDArray[numpy.float32]: ...
     def take_action(self, joint_names:typing.Sequence[builtins.str], action:numpy.typing.NDArray[numpy.float32]) -> None: ...
 
 class PyModelProvider:

{kinfer-0.4.2.dist-info → kinfer-0.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: kinfer
-Version: 0.4.2
+Version: 0.4.3
 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
@@ -21,7 +21,7 @@ Provides-Extra: pytorch
 Requires-Dist: torch; extra == "pytorch"
 Provides-Extra: jax
 Requires-Dist: tensorflow; extra == "jax"
-Requires-Dist: tf2onnx; extra == "jax"
+Requires-Dist: tf2onnx>=1.16.0; extra == "jax"
 Requires-Dist: jax; extra == "jax"
 Requires-Dist: equinox; extra == "jax"
 Requires-Dist: numpy<2; extra == "jax"
@@ -34,7 +34,7 @@ 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: tf2onnx>=1.16.0; extra == "all"
 Requires-Dist: jax; extra == "all"
 Requires-Dist: equinox; extra == "all"
 Requires-Dist: numpy<2; extra == "all"

{kinfer-0.4.2.dist-info → kinfer-0.4.3.dist-info}/RECORD

@@ -1,26 +1,26 @@
-kinfer/rust_bindings.cpython-311-darwin.so,sha256=fHPKAukypHHVpNxF5OaE_l8IdxAliPUkGXVEhrO7yh8,1933008
+kinfer/rust_bindings.cpython-311-darwin.so,sha256=ZUWCReCTDKhzqBteg3k4nJ-15ubYFMCSBYQdADqcvd0,1953904
 kinfer/requirements.txt,sha256=j08HO4ptA5afuj99j8FlAP2qla5Zf4_OiEBtgAmF7Jg,90
 kinfer/__init__.py,sha256=i5da6ND827Cgn8PFKzDCmEBk14ptQasLa_9fdof4Y9c,398
 kinfer/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-kinfer/rust_bindings.pyi,sha256=9u0_TuHwV9xQH3O5OKbYfUZ_Uj6bsUwEOW0bfZujRMQ,1709
+kinfer/rust_bindings.pyi,sha256=lyBw-t_9ONYwVSaUTeeYDWpBbhKod1wf9Dvc4MgnC1g,1776
 kinfer/rust/Cargo.toml,sha256=uxjvuM1Sm82UgCdHW3R6PbbiJjXvfE4VCLvjGPqp2mY,606
-kinfer/rust/src/runtime.rs,sha256=eDflFSIa4IxOmG3SInwrHrYc02UzEANZMkK-liA7pVA,3407
+kinfer/rust/src/runtime.rs,sha256=l9wG0XbZmpHSqho3SfGuW54P7YXs8eXZ365kOj7UqUE,3321
 kinfer/rust/src/lib.rs,sha256=Z3dWdhKhqhWqPJee6vWde4ptqquOYW6W9wB0wyaKCyk,71
-kinfer/rust/src/model.rs,sha256=NqxWjSDivkN2mLAGCqytjfYHlx8dpP4zqoatx7pRWg4,11419
+kinfer/rust/src/model.rs,sha256=zY_XmkCZOq8VhVE_LpgTxKDrV-FceY8tnkkqS1tFa8c,11890
 kinfer/common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kinfer/common/types.py,sha256=jiuCfoTEumOJm4VZ03VoFVIL5PMJ1tCG7b2Kr66qJQA,176
 kinfer/export/serialize.py,sha256=7rcSrXawVnJinxeJp26wovUwoR04BnJ5CZAV9vwKUUo,3474
 kinfer/export/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 kinfer/export/jax.py,sha256=gZWWUJKb3MPEtfNYeKfPnypeK0elJwGuDYfpCW0GqNY,1413
 kinfer/export/pytorch.py,sha256=rwHGOTcd08IeL5-HZ_7JVUMPMHleiHI8-TkShAoRrYM,1612
-kinfer/export/common.py,sha256=QWcLXY-jzZCbrVPWN7gFGLE5RbnZHLhuuvoH9agEvwk,1244
+kinfer/export/common.py,sha256=AYN_gwmxmSlZ7BuswGUKYcY8uVUCh9UyXnWw-k8RVys,1290
 kinfer/rust_bindings/Cargo.toml,sha256=i1RGB9VNd9Q4FJ6gGwjZJQYo8DBBvpVWf3GJ95EfVgM,637
 kinfer/rust_bindings/pyproject.toml,sha256=jLcJuHCnQRh9HWR_R7a9qLHwj6LMBgnHyeKK_DruO1Y,135
-kinfer/rust_bindings/rust_bindings.pyi,sha256=9u0_TuHwV9xQH3O5OKbYfUZ_Uj6bsUwEOW0bfZujRMQ,1709
-kinfer/rust_bindings/src/lib.rs,sha256=ydTTyx68d1fZDzc6kjFpU9ePDQa9IBD8Gf2rPy-kutI,10840
+kinfer/rust_bindings/rust_bindings.pyi,sha256=lyBw-t_9ONYwVSaUTeeYDWpBbhKod1wf9Dvc4MgnC1g,1776
+kinfer/rust_bindings/src/lib.rs,sha256=pjsRej2BxTsecG9Y7GdBpW6GGR-pgbBk3yrKqEEQ3K0,11547
 kinfer/rust_bindings/src/bin/stub_gen.rs,sha256=hhoVGnaSfazbSfj5a4x6mPicGPOgWQAfsDmiPej0B6Y,133
-kinfer-0.4.2.dist-info/RECORD,,
-kinfer-0.4.2.dist-info/WHEEL,sha256=rtKwvZSAzWV03G3Ircwq_TRBlj2DQz4ocNXl0bd9DbU,136
-kinfer-0.4.2.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
-kinfer-0.4.2.dist-info/METADATA,sha256=utM93j1dUATyhQF8h2ittLqKCOwCAt_M-pnJ68dODh4,1745
-kinfer-0.4.2.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069
+kinfer-0.4.3.dist-info/RECORD,,
+kinfer-0.4.3.dist-info/WHEEL,sha256=3lrG374qykB8NIZqJ4ApcNfFLoAGB9tcfdD4_UMfO40,136
+kinfer-0.4.3.dist-info/top_level.txt,sha256=6mY_t3PYr3Dm0dpqMk80uSnArbvGfCFkxOh1QWtgDEo,7
+kinfer-0.4.3.dist-info/METADATA,sha256=DZKaMDT-DToGkw_Ge_nSYBoo3ZAzxatlv6sR0x9kv-A,1761
+kinfer-0.4.3.dist-info/licenses/LICENSE,sha256=Qw-Z0XTwS-diSW91e_jLeBPX9zZbAatOJTBLdPHPaC0,1069

{kinfer-0.4.2.dist-info → kinfer-0.4.3.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.4.0)
+Generator: setuptools (80.7.1)
 Root-Is-Purelib: false
 Tag: cp311-cp311-macosx_11_0_arm64
 Generator: delocate 0.13.0