stem_splitter_core/
pipeline.rs

1use 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}