From 1427e3eecd0dc035ae8f2215c09d6f5bfba12596 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 18 Jul 2023 23:05:46 -0400 Subject: [PATCH] Handle zip files with .index files --- Cargo.lock | 8 + Cargo.toml | 1 + crates/lib/files/archives/Cargo.toml | 25 ++++ crates/lib/files/archives/src/lib.rs | 21 +++ .../files/archives/src/list_archive_files.rs | 31 ++++ .../rvc_v2/extract_rvc_files.rs | 96 ++++++++++++ .../job_types/voice_conversion/rvc_v2/mod.rs | 1 + .../rvc_v2/process_rvc_v2_model.rs | 141 ++++-------------- 8 files changed, 210 insertions(+), 114 deletions(-) create mode 100644 crates/lib/files/archives/Cargo.toml create mode 100644 crates/lib/files/archives/src/lib.rs create mode 100644 crates/lib/files/archives/src/list_archive_files.rs create mode 100644 crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/extract_rvc_files.rs diff --git a/Cargo.lock b/Cargo.lock index f80311b30f..c136ee19d8 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -883,6 +883,14 @@ dependencies = [ "x11rb", ] +[[package]] +name = "archives" +version = "0.0.1" +dependencies = [ + "errors", + "zip", +] + [[package]] name = "arrayref" version = "0.3.6" diff --git a/Cargo.toml b/Cargo.toml index 39b9d4e320..5a345dd197 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -30,6 +30,7 @@ members = [ "crates/lib/deprecated/http_server_common", # TODO(bt, 2023-03-20): Move useful pieces to actix_helpers "crates/lib/envvar", "crates/lib/errors", + "crates/lib/files/archives", "crates/lib/files/filesys", "crates/lib/files/hashing", "crates/lib/files/media", diff --git a/crates/lib/files/archives/Cargo.toml b/crates/lib/files/archives/Cargo.toml new file mode 100644 index 0000000000..5651decbb2 --- /dev/null +++ b/crates/lib/files/archives/Cargo.toml @@ -0,0 +1,25 @@ +[package] +name = "archives" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +# NB: Public "archives" crate does not exist, so it's safe to use for now. +# We'll use this for zip, tgz, rar, etc. file types and related utilities. +name = "archives" +path = "src/lib.rs" + +[dependencies] +# Internal +errors = { path = "../../errors" } + +# External +zip = "0.6.6" + +[dev-dependencies] +# None yet diff --git a/crates/lib/files/archives/src/lib.rs b/crates/lib/files/archives/src/lib.rs new file mode 100644 index 0000000000..c43bc2d8f6 --- /dev/null +++ b/crates/lib/files/archives/src/lib.rs @@ -0,0 +1,21 @@ +//! archives +//! +//! The purpose of this library is to collect common utilities for +//! zip, tar.gz, rar, etc. files. +//! + +// Never allow these +#![forbid(private_in_public)] +#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things + +// Okay to toggle +#![forbid(unreachable_patterns)] +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +pub mod list_archive_files; diff --git a/crates/lib/files/archives/src/list_archive_files.rs b/crates/lib/files/archives/src/list_archive_files.rs new file mode 100644 index 0000000000..46af211040 --- /dev/null +++ b/crates/lib/files/archives/src/list_archive_files.rs @@ -0,0 +1,31 @@ +use std::path::Path; +use errors::{anyhow, AnyhowResult}; + +// NB: Untested +pub fn list_archive_files>(file_path: P) -> AnyhowResult> { + let file = std::fs::File::open(&file_path)?; + let reader = std::io::BufReader::new(file); + let mut archive = zip::ZipArchive::new(reader)?; + + let mut entries = Vec::with_capacity(archive.len()); + + for i in 0..archive.len() { + let file = archive.by_index(i)?; + + let _inner_path = match file.enclosed_name() { + Some(path) => path, + None => { + return Err(anyhow!("Entry {} has a suspicious path", file.name())); + } + }; + + if (*file.name()).ends_with('/') { + continue; // Directory + } + + //entries.push(format!("{}", _inner_path.display())); + entries.push(file.name().to_string()); + } + + Ok(entries) +} diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/extract_rvc_files.rs b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/extract_rvc_files.rs new file mode 100644 index 0000000000..08e78263b8 --- /dev/null +++ b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/extract_rvc_files.rs @@ -0,0 +1,96 @@ +use errors::AnyhowResult; +use std::path::{Path, PathBuf}; +use log::warn; +use tempdir::TempDir; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use mimetypes::mimetype_for_file::get_mimetype_for_file; + +pub enum DownloadedRvcFile { + /// Error - wrong file type was uploaded + InvalidModel, + + /// Only the model file (.pth) was uploaded + ModelFileOnly { + model_file: PathBuf, + }, + + /// Archive containing both files was uploaded. + ModelAndIndexFile { + model_file: PathBuf, + index_file: PathBuf, + }, +} + +pub fn extract_rvc_files(download_file: &Path, temp_dir: &TempDir) -> AnyhowResult { + let maybe_mimetype = get_mimetype_for_file(download_file)?; + + if maybe_mimetype != Some("application/zip") { + warn!("File must be an application/zip (either .pth or .zip)"); + return Ok(DownloadedRvcFile::InvalidModel); + } + + let file = std::fs::File::open(download_file)?; + let reader = std::io::BufReader::new(file); + let mut archive = zip::ZipArchive::new(reader)?; + + if archive.len() > 50 { + // If there are hundreds of entries, this is likely a .pth file rather than .zip file. + // The fact that it looks like a zip file is a consequence of the .pth serialization format + // being the same as zip archives. + // In any case, we'll verify that the model is valid downstream of this call. + return Ok(DownloadedRvcFile::ModelFileOnly { model_file: download_file.to_path_buf() }); + } + + if archive.len() > 3 { + // There's something suspicious if the model files has more than 3 entries. + // It should have the .pth file, a .index file, and _maybe_ a .txt file for credits. + warn!("File has wrong number of entries to be a valid model: {}", archive.len()); + return Ok(DownloadedRvcFile::InvalidModel); + } + + let mut maybe_path_to_model = None; + let mut maybe_path_to_index = None; + + for i in 0..archive.len() { + let mut file = archive.by_index(i)?; + let filename = file.name().to_lowercase(); + + let temp_file_path; + + if filename.ends_with(".pth") { + temp_file_path = temp_dir.path().join("model.pth"); + maybe_path_to_model = Some(temp_file_path.clone()); + } else if filename.ends_with(".index") || filename.ends_with(".idx") { + temp_file_path = temp_dir.path().join("model.index"); + maybe_path_to_index = Some(temp_file_path.clone()); + } else { + continue; + } + + let mut outfile = std::fs::File::create(&temp_file_path)?; + std::io::copy(&mut file, &mut outfile)?; + } + + let path_to_model = match maybe_path_to_model { + Some(path_to_model) => path_to_model, + None => { + // It isn't valid to not have a model file. + if let Some(path_to_index) = maybe_path_to_index { + safe_delete_temp_file(&path_to_index); + } + warn!("Archive did not have a model file within."); + return Ok(DownloadedRvcFile::InvalidModel); + } + }; + + if let Some(path_to_index) = maybe_path_to_index { + Ok(DownloadedRvcFile::ModelAndIndexFile { + model_file: path_to_model, + index_file: path_to_index, + }) + } else { + Ok(DownloadedRvcFile::ModelFileOnly { + model_file: path_to_model, + }) + } +} diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/mod.rs b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/mod.rs index 719374f37d..6b2cda23c8 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/mod.rs +++ b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/mod.rs @@ -1,2 +1,3 @@ +pub mod extract_rvc_files; pub mod process_rvc_v2_model; pub mod rvc_v2_model_check_command; \ No newline at end of file diff --git a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs index fe0f35263f..38f3be70bf 100644 --- a/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs +++ b/crates/service/job/download_job/src/job_types/voice_conversion/rvc_v2/process_rvc_v2_model.rs @@ -5,6 +5,7 @@ use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_d use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; use crate::JobState; use crate::job_loop::job_results::JobResults; +use crate::job_types::voice_conversion::rvc_v2::extract_rvc_files::{DownloadedRvcFile, extract_rvc_files}; use crate::job_types::voice_conversion::rvc_v2::rvc_v2_model_check_command::CheckArgs; use crockford::crockford_entropy_lower; use enums::by_table::voice_conversion_models::voice_conversion_model_type::VoiceConversionModelType; @@ -13,7 +14,6 @@ use errors::AnyhowResult; use filesys::file_size::file_size; use jobs_common::redis_job_status_logger::RedisJobStatusLogger; use log::{error, info, warn}; -use mimetypes::mimetype_for_file::get_mimetype_for_file; use mysql_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; use mysql_queries::queries::voice_conversion::models::insert_voice_conversion_model_from_download_job::{insert_voice_conversion_model_from_download_job, InsertVoiceConversionModelArgs}; use std::path::PathBuf; @@ -28,137 +28,50 @@ pub async fn process_rvc_v2_model<'a, 'b>( redis_logger: &'a mut RedisJobStatusLogger<'b>, ) -> AnyhowResult { + info!("Processing model file..."); + redis_logger.log_status("checking rvc (v2) model")?; + // ==================== DETERMINE FILE CONTENTS ==================== // let original_download_file_path = PathBuf::from(download_filename.clone()); - let maybe_mimetype = get_mimetype_for_file(&original_download_file_path)?; + let download_contents = extract_rvc_files(&original_download_file_path, temp_dir)?; - for _ in 0..20 { - info!("Mimetype: {:?}", &maybe_mimetype); - } + let original_model_file_path; + let maybe_original_model_index_file_path: Option; // TODO RENAME MAYBE - // The mimetype MUST be application/zip! - // There are two cases we support here: - // 1. If it's a .pth file (pytorch pickle), it's zip compressed, but it's 100% model data - // (hundreds of files for weights) - // 2. If it's a .zip file, it should have a .pt file and possibly an optional .index file - // we'll allow up to one extra file (txt or readme) - if maybe_mimetype == Some("application/zip") { - - info!("\n========================================\n"); - - let file = std::fs::File::open(&original_download_file_path)?; - let reader = std::io::BufReader::new(file); - let mut archive = zip::ZipArchive::new(reader)?; - - // https://drive.google.com/file/d/1YP8_OSBrtwz1Sf9gCw-9daizpyBVMtGS/view?usp=drive_link - // Entry 0 is a file with name "added_IVF293_Flat_nprobe_1_TwilightSparkle_v2.index" (36187059 bytes) - // Entry 1 is a file with name "TwilightSparkle.pth" (55226951 bytes) - - // https://drive.google.com/file/d/1TThj37azqeNRfnQdNN0XrzCHg03xrbhN/view?usp=sharing - // (so-vits-svc) checkpoint_best_legacy_500.pt - // Entry 0 is a file with name "archive/data.pkl" (71294 bytes) - // Entry 1 is a file with name "archive/data/0" (3072 bytes) - // Entry 2 is a file with name "archive/data/1" (516096 bytes) Entry 3 is a file with name "archive/data/10" (2097152 bytes) Entry 4 is a file with name "archive/data/100" (2359296 bytes) - // Entry 5 is a file with name "archive/data/101" (3072 bytes) - // Entry 6 is a file with name "archive/data/102" (2359296 bytes) - // Entry 7 is a file with name "archive/data/103" (3072 bytes) - // ... - // Entry 222 is a file with name "archive/data/97" (3072 bytes) - // Entry 223 is a file with name "archive/data/98" (2359296 bytes) - // Entry 224 is a file with name "archive/data/99" (3072 bytes) - // Entry 225 is a file with name "archive/version" (2 bytes) - - // https://drive.google.com/file/d/13Jo5wxCogcOM_qB5j2_lnNrnoPbNziJM/view?usp=sharing - // AnneBoonchuy.pth - // Entry 0 is a file with name "AnneBoonchuy/data.pkl" (61476 bytes) - // Entry 1 is a file with name "AnneBoonchuy/data/0" (294912 bytes) - // Entry 2 is a file with name "AnneBoonchuy/data/1" (384 bytes) - // Entry 3 is a file with name "AnneBoonchuy/data/10" (384 bytes) - // Entry 4 is a file with name "AnneBoonchuy/data/100" (384 bytes) - // Entry 5 is a file with name "AnneBoonchuy/data/101" (384 bytes) - // ... - - // https://drive.google.com/file/d/1EuUgaYFsSO-CVrEg-CeN5nN6lIb7S9UR/view?usp=sharing - // AnneBoonchuy.index - // - - - for i in 0..archive.len() { - let file = archive.by_index(i).unwrap(); - - let outpath = match file.enclosed_name() { - Some(path) => path, - None => { - println!("Entry {} has a suspicious path", file.name()); - continue; - } - }; - - { - let comment = file.comment(); - if !comment.is_empty() { - println!("Entry {i} comment: {comment}"); - } - } - - if (*file.name()).ends_with('/') { - println!( - "Entry {} is a directory with name \"{}\"", - i, - outpath.display() - ); - } else { - println!( - "Entry {} is a file with name \"{}\" ({} bytes)", - i, - outpath.display(), - file.size() - ); - } + match download_contents { + DownloadedRvcFile::InvalidModel => { + return Err(anyhow!("download did not contain a valid model payload")); + } + DownloadedRvcFile::ModelFileOnly { model_file } => { + info!("RVC download only contains a model (no index file)."); + original_model_file_path = model_file; + maybe_original_model_index_file_path = None; + } + DownloadedRvcFile::ModelAndIndexFile { model_file, index_file } => { + info!("RVC download contains both a model and an index file."); + original_model_file_path = model_file; + maybe_original_model_index_file_path = Some(index_file); } - - info!("\n========================================\n"); } - // 1) Tom_Petty_TalkNet.zip - // https://drive.google.com/file/d/1hlCZ2BJJWVEd_9Fb9JPhrHX9IyeGOckK/view?usp=drive_link - // Some("application/zip") - - // 2) BartSimpson.pth - // https://drive.google.com/file/d/1oShUemIvi8h8KOxqgB0R2umNbCC8Y3rn/view?usp=drive_link - // Some("application/zip") -- ugh - // - // - // ==================== RUN MODEL CHECK ==================== // - // TODO: Unzip bundled index files. - // TODO: Handle models without an index file!? (Maybe we'll be lazy and _not_.) - info!("Checking that model is valid..."); - redis_logger.log_status("checking rvc (v2) model")?; - - let original_model_file_path = PathBuf::from("/home/bt/models/rvc_v2/AnneBoonchuy.pth"); - //let original_model_index_file_path = PathBuf::from("/home/bt/models/rvc_v2/added_IVF119_Flat_nprobe_1_AnneBoonchuy_v2.index"); - let original_model_index_file_path : Option = None; - - //let input_wav_path = PathBuf::from("input.wav"); // NB: Bundled with repo let output_wav_path = temp_dir.path().join("output.wav"); let model_check_result = job_state.sidecar_configs.rvc_v2_model_check_command.execute_check(CheckArgs { model_path: &original_model_file_path, - maybe_model_index_path: original_model_index_file_path.as_deref(), - maybe_input_path: None, //Some(&input_wav_path), + maybe_model_index_path: maybe_original_model_index_file_path.as_deref(), + maybe_input_path: None, output_path: &output_wav_path, - //device: Device::Cuda, }); if let Err(e) = model_check_result { safe_delete_temp_file(&original_model_file_path); - safe_delete_possible_temp_file(original_model_index_file_path.as_deref()); + safe_delete_possible_temp_file(maybe_original_model_index_file_path.as_deref()); safe_delete_temp_file(&output_wav_path); safe_delete_temp_directory(&temp_dir); return Err(anyhow!("model check error: {:?}", e)); @@ -189,7 +102,7 @@ pub async fn process_rvc_v2_model<'a, 'b>( if let Err(err) = job_state.bucket_client.upload_filename(&model_bucket_path, &original_model_file_path).await { error!("Problem uploading model file: {:?}", err); safe_delete_temp_file(&original_model_file_path); - safe_delete_possible_temp_file(original_model_index_file_path.as_deref()); + safe_delete_possible_temp_file(maybe_original_model_index_file_path.as_deref()); safe_delete_temp_file(&output_wav_path); safe_delete_temp_directory(&temp_dir); return Err(err); @@ -197,7 +110,7 @@ pub async fn process_rvc_v2_model<'a, 'b>( // ==================== UPLOAD ORIGINAL MODEL INDEX FILE ==================== // - if let Some(original_index_file_path) = original_model_index_file_path.as_deref() { + if let Some(original_index_file_path) = maybe_original_model_index_file_path.as_deref() { let model_index_bucket_path = job_state.bucket_path_unifier.rvc_v2_model_index_path(&private_bucket_hash); info!("Destination bucket path (index): {:?}", &model_index_bucket_path); @@ -217,7 +130,7 @@ pub async fn process_rvc_v2_model<'a, 'b>( // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... info!("Done uploading; deleting temporary files and paths..."); safe_delete_temp_file(&original_model_file_path); - safe_delete_possible_temp_file(original_model_index_file_path.as_deref()); + safe_delete_possible_temp_file(maybe_original_model_index_file_path.as_deref()); safe_delete_temp_file(&output_wav_path); safe_delete_temp_directory(&temp_dir); @@ -234,7 +147,7 @@ pub async fn process_rvc_v2_model<'a, 'b>( creator_user_token: &job.creator_user_token, creator_ip_address: &job.creator_ip_address, creator_set_visibility: Visibility::Public, // TODO: All models default to public at start - has_index_file: original_model_index_file_path.is_some(), + has_index_file: maybe_original_model_index_file_path.is_some(), private_bucket_hash: &private_bucket_hash, private_bucket_object_name: "", // TODO: This should go away. mysql_pool: &job_state.mysql_pool,