diff --git a/.env b/.env index 06dd83f350..e7f1c54d72 100644 --- a/.env +++ b/.env @@ -27,6 +27,7 @@ W2L_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/Wav2Lip #W2L_MODEL_CHECKPOINT=/home/bt/models/wav2lip/s3fd-619a316812.pth W2L_MODEL_CHECKPOINT=/home/bt/models/wav2lip/wav2lip_gan.pth W2L_MODEL_FILENAME=wav2lip_gan.pth +W2L_END_BUMP_FILENAME=vocodes-short-end-bump.mp4 W2L_DOWNLOAD_SCRIPT_NAME=vocodes_process_upload.py W2L_INFERENCE_SCRIPT_NAME=vocodes_process_inference.py diff --git a/models/Wav2Lip/vocodes_process_inference.py b/models/Wav2Lip/vocodes_process_inference.py index 37e7329834..2348a4b976 100644 --- a/models/Wav2Lip/vocodes_process_inference.py +++ b/models/Wav2Lip/vocodes_process_inference.py @@ -48,6 +48,9 @@ parser.add_argument('--cached_faces_filename', type=str, parser.add_argument('--audio_filename', type=str, help='Filepath of video/audio file to use as raw audio source', required=True) +parser.add_argument('--end_bump_file', type=str, + help='Video file to concatenate at the end') + parser.add_argument('--output_video_filename', type=str, help='Output filename for the final video', required=True) @@ -95,10 +98,6 @@ parser.add_argument('--audio_start_pad_millis', default=0, type=int, parser.add_argument('--audio_end_pad_millis', default=0, type=int, help='Seconds to pad the end of the audio') -parser.add_argument('--end_bump_file', type=str, - help='Video file to concatenate at the end') - - # Purely for debugging on the host machine: parser.add_argument('--preserve_tempdir', default=False, action='store_true', help='Keep the tempdir arround for debugging') diff --git a/src/buckets/bucket_path_unifier.rs b/src/buckets/bucket_path_unifier.rs index e47ea2e0a1..88dec45496 100644 --- a/src/buckets/bucket_path_unifier.rs +++ b/src/buckets/bucket_path_unifier.rs @@ -10,6 +10,7 @@ pub struct BucketPathUnifier { pub tts_inference_output_root: PathBuf, pub w2l_inference_output_root: PathBuf, pub w2l_model_root: PathBuf, + pub w2l_end_bump_root: PathBuf, } impl BucketPathUnifier { @@ -29,18 +30,20 @@ impl BucketPathUnifier { tts_inference_output_root: PathBuf::from("/tts_inference_output"), w2l_inference_output_root: PathBuf::from("/w2l_inference_output"), w2l_model_root: PathBuf::from("/w2l_pretrained_models"), + w2l_end_bump_root: PathBuf::from("/w2l_end_bumps"), } } - pub fn end_bump_video_for_w2l_path(&self, end_bump_filename: &str) -> String { - "".to_string() - } - // W2L pretrained models. There are only two. pub fn w2l_pretrained_models_path(&self, w2l_model_name: &str) -> PathBuf { self.w2l_model_root.join(w2l_model_name) } + // W2L "end bumps" are videos added at the end. + pub fn end_bump_video_for_w2l_path(&self, end_bump_filename: &str) -> PathBuf { + self.w2l_end_bump_root.join(end_bump_filename) + } + // The video or images uploaded as templates // eg. /user_uploaded_w2l_templates/1/5/1/151a[...60]... pub fn media_templates_for_w2l_path(&self, template_file_hash: &str) -> PathBuf { @@ -79,6 +82,7 @@ impl BucketPathUnifier { let video_filename = format!("vocodes_video_{}.mp4", video_filename); let hashed_path = Self::hashed_directory_path(token); + let hashed_path = hashed_path.to_lowercase(); return self.w2l_inference_output_root .join(hashed_path) @@ -119,6 +123,7 @@ mod tests { tts_inference_output_root: PathBuf::from("/test_path_tts_output"), w2l_inference_output_root: PathBuf::from("/test_path_w2l_output"), w2l_model_root: PathBuf::from("/test_path_w2l_pretrained_models"), + w2l_end_bump_root: PathBuf::from("/test_path_w2l_end_bumps"), } } @@ -129,6 +134,13 @@ mod tests { "/test_path_w2l_pretrained_models/model.pth"); } + #[test] + fn test_end_bump_video_for_w2l_path() { + let paths = get_instance(); + assert_eq!(paths.end_bump_video_for_w2l_path("logo.mp4").to_str().unwrap(), + "/test_path_w2l_end_bumps/logo.mp4"); + } + #[test] fn test_user_audio_for_w2l_inference_path() { let paths = get_instance(); @@ -153,12 +165,19 @@ mod tests { #[test] fn test_w2l_inference_video_output_path() { let paths = get_instance(); + + // Case 1: Tokens without a "token type" assert_eq!(paths.w2l_inference_video_output_path("foobar").to_str().unwrap(), "/test_path_w2l_output/f/o/o/vocodes_video_foobar.mp4"); - // Note: it also removes the token type from dir path and handles the colon: + // Case 2: Tokens with a "token type" + // Note: that it removes the token type from dir path and handles the colon: assert_eq!(paths.w2l_inference_video_output_path("type:abcdef").to_str().unwrap(), "/test_path_w2l_output/a/b/c/vocodes_video_typeabcdef.mp4"); + + // It also handles capitalization + assert_eq!(paths.w2l_inference_video_output_path("TYPE:ABCDEF").to_str().unwrap(), + "/test_path_w2l_output/a/b/c/vocodes_video_TYPEABCDEF.mp4"); } #[test] diff --git a/src/script_execution/wav2lip_inference_command.rs b/src/script_execution/wav2lip_inference_command.rs index 346425eae4..bfe83a6be9 100644 --- a/src/script_execution/wav2lip_inference_command.rs +++ b/src/script_execution/wav2lip_inference_command.rs @@ -27,16 +27,18 @@ impl Wav2LipInferenceCommand { } } - pub fn execute>(&self, - checkpoint_path: P, - audio_filename: P, - media_template_filename: P, - cached_faces_filename: P, - output_metadata_filename: P, - output_video_filename: P, - is_image: bool, - spawn_process: bool) -> AnyhowResult<()> - { + pub fn execute>( + &self, + checkpoint_path: P, + audio_filename: P, + end_bump_video_filename: P, + media_template_filename: P, + cached_faces_filename: P, + output_metadata_filename: P, + output_video_filename: P, + is_image: bool, + spawn_process: bool + ) -> AnyhowResult<()> { let mut command = String::new(); command.push_str("echo 'test'"); @@ -53,6 +55,8 @@ impl Wav2LipInferenceCommand { command.push_str(&media_template_filename.as_ref().display().to_string()); command.push_str(" --audio_filename "); command.push_str(&audio_filename.as_ref().display().to_string()); + command.push_str(" --end_bump_file "); + command.push_str(&end_bump_video_filename.as_ref().display().to_string()); command.push_str(" --cached_faces_filename "); command.push_str(&cached_faces_filename.as_ref().display().to_string()); command.push_str(" --output_video_filename "); diff --git a/src/util/semi_persistent_cache_dir.rs b/src/util/semi_persistent_cache_dir.rs index 6e3bd89875..30f3211858 100644 --- a/src/util/semi_persistent_cache_dir.rs +++ b/src/util/semi_persistent_cache_dir.rs @@ -8,6 +8,7 @@ pub struct SemiPersistentCacheDir { cache_root: PathBuf, tts_model_root: PathBuf, w2l_model_root: PathBuf, + w2l_end_bump_root: PathBuf, w2l_face_templates_root: PathBuf, w2l_templates_media_root: PathBuf, video_asset_root: PathBuf, // end bump, etc. @@ -28,6 +29,7 @@ impl SemiPersistentCacheDir { cache_root: cache_root.clone(), tts_model_root: cache_root.join("tts/models/"), w2l_model_root: cache_root.join("w2l/models/"), + w2l_end_bump_root: cache_root.join("w2l/end_bumps/"), w2l_face_templates_root: cache_root.join("w2l/face_templates/"), w2l_templates_media_root: cache_root.join("w2l/template_media/"), video_asset_root: cache_root.join("static_video_assets/"), @@ -58,6 +60,21 @@ impl SemiPersistentCacheDir { Ok(()) } + // ==================== W2L END BUMPS (short video logos) ==================== + + pub fn w2l_end_bump_path(&self, end_bump_filename: &str) -> PathBuf { + self.w2l_end_bump_root.join(end_bump_filename) + } + + pub fn w2l_end_bump_directory(&self) -> &Path { + &self.w2l_end_bump_root + } + + pub fn create_w2l_end_bump_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.w2l_end_bump_directory())?; + Ok(()) + } + // ==================== W2L MEDIA ==================== /// We cache W2L media here. diff --git a/src/w2l_inference_job.rs b/src/w2l_inference_job.rs index aa93e1874b..c43d2debf7 100644 --- a/src/w2l_inference_job.rs +++ b/src/w2l_inference_job.rs @@ -89,6 +89,7 @@ struct Inferencer { // Command to run pub inference_script: String, pub w2l_model_filename: String, + pub w2l_end_bump_filename: String, } #[tokio::main] @@ -179,13 +180,16 @@ async fn main() -> AnyhowResult<()> { info!("Creating pod semi-persistent cache dirs..."); semi_persistent_cache.create_w2l_model_path()?; + semi_persistent_cache.create_w2l_end_bump_path()?; semi_persistent_cache.create_w2l_face_template_path()?; semi_persistent_cache.create_w2l_template_media_path()?; - semi_persistent_cache.create_w2l_model_path()?; let w2l_model_filename = easyenv::get_env_string_or_default( "W2L_MODEL_FILENAME", "wav2lip_gan.pth"); + let w2l_end_bump_filename = easyenv::get_env_string_or_default( + "W2L_END_BUMP_FILENAME", "vocodes-short-end-bump.mp4"); + let inferencer = Inferencer { download_temp_directory: temp_directory, mysql_pool, @@ -200,6 +204,7 @@ async fn main() -> AnyhowResult<()> { bucket_path_unifier: BucketPathUnifier::default_paths(), semi_persistent_cache, w2l_model_filename, + w2l_end_bump_filename, }; main_loop(inferencer).await; @@ -314,7 +319,7 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An let model_fs_path = inferencer.semi_persistent_cache.w2l_model_path(&model_filename); if !model_fs_path.exists() { - info!("Model file does not exist: {:?}", &model_fs_path); + warn!("Model file does not exist: {:?}", &model_fs_path); let model_object_path = inferencer.bucket_path_unifier .w2l_pretrained_models_path(&model_filename); @@ -329,6 +334,27 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An info!("Downloaded model from bucket!"); } + // ==================== CONFIRM OR DOWNLOAD W2L END BUMP ==================== // + + 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); + + if !end_bump_fs_path.exists() { + warn!("End bump file does not exist: {:?}", &end_bump_fs_path); + + 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 { @@ -442,6 +468,7 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An inferencer.w2l_inference.execute( &model_fs_path, &audio_fs_path, + &end_bump_fs_path, &template_media_fs_path, &face_template_fs_path, &output_metadata_fs_path,