Handle zip files with .index files

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"

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",

crates/lib/files/archives/Cargo.toml

@@ -0,0 +1,25 @@
+[package]
+name = "archives"
+edition = "2018"
+version = "0.0.1"
+authors = [
+    "Brandon Thomas <bt@brand.io>",
+    "Brandon Thomas <echelon@gmail.com>",
+]
+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

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;

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<P: AsRef<Path>>(file_path: P) -> AnyhowResult<Vec<String>> {
+  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)
+}

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<DownloadedRvcFile> {
+  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,
+    })
+  }
+}

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;

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<JobResults> {
 
+  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<PathBuf>; // 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<PathBuf> = 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,