Update generate_multitalk.py

generate_multitalk.py

@@ -23,7 +23,6 @@ from wan.utils.utils import cache_image, cache_video, str2bool
 from wan.utils.multitalk_utils import save_video_ffmpeg
 
 from transformers import Wav2Vec2FeatureExtractor
-from transformers import Wav2Vec2ForCTC
 from src.audio_analysis.wav2vec2 import Wav2Vec2Model
 
 import librosa
@@ -215,7 +214,7 @@ def _parse_args():
     return args
 
 def custom_init(device, wav2vec):
-    audio_encoder = Wav2Vec2Model.from_pretrained(args.wav2vec_dir, attn_implementation="eager").to(device)
+    audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec, attn_implementation="eager").to(device)
     audio_encoder.freeze_feature_extractor()
     wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec, local_files_only=True)
     return wav2vec_feature_extractor, audio_encoder
@@ -373,50 +372,50 @@ def generate(args):
 
     assert args.task == "multitalk-14B", 'You should choose multitalk in args.task.'
 
-    # Initialize a placeholder for all processes
-    input_data = None
-    
-    # Let only the main process prepare the data
+    # Let only the main process (rank 0) prepare the audio embeddings and overwrite the input JSON file.
     if rank == 0:
         with open(args.input_json, 'r', encoding='utf-8') as f:
             input_data = json.load(f)
 
-            wav2vec_feature_extractor, audio_encoder= custom_init('cpu', args.wav2vec_dir)
-            args.audio_save_dir = os.path.join(args.audio_save_dir, input_data['cond_image'].split('/')[-1].split('.')[0])
-            os.makedirs(args.audio_save_dir,exist_ok=True)
-
-            if len(input_data['cond_audio'])==2:
-                new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
-                audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
-                audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
-                emb1_path = os.path.join(args.audio_save_dir, '1.pt')
-                emb2_path = os.path.join(args.audio_save_dir, '2.pt')
-                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
-                sf.write(sum_audio, sum_human_speechs, 16000)
-                torch.save(audio_embedding_1, emb1_path)
-                torch.save(audio_embedding_2, emb2_path)
-                input_data['cond_audio']['person1'] = emb1_path
-                input_data['cond_audio']['person2'] = emb2_path
-                input_data['video_audio'] = sum_audio
-            elif len(input_data['cond_audio'])==1:
-                human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
-                audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder)
-                emb_path = os.path.join(args.audio_save_dir, '1.pt')
-                sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
-                sf.write(sum_audio, human_speech, 16000)
-                torch.save(audio_embedding, emb_path)
-                input_data['cond_audio']['person1'] = emb_path
-                input_data['video_audio'] = sum_audio
-
-    # Broadcast the data from rank 0 to all other processes
+        wav2vec_feature_extractor, audio_encoder = custom_init('cpu', args.wav2vec_dir)
+        args.audio_save_dir = os.path.join(args.audio_save_dir, input_data['cond_image'].split('/')[-1].split('.')[0])
+        os.makedirs(args.audio_save_dir, exist_ok=True)
+
+        if len(input_data['cond_audio']) == 2:
+            new_human_speech1, new_human_speech2, sum_human_speechs = audio_prepare_multi(input_data['cond_audio']['person1'], input_data['cond_audio']['person2'], input_data['audio_type'])
+            audio_embedding_1 = get_embedding(new_human_speech1, wav2vec_feature_extractor, audio_encoder)
+            audio_embedding_2 = get_embedding(new_human_speech2, wav2vec_feature_extractor, audio_encoder)
+            emb1_path = os.path.join(args.audio_save_dir, '1.pt')
+            emb2_path = os.path.join(args.audio_save_dir, '2.pt')
+            sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+            sf.write(sum_audio, sum_human_speechs, 16000)
+            torch.save(audio_embedding_1, emb1_path)
+            torch.save(audio_embedding_2, emb2_path)
+            input_data['cond_audio']['person1'] = emb1_path
+            input_data['cond_audio']['person2'] = emb2_path
+            input_data['video_audio'] = sum_audio
+        elif len(input_data['cond_audio']) == 1:
+            human_speech = audio_prepare_single(input_data['cond_audio']['person1'])
+            audio_embedding = get_embedding(human_speech, wav2vec_feature_extractor, audio_encoder)
+            emb_path = os.path.join(args.audio_save_dir, '1.pt')
+            sum_audio = os.path.join(args.audio_save_dir, 'sum.wav')
+            sf.write(sum_audio, human_speech, 16000)
+            torch.save(audio_embedding, emb_path)
+            input_data['cond_audio']['person1'] = emb_path
+            input_data['video_audio'] = sum_audio
+
+        # Overwrite the temporary JSON file with the updated paths
+        with open(args.input_json, 'w', encoding='utf-8') as f:
+            json.dump(input_data, f, indent=4)
+
+    # Barrier to ensure rank 0 has finished writing all files (embeddings AND the json)
     if dist.is_initialized():
-        objects_to_broadcast = [input_data] if rank == 0 else [None]
-        dist.broadcast_object_list(objects_to_broadcast, src=0)
-        input_data = objects_to_broadcast[0]
-
-        # Wait for all file I/O to be complete before proceeding
         dist.barrier()
 
+    # Now, ALL processes read the (now updated) JSON file, ensuring data consistency.
+    with open(args.input_json, 'r', encoding='utf-8') as f:
+        input_data = json.load(f)
+
     logging.info("Creating MultiTalk pipeline.")
     wan_i2v = wan.MultiTalkPipeline(
         config=cfg,
@@ -461,7 +460,7 @@ def generate(args):
             args.save_file = f"{args.task}_{args.size.replace('*','x') if sys.platform=='win32' else args.size}_{args.ulysses_size}_{args.ring_size}_{formatted_prompt}_{formatted_time}"
 
         logging.info(f"Saving generated video to {args.save_file}.mp4")
-        save_video_ffmpeg(video, args.save_file, [input_data['video_audio']])
+        save__video_ffmpeg(video, args.save_file, [input_data['video_audio']])
 
     logging.info("Finished.")