stem_splitter_core/
pipeline.rs1use crate::{
2 audio::{read_audio, write_audio},
3 model::{PythonModel, StemModel},
4 types::{AudioData, SplitConfig, StemResult},
5};
6use anyhow::{Context, Result};
7use std::path::{Path, PathBuf};
8
9pub fn split_file(path: &str, config: SplitConfig) -> Result<StemResult> {
10 std::fs::metadata(path).with_context(|| format!("File does not exist: {}", path))?;
11
12 let model = PythonModel::new();
13
14 let audio = read_audio(path)?;
15 let tmp_output_dir = PathBuf::from("./tmp");
16 let result = model
17 .separate(&audio.samples, audio.channels, &tmp_output_dir)
18 .with_context(|| format!("Failed to separate stems for file: {}", path))?;
19
20 let file_stem = Path::new(path)
21 .file_stem()
22 .and_then(|s| s.to_str())
23 .unwrap_or("output");
24
25 let base_path = PathBuf::from(&config.output_dir).join(file_stem);
26 save_stems(&result, base_path.to_str().unwrap(), audio.sample_rate)?;
27
28 Ok(result)
29}
30
31fn save_stems(result: &StemResult, base_path: &str, sample_rate: u32) -> Result<()> {
32 let make_path = |stem: &str| format!("{base_path}_{stem}.wav");
33
34 let stems = [
35 ("vocals", &result.vocals),
36 ("drums", &result.drums),
37 ("bass", &result.bass),
38 ("other", &result.other),
39 ];
40
41 for (name, samples) in stems {
42 write_audio(
43 &make_path(name),
44 &AudioData {
45 samples: samples.clone(),
46 sample_rate,
47 channels: 1,
48 },
49 )?;
50 }
51
52 Ok(())
53}