From a0919fffee3c59c613c214ec7f1e84613bf19a36 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 1 Jun 2021 14:23:16 -0400 Subject: [PATCH] overhaul tts_job --- .../up.sql | 4 +- src/common_queries/firehose_publisher.rs | 15 ++ src/job_queries/tts_inference_job_queries.rs | 13 +- .../tacotron_inference_command.rs | 5 +- src/tts_inference_job.rs | 233 +++++++----------- 5 files changed, 116 insertions(+), 154 deletions(-) diff --git a/migrations/2021-05-15-212918_inference_results/up.sql b/migrations/2021-05-15-212918_inference_results/up.sql index b8b405c918..609d35718e 100644 --- a/migrations/2021-05-15-212918_inference_results/up.sql +++ b/migrations/2021-05-15-212918_inference_results/up.sql @@ -38,6 +38,9 @@ CREATE TABLE tts_results ( -- Where the wav, spectrogram, and etc. are located. public_bucket_hash CHAR(64) NOT NULL, + -- Where the wav, spectrogram, and etc. are located. + public_bucket_wav_audio_path VARCHAR(255) NOT NULL, + public_bucket_spectrogram_path VARCHAR(255) NOT NULL, -- Where the wav, spectrogram, and etc. are located. inference_public_bucket_uuid VARCHAR(36) NOT NULL, @@ -45,7 +48,6 @@ CREATE TABLE tts_results ( -- ========== METADATA ========== file_size_bytes INT(10) NOT NULL DEFAULT 0, - mime_type VARCHAR(32) NOT NULL DEFAULT '', duration_millis INT(10) NOT NULL DEFAULT 0, -- ========== MODERATION DETAILS ========== diff --git a/src/common_queries/firehose_publisher.rs b/src/common_queries/firehose_publisher.rs index c4de611150..98e4e00885 100644 --- a/src/common_queries/firehose_publisher.rs +++ b/src/common_queries/firehose_publisher.rs @@ -103,6 +103,21 @@ impl FirehosePublisher { Ok(()) } + pub async fn tts_inference_finished( + &self, + maybe_user_token: Option<&str>, + model_token: &str, + result_token: &str + ) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::TtsInferenceCompleted, + maybe_user_token, + Some(model_token), + Some(result_token) + ).await?; + Ok(()) + } + pub async fn enqueue_w2l_template_upload(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { let _record_id = self.insert( FirehoseEvent::W2lTemplateUploadStarted, diff --git a/src/job_queries/tts_inference_job_queries.rs b/src/job_queries/tts_inference_job_queries.rs index 6f12a5649d..1f201a78e9 100644 --- a/src/job_queries/tts_inference_job_queries.rs +++ b/src/job_queries/tts_inference_job_queries.rs @@ -135,8 +135,8 @@ pub async fn insert_tts_result>( pool: &MySqlPool, job: &TtsInferenceJobRecord, bucket_audio_results_path: P, + bucket_spectrogram_results_path: P, file_size_bytes: u64, - maybe_mime_type: Option<&str>, duration_millis: u64 ) -> AnyhowResult<(u64, String)> { @@ -147,6 +147,11 @@ pub async fn insert_tts_result>( .display() .to_string(); + let bucket_spectrogram_result_path = &bucket_audio_results_path + .as_ref() + .display() + .to_string(); + let query_result = sqlx::query!( r#" INSERT INTO tts_results @@ -160,10 +165,10 @@ SET creator_set_visibility = 'public', public_bucket_hash = 'TODO', - public_bucket_audio_path = ?, + public_bucket_wav_audio_path = ?, + public_bucket_spectrogram_path = ?, file_size_bytes = ?, - mime_type = ?, duration_millis = ? "#, inference_result_token, @@ -174,9 +179,9 @@ SET job.creator_ip_address.clone(), bucket_audio_result_path, + bucket_spectrogram_result_path, file_size_bytes, - maybe_mime_type.unwrap_or(""), duration_millis ) .execute(pool) diff --git a/src/script_execution/tacotron_inference_command.rs b/src/script_execution/tacotron_inference_command.rs index fc22274dc1..7a9b3c0344 100644 --- a/src/script_execution/tacotron_inference_command.rs +++ b/src/script_execution/tacotron_inference_command.rs @@ -29,8 +29,9 @@ impl TacotronInferenceCommand { synthesizer_checkpoint_path: P, vocoder_checkpoint_path: P, input_text_filename: P, - output_metadata_filename: P, output_audio_filename: P, + output_spectrogram_filename: P, + output_metadata_filename: P, spawn_process: bool ) -> AnyhowResult<()> { let mut command = String::new(); @@ -51,6 +52,8 @@ impl TacotronInferenceCommand { command.push_str(&input_text_filename.as_ref().display().to_string()); command.push_str(" --output_audio_filename "); command.push_str(&output_audio_filename.as_ref().display().to_string()); + command.push_str(" --output_spectrogram_filename "); + command.push_str(&output_spectrogram_filename.as_ref().display().to_string()); command.push_str(" --output_metadata_filename "); command.push_str(&output_metadata_filename.as_ref().display().to_string()); diff --git a/src/tts_inference_job.rs b/src/tts_inference_job.rs index 6d56196b8d..9ed80a1941 100644 --- a/src/tts_inference_job.rs +++ b/src/tts_inference_job.rs @@ -244,7 +244,7 @@ async fn main_loop(inferencer: Inferencer) { } } -async fn process_jobs(inferencer: &Inferencer, jobs: Vec) -> AnyhowResult<()> { +async fn process_jobs(inferencer: &Inferencer, jobs: Vec) -> AnyhowResult<()> { for job in jobs.into_iter() { let result = process_job(inferencer, &job).await; match result { @@ -252,7 +252,7 @@ async fn process_jobs(inferencer: &Inferencer, jobs: Vec) Err(e) => { warn!("Failure to process job: {:?}", e); let failure_reason = ""; - let _r = mark_w2l_inference_job_failure( + let _r = mark_tts_inference_job_failure( &inferencer.mysql_pool, &job, failure_reason) @@ -278,20 +278,21 @@ fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { Ok(serde_json::from_str(&buffer)?) } -async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> AnyhowResult<()> { +async fn process_job(inferencer: &Inferencer, job: &TtsInferenceJobRecord) -> AnyhowResult<()> { // TODO 1. Mark processing - // TODO 2. Check if vocoder model is downloaded / download to stable location + // TODO 2. Check if vocoder model is downloaded / download to stable location (DONE) - // TODO 3. Query model by token. - // TODO 4. Check if model is downloaded, otherwise download to stable location + // TODO 3. Query model by token. (DONE) + // TODO 4. Check if model is downloaded, otherwise download to stable location (DONE) - // TODO 5. Process Inference + // TODO 5. Write text to file + // TODO 6. Process Inference - // TODO 6. Upload Result - // TODO 7. Save record - // TODO 8. Mark job done + // TODO 7. Upload Result + // TODO 8. Save record + // TODO 9. Mark job done // ==================== CONFIRM OR DOWNLOAD TTS VOCODER MODEL ==================== // @@ -302,212 +303,148 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An if !tts_vocoder_model_fs_path.exists() { warn!("Vocoder model file does not exist: {:?}", &tts_vocoder_model_fs_path); - let model_object_path = inferencer.bucket_path_unifier - .w2l_pretrained_models_path(&model_filename); + let tts_vocoder_model_object_path = inferencer.bucket_path_unifier + .tts_pretrained_vocoders_path(&tts_vocoder_model_filename); - info!("Download from bucket path: {:?}", &model_object_path); + info!("Download vocoder from bucket path: {:?}", &tts_vocoder_model_object_path); inferencer.private_bucket_client.download_file_to_disk( - &model_object_path, - &model_fs_path + &tts_vocoder_model_object_path, + &tts_vocoder_model_fs_path ).await?; - info!("Downloaded model from bucket!"); + info!("Downloaded tts vocoder model from bucket!"); } - // ==================== CONFIRM OR DOWNLOAD W2L END BUMP ==================== // + // ==================== LOOK UP TTS SYNTHESIZER RECORD (WHICH CONTAINS ITS BUCKET PATH) ==================== // - let end_bump_filename = inferencer.w2l_end_bump_filename.clone(); - let end_bump_fs_path = inferencer.semi_persistent_cache.w2l_end_bump_path(&end_bump_filename); + info!("Looking up TTS model by token: {}", &job.model_token); - if !end_bump_fs_path.exists() { - warn!("End bump file does not exist: {:?}", &end_bump_fs_path); + let query_result = get_tts_model_by_token( + &inferencer.mysql_pool, + &job.model_token).await?; - let end_bump_object_path = inferencer.bucket_path_unifier - .end_bump_video_for_w2l_path(&end_bump_filename); - - info!("Download from bucket path: {:?}", &end_bump_object_path); - - inferencer.private_bucket_client.download_file_to_disk( - &end_bump_object_path, - &end_bump_fs_path - ).await?; - - info!("Downloaded end bump from bucket!"); - } - - // ==================== LOOK UP TEMPLATE RECORD ==================== // - - let template_token = match &job.maybe_w2l_template_token { - Some(token) => token.to_string(), + let tts_model = match query_result { + Some(model) => model, None => { - warn!("non-template token based inference not yet supported"); - return Err(anyhow!("non-template token based inference not yet supported")) + warn!("TTS model not found: {}", &job.model_token); + return Err(anyhow!("Model not found!")) }, }; - info!("Looking up w2l template by token: {}", &template_token); + // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // - let query_result = get_w2l_template_by_token(&inferencer.mysql_pool, &template_token).await?; + // TODO: Let's just put paths in the db + // TODO: We'll probably need to LRU cache these. - let w2l_template = match query_result { - Some(template) => template, - None => { - warn!("W2L Template not found: {}", &template_token); - return Err(anyhow!("Template not found!")) - }, - }; + let tts_synthesizer_fs_path = inferencer.semi_persistent_cache.tts_synthesizer_model_path( + &tts_model.model_token); - // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE AUDIO OR VIDEO ==================== // + if !tts_synthesizer_fs_path.exists() { + info!("TTS synthesizer model file does not exist: {:?}", &tts_synthesizer_fs_path); - // Template is based on the `private_bucket_hash`: - // - private_bucket_hash: 1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 - // - private_bucket_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 - // - private_bucket_cached_faces_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_detected_faces.pickle - // - maybe_public_bucket_preview_image_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_preview.webp + let tts_synthesizer_object_path = inferencer.bucket_path_unifier + .tts_synthesizer_path(&tts_model.private_bucket_hash); - let template_media_fs_path = inferencer.semi_persistent_cache.w2l_template_media_path( - &w2l_template.private_bucket_hash); - - if !template_media_fs_path.exists() { - info!("W2L template media file does not exist: {:?}", &template_media_fs_path); - - let template_media_object_path = inferencer.bucket_path_unifier - .media_templates_for_w2l_path(&w2l_template.private_bucket_hash); - - info!("Download from template media path: {:?}", &template_media_object_path); + info!("Download from template media path: {:?}", &tts_synthesizer_object_path); inferencer.private_bucket_client.download_file_to_disk( - &template_media_object_path, - &template_media_fs_path + &tts_synthesizer_object_path, + &tts_synthesizer_fs_path ).await?; info!("Downloaded template media from bucket!"); } - // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE FACE ==================== // + // ==================== WRITE TEXT TO FILE ==================== // - // Template is based on the `private_bucket_hash`: - // - private_bucket_hash: 1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 - // - private_bucket_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2 - // - private_bucket_cached_faces_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_detected_faces.pickle - // - maybe_public_bucket_preview_image_object_name: /user_uploaded_w2l_templates/1/5/1/1519edf86e6975fdcd0a56a5953d84948db79f2b9ce588818d7fa544d5cb38b2_preview.webp - - let face_template_fs_path = inferencer.semi_persistent_cache.w2l_face_template_path( - &w2l_template.private_bucket_hash); - - if !face_template_fs_path.exists() { - info!("W2L face template file does not exist: {:?}", &face_template_fs_path); - - let face_template_object_path = inferencer.bucket_path_unifier - .precomputed_faces_for_w2l_path(&w2l_template.private_bucket_hash); - - info!("Download from face template path: {:?}", &face_template_object_path); - - inferencer.private_bucket_client.download_file_to_disk( - &face_template_object_path, - &face_template_fs_path - ).await?; - - info!("Downloaded face template from bucket!"); - } - - // ==================== DOWNLOAD USER AUDIO ==================== // - - let temp_dir = format!("temp_{}", job.id); + let temp_dir = format!("temp_tts_inference_{}", job.id); let temp_dir = TempDir::new(&temp_dir)?; // NB: Exists until it goes out of scope. - let audio_bucket_hash = match &job.maybe_public_audio_bucket_hash { - Some(l) => l.clone(), - None => { - warn!("Only W2L jobs with user-uploaded audio are supported right now"); - return Err(anyhow!("Only W2L jobs with user-uploaded audio are supported right now")) - }, - }; - - let audio_fs_path = temp_dir.path().join(&audio_bucket_hash); - - let audio_object_path = inferencer.bucket_path_unifier - .user_audio_for_w2l_inference_path(&audio_bucket_hash); - - inferencer.private_bucket_client.download_file_to_disk( - &audio_object_path, - &audio_fs_path - ).await?; - + let text_input_fs_path = None; // ==================== RUN INFERENCE ==================== // - let output_video_fs_path = temp_dir.path().join("output.mp4"); + let output_audio_fs_path = temp_dir.path().join("output.mp4"); let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + let output_spectrogram_fs_path = temp_dir.path().join("spectrogram.json"); let is_image = w2l_template.template_type.contains("image"); - info!("Is image? {}", is_image); - info!("Running W2L inference..."); + info!("Running TTS inference..."); - inferencer.w2l_inference.execute( - &model_fs_path, - &audio_fs_path, - &end_bump_fs_path, - &template_media_fs_path, - &face_template_fs_path, + info!("Expected output audio filename: {:?}", &output_audio_fs_path); + info!("Expected output spectrogram filename: {:?}", &output_spectrogram_fs_path); + info!("Expected output metadata filename: {:?}", &output_metadata_fs_path); + + inferencer.tts_inference.execute( + &tts_synthesizer_fs_path, + &tts_vocoder_model_fs_path, + &text_input_fs_path, + &output_audio_fs_path, + &output_spectrogram_fs_path, &output_metadata_fs_path, - &output_video_fs_path, false, - false )?; - info!("Output filename: {:?}", &output_video_fs_path); - // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // info!("Checking that output files exist..."); - check_file_exists(&output_video_fs_path)?; + check_file_exists(&output_audio_fs_path)?; + check_file_exists(&output_spectrogram_fs_path)?; check_file_exists(&output_metadata_fs_path)?; let file_metadata = read_metadata_file(&output_metadata_fs_path)?; - // ==================== UPLOAD TO BUCKETS ==================== // + // ==================== UPLOAD AUDIO TO BUCKET ==================== // - let result_object_path = inferencer.bucket_path_unifier.w2l_inference_video_output_path( + let audio_result_object_path = inferencer.bucket_path_unifier.tts_inference_wav_audio_output_path( &job.inference_job_token); - info!("Image/video destination bucket path: {:?}", &result_object_path); + info!("Audio destination bucket path: {:?}", &audio_result_object_path); - info!("Uploading image/video..."); - - let original_mime_type = file_metadata.mimetype - .as_deref() - .unwrap_or("application/octet-stream"); + info!("Uploading audio..."); inferencer.public_bucket_client.upload_filename_with_content_type( - &result_object_path, - &output_video_fs_path, - original_mime_type) + &audio_result_object_path, + &output_audio_fs_path, + "audio/wav") + .await?; + + // ==================== UPLOAD SPECTROGRAM TO BUCKETS ==================== // + + let spectrogram_result_object_path = inferencer.bucket_path_unifier.tts_inference_spectrogram_output_path( + &job.inference_job_token); + + info!("Spectrogram destination bucket path: {:?}", &spectrogram_result_object_path); + + info!("Uploading spectrogram..."); + + inferencer.public_bucket_client.upload_filename_with_content_type( + &spectrogram_result_object_path, + &output_spectrogram_fs_path, + "application/json") .await?; // ==================== SAVE RECORDS ==================== // - info!("Saving w2l inference record..."); - let (id, inference_result_token) = insert_w2l_result( + info!("Saving tts inference record..."); + let (id, inference_result_token) = insert_tts_result( &inferencer.mysql_pool, job, - &result_object_path, + &audio_result_object_path, + &spectrogram_result_object_path, file_metadata.file_size_bytes, - file_metadata.mimetype.as_deref(), - file_metadata.width, - file_metadata.height, file_metadata.duration_millis.unwrap_or(0)) .await?; info!("Marking job complete..."); - mark_w2l_inference_job_done(&inferencer.mysql_pool, job, true).await?; + mark_tts_inference_job_done(&inferencer.mysql_pool, job, true).await?; - inferencer.firehose_publisher.w2l_inference_finished( + inferencer.firehose_publisher.tts_inference_finished( job.maybe_creator_user_token.as_deref(), - &job.inference_job_token, + &tts_model.model_token, &inference_result_token) .await .map_err(|e| {