diff --git a/lightx2v/models/runners/wan/wan_infinitetalk_runner.py b/lightx2v/models/runners/wan/wan_infinitetalk_runner.py index 857d7b7a7..8a1894afa 100644 --- a/lightx2v/models/runners/wan/wan_infinitetalk_runner.py +++ b/lightx2v/models/runners/wan/wan_infinitetalk_runner.py @@ -120,6 +120,7 @@ def __init__(self, config): self.cond_frame_cache = {} self.cond_video_fps = None self.cond_video_frame_count = None + self.stream_saved_video_needs_audio_remux = False def init_scheduler(self): self.scheduler = InfiniteTalkScheduler(self.config) @@ -738,6 +739,7 @@ def _build_clip_context(self, cond_image): def _run_dit_clip(self, dit_inputs): infer_steps = self.scheduler.infer_steps for step_index in range(infer_steps): + self.check_stop() logger.info(f"==> step_index: {step_index + 1} / {infer_steps}") with ProfilingContext4DebugL1("step_pre"): self.scheduler.step_pre(step_index) @@ -957,8 +959,10 @@ def run_main(self): for segment_idx in range(self.video_segment_num): logger.info(f"start InfiniteTalk segment {segment_idx + 1}/{self.video_segment_num}") with ProfilingContext4DebugL1(f"segment end2end {segment_idx + 1}/{self.video_segment_num}"): + self.check_stop() self.init_run_segment(segment_idx) latents = self.run_segment(segment_idx) + self.check_stop() self.end_run_segment(segment_idx, latents) if self.stream_save_video: @@ -971,6 +975,7 @@ def run_main(self): def process_images_after_vae_decoder(self): if self.stream_save_video: if self.input_info.save_result_path is not None and is_main_process(): + self.stream_saved_video_needs_audio_remux = True logger.info(f"Video saved to {self.input_info.save_result_path}") return {"video": None} @@ -1086,6 +1091,19 @@ def end_run(self): if self.va_controller is not None: self.va_controller.clear() self.va_controller = None + if self.stream_saved_video_needs_audio_remux: + out_path = self.input_info.save_result_path + mux_audio = self._resolve_mux_audio_path() + try: + if not mux_audio or not os.path.isfile(mux_audio): + audio_input = getattr(self.input_info, "audio_path", None) or self.config.get("audio_path", "") + raise FileNotFoundError(f"InfiniteTalk mux audio is unavailable for audio input: {audio_input}") + if not os.path.isfile(out_path): + raise FileNotFoundError(f"InfiniteTalk stream video is unavailable for audio mux: {out_path}") + logger.info(f"Muxing InfiniteTalk stream audio {mux_audio} into {out_path}") + self._mux_audio(out_path, mux_audio) + finally: + self.stream_saved_video_needs_audio_remux = False self._remove_video_audio_path() self._clear_cond_frame_source() self._remove_cond_video_temp_path() @@ -1106,6 +1124,7 @@ def run_pipeline(self, input_info): if GET_RECORDER_MODE(): monitor_cli.lightx2v_worker_request_count.inc() self.input_info = input_info + self.stream_saved_video_needs_audio_remux = False try: self._prepare_input_data() if self.reuse: