From 8fdcb9a85c8d77ca550311e864c7828c034b1dec Mon Sep 17 00:00:00 2001
From: Paul Sajna <hello@paulsajna.com>
Date: Tue, 6 Jun 2023 21:01:50 -0700
Subject: [PATCH] audiowmark integration

---
 .../src/util => lib/jobs_common/src}/audiowmark.rs       | 0
 crates/lib/jobs_common/src/lib.rs                        | 3 ++-
 .../users/user_sessions/get_user_session_by_token.rs     | 2 +-
 .../user_sessions/get_user_session_by_token_light.rs     | 2 +-
 .../src/job/job_types/vc/so_vits_svc/process_job.rs      | 9 ++++++++-
 .../src/job_steps/process_single_job.rs                  | 7 ++++++-
 6 files changed, 18 insertions(+), 5 deletions(-)
 rename crates/{service/job/inference_job/src/util => lib/jobs_common/src}/audiowmark.rs (100%)

diff --git a/crates/service/job/inference_job/src/util/audiowmark.rs b/crates/lib/jobs_common/src/audiowmark.rs
similarity index 100%
rename from crates/service/job/inference_job/src/util/audiowmark.rs
rename to crates/lib/jobs_common/src/audiowmark.rs
diff --git a/crates/lib/jobs_common/src/lib.rs b/crates/lib/jobs_common/src/lib.rs
index 1e4864a7f5..be00f1725d 100644
--- a/crates/lib/jobs_common/src/lib.rs
+++ b/crates/lib/jobs_common/src/lib.rs
@@ -3,4 +3,5 @@
 pub mod job_progress_reporter;
 pub mod noop_logger;
 pub mod redis_job_status_logger;
-pub mod semi_persistent_cache_dir;
\ No newline at end of file
+pub mod semi_persistent_cache_dir;
+pub mod audiowmark;
diff --git a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs
index 0fe46a9549..e88adf1211 100644
--- a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs
+++ b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs
@@ -202,7 +202,7 @@ WHERE user_sessions.token = ?
     },
     Err(err) => {
       match err {
-        RowNotFound => {
+        sqlx::Error::RowNotFound => {
           warn!("Valid cookie; invalid session: {}", session_token);
           Ok(None)
         },
diff --git a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token_light.rs b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token_light.rs
index 62a4c69bdd..43bfdd237e 100644
--- a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token_light.rs
+++ b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token_light.rs
@@ -38,7 +38,7 @@ AND deleted_at IS NULL
     Ok(session_record) => Ok(Some(session_record)),
     Err(err) => {
       match err {
-        RowNotFound => {
+        sqlx::Error::RowNotFound => {
           warn!("Valid cookie; invalid session: {}", session_token);
           Ok(None)
         },
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 71c9e169eb..e2d598238d 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
@@ -9,6 +9,7 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 use crate::job::job_types::vc::so_vits_svc::so_vits_svc_inference_command::{Device, InferenceArgs};
 use crate::job_dependencies::JobDependencies;
 use crate::util::maybe_download_file_from_bucket::maybe_download_file_from_bucket;
+use jobs_common::audiowmark::audiowmark_add;
 use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
 use errors::AnyhowResult;
 use filesys::create_dir_all_if_missing::create_dir_all_if_missing;
@@ -252,6 +253,11 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
     file_size_bytes,
   };
 
+  // ==================== Watermark the result audio ================== //
+  let output_audio_watermarked_fs_path = work_temp_dir.path().join("output_watermarked.wav");
+  let _watermark_process_stdout = audiowmark_add(output_audio_fs_path, 40, b"FAKE YOU DOT COM", &output_audio_watermarked_fs_path).map_err(|e| ProcessSingleJobError::Other(e.into()); 
+  
+
   //safe_delete_temp_file(&output_metadata_fs_path);
 
   // ==================== UPLOAD AUDIO TO BUCKET ==================== //
@@ -269,11 +275,12 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result<JobSuccess
 
   args.job_dependencies.public_bucket_client.upload_filename_with_content_type(
     &result_bucket_object_pathbuf,
-    &output_audio_fs_path,
+    &output_audio_watermarked_fs_path,
     "audio/wav")
       .await
       .map_err(|e| ProcessSingleJobError::Other(e))?;
 
+  safe_delete_temp_file(&output_audio_watermarked_fs_path);
   safe_delete_temp_file(&output_audio_fs_path);
 
 //  // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== //
diff --git a/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs b/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs
index 75096d2080..2624cd6c8f 100644
--- a/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs
+++ b/crates/service/job/tts_inference_job/src/job_steps/process_single_job.rs
@@ -30,6 +30,7 @@ use tempdir::TempDir;
 use tts_common::clean_symbols::clean_symbols;
 use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic;
 use tts_common::text_pipelines::text_pipeline_type::TextPipelineType;
+use jobs_common::audiowmark::audiowmark_add;
 
 /// Text starting with this will be treated as a test request.
 /// This allows the request to bypass the model cache and query the latest TTS model.
@@ -349,6 +350,9 @@ pub async fn process_single_job(
 
   safe_delete_temp_file(&output_metadata_fs_path);
 
+  // =================== Watermark Audio ============================ //
+  let output_audio_watermarked_fs_path = temp_dir.path().join("output_watermarked.wav");
+  let _audiowmark_stdout = audiowmark_add(&output_audio_fs_path, 40, *b"FAKE YOU DOT COM", &output_audio_watermarked_fs_path).map_err(|e| ProcessSingleJobError::Other(e.into()))?; 
   // ==================== UPLOAD AUDIO TO BUCKET ==================== //
 
   job_progress_reporter.log_status("uploading result")
@@ -363,12 +367,13 @@ pub async fn process_single_job(
 
   job_args.public_bucket_client.upload_filename_with_content_type(
     &audio_result_object_path,
-    &output_audio_fs_path,
+    &output_audio_watermarked_fs_path,
     "audio/wav")
       .await
       .map_err(|e| ProcessSingleJobError::Other(e))?;
 
   safe_delete_temp_file(&output_audio_fs_path);
+  safe_delete_temp_file(&output_audio_watermarked_fs_path);
 
   // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== //