Update generate_multitalk.py

generate_multitalk.py

@@ -207,14 +207,14 @@ def _parse_args():
         help="Norm threshold used in adaptive projected guidance (APG)."
     )
 
-    
+
     args = parser.parse_args()
 
     _validate_args(args)
 
     return args
 
-def custom_init(device, wav2vec):    
+def custom_init(device, wav2vec):
     audio_encoder = Wav2Vec2Model.from_pretrained(args.wav2vec_dir, attn_implementation="eager").to(device)
     audio_encoder.freeze_feature_extractor()
     wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec, local_files_only=True)
@@ -244,7 +244,7 @@ def audio_prepare_multi(left_path, right_path, audio_type, sample_rate=16000):
         new_human_speech1 = human_speech_array1
         new_human_speech2 = human_speech_array2
     elif audio_type=='add':
-        new_human_speech1 = np.concatenate([human_speech_array1[: human_speech_array1.shape[0]], np.zeros(human_speech_array2.shape[0])]) 
+        new_human_speech1 = np.concatenate([human_speech_array1[: human_speech_array1.shape[0]], np.zeros(human_speech_array2.shape[0])])
         new_human_speech2 = np.concatenate([np.zeros(human_speech_array1.shape[0]), human_speech_array2[:human_speech_array2.shape[0]]])
     sum_human_speechs = new_human_speech1 + new_human_speech2
     return new_human_speech1, new_human_speech2, sum_human_speechs
@@ -359,21 +359,6 @@ def generate(args):
             ulysses_degree=args.ulysses_size,
         )
 
-    # TODO: use prompt refine
-    # if args.use_prompt_extend:
-    #     if args.prompt_extend_method == "dashscope":
-    #         prompt_expander = DashScopePromptExpander(
-    #             model_name=args.prompt_extend_model,
-    #             is_vl="i2v" in args.task or "flf2v" in args.task)
-    #     elif args.prompt_extend_method == "local_qwen":
-    #         prompt_expander = QwenPromptExpander(
-    #             model_name=args.prompt_extend_model,
-    #             is_vl="i2v" in args.task,
-    #             device=rank)
-    #     else:
-    #         raise NotImplementedError(
-    #             f"Unsupport prompt_extend_method: {args.prompt_extend_method}")
-
     cfg = WAN_CONFIGS[args.task]
     if args.ulysses_size > 1:
         assert cfg.num_heads % args.ulysses_size == 0, f"`{cfg.num_heads=}` cannot be divided evenly by `{args.ulysses_size=}`."
@@ -387,66 +372,50 @@ def generate(args):
         args.base_seed = base_seed[0]
 
     assert args.task == "multitalk-14B", 'You should choose multitalk in args.task.'
-    
-
-    # TODO: add prompt refine
-    # img = Image.open(args.image).convert("RGB")
-    # if args.use_prompt_extend:
-    #     logging.info("Extending prompt ...")
-    #     if rank == 0:
-    #         prompt_output = prompt_expander(
-    #             args.prompt,
-    #             tar_lang=args.prompt_extend_target_lang,
-    #             image=img,
-    #             seed=args.base_seed)
-    #         if prompt_output.status == False:
-    #             logging.info(
-    #                 f"Extending prompt failed: {prompt_output.message}")
-    #             logging.info("Falling back to original prompt.")
-    #             input_prompt = args.prompt
-    #         else:
-    #             input_prompt = prompt_output.prompt
-    #         input_prompt = [input_prompt]
-    #     else:
-    #         input_prompt = [None]
-    #     if dist.is_initialized():
-    #         dist.broadcast_object_list(input_prompt, src=0)
-    #     args.prompt = input_prompt[0]
-    #     logging.info(f"Extended prompt: {args.prompt}")
-
-    # read input files
 
+    # Initialize a placeholder for all processes
+    input_data = None
     
+    # Let only the main process prepare the data
+    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
+    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]
 
-    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
+        # Wait for all file I/O to be complete before proceeding
+        dist.barrier()
 
     logging.info("Creating MultiTalk pipeline.")
     wan_i2v = wan.MultiTalkPipeline(
@@ -455,8 +424,8 @@ def generate(args):
         device_id=device,
         rank=rank,
         t5_fsdp=args.t5_fsdp,
-        dit_fsdp=args.dit_fsdp, 
-        use_usp=(args.ulysses_size > 1 or args.ring_size > 1),  
+        dit_fsdp=args.dit_fsdp,
+        use_usp=(args.ulysses_size > 1 or args.ring_size > 1),
         t5_cpu=args.t5_cpu
     )
 
@@ -465,7 +434,7 @@ def generate(args):
         wan_i2v.enable_vram_management(
             num_persistent_param_in_dit=args.num_persistent_param_in_dit
         )
-    
+
     logging.info("Generating video ...")
     video = wan_i2v.generate(
         input_data,
@@ -481,19 +450,19 @@ def generate(args):
         max_frames_num=args.frame_num if args.mode == 'clip' else 1000,
         extra_args=args,
         )
-    
+
 
     if rank == 0:
-        
+
         if args.save_file is None:
-            formatted_time = datetime.now().strftime("%Y%m%d_%H%M%S")
+            formatted_time = datetime..now().strftime("%Y%m%d_%H%M%S")
             formatted_prompt = input_data['prompt'].replace(" ", "_").replace("/",
                                                                         "_")[:50]
             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']])
-        
+
     logging.info("Finished.")
 
     if torch.cuda.is_available():
@@ -506,4 +475,4 @@ def generate(args):
 
 if __name__ == "__main__":
     args = _parse_args()
-    generate(args)
+    generate(args)