stem_splitter_core/
pipeline.rs

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

stem_splitter_core/
pipeline.rs
