beat_this/runtime/
rten.rs

1use std::collections::HashMap;
2use std::path::Path;
3
4use anyhow::{anyhow, Result};
5use rten::Value as RtenValue;
6use rten::{Model as RtenGraph, NodeId};
7use rten_tensor::{AsView, Layout};
8
9use super::{Model, Runtime, Tensor};
10
11/// Pure-Rust ONNX inference runtime backed by rten.
12///
13/// Uses Rayon internally for multi-threaded inference (defaults to physical
14/// core count). No configuration needed — just load and run.
15pub struct RtenRuntime;
16
17impl Runtime for RtenRuntime {
18    type Model = RtenModel;
19
20    fn load_model(&self, path: &Path) -> Result<RtenModel> {
21        let model = RtenGraph::load_file(path)?;
22
23        // Build name→NodeId map for inputs
24        let input_map: HashMap<String, NodeId> = model
25            .input_ids()
26            .iter()
27            .filter_map(|&id| {
28                let info = model.node_info(id)?;
29                let name = info.name()?;
30                Some((name.to_string(), id))
31            })
32            .collect();
33
34        // Build NodeId→name map for outputs (reverse lookup when returning results)
35        let output_names: Vec<(NodeId, String)> = model
36            .output_ids()
37            .iter()
38            .filter_map(|&id| {
39                let info = model.node_info(id)?;
40                let name = info.name()?;
41                Some((id, name.to_string()))
42            })
43            .collect();
44
45        let output_ids: Vec<NodeId> = model.output_ids().to_vec();
46
47        Ok(RtenModel {
48            model,
49            input_map,
50            output_names,
51            output_ids,
52        })
53    }
54}
55
56pub struct RtenModel {
57    model: RtenGraph,
58    /// "mel_spectrogram" → NodeId
59    input_map: HashMap<String, NodeId>,
60    /// [(NodeId, "beat"), (NodeId, "downbeat")]
61    output_names: Vec<(NodeId, String)>,
62    /// Ordered output node IDs for model.run()
63    output_ids: Vec<NodeId>,
64}
65
66impl Model for RtenModel {
67    fn run(&mut self, inputs: &[(&str, &Tensor)]) -> Result<HashMap<String, Tensor>> {
68        // Convert named inputs to (NodeId, Value) pairs
69        let rten_inputs: Vec<(NodeId, RtenValue)> = inputs
70            .iter()
71            .map(|(name, tensor)| {
72                let node_id = self
73                    .input_map
74                    .get(*name)
75                    .ok_or_else(|| anyhow!("rten: unknown input name '{}'", name))?;
76                let value = RtenValue::from_shape(tensor.shape.as_slice(), tensor.data.clone())
77                    .map_err(|e| {
78                        anyhow!("rten: failed to create input tensor '{}': {}", name, e)
79                    })?;
80                Ok((*node_id, value))
81            })
82            .collect::<Result<Vec<_>>>()?;
83
84        // model.run takes Vec<(NodeId, ValueOrView)> — convert via (&val).into()
85        let inputs_with_views: Vec<_> = rten_inputs
86            .iter()
87            .map(|(id, val)| (*id, val.into()))
88            .collect();
89
90        let outputs = self.model.run(inputs_with_views, &self.output_ids, None)?;
91
92        // Convert outputs to named Tensor map
93        let mut result = HashMap::new();
94        for (&id, value) in self.output_ids.iter().zip(outputs) {
95            let name = self
96                .output_names
97                .iter()
98                .find(|(nid, _)| *nid == id)
99                .map(|(_, n)| n.clone())
100                .unwrap_or_else(|| format!("output_{:?}", id));
101
102            let rten_tensor = value
103                .into_tensor::<f32>()
104                .ok_or_else(|| anyhow!("rten: output '{}' is not f32", name))?;
105            let shape: Vec<usize> = rten_tensor.shape().to_vec();
106            let data: Vec<f32> = rten_tensor.to_vec();
107
108            result.insert(name, Tensor { shape, data });
109        }
110
111        Ok(result)
112    }
113}
beat_this/runtime/rten.rs

beat_this/runtime/
rten.rs
