Skip to content

feat: add deletion policy to allow library conversion #42

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
acolombier wants to merge 1 commit into
base: main
Choose a base branch
from
Open

feat: add deletion policy to allow library conversion #42

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions Dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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" ]
116 changes: 113 additions & 3 deletions cli/src/cli.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,4 @@
use std::path::PathBuf;
use std::path::{Component, PathBuf};

use clap::{builder::ValueParser, value_parser, ArgAction, Parser, Subcommand};
use stemgen::{
Expand Down Expand Up @@ -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<PathBuf, Box<dyn std::error::Error>>
{
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<Component> = 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<dyn std::error::Error>> {
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<DeleteOriginal, String> {
value.try_into()
}

impl TryFrom<&str> for DeleteOriginal {
type Error = String;

fn try_from(value: &str) -> Result<Self, Self::Error> {
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)]
Expand All @@ -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)]
Expand Down Expand Up @@ -134,12 +224,32 @@ pub fn prepare_ffmpeg(ctx: &Cli) -> Result<(), Box<dyn std::error::Error>> {

#[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"));
}
}
178 changes: 100 additions & 78 deletions cli/src/generate.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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};
Expand Down Expand Up @@ -81,100 +81,122 @@ pub fn generate(ctx: &Cli, command: &GenerateArgs) -> Result<bool, Box<dyn std::
ctx.ext
);
let output_file = command.output.join(output_filename);
if matches!((canonicalize(file), canonicalize(&output_file)), (Ok(file), Ok(output_file)) if output_file == file) {
continue;
}
if output_file.exists() {
if !ctx.force {
eprintln!(
"Cannot proceed with {}: stem file already exist in output directory!",
output_file.display()
);
has_failure |= true;
if let Err(err) = command.delete_original.proceed(&file, &output_file) {
eprintln!(
"Cannot replace original file {}: {}",
file.display(),
err
);
}
continue;
}
std::fs::remove_file(&output_file)?;
}
let mut input = Track::new(file)?;
let mut nistem = if command.preserved_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("#>-"),
);
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<f32> = vec![0f32; 343980 * 2];
let mut original_packets = Vec::with_capacity(512);
let mut original_buffer: Vec<f32> = 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<f32> = vec![0f32; 343980 * 2];
let mut original_packets = Vec::with_capacity(512);
let mut original_buffer: Vec<f32> = 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)
}
Expand All @@ -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()
Expand All @@ -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()
Expand Down
5 changes: 3 additions & 2 deletions cli/src/main.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -39,7 +39,7 @@ mod tests {
};

use crate::{
cli::{Commands, CreateArgs, GenerateArgs}, Cli
cli::{Commands, CreateArgs, DeleteOriginal, GenerateArgs}, Cli
};

#[test]
Expand Down Expand Up @@ -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,
Expand Down
Loading