Skip to main content

supertonic_core/
style.rs

1use ndarray::Array3;
2use serde::{Deserialize, Serialize};
3use std::fs::File;
4use std::io::BufReader;
5use std::path::Path;
6
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct VoiceStyleData {
9    pub style_ttl: StyleComponent,
10    pub style_dp: StyleComponent,
11}
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct StyleComponent {
15    pub data: Vec<Vec<Vec<f32>>>,
16    pub dims: Vec<usize>,
17    #[serde(rename = "type")]
18    pub dtype: String,
19}
20
21#[derive(Debug, Clone)]
22pub struct Style {
23    pub ttl: Array3<f32>,
24    pub dp: Array3<f32>,
25}
26
27pub fn load_voice_style(paths: &[String]) -> Result<Style, anyhow::Error> {
28    let bsz = paths.len();
29    let first_file = File::open(&paths[0])?;
30    let first_reader = BufReader::new(first_file);
31    let first_data: VoiceStyleData = serde_json::from_reader(first_reader)?;
32
33    let ttl_dims = &first_data.style_ttl.dims;
34    let dp_dims = &first_data.style_dp.dims;
35
36    let ttl_dim1 = ttl_dims[1];
37    let ttl_dim2 = ttl_dims[2];
38    let dp_dim1 = dp_dims[1];
39    let dp_dim2 = dp_dims[2];
40
41    let ttl_size = bsz * ttl_dim1 * ttl_dim2;
42    let dp_size = bsz * dp_dim1 * dp_dim2;
43    let mut ttl_flat = vec![0.0f32; ttl_size];
44    let mut dp_flat = vec![0.0f32; dp_size];
45
46    for (i, path) in paths.iter().enumerate() {
47        let file = File::open(path)?;
48        let reader = BufReader::new(file);
49        let data: VoiceStyleData = serde_json::from_reader(reader)?;
50
51        let ttl_offset = i * ttl_dim1 * ttl_dim2;
52        let mut idx = 0;
53        for batch in &data.style_ttl.data {
54            for row in batch {
55                for &val in row {
56                    ttl_flat[ttl_offset + idx] = val;
57                    idx += 1;
58                }
59            }
60        }
61
62        let dp_offset = i * dp_dim1 * dp_dim2;
63        idx = 0;
64        for batch in &data.style_dp.data {
65            for row in batch {
66                for &val in row {
67                    dp_flat[dp_offset + idx] = val;
68                    idx += 1;
69                }
70            }
71        }
72    }
73
74    let ttl = Array3::from_shape_vec((bsz, ttl_dim1, ttl_dim2), ttl_flat)?;
75    let dp = Array3::from_shape_vec((bsz, dp_dim1, dp_dim2), dp_flat)?;
76
77    Ok(Style { ttl, dp })
78}
79
80impl Style {
81    pub fn load<P: AsRef<Path>>(path: P) -> Result<Self, anyhow::Error> {
82        let path_str = path.as_ref().to_string_lossy().to_string();
83        load_voice_style(&[path_str])
84    }
85}