diff --git a/models/Wav2Lip/vocodes_process_inference.py b/models/Wav2Lip/vocodes_process_inference.py index 7fd4a2032d..1e0b0f5a85 100644 --- a/models/Wav2Lip/vocodes_process_inference.py +++ b/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 diff --git a/models/Wav2Lip/vocodes_process_upload.py b/models/Wav2Lip/vocodes_process_upload.py index dc4be013f0..5261ffd0e1 100644 --- a/models/Wav2Lip/vocodes_process_upload.py +++ b/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): diff --git a/src/buckets/bucket_path_unifier.rs b/src/buckets/bucket_path_unifier.rs index 8410901a2a..b314715a99 100644 --- a/src/buckets/bucket_path_unifier.rs +++ b/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(); diff --git a/src/util/semi_persistent_cache_dir.rs b/src/util/semi_persistent_cache_dir.rs index a3e226c863..fd03496b03 100644 --- a/src/util/semi_persistent_cache_dir.rs +++ b/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. diff --git a/src/w2l_inference_job.rs b/src/w2l_inference_job.rs index eb57547605..989c9575f6 100644 --- a/src/w2l_inference_job.rs +++ b/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?;