diff --git a/Dockerfile b/Dockerfile index e9de7f0..462cef9 100644 --- a/Dockerfile +++ b/Dockerfile @@ -53,4 +53,6 @@ RUN --mount=type=cache,id=final,target=/var/cache/apt,sharing=locked \ rm -rf /var/lib/{dpkg,cache,log}/ COPY --from=builder --chown=1000:1000 /build/target/release/stemgen /usr/bin/stemgen COPY --from=builder --chown=1000:1000 /build/target/release/libonnxruntime_providers*.so /usr/lib +RUN useradd -m -u 1000 user +USER user ENTRYPOINT [ "/usr/bin/stemgen" ] diff --git a/cli/src/cli.rs b/cli/src/cli.rs index cc1be4a..4aaa44a 100644 --- a/cli/src/cli.rs +++ b/cli/src/cli.rs @@ -1,4 +1,4 @@ -use std::path::PathBuf; +use std::path::{Component, PathBuf}; use clap::{builder::ValueParser, value_parser, ArgAction, Parser, Subcommand}; use stemgen::{ @@ -87,6 +87,94 @@ pub struct CreateArgs { pub copy_id3tags_from_mastered: bool, } +#[derive(Debug, Parser, Default, Copy, Clone)] +pub enum DeleteOriginal { + #[default] + No, + #[cfg(unix)] + Symlink, + Yes +} + +fn diff_paths(old: &PathBuf, new: &PathBuf) -> Result> +{ + let mut ita = new.parent().ok_or("expected a parented target")?.components(); + let mut itb = old.parent().ok_or("expected a parented target")?.components(); + let mut comps: Vec = vec![]; + + // ./foo and foo are the same + if let Some(Component::CurDir) = ita.clone().next() { + ita.next(); + } + if let Some(Component::CurDir) = itb.clone().next() { + itb.next(); + } + + loop { + match (ita.next(), itb.next()) { + (None, None) => break, + (Some(a), None) => { + comps.push(a); + comps.extend(ita.by_ref()); + break; + } + (None, _) => comps.push(Component::ParentDir), + (Some(a), Some(b)) if comps.is_empty() && a == b => (), + (Some(a), Some(b)) if b == Component::CurDir => comps.push(a), + (Some(_), Some(b)) if b == Component::ParentDir => return Err("unexpected parent dir".into()), + (Some(a), Some(_)) => { + comps.push(Component::ParentDir); + for _ in itb { + comps.push(Component::ParentDir); + } + comps.push(a); + comps.extend(ita.by_ref()); + break; + } + } + } + let rel: PathBuf = comps.iter().map(|c| c.as_os_str()).collect(); + let filename = new.file_name().ok_or("missing filename in target")?; + Ok(rel.join(filename)) +} + +impl DeleteOriginal { + pub(crate) fn proceed(&self, original: &PathBuf, new: &PathBuf) -> Result<(), Box> { + match self { + DeleteOriginal::No => Ok(()), + #[cfg(unix)] + DeleteOriginal::Symlink => { + use std::fs::canonicalize; + + let new = canonicalize(new)?; + let original = canonicalize(original)?; + std::fs::remove_file(&original)?; + + std::os::unix::fs::symlink(diff_paths(&original, &new)?, original) + }, + DeleteOriginal::Yes => + std::fs::remove_file(original), + }.map_err(|e|e.into()) + } +} + +fn parse_deleteoriginal(value: &str) -> Result { + value.try_into() +} + +impl TryFrom<&str> for DeleteOriginal { + type Error = String; + + fn try_from(value: &str) -> Result { + match value.to_ascii_lowercase().as_str() { + "no" => Ok(Self::No), + "symlink" => Ok(Self::Symlink), + "yes" => Ok(Self::Yes), + _ => Err("unsupported value".to_owned()), + } + } +} + #[derive(Debug, Parser, Default)] pub struct GenerateArgs { #[arg(num_args = 1.., value_name = "FILES", help = "path(s) to a file supported by the FFmpeg codec available on your machine. Advanced glob pattern can be used such as '~/Music/**/*.mp3'", required = true)] @@ -105,7 +193,9 @@ pub struct GenerateArgs { )] pub thread: usize, #[arg(long, default_value_t = false)] - pub preserved_original_as_master: bool, + pub preserve_original_as_master: bool, + #[arg(long, value_enum, value_parser = ValueParser::new(parse_deleteoriginal), default_value = "no")] + pub delete_original: DeleteOriginal, } #[derive(Debug, Subcommand)] @@ -134,12 +224,32 @@ pub fn prepare_ffmpeg(ctx: &Cli) -> Result<(), Box> { #[cfg(test)] mod tests { + use std::path::PathBuf; + use clap::CommandFactory; - use crate::Cli; + use crate::{cli::diff_paths, Cli}; #[test] fn verify_cmd() { Cli::command().debug_assert(); } + + #[test] + fn test_can_resolve_relative_to_original(){ + let original = PathBuf::from("../Manau - La Tribu de Dana.mp3"); + let new = PathBuf::from("../Manau - La Tribu de Dana.stem.mp4"); + + assert_eq!(diff_paths(&original, &new).unwrap(), PathBuf::from("Manau - La Tribu de Dana.stem.mp4")); + + let original = PathBuf::from("./Manau - La Tribu de Dana.mp3"); + let new = PathBuf::from("../Manau - La Tribu de Dana.stem.mp4"); + + assert_eq!(diff_paths(&original, &new).unwrap(), PathBuf::from("../Manau - La Tribu de Dana.stem.mp4")); + + let original = PathBuf::from("/mnt/Stereo/Manau - La Tribu de Dana.mp3"); + let new = PathBuf::from("/mnt/Stem/Manau - La Tribu de Dana.stem.mp4"); + + assert_eq!(diff_paths(&original, &new).unwrap(), PathBuf::from("../Stem/Manau - La Tribu de Dana.stem.mp4")); + } } diff --git a/cli/src/generate.rs b/cli/src/generate.rs index dd697cd..56c284f 100644 --- a/cli/src/generate.rs +++ b/cli/src/generate.rs @@ -1,4 +1,4 @@ -use std::{ffi::OsStr, path::PathBuf}; +use std::{ffi::OsStr, fs::canonicalize, path::PathBuf}; use glob::glob; use indicatif::{ProgressBar, ProgressStyle}; @@ -81,6 +81,9 @@ pub fn generate(ctx: &Cli, command: &GenerateArgs) -> Result Result-"), - ); + let mut result = ||{ + let mut input = Track::new(file)?; + let mut nistem = if command.preserve_original_as_master { + NIStem::new_with_preserved_original(&output_file, input.args(), ctx)? + } else { + NIStem::new_with_consistent_streams(&output_file, ctx)? + }; + nistem.clone(file)?; + let mut read = 0; + let pb = ProgressBar::new(2 * input.total() as u64); + pb.set_style( + ProgressStyle::with_template( + &format!("{{spinner:.green}} {} [{{wide_bar:.cyan/blue}}] [{{elapsed_precise}}] {{percent}}% ({{eta}})", filename.display()), + ) + .unwrap() + .progress_chars("#>-"), + ); - loop { - let mut buf: Vec = vec![0f32; 343980 * 2]; - let mut original_packets = Vec::with_capacity(512); - let mut original_buffer: Vec = Vec::with_capacity(512); - - let (data, eof) = loop { - let size = input.read( - if matches!(nistem, NIStem::PreservedMaster(..)) { - Some(&mut original_packets) - } else { - None - }, - &mut buf, - )?; - read += size; - if matches!(nistem, NIStem::ConsistentStream(..)) { - original_buffer.extend(buf[..size].to_vec()); - } - if let Some(mut data) = demucs.send(&buf[..size])? { + loop { + let mut buf: Vec = vec![0f32; 343980 * 2]; + let mut original_packets = Vec::with_capacity(512); + let mut original_buffer: Vec = Vec::with_capacity(512); + + let (data, eof) = loop { + let size = input.read( + if matches!(nistem, NIStem::PreservedMaster(..)) { + Some(&mut original_packets) + } else { + None + }, + &mut buf, + )?; + read += size; if matches!(nistem, NIStem::ConsistentStream(..)) { - data.insert(0, original_buffer); + original_buffer.extend(buf[..size].to_vec()); } - break (data, false) - } - if size != buf.len() { - let mut data = demucs.flush()?; - if matches!(nistem, NIStem::ConsistentStream(..)) { - data.insert(0, original_buffer); + if let Some(mut data) = demucs.send(&buf[..size])? { + if matches!(nistem, NIStem::ConsistentStream(..)) { + data.insert(0, original_buffer); + } + break (data, false) + } + if size != buf.len() { + let mut data = demucs.flush()?; + if matches!(nistem, NIStem::ConsistentStream(..)) { + data.insert(0, original_buffer); + } + break (data, true); } - break (data, true); + }; + pb.set_position(read as u64 / sample_rate); + match nistem { + NIStem::PreservedMaster(..) => nistem.write_preserved(original_packets, data)?, + NIStem::ConsistentStream(..) => nistem.write_consistent(data)?, } - }; - pb.set_position(read as u64 / sample_rate); - match nistem { - NIStem::PreservedMaster(..) => nistem.write_preserved(original_packets, data)?, - NIStem::ConsistentStream(..) => nistem.write_consistent(data)?, - } - if eof { - break; + if eof { + break; + } } - } + nistem.flush(nistem::Atom { + stems: [ + nistem::AtomStem { + color: ctx.drum_stem_color.to_owned(), + name: ctx.drum_stem_label.to_owned(), + }, + nistem::AtomStem { + color: ctx.bass_stem_color.to_owned(), + name: ctx.bass_stem_label.to_owned(), + }, + nistem::AtomStem { + color: ctx.other_stem_color.to_owned(), + name: ctx.other_stem_label.to_owned(), + }, + nistem::AtomStem { + color: ctx.vocal_stem_color.to_owned(), + name: ctx.vocal_stem_label.to_owned(), + }, + ], + version: 1, + ..Default::default() + })?; + pb.finish_with_message(format!("generated {}", filename.display())); + command.delete_original.proceed(&file, &output_file) + }; - pb.finish_with_message(format!("downloaded {}", filename.display())); - nistem.flush(nistem::Atom { - stems: [ - nistem::AtomStem { - color: ctx.drum_stem_color.to_owned(), - name: ctx.drum_stem_label.to_owned(), - }, - nistem::AtomStem { - color: ctx.bass_stem_color.to_owned(), - name: ctx.bass_stem_label.to_owned(), - }, - nistem::AtomStem { - color: ctx.other_stem_color.to_owned(), - name: ctx.other_stem_label.to_owned(), - }, - nistem::AtomStem { - color: ctx.vocal_stem_color.to_owned(), - name: ctx.vocal_stem_label.to_owned(), - }, - ], - version: 1, - ..Default::default() - })?; + if let Err(err) = result() { + eprintln!( + "Cannot proceed with {}: {}", + file.display(), + err.to_string() + ); + has_failure |= true; + continue; + } } Ok(has_failure) } @@ -199,7 +221,7 @@ mod tests { command: Commands::Generate(GenerateArgs { files: vec!["../testdata/Oddchap - Sound 104.mp3".into()], output: "..".into(), - preserved_original_as_master: false, + preserve_original_as_master: false, ..Default::default() }), ..Default::default() @@ -223,7 +245,7 @@ mod tests { command: Commands::Generate(GenerateArgs { files: vec!["../**/*.mp3".into()], output: "..".into(), - preserved_original_as_master: false, + preserve_original_as_master: false, ..Default::default() }), ..Default::default() diff --git a/cli/src/main.rs b/cli/src/main.rs index 9e6e428..acb8c1e 100644 --- a/cli/src/main.rs +++ b/cli/src/main.rs @@ -39,7 +39,7 @@ mod tests { }; use crate::{ - cli::{Commands, CreateArgs, GenerateArgs}, Cli + cli::{Commands, CreateArgs, DeleteOriginal, GenerateArgs}, Cli }; #[test] @@ -67,9 +67,10 @@ mod tests { files, output, device: Device::CPU, + delete_original: DeleteOriginal::No, model: Model::Url(model_url), thread: 4, - preserved_original_as_master: false + preserve_original_as_master: false }), drum_stem_label, bass_stem_label,