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}