download template media, start work on process script

models/Wav2Lip/vocodes_process_inference.py

@@ -28,8 +28,22 @@ 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('--image_or_video_type', type=str,
+#                    help='either "image" or "video', required=True)
+
+# This is a little more deliberate than "static" and causes us to pick an FPS, etc.
+# This is useful when the filename doesn't have an extension.
+parser.add_argument('--is_image', default=False, action='store_true',
+                    help='Denote that the input is a single-frame image, not a video')
+
+parser.add_argument('--cached_faces_filename', type=str,
+                    help='Filename for the cached faces file to use', required=True)
+
+#parser.add_argument('--face', type=str,
+#                    help='Filepath of video/image that contains faces to use', required=True)
 parser.add_argument('--audio', type=str,
                     help='Filepath of video/audio file to use as raw audio source', required=True)
 parser.add_argument('--outfile', type=str, help='Video path to save result. See default for an e.g.',
@@ -76,93 +90,92 @@ parser.add_argument('--audio_end_pad_millis', default=0, type=int,
 parser.add_argument('--end_bump_file', type=str,
                     help='Video file to concatenate at the end')
 
-# This is a little more deliberate than "static" and causes us to pick an FPS, etc.
-# This is useful when the filename doesn't have an extension.
-parser.add_argument('--is_image', default=False, action='store_true',
-                    help='Denote that the input is a single-frame image, not a video')
 
 # Purely for debugging on the host machine:
 parser.add_argument('--preserve_tempdir', default=False, action='store_true',
                     help='Keep the tempdir arround for debugging')
 
+#parser.add_argument('--frame_width', type=int, help='Width of the frame')
+#parser.add_argument('--frame_height', type=int, help='Height of the frame')
+
 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']:
-    args.static = True
-elif args.is_image:
-    args.static = True
+#if os.path.isfile(args.face) and os.path.splitext(args.face)[1] in ['.jpg', '.png', '.jpeg']:
+#    args.static = True
+#elif args.is_image:
+#    args.static = True
 
-def get_smoothened_boxes(boxes, T):
-    for i in range(len(boxes)):
-        if i + T > len(boxes):
-            window = boxes[len(boxes) - T:]
-        else:
-            window = boxes[i : i + T]
-        boxes[i] = np.mean(window, axis=0)
-    return boxes
+#def get_smoothened_boxes(boxes, T):
+#    for i in range(len(boxes)):
+#        if i + T > len(boxes):
+#            window = boxes[len(boxes) - T:]
+#        else:
+#            window = boxes[i : i + T]
+#        boxes[i] = np.mean(window, axis=0)
+#    return boxes
 
-def face_detect(images):
-    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
-                                            flip_input=False, device=device)
+#def face_detect(images):
+#    detector = face_detection.FaceAlignment(face_detection.LandmarksType._2D,
+#                                            flip_input=False, device=device)
+#
+#    batch_size = args.face_det_batch_size
+#
+#    while 1:
+#        predictions = []
+#        try:
+#            for i in tqdm(range(0, len(images), batch_size)):
+#                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
+#        except RuntimeError:
+#            if batch_size == 1:
+#                raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
+#            batch_size //= 2
+#            print('Recovering from OOM error; New batch size: {}'.format(batch_size))
+#            continue
+#        break
+#
+#    results = []
+#    pady1, pady2, padx1, padx2 = args.pads
+#    for rect, image in zip(predictions, images):
+#        if rect is None:
+#            cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
+#            raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
+#
+#        y1 = max(0, rect[1] - pady1)
+#        y2 = min(image.shape[0], rect[3] + pady2)
+#        x1 = max(0, rect[0] - padx1)
+#        x2 = min(image.shape[1], rect[2] + padx2)
+#
+#        results.append([x1, y1, x2, y2])
+#
+#    boxes = np.array(results)
+#    if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
+#    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
+#
+#    del detector
+#    return results
 
-    batch_size = args.face_det_batch_size
-
-    while 1:
-        predictions = []
-        try:
-            for i in tqdm(range(0, len(images), batch_size)):
-                predictions.extend(detector.get_detections_for_batch(np.array(images[i:i + batch_size])))
-        except RuntimeError:
-            if batch_size == 1:
-                raise RuntimeError('Image too big to run face detection on GPU. Please use the --resize_factor argument')
-            batch_size //= 2
-            print('Recovering from OOM error; New batch size: {}'.format(batch_size))
-            continue
-        break
-
-    results = []
-    pady1, pady2, padx1, padx2 = args.pads
-    for rect, image in zip(predictions, images):
-        if rect is None:
-            cv2.imwrite('temp/faulty_frame.jpg', image) # check this frame where the face was not detected.
-            raise ValueError('Face not detected! Ensure the video contains a face in all the frames.')
-
-        y1 = max(0, rect[1] - pady1)
-        y2 = min(image.shape[0], rect[3] + pady2)
-        x1 = max(0, rect[0] - padx1)
-        x2 = min(image.shape[1], rect[2] + padx2)
-
-        results.append([x1, y1, x2, y2])
-
-    boxes = np.array(results)
-    if not args.nosmooth: boxes = get_smoothened_boxes(boxes, T=5)
-    results = [[image[y1: y2, x1:x2], (y1, y2, x1, x2)] for image, (x1, y1, x2, y2) in zip(images, boxes)]
-
-    del detector
-    return results
-
-def detect_faces_in_frames(frames, video_faces_pickle_file):
-    t0 = datetime.datetime.now()
-    if args.box[0] == -1:
-        if not args.static:
-            face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
-        else:
-            face_det_results = face_detect([frames[0]])
-    else:
-        print('Using the specified bounding box instead of face detection...')
-        y1, y2, x1, x2 = args.box
-        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
-    t1 = datetime.datetime.now()
-
-    d = t1 - t0
-    print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds()))
-
-    with open(video_faces_pickle_file, 'wb') as f:
-        print('Saving pickle file: {}'.format(video_faces_pickle_file))
-        pickle.dump(face_det_results, f)
-
-    return face_det_results
+#def detect_faces_in_frames(frames, video_faces_pickle_file):
+#    t0 = datetime.datetime.now()
+#    if args.box[0] == -1:
+#        if not args.static:
+#            face_det_results = face_detect(frames) # BGR2RGB for CNN face detection
+#        else:
+#            face_det_results = face_detect([frames[0]])
+#    else:
+#        print('Using the specified bounding box instead of face detection...')
+#        y1, y2, x1, x2 = args.box
+#        face_det_results = [[f[y1: y2, x1:x2], (y1, y2, x1, x2)] for f in frames]
+#    t1 = datetime.datetime.now()
+#
+#    d = t1 - t0
+#    print('Total seconds to detect faces in {} frames: {}'.format(len(frames), d.total_seconds()))
+#
+#    with open(video_faces_pickle_file, 'wb') as f:
+#        print('Saving pickle file: {}'.format(video_faces_pickle_file))
+#        pickle.dump(face_det_results, f)
+#
+#    return face_det_results
 
 #def datagen(frames, mels):
 def datagen(frames, face_det_results, mels):
@@ -308,8 +321,9 @@ def maybe_concatenate_end_bump(tempdir, args, frame_w, frame_h):
 
 
 def main(tempdir):
-    video_faces_pickle_file = args.face + '.faces'
-    print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True)
+    video_faces_pickle_file = args.cached_faces_filename
+    #video_faces_pickle_file = args.face + '.faces'
+    #print('Video faces pickle file: {}'.format(video_faces_pickle_file), flush=True)
 
     frame_w = 0
     frame_h = 0
@@ -351,9 +365,13 @@ def main(tempdir):
 
         frame_h, frame_w = full_frames[0].shape[:-1]
 
+    #frame_w = args.frame_width
+    #frame_h = args.frame_height
+
     print("Number of frames available for inference: "+str(len(full_frames)), flush=True)
     print("Frame dimensions: {}x{}".format(frame_w, frame_h), flush=True)
 
+    # TODO(bt): This could be more efficient
     if not args.audio.endswith('.wav'):
         print('Extracting raw audio...')
         temp_wav_filename = os.path.join(tempdir, 'temp.wav') # previously 'temp/temp.wav'
@@ -365,6 +383,7 @@ def main(tempdir):
 
     maybe_pad_audio_file(tempdir, args)
 
+    # TODO(bt): Save the spectrogram?
     wav = audio.load_wav(args.audio, 16000)
     mel = audio.melspectrogram(wav)
     print(mel.shape)
@@ -386,26 +405,32 @@ def main(tempdir):
     print("Length of mel chunks: {}".format(len(mel_chunks)), flush=True)
 
     if len(mel_chunks) > 2000:
-        # TODO: Catch that the file is too long earlier on.
+        # TODO(bt): Catch that the file is too long earlier on.
         raise Exception("Too many mel chunks: {}".format(len(mel_chunks)))
 
     full_frames = full_frames[:len(mel_chunks)]
 
-    # TODO: Do this offline.
-    print('Face detection time...', flush=True)
-    if not os.path.isfile(video_faces_pickle_file):
-        print('Detecting faces...', flush=True)
-        face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
-    else:
-        with open(video_faces_pickle_file, 'rb') as f:
-            face_det_results = pickle.load(f)
-        if len(face_det_results) < len(full_frames):
-            print("Face detect results have {} frames, but we need {}".format(len(face_det_results), len(full_frames)))
-            print('Detecting faces again...', flush=True)
-            face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
-        else:
-            print("We don't need to find faces! Woo!", flush=True)
+    #print('Face detection time...', flush=True)
+    #if not os.path.isfile(video_faces_pickle_file):
+    #    print('Detecting faces...', flush=True)
+    #    face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
+    #else:
+    #    with open(video_faces_pickle_file, 'rb') as f:
+    #        face_det_results = pickle.load(f)
+    #    if len(face_det_results) < len(full_frames):
+    #        print("Face detect results have {} frames, but we need {}".format(len(face_det_results), len(full_frames)))
+    #        print('Detecting faces again...', flush=True)
+    #        face_det_results = detect_faces_in_frames(full_frames, video_faces_pickle_file)
+    #    else:
+    #        print("We don't need to find faces! Woo!", flush=True)
 
+
+    # Load face detection results from pickle file
+    video_faces_pickle_file = args.cached_faces_filename
+    with open(video_faces_pickle_file, 'rb') as f:
+        face_det_results = pickle.load(f)
+
+    #face_det_results = pickle.load(args.cached_faces_filename)
     face_det_results = face_det_results[:len(mel_chunks)]
 
     batch_size = args.wav2lip_batch_size

models/Wav2Lip/vocodes_process_upload.py

@@ -243,26 +243,26 @@ mel_step_size = 16
 device = 'cuda' if torch.cuda.is_available() else 'cpu'
 print('Using {} for inference.'.format(device))
 
-def _load(checkpoint_path):
-    if device == 'cuda':
-        checkpoint = torch.load(checkpoint_path)
-    else:
-        checkpoint = torch.load(checkpoint_path,
-                                map_location=lambda storage, loc: storage)
-    return checkpoint
+#def _load(checkpoint_path):
+#    if device == 'cuda':
+#        checkpoint = torch.load(checkpoint_path)
+#    else:
+#        checkpoint = torch.load(checkpoint_path,
+#                                map_location=lambda storage, loc: storage)
+#    return checkpoint
 
-def load_model(path):
-    model = Wav2Lip()
-    print("Load checkpoint from: {}".format(path), flush=True)
-    checkpoint = _load(path)
-    s = checkpoint["state_dict"]
-    new_s = {}
-    for k, v in s.items():
-        new_s[k.replace('module.', '')] = v
-    model.load_state_dict(new_s)
-
-    model = model.to(device)
-    return model.eval()
+#def load_model(path):
+#    model = Wav2Lip()
+#    print("Load checkpoint from: {}".format(path), flush=True)
+#    checkpoint = _load(path)
+#    s = checkpoint["state_dict"]
+#    new_s = {}
+#    for k, v in s.items():
+#        new_s[k.replace('module.', '')] = v
+#    model.load_state_dict(new_s)
+#
+#    model = model.to(device)
+#    return model.eval()
 
 
 def main(tempdir):

src/buckets/bucket_path_unifier.rs

@@ -41,6 +41,15 @@ impl BucketPathUnifier {
     self.w2l_model_root.join(w2l_model_name)
   }
 
+  // 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 {
+    let hashed_path = Self::hashed_directory_path(template_file_hash);
+    self.user_uploaded_w2l_templates_root
+      .join(hashed_path)
+      .join(template_file_hash)
+  }
+
   // These share the same directory as the uploaded w2l template media.
   // eg. /user_uploaded_w2l_templates/1/5/1/151a[...60]_detected_faces.pickle
   pub fn precomputed_faces_for_w2l_path(&self, template_file_hash: &str) -> PathBuf {
@@ -110,6 +119,13 @@ mod tests {
                "/test_path_w2l_audio/f/o/o/foobar");
   }
 
+  #[test]
+  fn test_media_templates_for_w2l_path() {
+    let paths = get_instance();
+    assert_eq!(paths.media_templates_for_w2l_path("foobar").to_str().unwrap(),
+               "/test_path_w2l_templates/f/o/o/foobar");
+  }
+
   #[test]
   fn test_precomputed_faces_for_w2l_path() {
     let paths = get_instance();

src/util/semi_persistent_cache_dir.rs

@@ -8,6 +8,7 @@ pub struct SemiPersistentCacheDir {
   tts_model_root: PathBuf,
   w2l_model_root: PathBuf,
   w2l_face_templates_root: PathBuf,
+  w2l_templates_media_root: PathBuf,
   video_asset_root: PathBuf, // end bump, etc.
 }
 
@@ -27,6 +28,7 @@ impl SemiPersistentCacheDir {
       tts_model_root: cache_root.join("tts/models/"),
       w2l_model_root: cache_root.join("w2l/models/"),
       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/"),
     }
   }
@@ -55,7 +57,24 @@ impl SemiPersistentCacheDir {
     Ok(())
   }
 
-  // ==================== W2L FACES ====================
+  // ==================== W2L MEDIA ====================
+
+  /// We cache W2L media here.
+  /// We'll likely need to LRU cache them.
+  pub fn w2l_template_media_path(&self, template_private_bucket_hash: &str) -> PathBuf {
+    self.w2l_templates_media_root.join(template_private_bucket_hash)
+  }
+
+  pub fn w2l_template_media_directory(&self) -> &Path {
+    &self.w2l_templates_media_root
+  }
+
+  pub fn create_w2l_template_media_path(&self) -> AnyhowResult<()> {
+    let _ = fs::create_dir_all(self.w2l_template_media_directory())?;
+    Ok(())
+  }
+
+  // ==================== W2L CACHED FACES ====================
 
   /// We cache W2L faces here.
   /// We'll likely need to LRU cache them.

src/w2l_inference_job.rs

@@ -179,6 +179,7 @@ async fn main() -> AnyhowResult<()> {
   info!("Creating pod semi-persistent cache dirs...");
   semi_persistent_cache.create_w2l_model_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(
@@ -327,13 +328,7 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An
     info!("Downloaded model from bucket!");
   }
 
-  // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE FACE ==================== //
-
-  // 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
+  // ==================== LOOK UP TEMPLATE RECORD ==================== //
 
   let template_token = match &job.maybe_w2l_template_token {
     Some(token) => token.to_string(),
@@ -355,19 +350,55 @@ async fn process_job(inferencer: &Inferencer, job: &W2lInferenceJobRecord) -> An
     },
   };
 
+  // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE AUDIO OR VIDEO ==================== //
+
+  // 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 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);
+
+    inferencer.private_bucket_client.download_file_to_disk(
+      &template_media_object_path,
+      &template_media_fs_path
+    ).await?;
+
+    info!("Downloaded template media from bucket!");
+  }
+
+
+  // ==================== CONFIRM OR DOWNLOAD W2L TEMPLATE FACE ==================== //
+
+  // 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 template_object_path = inferencer.bucket_path_unifier
+    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: {:?}", &template_object_path);
+    info!("Download from face template path: {:?}", &face_template_object_path);
 
     inferencer.private_bucket_client.download_file_to_disk(
-      &template_object_path,
+      &face_template_object_path,
       &face_template_fs_path
     ).await?;