Preprocess W2L uploads (but glitchy)

.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

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)

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,

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...')

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<u64>
 {
   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;

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,

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);