fix: honor cancel signal in InfiniteTalk runner

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: