diff --git a/.env b/.env index b978eb5bf5..4f9feebeb5 100644 --- a/.env +++ b/.env @@ -17,3 +17,8 @@ TTS_DOWNLOAD_BUCKET_ROOT=/user_uploaded_tts_models W2L_DOWNLOAD_BUCKET_NAME=dev-vocodes-private-uploads W2L_DOWNLOAD_BUCKET_ROOT=/user_uploaded_w2l_templates + +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_SCRIPT_NAME=vocodes_process_upload.py \ No newline at end of file diff --git a/debug/curl/upload_w2l_image_gdrive.py b/debug/curl/upload_w2l_image_gdrive.py index aaa9ec65fb..aadb14da89 100755 --- a/debug/curl/upload_w2l_image_gdrive.py +++ b/debug/curl/upload_w2l_image_gdrive.py @@ -32,10 +32,15 @@ print("===== Upload W2L =====") cookies = { 'session': session_cookie } +# mario.zip +#download_url = 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing' +# elon_integration_test.jpg +download_url = 'https://drive.google.com/file/d/1yoBSb6nwFUZVI4CD-nM2BSofGdBjC8Nn/view?usp=sharing' + payload = { 'idempotency_token': str(uuid.uuid4()), 'title': 'this is a new image template', - 'download_url': 'https://drive.google.com/file/d/15-tkgblTZpa0ifmvvyr_kgWy2cguTpAe/view?usp=sharing', + 'download_url': download_url, } r = requests.post(upload_url, cookies=cookies, json=payload) diff --git a/migrations/2021-05-15-071601_assets/up.sql b/migrations/2021-05-15-071601_assets/up.sql index c680f0cd16..0f1fcf6241 100644 --- a/migrations/2021-05-15-071601_assets/up.sql +++ b/migrations/2021-05-15-071601_assets/up.sql @@ -192,6 +192,8 @@ CREATE TABLE w2l_templates ( private_bucket_hash CHAR(64) NOT NULL, -- The "full url" version of the path private_bucket_object_name VARCHAR(255) NOT NULL, + -- The "full url" of the cached faces object + private_bucket_cached_faces_object_name VARCHAR(255) NOT NULL, -- For the thumbnail we show. public_bucket_hash CHAR(64) NOT NULL, diff --git a/models/Wav2Lip/vocodes_process_upload.py b/models/Wav2Lip/vocodes_process_upload.py index cd1cda16b5..848f51e11b 100644 --- a/models/Wav2Lip/vocodes_process_upload.py +++ b/models/Wav2Lip/vocodes_process_upload.py @@ -50,9 +50,12 @@ parser = argparse.ArgumentParser(description='Inference code to lip-sync videos parser.add_argument('--checkpoint_path', type=str, help='Name of saved checkpoint to load weights from', required=True) -parser.add_argument('--face', type=str, +parser.add_argument('--image_or_video_filename', type=str, help='Filepath of video/image that contains faces to use', required=True) +parser.add_argument('--output_cached_faces_filename', type=str, + help='Output filename for the cached faces file', required=True) + # NB: Not needed for processing the upload file for faces: # # parser.add_argument('--audio', type=str, @@ -114,7 +117,8 @@ parser.add_argument('--preserve_tempdir', default=False, action='store_true', args = parser.parse_args() args.img_size = 96 -if os.path.isfile(args.face) and os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg']: +if os.path.isfile(args.image_or_video_filename) \ + and os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg']: args.static = True elif args.is_image: args.static = True @@ -258,23 +262,23 @@ def load_model(path): def main(tempdir): - video_faces_pickle_file = args.face + '.faces.pickle' + video_faces_pickle_file = args.output_cached_faces_filename print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True) frame_w = 0 frame_h = 0 - if not os.path.isfile(args.face): - raise ValueError('--face argument must be a valid path to video/image file') + if not os.path.isfile(args.image_or_video_filename): + raise ValueError('--image_or_video_filename argument must be a valid path to video/image file') - elif os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg'] \ + elif os.path.splitext(args.image_or_video_filename)[1] in ['.jpg', '.png', '.jpeg'] \ or args.is_image: - full_frames = [cv2.imread(args.face)] + full_frames = [cv2.imread(args.image_or_video_filename)] fps = args.fps frame_h, frame_w = full_frames[0].shape[:-1] else: - video_stream = cv2.VideoCapture(args.face) + video_stream = cv2.VideoCapture(args.image_or_video_filename) fps = video_stream.get(cv2.CAP_PROP_FPS) print('Reading video frames...') diff --git a/src/job_queries/w2l_download_job_queries.rs b/src/job_queries/w2l_download_job_queries.rs index 7acfe7a262..0141f95b3d 100644 --- a/src/job_queries/w2l_download_job_queries.rs +++ b/src/job_queries/w2l_download_job_queries.rs @@ -113,7 +113,8 @@ pub async fn insert_w2l_template(pool: &MySqlPool, template_type: &str, // TODO: ENUM! job: &W2lTemplateUploadJobRecord, private_bucket_hash: &str, - private_bucket_object_name: &str) + private_bucket_object_name: &str, + private_bucket_cached_faces_object_name: &str) -> AnyhowResult { let model_token = random_token(32); @@ -133,7 +134,8 @@ SET creator_ip_address = ?, original_download_url = ?, private_bucket_hash = ?, - private_bucket_object_name = ? + private_bucket_object_name = ?, + private_bucket_cached_faces_object_name = ? "#, model_token, template_type, @@ -143,7 +145,8 @@ SET job.creator_ip_address.clone(), job.download_url.clone(), private_bucket_hash.to_string(), - private_bucket_object_name.to_string() + private_bucket_object_name.to_string(), + private_bucket_cached_faces_object_name.to_string() ) .execute(pool) .await; diff --git a/src/script_execution/wav2lip_process_upload_command.rs b/src/script_execution/wav2lip_process_upload_command.rs index 73df5307d9..d5eec7eba7 100644 --- a/src/script_execution/wav2lip_process_upload_command.rs +++ b/src/script_execution/wav2lip_process_upload_command.rs @@ -27,11 +27,9 @@ impl Wav2LipPreprocessClient { } pub fn execute(&self, - audio_filename: &str, image_or_video_filename: &str, - output_video_filename: &str, + output_cached_faces_filename: &str, is_image: bool, - disable_end_bump: bool, spawn_process: bool) -> AnyhowResult<()> { let mut command = String::new(); @@ -46,12 +44,10 @@ impl Wav2LipPreprocessClient { command.push_str(&self.script_name); command.push_str(" --checkpoint_path "); command.push_str(&self.checkpoint_path); - command.push_str(" --face "); + command.push_str(" --image_or_video_filename "); command.push_str(image_or_video_filename); - //command.push_str(" --audio "); - //command.push_str(audio_filename); - //command.push_str(" --outfile "); - //command.push_str(output_video_filename); + command.push_str(" --output_cached_faces_filename"); + command.push_str(output_cached_faces_filename); if is_image { command.push_str(" --is_image "); @@ -74,14 +70,14 @@ impl Wav2LipPreprocessClient { .write(true) .create(true) .truncate(true) - .open("/tmp/stdout.txt")?; + .open("/tmp/wav2lip_upload_stdout.txt")?; let stderr_file = OpenOptions::new() .read(true) .write(true) .create(true) .truncate(true) - .open("/tmp/stderr.txt")?; + .open("/tmp/wav2lip_upload_stderr.txt")?; let mut p = Popen::create(&command_parts, PopenConfig { //stdout: Redirection::Pipe, diff --git a/src/w2l_download_job.rs b/src/w2l_download_job.rs index 39f136e204..a7f06e0ee2 100644 --- a/src/w2l_download_job.rs +++ b/src/w2l_download_job.rs @@ -35,6 +35,7 @@ use std::path::{PathBuf, Path}; use std::process::Command; use std::time::Duration; use tempdir::TempDir; +use crate::script_execution::wav2lip_process_upload_command::Wav2LipPreprocessClient; // Buckets const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; @@ -43,6 +44,11 @@ const ENV_REGION_NAME : &'static str = "REGION_NAME"; const ENV_BUCKET_NAME : &'static str = "W2L_DOWNLOAD_BUCKET_NAME"; const ENV_BUCKET_ROOT : &'static str = "W2L_DOWNLOAD_BUCKET_ROOT"; +// Python code +const ENV_CODE_DIRECTORY : &'static str = "W2L_CODE_DIRECTORY"; +const ENV_MODEL_CHECKPOINT : &'static str = "W2L_MODEL_CHECKPOINT"; +const ENV_SCRIPT_NAME : &'static str = "W2L_SCRIPT_NAME"; + const DEFAULT_RUST_LOG: &'static str = "debug,actix_web=info"; const DEFAULT_TEMP_DIR: &'static str = "/tmp"; @@ -51,6 +57,7 @@ struct Downloader { pub mysql_pool: MySqlPool, pub bucket_client: BucketClient, pub google_drive_downloader: GoogleDriveDownloadCommand, + pub w2l_processor: Wav2LipPreprocessClient, // Command to run pub download_script: String, // Root to store W2L templates @@ -87,6 +94,16 @@ async fn main() -> AnyhowResult<()> { None, )?; + let py_code_directory = easyenv::get_env_string_required(ENV_CODE_DIRECTORY)?; + let py_script_name = easyenv::get_env_string_required(ENV_SCRIPT_NAME)?; + let py_model_checkpoint = easyenv::get_env_string_required(ENV_MODEL_CHECKPOINT)?; + + let w2l_preprecess_command = Wav2LipPreprocessClient::new( + &py_code_directory, + &py_script_name, + &py_model_checkpoint, + ); + let temp_directory = easyenv::get_env_string_or_default( "DOWNLOAD_TEMP_DIR", DEFAULT_TEMP_DIR); @@ -119,6 +136,7 @@ async fn main() -> AnyhowResult<()> { bucket_client, download_script, google_drive_downloader, + w2l_processor: w2l_preprecess_command, bucket_root_w2l_template_uploads: bucket_root.to_string(), }; @@ -211,23 +229,52 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord) .map(|c| c.to_string()) .unwrap_or("".to_string()); + // ==================== DOWNLOAD FILE ==================== // + info!("Calling downloader..."); let download_filename = downloader.google_drive_downloader .download_file(&download_url, &temp_dir).await?; + let file_path = PathBuf::from(&download_filename); + + // ==================== PROCESS FACES ==================== // + + let cached_faces_filename = format!("{}_detected_faces.pickle", &download_filename); + let is_image = false; // TODO: Don't always treat as video. + let spawn_process = true; + + downloader.w2l_processor.execute( + &download_filename, + &cached_faces_filename, + is_image, + spawn_process)?; + + + // ==================== UPLOAD TO BUCKET ==================== // + let private_bucket_hash = get_file_hash(&download_filename)?; info!("File hash: {}", private_bucket_hash); - // NB: /.../a/b/c/d/abcdefg.bin - let object_name = hash_to_bucket_path( + // Full path to video/image + let full_object_path = hash_to_bucket_path( &private_bucket_hash, Some(&downloader.bucket_root_w2l_template_uploads))?; - info!("Destination bucket path: {}", object_name); + info!("Image/video destination bucket path: {}", full_object_path); - let file_path = PathBuf::from(download_filename); - downloader.bucket_client.upload_filename(&object_name, &file_path).await?; + // Full path to cached faces + let full_object_path_cached_faces = format!("{}_detected_faces.pickle", full_object_path); + + info!("Cached faces destination bucket path: {}", full_object_path_cached_faces); + + info!("Uploading image/video..."); + downloader.bucket_client.upload_filename(&full_object_path, &file_path).await?; + + let cached_faces_path = PathBuf::from(&cached_faces_filename); + + info!("Uploading cached faces..."); + downloader.bucket_client.upload_filename(&full_object_path_cached_faces, &cached_faces_path).await?; // TODO: let template_type = "image"; @@ -238,7 +285,8 @@ async fn process_job(downloader: &Downloader, job: &W2lTemplateUploadJobRecord) template_type, job, &private_bucket_hash, - &object_name) + &full_object_path, + &full_object_path_cached_faces) .await?; info!("Saved model record: {}", id);