From b8b10f6fc157bb80a4564bd2e2e49a80a3f4a2e0 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Thu, 18 May 2023 12:17:08 -0700
Subject: [PATCH] f0-method and frequency transpose for job

---
 .../src/payloads/generic_inference_args.rs    | 12 +++----
 .../job_types/vc/so_vits_svc/process_job.rs   | 31 ++++++++++++++-----
 .../so_vits_svc_inference_command.rs          | 25 +++++++++++++++
 3 files changed, 55 insertions(+), 13 deletions(-)

diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args.rs
index 72b12a44f2..e7236d6041 100644
--- a/crates/lib/mysql_queries/src/payloads/generic_inference_args.rs
+++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args.rs
@@ -58,16 +58,16 @@ pub enum PolymorphicInferenceArgs {
     #[serde(skip_serializing_if = "Option::is_none")]
     auto_predict_f0: Option<bool>,
 
-    /// Argument for so-vits-svc
-    /// The python model defaults to true, but that sounds awful,
-    /// so we default to false unless specified.
+    /// Argument for so-vits-svc (-fm / --f0-method)
+    /// Crepe, dio, harvest, etc.
+    /// If unspecified, the model defaults to crepe
     #[serde(rename = "fm")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes.
     #[serde(skip_serializing_if = "Option::is_none")]
     override_f0_method: Option<FundamentalFrequencyMethodForJob>,
 
-    /// Argument for so-vits-svc
-    /// The python model defaults to true, but that sounds awful,
-    /// so we default to false unless specified.
+    /// Argument for so-vits-svc (-t / --transpose)
+    /// Pitch control.
+    /// If unspecified, the model uses "0".
     #[serde(rename = "t")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes.
     #[serde(skip_serializing_if = "Option::is_none")]
     transpose: Option<i32>,
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs
index 5fb962d84c..140c43ce46 100644
--- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs
+++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs
@@ -155,22 +155,37 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
   //info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path);
 
   // TODO: Limit output length for premium.
-  // NB: Tacotron operates on decoder steps. 1000 steps is the default and correlates to
-  //  roughly 12 seconds max. Here we map seconds to decoder steps.
-  //let max_decoder_steps = seconds_to_decoder_steps(job.max_duration_seconds);
 
-  // If not specified by the user, turn off auto prediction. It sounds awful.
-  let auto_predict_f0 = job.maybe_inference_args
+  let maybe_args = job.maybe_inference_args
       .as_ref()
       .map(|args| args.args.as_ref())
-      .flatten()
+      .flatten();
+
+
+  // If not specified by the user, turn off auto prediction. It sounds awful.
+  let auto_predict_f0 = maybe_args
       .map(|args| match args {
         PolymorphicInferenceArgs::Tts { .. } => None,
-        PolymorphicInferenceArgs::Vc { auto_predict_f0 } => *auto_predict_f0,
+        PolymorphicInferenceArgs::Vc { auto_predict_f0, .. } => *auto_predict_f0,
       })
       .flatten()
       .unwrap_or(false);
 
+  let maybe_transpose = maybe_args
+      .map(|args| match args {
+        PolymorphicInferenceArgs::Tts { .. } => None,
+        PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose,
+      })
+      .flatten();
+
+  let maybe_override_f0_method = maybe_args
+      .map(|args| match args {
+        PolymorphicInferenceArgs::Tts { .. } => None,
+        PolymorphicInferenceArgs::Vc { override_f0_method, .. } =>
+          *override_f0_method,
+      })
+      .flatten();
+
   // ==================== RUN INFERENCE SCRIPT ==================== //
 
   let inference_start_time = Instant::now();
@@ -186,6 +201,8 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
         maybe_config_path: None,
         device: Device::Cuda,
         auto_predict_f0,
+        maybe_override_f0_method,
+        maybe_transpose,
       });
 
   let inference_duration = Instant::now().duration_since(inference_start_time);
diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
index 6ccb4be258..26b5a88e49 100644
--- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
+++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/so_vits_svc_inference_command.rs
@@ -8,6 +8,7 @@ use std::fs::OpenOptions;
 use std::path::{Path, PathBuf};
 use std::time::Duration;
 use subprocess::{ExitStatus, Popen, PopenConfig, Redirection};
+use mysql_queries::payloads::generic_inference_args::FundamentalFrequencyMethodForJob;
 use subprocess_common::docker_options::{DockerEnvVar, DockerFilesystemMount, DockerGpu, DockerOptions};
 
 /// This command is used to check tacotron for being a real model
@@ -71,6 +72,12 @@ pub struct InferenceArgs<P: AsRef<Path>> {
   /// This sounds better when left off, but it defaults to *ON* if not specified.
   pub auto_predict_f0: bool,
 
+  /// --f0-method: f0 prediction method
+  pub maybe_override_f0_method: Option<FundamentalFrequencyMethodForJob>,
+
+  /// --transpose: pitch adjustment
+  pub maybe_transpose: Option<i32>,
+
   /// --device: cpu or cuda
   pub device: Device,
 }
@@ -236,6 +243,24 @@ impl SoVitsSvcInferenceCommand {
     command.push_str(if args.auto_predict_f0 { "true" } else { "false" });
     command.push_str(" ");
 
+    if let Some(transpose) = args.maybe_transpose {
+      let value = transpose.to_string();
+      command.push_str(" --transpose ");
+      command.push_str(&value);
+      command.push_str(" ");
+    }
+
+    if let Some(f0_method) = args.maybe_override_f0_method {
+      let method = match f0_method {
+        FundamentalFrequencyMethodForJob::Crepe => "crepe",
+        FundamentalFrequencyMethodForJob::Dio => "dio",
+        FundamentalFrequencyMethodForJob::Harvest => "harvest",
+      };
+      command.push_str(" --f0-method ");
+      command.push_str(&method);
+      command.push_str(" ");
+    }
+
     let device = match args.device {
       Device::Cuda => "cuda",
       Device::Cpu => "cpu",