From c8b79bfb40b049e6e78958190f9570d9738ea974 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 9 Oct 2023 21:44:14 -0400 Subject: [PATCH] cleanup logic --- .../src/seeding/zero_shot_tts.rs | 1 - .../job_types/tts/process_single_tts_job.rs | 177 +++++++++--------- .../job/job_types/tts/vall_e_x/process_job.rs | 108 ++++------- .../job_types/tts/vall_e_x/validate_job.rs | 12 +- 4 files changed, 133 insertions(+), 165 deletions(-) diff --git a/crates/cli/dev_database_seed/src/seeding/zero_shot_tts.rs b/crates/cli/dev_database_seed/src/seeding/zero_shot_tts.rs index 637e361888..e8e2c8ba8e 100644 --- a/crates/cli/dev_database_seed/src/seeding/zero_shot_tts.rs +++ b/crates/cli/dev_database_seed/src/seeding/zero_shot_tts.rs @@ -51,7 +51,6 @@ async fn create_voice_records( maybe_creator_user_token: Some(creator_user_token.as_str()), creator_ip_address: "127.0.0.1", creator_set_visibility: &Default::default(), - maybe_mod_user_token: None, mysql_pool, }).await?; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs index 58bcffa96f..3f0fb5c096 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs @@ -8,10 +8,10 @@ use tokens::tokens::tts_models::TtsModelToken; use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::tts::{ tacotron2_v2_early_fakeyou, vits }; +use crate::job::job_types::tts::{tacotron2_v2_early_fakeyou, vits}; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::process_job::ProcessJobArgs; -use crate::job::job_types::tts::vits::process_job::VitsProcessJobArgs; use crate::job::job_types::tts::vall_e_x::process_job::VALLEXProcessJobArgs; +use crate::job::job_types::tts::vits::process_job::VitsProcessJobArgs; use crate::job_dependencies::JobDependencies; use super::vall_e_x; @@ -21,97 +21,94 @@ pub async fn process_single_tts_job( job: &AvailableInferenceJob ) -> Result { - let tts_model_token = match job.maybe_model_token.as_deref() { - None => { - return Err(ProcessSingleJobError::Other(anyhow!("no model token on job"))); - } - Some(model_token) => TtsModelToken::new_from_str(model_token), - }; + // TODO: Move common checks for slurs, etc. here. + let raw_inference_text = job.maybe_raw_inference_text + .as_deref() + .ok_or(ProcessSingleJobError::Other(anyhow!("no inference text")))?; - // TODO: Move common checks for slurs, etc. here. - let raw_inference_text = job.maybe_raw_inference_text - .as_deref() - .ok_or(ProcessSingleJobError::Other(anyhow!("no inference text")))?; - - // TODO: Interrogate cache (which also depends on other flags) - let maybe_tts_model = get_tts_model_for_inference_improved( - &job_dependencies.mysql_pool, - tts_model_token.as_str() - ).await.map_err(|err| ProcessSingleJobError::Other(anyhow!("database error: {:?}", err)))?; - - // TODO: Attempt to grab job lock - // // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // - // - // info!("Attempting to grab lock for job: {}", job.inference_job_token); - // - // let lock_acquired = - // mark_tts_inference_job_pending_and_grab_lock(&job_args.mysql_pool, job.id) - // .await - // .map_err(|e| ProcessSingleJobError::Other(e))?; - // - // if !lock_acquired { - // warn!("Could not acquire job lock for: {:?}", &job.id); - // let duration = start.elapsed(); - // - // since_creation_span.set_attribute("status", "failure"); - // since_creation_span.set_duration(duration); - // - // job_iteration_span.set_attribute("status", "failure"); - // job_iteration_span.set_duration(duration); - // - // return Ok((since_creation_span, job_iteration_span)); - // } - // - // info!("Lock acquired for job: {}", job.inference_job_token); - - // Look for the Zeroshot job type - // let job_success_result = match job.maybe_model_type { - // Some(InferenceModelType::VallEX) => { - - // } - // Some(model_type) => return Err(ProcessSingleJobError::Other(anyhow!("wrong model type: {:?}", model_type))), - // None => return Err(ProcessSingleJobError::Other(anyhow!("no model type in record"))), - // }; - - if let Some(model_type) = job.maybe_model_type { - if model_type == InferenceModelType::VallEX { - let job_success_result = vall_e_x::process_job::process_job(VALLEXProcessJobArgs { - job_dependencies, - job, - }).await?; - Ok::(job_success_result); - } - // do nothing try regular tts else other zero shot models + match job.maybe_model_type { + Some(InferenceModelType::VallEX) => { + // Zero-shot TTS does not need a fine-tuned model token. + dispatch_zero_shot_model( + job_dependencies, + job, + &raw_inference_text + ).await + }, + Some(InferenceModelType::Tacotron2 | InferenceModelType::Vits) => { + // All other TTS types require a fine-tuned TTS database record. + dispatch_fine_tuned_weights_model( + job_dependencies, + job, + &raw_inference_text + ).await } + Some(other_model_type) => { + Err(ProcessSingleJobError::Other(anyhow!("wrong model type for TTS: {:?}", other_model_type))) + } + None => { + Err(ProcessSingleJobError::Other(anyhow!("tts model type not set"))) + } + } +} - // lets run the regular tts - let tts_model = match maybe_tts_model { - None => { - return Err( - ProcessSingleJobError::Other(anyhow!("tts model not found: {:?}", tts_model_token)) - ); - } - Some(model) => model, - }; +async fn dispatch_zero_shot_model( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob, + raw_inference_text: &str, +) -> Result { + match job.maybe_model_type { + Some(InferenceModelType::VallEX) => { + vall_e_x::process_job::process_job(VALLEXProcessJobArgs { + job_dependencies, + job, + }).await + } + _ => { + Err(ProcessSingleJobError::Other(anyhow!("not a zero-shot model"))) + } + } +} - let job_success_result: JobSuccessResult = match tts_model.tts_model_type { - TtsModelType::Tacotron2 => { - tacotron2_v2_early_fakeyou::process_job::process_job(ProcessJobArgs { - job_dependencies, - job, - tts_model: &tts_model, - raw_inference_text, - }).await? - } - TtsModelType::Vits => { - vits::process_job::process_job(VitsProcessJobArgs { - job_dependencies, - job, - tts_model: &tts_model, - raw_inference_text, - }).await? - } - }; - Ok(job_success_result) +async fn dispatch_fine_tuned_weights_model( + job_dependencies: &JobDependencies, + job: &AvailableInferenceJob, + raw_inference_text: &str, +) -> Result { + let tts_model_token = job.maybe_model_token + .as_deref() + .ok_or(ProcessSingleJobError::Other(anyhow!("no model token on job")))?; + + // TODO(bt,2023-10-09): Interrogate TTS model cache before querying database. + let maybe_tts_model = get_tts_model_for_inference_improved( + &job_dependencies.mysql_pool, + tts_model_token, + ).await.map_err(|err| ProcessSingleJobError::Other(anyhow!("database error: {:?}", err)))?; + + let tts_model = match maybe_tts_model { + None => { + return Err(ProcessSingleJobError::Other(anyhow!("tts model not found: {:?}", tts_model_token))); + }, + Some(tts_model) => tts_model, + }; + + match tts_model.tts_model_type { + TtsModelType::Tacotron2 => { + tacotron2_v2_early_fakeyou::process_job::process_job(ProcessJobArgs { + job_dependencies, + job, + tts_model: &tts_model, + raw_inference_text, + }).await + } + TtsModelType::Vits => { + vits::process_job::process_job(VitsProcessJobArgs { + job_dependencies, + job, + tts_model: &tts_model, + raw_inference_text, + }).await + } + } } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs index 3230a6efe2..4f59c8a509 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job.rs @@ -1,93 +1,65 @@ -use std::fs::read_to_string; -use std::time::Instant; - use anyhow::anyhow; -use log::{error, info, warn}; +use log::{error, info}; -use buckets::public::media_files::original_file::MediaFileBucketPath; -use container_common::filesystem::check_file_exists::check_file_exists; -use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; -use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; -use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; -use filesys::file_size::file_size; -use hashing::sha256::sha256_hash_file::sha256_hash_file; -use mimetypes::mimetype_for_file::get_mimetype_for_file; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; -use mysql_queries::queries::media_files::insert_media_file_from_face_animation::{insert_media_file_from_face_animation, InsertArgs}; -use tokens::users::user::UserToken; -use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; +use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; - -//use crate::job::job_types::lipsync::sad_talker::download_audio_file::download_audio_file; -//use crate::job::job_types::lipsync::sad_talker::sad_talker_inference_command::InferenceArgs; - -use crate::job::job_types::tts::vall_e_x::validate_job::validate_job; - use crate::job_dependencies::JobDependencies; - // This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked pub struct VALLEXProcessJobArgs<'a> { - pub job_dependencies: &'a JobDependencies, - pub job: &'a AvailableInferenceJob, - } + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} - pub async fn process_job(args: VALLEXProcessJobArgs<'_>) -> Result { +pub async fn process_job(args: VALLEXProcessJobArgs<'_>) -> Result { - // get job args and dependencies and execute them with the inputs - let job = args.job; - let deps = args.job_dependencies; + // get job args and dependencies and execute them with the inputs + let job = args.job; + let deps = args.job_dependencies; - let mut job_progress_reporter = args.job_dependencies + let job_progress_reporter = args.job_dependencies .job_progress_reporter .new_generic_inference(job.inference_job_token.as_str()) .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; - // validate the inputs on the job shouldn't this happen at json - //let job_args = validate_job(job)?; + // validate the inputs on the job shouldn't this happen at json + //let job_args = validate_job(job)?; + // Need to download the models + info!("Download models (if not present)..."); - // Need to download the models - info!("Download models (if not present)..."); - let mut i : usize = 0; + for downloader in deps.job_type_details.vall_e_x.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads, + ).await; - for downloader in deps.job_type_details.vall_e_x.downloaders.all_downloaders() { - - // Temporary debugging - info!("Downloader {}", i); - i = i + 1; - - let result = downloader.download_if_not_on_filesystem( - &args.job_dependencies.private_bucket_client, - &args.job_dependencies.fs.scoped_temp_dir_creator_for_downloads, - ).await; - - if let Err(e) = result { - error!("could not download: {:?}", e); - return Err(ProcessSingleJobError::from_anyhow_error(e)) - } + if let Err(e) = result { + error!("could not download: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) } - - // Download embeddings files - - // rinr inference - - // ==================== TEMP DIR ==================== // - // create a temp dir - - // run inference - - // upload audio to bucket - - // deletetemp dir files - - // save result ? - - - } + // Download embeddings files + + // rinr inference + + // ==================== TEMP DIR ==================== // + // create a temp dir + + // run inference + + // upload audio to bucket + + // deletetemp dir files + + // save result ? + + Err(ProcessSingleJobError::InvalidJob(anyhow!("this job flow is not yet implemented"))) +} + #[cfg(test)] mod test { diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/validate_job.rs index 81177b61ae..ebc70db63a 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/validate_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/validate_job.rs @@ -1,9 +1,9 @@ +use anyhow::anyhow; + use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; -use crate::job::{job_types::lipsync::sad_talker::validate_job::JobArgs, job_loop::process_single_job_error::ProcessSingleJobError}; +use crate::job::{job_loop::process_single_job_error::ProcessSingleJobError, job_types::lipsync::sad_talker::validate_job::JobArgs}; - - -pub fn validate_job(job: &AvailableInferenceJob) -> Result { - -} \ No newline at end of file +pub fn validate_job(job: &AvailableInferenceJob) -> Result { + Err(ProcessSingleJobError::InvalidJob(anyhow!("job not yet written"))) +}