From 4519bd8fb783fddc8992462ba6b71c1048d829bc Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Wed, 8 Jul 2026 15:06:20 +0800
Subject: [PATCH 1/9] Fix pause set input_info=None error

---
 lightx2v/models/runners/base_runner.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/lightx2v/models/runners/base_runner.py b/lightx2v/models/runners/base_runner.py
index 2a427d004..0825b6ce0 100755
--- a/lightx2v/models/runners/base_runner.py
+++ b/lightx2v/models/runners/base_runner.py
@@ -337,8 +337,4 @@ def check_stop(self):
                 print(f"end_run failed: {e}")
             raise Exception(f"find rank: {rank} stop_signal, stop running, it's an expected behavior")
         if paused == 1:
-            try:
-                self.end_run()
-            except Exception as e:
-                print(f"end_run failed: {e}")
             raise Exception(f"find rank: {rank} pause_signal, pause running, it's an expected behavior")

From e18d733970efe9aec0057d6a94fb1a36c9425f41 Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Wed, 8 Jul 2026 17:03:49 +0800
Subject: [PATCH 2/9] stream audio pipeline ok

---
 .../models/runners/wan/wan_audio_runner.py    | 113 +++++++-
 lightx2v/utils/va_controller.py               |  26 +-
 lightx2v/utils/va_reader_omni.py              | 261 ++++++++++++++++++
 3 files changed, 396 insertions(+), 4 deletions(-)

diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index ac7067a3c..d54ef046a 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -1275,8 +1275,10 @@ def _encode_audio_for_ar(self):
         self.inputs["audio_encoder_output"] = torch.stack(features_list, dim=0)
         self.inputs["previmg_encoder_output"] = {"prev_latents": None, "prev_mask": None, "prev_len": 0}
 
-    def init_run_segment(self, segment_idx):
+    def init_run_segment(self, segment_idx, origin_audio=None, latent_audio=None):
         self.segment_idx = segment_idx
+        if origin_audio is not None:
+            self._encode_audio_for_each_chunk(segment_idx, origin_audio, latent_audio)
 
     @ProfilingContext4DebugL1(
         "Prefill reference kv",
@@ -1340,9 +1342,8 @@ def init_kv_cache_manager(self):
                 torch.distributed.barrier(group=sp_group)
 
     @ProfilingContext4DebugL2("Run DiT")
-    def run_main(self):
+    def run_main_with_fixed_audio(self):
         try:
-            self.init_run()
             if self.config.get("compile", False) and hasattr(self.model, "comple"):
                 self.model.select_graph_for_compile(self.input_info)
             logger.info("start ar audio generation")
@@ -1388,6 +1389,7 @@ def run_main(self):
                             self.check_stop()
                             self.init_run_segment(segment_idx)
                             segment_latents = self.run_segment(segment_idx)
+                            logger.debug(f"segment_latents shape: {segment_latents.shape}")
 
                         vae_decoder.submit(self.decode_segment_latents, segment_idx, segment_latents)
                     segment_videos = vae_decoder.finish()
@@ -1409,3 +1411,108 @@ def run_main(self):
             if getattr(self, "va_controller", None) is not None:
                 self.va_controller.clear()
                 self.va_controller = None
+
+    # origin_audio: (4 * latent_per_chunk / fps * sr), eg. (4 * 3 / 16 * 16000)
+    # latent_audio: (latent_per_chunk, audio_window * sr) eg. (3, 1.0 * 16000)
+    def _encode_audio_for_each_chunk(self, segment_idx: int, origin_audio: np.ndarray, latent_audio: np.ndarray) -> torch.Tensor:
+        self.input_info.seed = self.input_info.seed + segment_idx
+        torch.manual_seed(self.input_info.seed)
+
+        # used for remux output video
+        end_idx = origin_audio.shape[0] // self._audio_processor.audio_frame_rate
+        # 第一段vae解码出来视频只有 (4 - 1) * latent_per_chunk 帧
+        if segment_idx == 0:
+            end_idx -= 1 * self.config["ar_config"]["num_frame_per_chunk"] * self._audio_processor.audio_frame_rate
+        origin_audio_tensor = torch.Tensor(origin_audio).float().unsqueeze(0)
+        self.segment = AudioSegment(origin_audio_tensor, 0, end_idx)
+
+        latent_audio_tensor = torch.Tensor(latent_audio).float().to(AI_DEVICE)
+        latent_length = latent_audio.shape[0]
+        audio_slice_list = self.audio_encoder.audio_feature_extractor(latent_audio_tensor, sampling_rate=self.config.get("audio_sr", 16000), return_tensors="pt").input_values.squeeze(0)
+
+        audio_slice_list = audio_slice_list.to(device=AI_DEVICE, dtype=GET_DTYPE())
+        audio_feat = self.audio_encoder.audio_feature_encoder(audio_slice_list, return_dict=True).last_hidden_state
+
+        audio_feat_per_latent = self.audio_sliding_processor.audio_feat_per_latent
+        left = audio_feat.shape[1] // 2 - audio_feat_per_latent // 2
+        assert left >= 0
+        audio_feat = audio_feat[:, left : left + audio_feat_per_latent]
+        audio_feat = audio_feat.reshape(1, latent_length, audio_feat_per_latent, -1).contiguous()
+        audio_feat = self.audio_adapter.forward_audio_proj(audio_feat, audio_feat.shape[1])
+
+        self.inputs["audio_encoder_output"] = audio_feat
+        self.inputs["previmg_encoder_output"] = {"prev_latents": None, "prev_mask": None, "prev_len": 0}
+
+    def run_main(self):
+        try:
+            self.va_controller = None
+            self.init_run()
+            logger.info(f"init va_recorder: {self.va_controller.recorder} and va_reader: {self.va_controller.reader}")
+
+            # fixed audio segments inputs
+            if self.va_controller.reader is None:
+                return self.run_main_with_fixed_audio()
+
+            self.va_controller.start()
+            if self.config.get("compile", False) and hasattr(self.model, "comple"):
+                self.model.select_graph_for_compile(self.input_info)
+            logger.info("start ar audio generation")
+
+            # steam audio input, video segment num is unlimited
+            self.video_segment_num = 1000000
+            segment_idx = 0
+            fail_count, max_fail_count = 0, 10
+            self.va_controller.before_control()
+
+            latent_shape = list(self.input_info.latent_shape)
+            latent_shape[1] = self.config["ar_config"]["num_frame_per_chunk"]
+            self.input_info.latent_shape = latent_shape
+            self.model.scheduler.prepare(seed=self.input_info.seed, latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
+            self._validate_prompt_travel_schedule()
+            self.init_kv_cache_manager()
+            self.prefill_reference_kv()
+
+            while True:
+                with ProfilingContext4DebugL1(f"stream segment get audio segment {segment_idx}"):
+                    control = self.va_controller.next_control()
+                    if control.action == "blank_to_voice":
+                        self.prev_video = control.data
+                    elif control.action == "wait":
+                        time.sleep(0.01)
+                        continue
+
+                    origin_audio, latent_audio, valid_duration = self.va_controller.reader.get_audio_segment()
+                    if origin_audio is None or latent_audio is None:
+                        fail_count += 1
+                        logger.warning(f"Failed to get audio chunk {fail_count} times")
+                        if fail_count > max_fail_count:
+                            raise Exception(f"Failed to get audio chunk {fail_count} times, stop reader")
+                        continue
+
+                with ProfilingContext4DebugL1(f"stream segment end2end {segment_idx}"):
+                    try:
+                        # reset pause signal
+                        self.pause_signal = False
+                        self.can_pause = valid_duration <= 1e-5
+                        self.init_run_segment(segment_idx, origin_audio, latent_audio)
+                        self.model.scheduler.prepare(seed=self.input_info.seed, latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
+                        self.check_stop()
+                        latents = self.run_segment(0)
+                        self.check_stop()
+                        self.gen_video = self.decode_segment_latents(segment_idx, latents)
+                        logger.debug(f"gen_video shape: {self.gen_video.shape}")
+                        self.check_stop()
+                        self.end_run_segment(segment_idx, valid_duration=valid_duration)
+                        segment_idx += 1
+                        fail_count = 0
+                    except Exception as e:
+                        if "pause_signal, pause running" in str(e):
+                            logger.warning(f"model infer audio pause: {e}, should continue")
+                        else:
+                            raise
+        finally:
+            if hasattr(self.model, "inputs"):
+                self.end_run()
+            if self.va_controller is not None:
+                self.va_controller.clear()
+                self.va_controller = None
diff --git a/lightx2v/utils/va_controller.py b/lightx2v/utils/va_controller.py
index c27346d85..7ff74b3f5 100644
--- a/lightx2v/utils/va_controller.py
+++ b/lightx2v/utils/va_controller.py
@@ -81,6 +81,13 @@ def init_base(self, config, input_info, has_vfi_model, has_vsr_model):
         max_end_idx = max(self.est_infer_end_idx, self.est_switch_image_end_idx, self.est_switch_action_end_idx)
         self.min_stay_queue_num = max_end_idx * 2 + 1
 
+        # for seko_talk_ar, the audio_sr is the same as the sample_rate
+        self.is_seko_talk_ar = config.get("model_cls") == "seko_talk_ar"
+        if self.is_seko_talk_ar:
+            self.latent_per_chunk = config["ar_config"]["num_frame_per_chunk"]
+            self.audio_window_secs = config.get("audio_window", 1.0)
+            self.look_ahead_secs = config.get("look_ahead", 0.0)
+
     def init_recorder(self):
         if not self.output_video_path or self.rank != self.target_recorder_rank:
             return
@@ -117,7 +124,24 @@ def init_reader(self, model_runner=None):
         prev_duration = self.prev_frame_length / self.target_fps
         omni_work_dir = os.getenv("OMNI_WORK_DIR", None)
         if omni_work_dir:
-            from lightx2v.utils.va_reader_omni import OmniVAReader
+            from lightx2v.utils.va_reader_omni import OmniVAReader, SekoAROmniVAReader
+
+            if self.is_seko_talk_ar:
+                self.reader = SekoAROmniVAReader(
+                    rank=self.rank,
+                    world_size=self.world_size,
+                    stream_url=self.audio_path["data"],
+                    sample_rate=self.audio_sr,
+                    latent_per_chunk=self.latent_per_chunk,
+                    audio_window_secs=self.audio_window_secs,
+                    look_ahead_secs=self.look_ahead_secs,
+                    video_fps=self.target_fps,
+                    target_rank=self.target_reader_rank,
+                    model_runner=model_runner,
+                    huoshan_tts_voice_type=self.audio_path.get("huoshan_tts_voice_type", None),
+                    stream_config=self.stream_config,
+                )
+                return
 
             self.reader = OmniVAReader(
                 rank=self.rank,
diff --git a/lightx2v/utils/va_reader_omni.py b/lightx2v/utils/va_reader_omni.py
index f89f2a557..0d6efa78f 100644
--- a/lightx2v/utils/va_reader_omni.py
+++ b/lightx2v/utils/va_reader_omni.py
@@ -559,6 +559,267 @@ def __del__(self):
         self.stop()
 
 
+class SekoAROmniVAReader(OmniVAReader):
+    def __init__(
+        self,
+        rank: int,
+        world_size: int,
+        stream_url: str,
+        latent_per_chunk: int = 3,
+        audio_window_secs: float = 1.0,
+        look_ahead_secs: float = 0.0,
+        video_fps: float = 16,
+        sample_rate: int = 16000,
+        audio_channels: int = 1,
+        target_rank: int = 0,
+        model_runner=None,
+        huoshan_tts_voice_type=None,
+        stream_config: dict = {},
+        **kwargs,
+    ):
+        self.rank = rank
+        self.world_size = world_size
+        self.stream_url = stream_url
+        self.sample_rate = sample_rate
+
+        self.audio_channels = audio_channels
+        self.latent_per_chunk = latent_per_chunk
+        self.audio_window_secs = audio_window_secs
+        self.look_ahead_secs = look_ahead_secs
+        self.video_fps = video_fps
+
+        self.target_rank = target_rank % self.world_size
+        self.origin_flag_tensor = torch.tensor([0], dtype=torch.int32).to(device="cuda")
+        self.latent_flag_tensor = torch.tensor([0], dtype=torch.int32).to(device="cuda")
+        self.valid_duration_tensor = torch.tensor([0], dtype=torch.float32).to(device="cuda")
+        self.immediate_switch_tensor = torch.tensor([0], dtype=torch.int32).to(device="cuda")
+
+        self.chat_adapter = None
+        self.model_runner = model_runner
+        self.huoshan_tts_voice_type = huoshan_tts_voice_type
+        self.stream_config = stream_config
+
+        assert self.audio_channels == 1, "Only mono audio is supported for OmniVAReader"
+        self.initAR()
+        logger.info(f"VAReader initialized for stream: {stream_url} target_rank: {self.target_rank}")
+
+    def initAR(self):
+        # ref: https://github.com/ModelTC/LightX2V/blob/main/lightx2v/models/input_encoders/hf/seko_audio/audio_adapter.py#L267-L280
+        # 假定 frame_per_latent = 4（每个latent包含4个frame）
+        # 左边界帧索引：[0, 1, 5,  9, 13, 17, ...]
+        # 右边界帧索引：[1, 5, 9, 13, 17, 21, ...]
+        # 中心时间戳（秒）：[(0+1)/2/fps, (1+5)/2/fps, (5+9)/2/fps, ...]，除第一个外，其余为(latent_idx*4-1)/fps
+        # 上边界时间戳（秒）:[1/fps, 5/fps, 9/fps, ...]+look_ahead_secs，即(latent_idx*4+1)/fps + look_ahead_secs
+        # 每一个 latent 对应的音频时间范围（假定每个chunk包含3个latent，即每3个latent取一次音频）：
+        #    [0.5/fps - audio_window_secs/2, 0.5/fps + audio_window_secs/2] 上界应小于 1/fps + look_ahead_secs
+        #    [3/fps - audio_window_secs/2, 3/fps + audio_window_secs/2]     上界应小于 5/fps + look_ahead_secs
+        # 取 [7/fps - audio_window_secs/2, 7/fps + audio_window_secs/2]     上界应小于 9/fps + look_ahead_secs
+        #    [11/fps - audio_window_secs/2, 11/fps + audio_window_secs/2]   上界应小于 13/fps + look_ahead_secs
+        #    [15/fps - audio_window_secs/2, 15/fps + audio_window_secs/2]   上界应小于 17/fps + look_ahead_secs
+        # 取 [19/fps - audio_window_secs/2, 19/fps + audio_window_secs/2]   上界应小于 21/fps + look_ahead_secs
+        #    [23/fps - audio_window_secs/2, 23/fps + audio_window_secs/2]   上界应小于 25/fps + look_ahead_secs
+        #    ...
+        # 每次取音频应当取 chunk 对应所有 latents 的音频块范围上界的最大值，每次按chunk取的音频段为：
+        #    [0, 9/fps + look_ahead_secs]
+        #    [9/fps + look_ahead_secs, 21/fps + look_ahead_secs]
+        #    [21/fps + look_ahead_secs, 33/fps + look_ahead_secs]
+        #    ...
+        # 需要根据这个恢复每个chunk内各个latent对应的音频段，用于audio encode 计算
+        # 这些音频对应的输出音频段应当按照中心坐标对齐：
+        #    [-3/fps, 9/fps],
+        #    [9/fps, 21/fps],
+        #    [21/fps, 33/fps],
+        #    ...
+        #  [后面-12/fps, (latent_per_chunk - 1 ) * 4 + 1  + chunk_idx * letent_per_chunk * 4]
+
+        # 第一次取 9/fps + look_ahead_secs 秒的音频
+        self.segment_duration = ((self.latent_per_chunk - 1) * 4 + 1) / self.video_fps + self.look_ahead_secs
+        # 后续每次取 12/fps 秒的音频
+        self.other_fetch_duration = self.latent_per_chunk * 4 / self.video_fps
+        self.chunk_idx = 0
+        self.last_chunk_audios = None
+
+        # 当前chunk的原音频，每个chunk合成的视频长度固定，对应音频长度为: 12/fps
+        self.origin_audio_tensor = torch.zeros(int(self.sample_rate * self.other_fetch_duration) * 2, dtype=torch.uint8, device="cuda")
+        # 当前chunk内各个latent对应的音频段，用于audio encode 计算
+        self.latent_audio_tensor = torch.zeros(self.latent_per_chunk * int(self.sample_rate * self.audio_window_secs) * 2, dtype=torch.uint8, device="cuda")
+
+    # 获取当前latent对应的边界索引帧序号
+    def get_latent_idxs(self, i, base_idx):
+        latent_idx = self.chunk_idx * self.latent_per_chunk + i
+        if latent_idx == 0:
+            center_secs = 0.5 / self.video_fps
+        else:
+            center_secs = (latent_idx * 4 - 1) / self.video_fps
+        # 当前 latent 音频窗口上界索引
+        right_idx = int((center_secs + self.audio_window_secs / 2) * self.sample_rate)
+        # 当前 latent 音频窗口下界索引
+        left_idx = right_idx - int(self.sample_rate * self.audio_window_secs)
+        # 当前 latent 音频段上界索引
+        look_ahead_secs = (latent_idx * 4 + 1) / self.video_fps + self.look_ahead_secs
+        look_ahead_idx = int(look_ahead_secs * self.sample_rate)
+        logger.debug(f"latent_idx: {latent_idx} center_secs: {center_secs} right_idx: {right_idx} left_idx: {left_idx} look_ahead_secs: {look_ahead_secs} look_ahead_idx: {look_ahead_idx}")
+
+        # 对齐当前缓存的所有音频的起始位置（只缓存了最多两个chunk的音频）
+        left_idx -= base_idx
+        right_idx -= base_idx
+        look_ahead_idx -= base_idx
+        logger.debug(f"adjusted idxs: {left_idx} {right_idx} {look_ahead_idx}")
+        return left_idx, right_idx, look_ahead_idx
+
+    # 获取当前chunk对应的原始音频的开始和结束索引帧
+    def get_chunk_origin_idxs(self, base_idx):
+        # end secs: 9/fps + chunk_idx * 12/fps
+        end_secs = ((self.latent_per_chunk - 1) * 4 + 1 + self.chunk_idx * self.latent_per_chunk * 4) / self.video_fps
+        start_secs = end_secs - (self.latent_per_chunk * 4) / self.video_fps
+        start_idx = int(start_secs * self.sample_rate)
+        end_idx = int(end_secs * self.sample_rate)
+        logger.debug(f"chunk_idx: {self.chunk_idx} start_idx: {start_idx} end_idx: {end_idx}")
+        # 对齐当前缓存的所有音频的起始位置（只缓存了最多两个chunk的音频）
+        start_idx -= base_idx
+        end_idx -= base_idx
+        logger.debug(f"adjusted idxs: {start_idx} {end_idx}")
+        return start_idx, end_idx
+
+    def prepare_ar_audios(self, merged_audio, origin_audios, latent_audios):
+        # 当前chunk音频的上一个chunk音频的开始时间，用于对齐音频初始位置
+        last_chunk_start_secs = 0
+        if self.chunk_idx > 1:
+            last_chunk_start_secs = self.segment_duration + (self.chunk_idx - 2) * self.other_fetch_duration
+        base_idx = int(last_chunk_start_secs * self.sample_rate)
+        audio_length = len(merged_audio)
+        logger.debug(
+            "chunk idx: ", self.chunk_idx, "start_secs: ", last_chunk_start_secs, "end_secs: ", last_chunk_start_secs + self.segment_duration, "base_idx: ", base_idx, "audio_length: ", audio_length
+        )
+
+        # 构造原始音频，用于合成最终视频
+        start_idx, end_idx = self.get_chunk_origin_idxs(base_idx)
+        pad_idx = max(-start_idx, 0)
+        real_start = max(start_idx, 0)
+        real_len = min(end_idx, audio_length) - real_start
+        logger.debug(f"----- center audios pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len} -----")
+        origin_audios[pad_idx : pad_idx + real_len] = merged_audio[real_start : real_start + real_len]
+
+        # 构造 latent 音频，用于 audio encode
+        for i in range(self.latent_per_chunk):
+            left_idx, right_idx, look_ahead_idx = self.get_latent_idxs(i, base_idx)
+            pad_idx = max(-left_idx, 0)
+            real_start = max(left_idx, 0)
+            real_len = min(look_ahead_idx, right_idx, audio_length) - real_start
+            logger.debug(f"----- {i} pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len} -----")
+            latent_audios[i, pad_idx : pad_idx + real_len] = merged_audio[real_start : real_start + real_len]
+
+    def prepare_audio_data(self, chat_audio_result):
+        sample_count = 0
+        audio = np.array([], dtype=np.int16)
+        origin_audios = np.zeros(int(self.sample_rate * self.other_fetch_duration), dtype=np.int16)
+        latent_audios = np.zeros((self.latent_per_chunk, int(self.sample_rate * self.audio_window_secs)), dtype=np.int16)
+
+        # convert chat audio result to mono and target sample rate
+        if chat_audio_result is not None:
+            audio_data, audio_info = chat_audio_result
+            audio, sample_count = self.convert_pcm_s16le_to_mono_resampled(audio_data, audio_info)
+        valid_duration = sample_count / self.sample_rate
+        if sample_count <= 0:
+            return origin_audios.tobytes(), latent_audios.tobytes(), valid_duration
+
+        if self.chunk_idx == 0:
+            expect_count = int(self.segment_duration * self.sample_rate)
+        else:
+            expect_count = int(self.other_fetch_duration * self.sample_rate)
+        assert sample_count <= expect_count, f"audio length {sample_count} > expect_count {expect_count}"
+
+        # pad 0 to the audio to make it the same length as expect_count
+        if sample_count < expect_count:
+            pad_count = expect_count - sample_count
+            logger.debug(f"pad {pad_count} samples to audio")
+            audio = np.pad(audio, (0, pad_count), mode="constant", constant_values=0)
+
+        # if is not the first segment, concat with previous segment
+        if self.last_chunk_audios is not None:
+            logger.debug(f"concat {self.last_chunk_audios.shape} with {audio.shape}")
+            merged_audio = np.concatenate([self.last_chunk_audios, audio])
+        else:
+            merged_audio = audio
+        self.prepare_ar_audios(merged_audio, origin_audios, latent_audios)
+        logger.debug(
+            f"chunk[{self.chunk_idx}] origin_audios: {origin_audios.shape} {origin_audios.dtype} {origin_audios.min()} {origin_audios.max()} latent_audios: {latent_audios.shape} {latent_audios.dtype} {latent_audios.min()} {latent_audios.max()} {sample_count}, valid_duration: {valid_duration}"
+        )
+
+        # update prev seg chunk
+        self.last_chunk_audios = audio
+        self.chunk_idx += 1
+        return origin_audios.tobytes(), latent_audios.tobytes(), valid_duration
+
+    def get_fetch_duration(self):
+        fetch_duration = self.segment_duration
+        # after immediate switch, reset prev seg chunk
+        if self.chat_adapter is not None and self.chat_adapter.reset_prev:
+            self.chat_adapter.reset_prev = False
+            self.chunk_idx = 0
+            self.last_chunk_audios = None
+            logger.warning(f"Reset prev seg chunk")
+        # first segment, fetch segment_duration, else fetch self.other_fetch_duration
+        if self.chunk_idx > 0 and self.last_chunk_audios is not None:
+            fetch_duration = self.other_fetch_duration
+        return fetch_duration
+
+    def get_audio_segment(self):
+        origin_audio = None
+        latent_audio = None
+        valid_duration = 0
+        try:
+            if self.rank == self.target_rank:
+                fetch_duration = self.get_fetch_duration()
+                # logger.info(f"Get segment, fetch_duration: {fetch_duration}")
+                if self.chat_adapter.status == "voice":
+                    audio_result = self.chat_adapter.get_audio(fetch_duration)
+                    origin_audio, latent_audio, valid_duration = self.prepare_audio_data(audio_result)
+                    # think all voice segments inferred, naturally switch to blank
+                    if audio_result is None:
+                        logger.info(f"Think all voice segments inferred, naturally switch to blank")
+                        self.chat_adapter.status = "blank"
+                        self.chunk_idx = 0
+                        self.last_chunk_audios = None
+                else:
+                    origin_audio, latent_audio, valid_duration = self.prepare_audio_data(None)
+
+            if self.world_size > 1:
+                origin_audio = self.braodcast_audio_data(origin_audio, self.origin_audio_tensor, self.origin_flag_tensor)
+                latent_audio = self.braodcast_audio_data(latent_audio, self.latent_audio_tensor, self.latent_flag_tensor)
+                valid_duration = self.braodcast_valid_duration(valid_duration)
+
+            origin_audio = self.bytes_to_ndarray(origin_audio)
+            latent_audio = self.bytes_to_ndarray(latent_audio)
+            if latent_audio is not None:
+                # latent_audio: (latent_per_chunk, audio_window * sr)
+                latent_audio = latent_audio.reshape(self.latent_per_chunk, -1)
+            return origin_audio, latent_audio, valid_duration
+
+        except Exception:
+            logger.warning(f"Failed to get voice segment: {traceback.format_exc()}")
+            return None, None, 0
+
+    def braodcast_audio_data(self, audio_data, audio_tensor, flag_tensor):
+        if self.rank == self.target_rank:
+            if audio_data is None:
+                flag_tensor.fill_(0)
+            else:
+                flag_tensor.fill_(1)
+                audio_tensor.copy_(torch.frombuffer(bytearray(audio_data), dtype=torch.uint8))
+                # logger.info(f"rank {self.rank} send audio_tensor: {self.audio_tensor.shape}")
+
+        dist.broadcast(flag_tensor, src=self.target_rank)
+        if flag_tensor.item() == 0:
+            return None
+
+        dist.broadcast(audio_tensor, src=self.target_rank)
+        if self.rank != self.target_rank:
+            # logger.info(f"rank {self.rank} recv audio_tensor: {self.audio_tensor.shape}")
+            audio_data = audio_tensor.cpu().numpy().tobytes()
+        return audio_data
+
+
 if __name__ == "__main__":
     WORLD_SIZE = int(os.environ.get("WORLD_SIZE", 1))
     RANK = int(os.environ.get("RANK", 0))

From f1e159140b8367a972ec5da5f4caa7252fff9b09 Mon Sep 17 00:00:00 2001
From: gushiqiao <975033167@qq.com>
Date: Thu, 9 Jul 2026 04:32:34 +0000
Subject: [PATCH 3/9] fix bugs

---
 .../networks/wan/infer/audio/pre_infer.py     | 11 ++-
 .../models/runners/wan/wan_audio_runner.py    | 74 ++++++++++++++++---
 2 files changed, 71 insertions(+), 14 deletions(-)

diff --git a/lightx2v/models/networks/wan/infer/audio/pre_infer.py b/lightx2v/models/networks/wan/infer/audio/pre_infer.py
index 32879b9d9..27a379220 100755
--- a/lightx2v/models/networks/wan/infer/audio/pre_infer.py
+++ b/lightx2v/models/networks/wan/infer/audio/pre_infer.py
@@ -280,9 +280,14 @@ def infer(self, weights, inputs, kv_start=0, kv_end=0):
 
         audio_encoder_output = inputs.get("audio_encoder_output")
         if audio_encoder_output is not None and not is_ref_prefill:
-            segment_idx = self.scheduler.seg_index
-            chunk_size = self.scheduler.chunk_size
-            audio_encoder_output = audio_encoder_output[:, segment_idx * chunk_size : segment_idx * chunk_size + grid_sizes_t]
+            if inputs.get("audio_encoder_output_is_chunk", False):
+                if audio_encoder_output.shape[1] < grid_sizes_t:
+                    raise ValueError(f"stream audio feature is too short: got {audio_encoder_output.shape[1]}, need {grid_sizes_t}")
+                audio_encoder_output = audio_encoder_output[:, :grid_sizes_t]
+            else:
+                segment_idx = self.scheduler.seg_index
+                chunk_size = self.scheduler.chunk_size
+                audio_encoder_output = audio_encoder_output[:, segment_idx * chunk_size : segment_idx * chunk_size + grid_sizes_t]
         elif is_ref_prefill:
             audio_encoder_output = None
 
diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index d54ef046a..57c768217 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -1194,7 +1194,7 @@ def _apply_prompt_travel_for_latent(self, latent_idx):
             f"start_latent={self.prompt_travel_segments[prompt_idx].start_latent}"
         )
 
-    def _validate_prompt_travel_schedule(self):
+    def _validate_prompt_travel_schedule(self, allow_open_ended=False):
         if not self.prompt_travel_segments:
             return
         latent_length = int(self.input_info.latent_shape[1])
@@ -1206,7 +1206,7 @@ def _validate_prompt_travel_schedule(self):
                 raise ValueError("invalid prompt travel schedule: start_latents must be non-decreasing")
             if segment.start_latent == prev_start_latent:
                 logger.warning(f"Prompt travel segment at frame {segment.start_frame} maps to duplicate latent {segment.start_latent}; the later segment will be selected at that latent.")
-            if segment.start_latent >= latent_length:
+            if not allow_open_ended and segment.start_latent >= latent_length:
                 raise ValueError(f"prompt travel start_frame {segment.start_frame} is outside latent length {latent_length}")
             prev_start_latent = segment.start_latent
 
@@ -1273,6 +1273,7 @@ def _encode_audio_for_ar(self):
             feat = self.audio_adapter.forward_audio_proj(feat, feat.shape[1])
             features_list.append(feat.squeeze(0))
         self.inputs["audio_encoder_output"] = torch.stack(features_list, dim=0)
+        self.inputs["audio_encoder_output_is_chunk"] = False
         self.inputs["previmg_encoder_output"] = {"prev_latents": None, "prev_mask": None, "prev_len": 0}
 
     def init_run_segment(self, segment_idx, origin_audio=None, latent_audio=None):
@@ -1304,7 +1305,11 @@ def run_segment(self, segment_idx=0):
         start = segment_idx * chunk_size
         end = start + chunk_size
         self.model.scheduler.set_timesteps(infer_steps, device=AI_DEVICE)
-        xt = self.model.scheduler.noise[:, start:end].to(AI_DEVICE)
+        chunk_noise = self._make_ar_chunk_noise(segment_idx)
+        if chunk_noise is None:
+            xt = self.model.scheduler.noise[:, start:end].to(AI_DEVICE)
+        else:
+            xt = chunk_noise.to(AI_DEVICE)
         for step_index in range(infer_steps):
             logger.info(f"==> chunk: {segment_idx + 1} / {self.video_segment_num}, step_index: {step_index + 1} / {infer_steps}")
             self.model.kv_cache_manager.current_step = step_index
@@ -1320,6 +1325,27 @@ def run_segment(self, segment_idx=0):
                 self.progress_callback((current_step / total_all_steps) * 100, 100)
         return xt
 
+    def run_stream_segment(self, segment_idx=0):
+        infer_steps = self.model.scheduler.infer_steps
+        self.model.scheduler.set_timesteps(infer_steps, device=AI_DEVICE)
+        chunk_noise = self._make_ar_chunk_noise(segment_idx)
+        if chunk_noise is None:
+            xt = self.model.scheduler.noise.to(AI_DEVICE)
+        else:
+            xt = chunk_noise.to(AI_DEVICE)
+        for step_index in range(infer_steps):
+            logger.info(f"==> stream chunk: {segment_idx + 1}, step_index: {step_index + 1} / {infer_steps}")
+            self.model.kv_cache_manager.current_step = step_index
+            with ProfilingContext4DebugL1("step_pre"):
+                self.model.scheduler.step_pre(segment_idx, step_index, xt)
+            with ProfilingContext4DebugL1("🚀 infer_main"):
+                self.model.infer(self.inputs)
+            with ProfilingContext4DebugL1("step_post"):
+                xt = self.model.scheduler.step_post(xt)
+            if self.progress_callback:
+                self.progress_callback(((step_index + 1) / infer_steps) * 100, 100)
+        return xt
+
     def decode_segment_latents(self, segment_idx: int, segment_latents: torch.Tensor) -> torch.Tensor:
         is_first = segment_idx == 0
         is_last = segment_idx == self.video_segment_num - 1
@@ -1353,6 +1379,7 @@ def run_main_with_fixed_audio(self):
             latent_shape[1] = self.inputs["audio_encoder_output"].shape[1]
             latent_shape[1] = (latent_shape[1] // self.config["ar_config"]["num_frame_per_chunk"]) * self.config["ar_config"]["num_frame_per_chunk"]
             self.input_info.latent_shape = latent_shape
+            self._enable_ar_chunk_noise(self.input_info.seed, latent_shape)
             self.model.scheduler.prepare(seed=self.input_info.seed, latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
             self.get_video_segment_num()
             self._validate_prompt_travel_schedule()
@@ -1415,14 +1442,15 @@ def run_main_with_fixed_audio(self):
     # origin_audio: (4 * latent_per_chunk / fps * sr), eg. (4 * 3 / 16 * 16000)
     # latent_audio: (latent_per_chunk, audio_window * sr) eg. (3, 1.0 * 16000)
     def _encode_audio_for_each_chunk(self, segment_idx: int, origin_audio: np.ndarray, latent_audio: np.ndarray) -> torch.Tensor:
-        self.input_info.seed = self.input_info.seed + segment_idx
-        torch.manual_seed(self.input_info.seed)
+        torch.manual_seed(self._ar_chunk_seed(segment_idx))
 
         # used for remux output video
         end_idx = origin_audio.shape[0] // self._audio_processor.audio_frame_rate
-        # 第一段vae解码出来视频只有 (4 - 1) * latent_per_chunk 帧
+        # The first cached VAE decode has no previous temporal cache, so it emits
+        # vae_stride_t - 1 fewer video frames than later chunks.
         if segment_idx == 0:
-            end_idx -= 1 * self.config["ar_config"]["num_frame_per_chunk"] * self._audio_processor.audio_frame_rate
+            vae_stride_t = int(self.config.get("vae_stride", [4])[0])
+            end_idx = max(end_idx - max(vae_stride_t - 1, 0), 0)
         origin_audio_tensor = torch.Tensor(origin_audio).float().unsqueeze(0)
         self.segment = AudioSegment(origin_audio_tensor, 0, end_idx)
 
@@ -1441,8 +1469,26 @@ def _encode_audio_for_each_chunk(self, segment_idx: int, origin_audio: np.ndarra
         audio_feat = self.audio_adapter.forward_audio_proj(audio_feat, audio_feat.shape[1])
 
         self.inputs["audio_encoder_output"] = audio_feat
+        self.inputs["audio_encoder_output_is_chunk"] = True
         self.inputs["previmg_encoder_output"] = {"prev_latents": None, "prev_mask": None, "prev_len": 0}
 
+    def _enable_ar_chunk_noise(self, base_seed: int, latent_shape):
+        self._ar_base_seed = int(base_seed)
+        self._ar_chunk_latent_shape = list(latent_shape)
+        self._ar_chunk_latent_shape[1] = int(self.config["ar_config"]["num_frame_per_chunk"])
+
+    def _ar_chunk_seed(self, segment_idx: int) -> int:
+        return int(getattr(self, "_ar_base_seed", self.input_info.seed)) + int(segment_idx)
+
+    def _make_ar_chunk_noise(self, segment_idx: int):
+        latent_shape = getattr(self, "_ar_chunk_latent_shape", None)
+        if latent_shape is None:
+            return None
+        seed = self._ar_chunk_seed(segment_idx)
+        torch.manual_seed(seed)
+        generator = torch.Generator("cpu").manual_seed(seed)
+        return torch.randn(*latent_shape, dtype=GET_DTYPE(), device="cpu", generator=generator)
+
     def run_main(self):
         try:
             self.va_controller = None
@@ -1464,11 +1510,15 @@ def run_main(self):
             fail_count, max_fail_count = 0, 10
             self.va_controller.before_control()
 
+            if self.config.get("ar_config", {}).get("local_attn_size", -1) == -1:
+                raise ValueError("seko_talk_ar stream requires ar_config.local_attn_size for rolling KV cache.")
+
             latent_shape = list(self.input_info.latent_shape)
             latent_shape[1] = self.config["ar_config"]["num_frame_per_chunk"]
             self.input_info.latent_shape = latent_shape
-            self.model.scheduler.prepare(seed=self.input_info.seed, latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
-            self._validate_prompt_travel_schedule()
+            self._enable_ar_chunk_noise(self.input_info.seed, latent_shape)
+            self.model.scheduler.prepare(seed=self._ar_chunk_seed(0), latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
+            self._validate_prompt_travel_schedule(allow_open_ended=True)
             self.init_kv_cache_manager()
             self.prefill_reference_kv()
 
@@ -1494,10 +1544,12 @@ def run_main(self):
                         # reset pause signal
                         self.pause_signal = False
                         self.can_pause = valid_duration <= 1e-5
+                        chunk_size = int(self.model.scheduler.chunk_size)
+                        self._apply_prompt_travel_for_latent(segment_idx * chunk_size)
                         self.init_run_segment(segment_idx, origin_audio, latent_audio)
-                        self.model.scheduler.prepare(seed=self.input_info.seed, latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
+                        self.model.scheduler.prepare(seed=self._ar_chunk_seed(segment_idx), latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
                         self.check_stop()
-                        latents = self.run_segment(0)
+                        latents = self.run_stream_segment(segment_idx)
                         self.check_stop()
                         self.gen_video = self.decode_segment_latents(segment_idx, latents)
                         logger.debug(f"gen_video shape: {self.gen_video.shape}")

From be94b081ee1573b7a09f8e88726db5d6cb6b5e39 Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Thu, 9 Jul 2026 15:41:47 +0800
Subject: [PATCH 4/9] Tidy code

---
 .../models/runners/wan/wan_audio_runner.py    |  5 ++-
 lightx2v/utils/va_reader_omni.py              | 37 +++++++++----------
 2 files changed, 21 insertions(+), 21 deletions(-)

diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index 57c768217..446f2693f 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -1334,7 +1334,7 @@ def run_stream_segment(self, segment_idx=0):
         else:
             xt = chunk_noise.to(AI_DEVICE)
         for step_index in range(infer_steps):
-            logger.info(f"==> stream chunk: {segment_idx + 1}, step_index: {step_index + 1} / {infer_steps}")
+            # logger.info(f"==> stream chunk: {segment_idx + 1}, step_index: {step_index + 1} / {infer_steps}")
             self.model.kv_cache_manager.current_step = step_index
             with ProfilingContext4DebugL1("step_pre"):
                 self.model.scheduler.step_pre(segment_idx, step_index, xt)
@@ -1451,6 +1451,7 @@ def _encode_audio_for_each_chunk(self, segment_idx: int, origin_audio: np.ndarra
         if segment_idx == 0:
             vae_stride_t = int(self.config.get("vae_stride", [4])[0])
             end_idx = max(end_idx - max(vae_stride_t - 1, 0), 0)
+            origin_audio = origin_audio[-end_idx * self._audio_processor.audio_frame_rate :]
         origin_audio_tensor = torch.Tensor(origin_audio).float().unsqueeze(0)
         self.segment = AudioSegment(origin_audio_tensor, 0, end_idx)
 
@@ -1552,7 +1553,7 @@ def run_main(self):
                         latents = self.run_stream_segment(segment_idx)
                         self.check_stop()
                         self.gen_video = self.decode_segment_latents(segment_idx, latents)
-                        logger.debug(f"gen_video shape: {self.gen_video.shape}")
+                        # logger.debug(f"gen_video shape: {self.gen_video.shape}")
                         self.check_stop()
                         self.end_run_segment(segment_idx, valid_duration=valid_duration)
                         segment_idx += 1
diff --git a/lightx2v/utils/va_reader_omni.py b/lightx2v/utils/va_reader_omni.py
index 0d6efa78f..3ab8c43f7 100644
--- a/lightx2v/utils/va_reader_omni.py
+++ b/lightx2v/utils/va_reader_omni.py
@@ -212,7 +212,7 @@ def recv_loop(self):
             try:
                 message = BSON.decode(message)
                 msg_type = message["type"]
-                logger.debug("Received message type: {}".format(msg_type))
+                # logger.debug("Received message type: {}".format(msg_type))
                 if msg_type == "AgentAudio":
                     audio = message["audio"]
                     if audio["type"] != "Pcm":
@@ -220,7 +220,7 @@ def recv_loop(self):
                         continue
                     pcm_data = audio["data"]
                     audio_info = AudioInfo(audio["info"])
-                    logger.debug("Received audio with duration: {}".format(audio_info.duration()))
+                    # logger.debug("Received audio with duration: {}".format(audio_info.duration()))
                     if self.audio_info is None:
                         self.audio_info = audio_info
                     else:
@@ -272,8 +272,8 @@ def get_audio(self, fetch_duration) -> (bytes, AudioInfo):
 
         # the actual sample count fetched
         sample_count = len(pcm_data) // (self.audio_info.channel_count * 2)
-        logger.debug("Fetched {} bytes audio".format(sample_count))
-        logger.debug("After fetch, there are {} bytes left".format(self.audio_buffer.current_size))
+        # logger.debug("Fetched {} bytes audio".format(sample_count))
+        # logger.debug("After fetch, there are {} bytes left".format(self.audio_buffer.current_size))
         audio_info = deepcopy(self.audio_info)
         audio_info.sample_count = sample_count
         return (pcm_data, audio_info)
@@ -658,13 +658,13 @@ def get_latent_idxs(self, i, base_idx):
         # 当前 latent 音频段上界索引
         look_ahead_secs = (latent_idx * 4 + 1) / self.video_fps + self.look_ahead_secs
         look_ahead_idx = int(look_ahead_secs * self.sample_rate)
-        logger.debug(f"latent_idx: {latent_idx} center_secs: {center_secs} right_idx: {right_idx} left_idx: {left_idx} look_ahead_secs: {look_ahead_secs} look_ahead_idx: {look_ahead_idx}")
+        # logger.debug(f"latent [{latent_idx}]: center_secs: {center_secs} ahead_sec: {look_ahead_secs} {left_idx} {right_idx} {look_ahead_idx}")
 
         # 对齐当前缓存的所有音频的起始位置（只缓存了最多两个chunk的音频）
         left_idx -= base_idx
         right_idx -= base_idx
         look_ahead_idx -= base_idx
-        logger.debug(f"adjusted idxs: {left_idx} {right_idx} {look_ahead_idx}")
+        # logger.debug(f"latent [{latent_idx}]: {left_idx} {right_idx} {look_ahead_idx}")
         return left_idx, right_idx, look_ahead_idx
 
     # 获取当前chunk对应的原始音频的开始和结束索引帧
@@ -674,11 +674,11 @@ def get_chunk_origin_idxs(self, base_idx):
         start_secs = end_secs - (self.latent_per_chunk * 4) / self.video_fps
         start_idx = int(start_secs * self.sample_rate)
         end_idx = int(end_secs * self.sample_rate)
-        logger.debug(f"chunk_idx: {self.chunk_idx} start_idx: {start_idx} end_idx: {end_idx}")
+        # logger.debug(f"chunk {self.chunk_idx} start_secs: {start_secs} end_secs: {end_secs} {start_idx} {end_idx}")
         # 对齐当前缓存的所有音频的起始位置（只缓存了最多两个chunk的音频）
         start_idx -= base_idx
         end_idx -= base_idx
-        logger.debug(f"adjusted idxs: {start_idx} {end_idx}")
+        # logger.debug(f"chunk {self.chunk_idx} adjusted: {start_idx} {end_idx}")
         return start_idx, end_idx
 
     def prepare_ar_audios(self, merged_audio, origin_audios, latent_audios):
@@ -688,16 +688,13 @@ def prepare_ar_audios(self, merged_audio, origin_audios, latent_audios):
             last_chunk_start_secs = self.segment_duration + (self.chunk_idx - 2) * self.other_fetch_duration
         base_idx = int(last_chunk_start_secs * self.sample_rate)
         audio_length = len(merged_audio)
-        logger.debug(
-            "chunk idx: ", self.chunk_idx, "start_secs: ", last_chunk_start_secs, "end_secs: ", last_chunk_start_secs + self.segment_duration, "base_idx: ", base_idx, "audio_length: ", audio_length
-        )
 
         # 构造原始音频，用于合成最终视频
         start_idx, end_idx = self.get_chunk_origin_idxs(base_idx)
         pad_idx = max(-start_idx, 0)
         real_start = max(start_idx, 0)
         real_len = min(end_idx, audio_length) - real_start
-        logger.debug(f"----- center audios pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len} -----")
+        # logger.debug(f"origin audios pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len}")
         origin_audios[pad_idx : pad_idx + real_len] = merged_audio[real_start : real_start + real_len]
 
         # 构造 latent 音频，用于 audio encode
@@ -706,7 +703,7 @@ def prepare_ar_audios(self, merged_audio, origin_audios, latent_audios):
             pad_idx = max(-left_idx, 0)
             real_start = max(left_idx, 0)
             real_len = min(look_ahead_idx, right_idx, audio_length) - real_start
-            logger.debug(f"----- {i} pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len} -----")
+            # logger.debug(f"latent audios {i} pad_idx: {pad_idx} real_start: {real_start} real_len: {real_len}")
             latent_audios[i, pad_idx : pad_idx + real_len] = merged_audio[real_start : real_start + real_len]
 
     def prepare_audio_data(self, chat_audio_result):
@@ -719,9 +716,9 @@ def prepare_audio_data(self, chat_audio_result):
         if chat_audio_result is not None:
             audio_data, audio_info = chat_audio_result
             audio, sample_count = self.convert_pcm_s16le_to_mono_resampled(audio_data, audio_info)
-        valid_duration = sample_count / self.sample_rate
+
         if sample_count <= 0:
-            return origin_audios.tobytes(), latent_audios.tobytes(), valid_duration
+            return origin_audios.tobytes(), latent_audios.tobytes(), 0
 
         if self.chunk_idx == 0:
             expect_count = int(self.segment_duration * self.sample_rate)
@@ -729,22 +726,24 @@ def prepare_audio_data(self, chat_audio_result):
             expect_count = int(self.other_fetch_duration * self.sample_rate)
         assert sample_count <= expect_count, f"audio length {sample_count} > expect_count {expect_count}"
 
+        valid_duration = self.other_fetch_duration
         # pad 0 to the audio to make it the same length as expect_count
         if sample_count < expect_count:
             pad_count = expect_count - sample_count
             logger.debug(f"pad {pad_count} samples to audio")
             audio = np.pad(audio, (0, pad_count), mode="constant", constant_values=0)
+            valid_duration -= pad_count / self.sample_rate
 
         # if is not the first segment, concat with previous segment
         if self.last_chunk_audios is not None:
-            logger.debug(f"concat {self.last_chunk_audios.shape} with {audio.shape}")
+            # logger.debug(f"concat {self.last_chunk_audios.shape} with {audio.shape}")
             merged_audio = np.concatenate([self.last_chunk_audios, audio])
         else:
             merged_audio = audio
         self.prepare_ar_audios(merged_audio, origin_audios, latent_audios)
-        logger.debug(
-            f"chunk[{self.chunk_idx}] origin_audios: {origin_audios.shape} {origin_audios.dtype} {origin_audios.min()} {origin_audios.max()} latent_audios: {latent_audios.shape} {latent_audios.dtype} {latent_audios.min()} {latent_audios.max()} {sample_count}, valid_duration: {valid_duration}"
-        )
+        # logger.debug(
+        #     f"chunk[{self.chunk_idx}] origin_audios: {origin_audios.shape} {origin_audios.dtype} {origin_audios.min()} {origin_audios.max()} latent_audios: {latent_audios.shape} {latent_audios.dtype} {latent_audios.min()} {latent_audios.max()} {sample_count}, valid_duration: {valid_duration}"
+        # )
 
         # update prev seg chunk
         self.last_chunk_audios = audio

From 722d3f5cfff3647e4cf513557baac5faf3de3bdb Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Thu, 9 Jul 2026 17:25:33 +0800
Subject: [PATCH 5/9] Fix no realtime output

---
 .../models/runners/wan/wan_audio_runner.py    | 30 +++++++++----------
 lightx2v/utils/va_controller.py               |  2 +-
 2 files changed, 15 insertions(+), 17 deletions(-)

diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index 446f2693f..500c66a81 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -1416,7 +1416,6 @@ def run_main_with_fixed_audio(self):
                             self.check_stop()
                             self.init_run_segment(segment_idx)
                             segment_latents = self.run_segment(segment_idx)
-                            logger.debug(f"segment_latents shape: {segment_latents.shape}")
 
                         vae_decoder.submit(self.decode_segment_latents, segment_idx, segment_latents)
                     segment_videos = vae_decoder.finish()
@@ -1524,21 +1523,20 @@ def run_main(self):
             self.prefill_reference_kv()
 
             while True:
-                with ProfilingContext4DebugL1(f"stream segment get audio segment {segment_idx}"):
-                    control = self.va_controller.next_control()
-                    if control.action == "blank_to_voice":
-                        self.prev_video = control.data
-                    elif control.action == "wait":
-                        time.sleep(0.01)
-                        continue
-
-                    origin_audio, latent_audio, valid_duration = self.va_controller.reader.get_audio_segment()
-                    if origin_audio is None or latent_audio is None:
-                        fail_count += 1
-                        logger.warning(f"Failed to get audio chunk {fail_count} times")
-                        if fail_count > max_fail_count:
-                            raise Exception(f"Failed to get audio chunk {fail_count} times, stop reader")
-                        continue
+                control = self.va_controller.next_control()
+                if control.action == "blank_to_voice":
+                    self.prev_video = control.data
+                elif control.action == "wait":
+                    time.sleep(0.01)
+                    continue
+
+                origin_audio, latent_audio, valid_duration = self.va_controller.reader.get_audio_segment()
+                if origin_audio is None or latent_audio is None:
+                    fail_count += 1
+                    logger.warning(f"Failed to get audio chunk {fail_count} times")
+                    if fail_count > max_fail_count:
+                        raise Exception(f"Failed to get audio chunk {fail_count} times, stop reader")
+                    continue
 
                 with ProfilingContext4DebugL1(f"stream segment end2end {segment_idx}"):
                     try:
diff --git a/lightx2v/utils/va_controller.py b/lightx2v/utils/va_controller.py
index 7ff74b3f5..6ad43abee 100644
--- a/lightx2v/utils/va_controller.py
+++ b/lightx2v/utils/va_controller.py
@@ -68,7 +68,7 @@ def init_base(self, config, input_info, has_vfi_model, has_vsr_model):
         # how many frames to publish stream as a batch
         self.slice_frame = config.get("slice_frame", self.prev_frame_length)
         # estimate the max infer seconds, for immediate switch with local omni
-        slice_interval = max(1, self.slice_frame / self.record_fps)
+        slice_interval = max(1e-5, self.slice_frame / self.record_fps)
 
         est_max_infer_secs = config.get("est_max_infer_secs", 0.6)
         est_max_switch_image_secs = config.get("est_max_switch_image_secs", 0)

From b8c8c27600626a7bfe80f4e5c5d0217fc7414dd2 Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Fri, 10 Jul 2026 13:09:04 +0800
Subject: [PATCH 6/9] Restore vae decoder with specific segment idx

---
 .../models/runners/wan/wan_audio_runner.py    | 32 +++++++++++++++++--
 lightx2v/utils/va_controller.py               |  3 ++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index 500c66a81..53ab16bdd 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -676,7 +676,11 @@ def end_run_segment(self, segment_idx, valid_duration=1e9):
             )
 
         if self.va_controller.recorder is not None:
-            self.va_controller.pub_livestream(video_seg, audio_seg, self.gen_video[:, :, :useful_length], valid_duration=valid_duration)
+            gen_video = self.gen_video[:, :, :useful_length]
+            # only save segment idx for seko_talk_ar
+            if self.config.get("model_cls") == "seko_talk_ar":
+                gen_video = torch.full((1, 1, useful_length), int(segment_idx), dtype=torch.int64, device=AI_DEVICE)
+            self.va_controller.pub_livestream(video_seg, audio_seg, gen_video, valid_duration=valid_duration)
         elif self.input_info.return_result_tensor:
             self.gen_video_final[self.segment.start_frame : self.segment.end_frame].copy_(video_seg)
             self.cut_audio_final[self.segment.start_frame * self._audio_processor.audio_frame_rate : self.segment.end_frame * self._audio_processor.audio_frame_rate].copy_(audio_seg)
@@ -1489,8 +1493,25 @@ def _make_ar_chunk_noise(self, segment_idx: int):
         generator = torch.Generator("cpu").manual_seed(seed)
         return torch.randn(*latent_shape, dtype=GET_DTYPE(), device="cpu", generator=generator)
 
+    def _save_vae_feat_map(self, segment_idx: int) -> None:
+        feat_map = self.vae_decoder.model._feat_map
+        self.vae_feat_maps[segment_idx] = [None if item is None else item.detach().clone() for item in feat_map]
+        min_segment_idx = segment_idx - int(self.config.get("ar_vae_cache_chunk", 7))
+        for key in list(self.vae_feat_maps.keys()):
+            if key < min_segment_idx:
+                self.vae_feat_maps.pop(key, None)
+
+    def _restore_vae_feat_map(self, segment_idx: int) -> None:
+        feat_map = self.vae_feat_maps.get(segment_idx)
+        if feat_map is None:
+            logger.warning(f"missing VAE feat_map for restoring segment_idx={segment_idx}")
+            return
+        self.vae_decoder.model._feat_map = list(feat_map)
+        logger.debug(f"restored VAE feat_map for segment_idx={segment_idx}")
+
     def run_main(self):
         try:
+            self.vae_feat_maps = {}
             self.va_controller = None
             self.init_run()
             logger.info(f"init va_recorder: {self.va_controller.recorder} and va_reader: {self.va_controller.reader}")
@@ -1525,7 +1546,12 @@ def run_main(self):
             while True:
                 control = self.va_controller.next_control()
                 if control.action == "blank_to_voice":
-                    self.prev_video = control.data
+                    restore_idx = int(control.data.view(-1)[0].item())
+                    # restor new idx must be less than current segment idx
+                    if restore_idx >= 0 and restore_idx + 1 < segment_idx:
+                        logger.warning(f"restore segment idx: {segment_idx} -> {restore_idx + 1}")
+                        self._restore_vae_feat_map(restore_idx)
+                        segment_idx = restore_idx + 1
                 elif control.action == "wait":
                     time.sleep(0.01)
                     continue
@@ -1554,6 +1580,7 @@ def run_main(self):
                         # logger.debug(f"gen_video shape: {self.gen_video.shape}")
                         self.check_stop()
                         self.end_run_segment(segment_idx, valid_duration=valid_duration)
+                        self._save_vae_feat_map(segment_idx)
                         segment_idx += 1
                         fail_count = 0
                     except Exception as e:
@@ -1562,6 +1589,7 @@ def run_main(self):
                         else:
                             raise
         finally:
+            self.vae_feat_maps = {}
             if hasattr(self.model, "inputs"):
                 self.end_run()
             if self.va_controller is not None:
diff --git a/lightx2v/utils/va_controller.py b/lightx2v/utils/va_controller.py
index 6ad43abee..27d948115 100644
--- a/lightx2v/utils/va_controller.py
+++ b/lightx2v/utils/va_controller.py
@@ -196,6 +196,9 @@ def before_control(self):
         if isinstance(self.reader, OmniVAReader):
             self.len_tensor = torch.tensor([0], dtype=torch.int32, device=AI_DEVICE)
             self.flag_tensor = torch.tensor([0], dtype=torch.int32, device=AI_DEVICE)
+            if self.is_seko_talk_ar:
+                self.prev_tensor = torch.full((1, 1, self.prev_frame_length), -1, dtype=torch.int64, device=AI_DEVICE)
+                return
             self.prev_tensor = torch.zeros((1, 3, self.prev_frame_length, self.tgt_h, self.tgt_w), dtype=torch.float, device=AI_DEVICE)
 
     def omni_reader_next_control(self):

From 98a685f2737bf0a8af374f9de11badc735c61e32 Mon Sep 17 00:00:00 2001
From: gushiqiao <975033167@qq.com>
Date: Fri, 10 Jul 2026 05:57:53 +0000
Subject: [PATCH 7/9] fix rope

---
 .../seko_talk/ar/seko_talk_ar_kv_dist.json    |  1 +
 .../wan/infer/audio/transformer_infer.py      | 36 +++++++++++++++----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/configs/seko_talk/ar/seko_talk_ar_kv_dist.json b/configs/seko_talk/ar/seko_talk_ar_kv_dist.json
index f4fefd50b..a5d157825 100644
--- a/configs/seko_talk/ar/seko_talk_ar_kv_dist.json
+++ b/configs/seko_talk/ar/seko_talk_ar_kv_dist.json
@@ -39,6 +39,7 @@
         "step_kv_cache": true,
         "sink_size": 2,
         "local_attn_size": 21,
+        "rope_position_mode": "local",
         "kv_offload": false,
         "async_vae_decode": false
     }
diff --git a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
index 5a416d5f4..415007c83 100755
--- a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
@@ -129,6 +129,8 @@ def __init__(self, config):
         super().__init__(config)
         self._setup_audio_post_adapter(config)
         self._audio_grid_meta_cache = {}
+        ar_config = config.get("ar_config", {})
+        self.rope_position_mode = ar_config.get("rope_position_mode", "local")
 
     def _audio_grid_meta(self, grid_sizes):
         key = (grid_sizes.data_ptr(), int(self.scheduler.seg_index), int(self.scheduler.step_index))
@@ -319,6 +321,7 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
         local_per_frame = num_new // frames if frames > 0 else 0
         cache_per_frame = local_per_frame if replicated_ref_prefill else local_per_frame * sp_world_size if seq_parallel else local_per_frame
         sink_tokens = self.kv_cache_manager.sink_size * cache_per_frame
+        use_local_cache_rope = self.rope_position_mode != "global"
 
         need_roll = self.kv_cache_manager.local_attn_size != -1 and current_end > global_end and cache_num_new + local_end > self.kv_cache_size
         if need_roll:
@@ -348,14 +351,18 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
                 shard_heads = self.num_heads // sp_world_size
                 h0 = sp_rank * shard_heads
                 h1 = h0 + shard_heads
-                kv_cache.store_kv(k_rope[:, h0:h1], v[:, h0:h1], local_start_idx, local_end_idx, self.block_idx)
+                k_to_store = k[:, h0:h1] if use_local_cache_rope else k_rope[:, h0:h1]
+                kv_cache.store_kv(k_to_store, v[:, h0:h1], local_start_idx, local_end_idx, self.block_idx)
             else:
-                start_frame = self.kv_cache_manager.ref_num_frames + segment_idx * frames
+                if use_local_cache_rope:
+                    start_frame = local_start_idx // max(cache_per_frame, 1)
+                else:
+                    start_frame = self.kv_cache_manager.ref_num_frames + segment_idx * frames
                 q_rope, k_rope = self._apply_rope_sp(q, k, grid_sizes, freqs, start_frame)
                 use_fp8_comm = self.config["parallel"].get("seq_p_fp8_comm", False)
                 use_fp4_comm = self.config["parallel"].get("seq_p_fp4_comm", False)
                 k_to_store = all2all_seq2head(
-                    k_rope,
+                    k if use_local_cache_rope else k_rope,
                     group=self.seq_p_group,
                     use_fp8_comm=use_fp8_comm,
                     use_fp4_comm=use_fp4_comm,
@@ -390,6 +397,19 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
             else:
                 attn_k = kv_cache.k_cache(self.block_idx, attn_start, local_end_idx)
                 attn_v = kv_cache.v_cache(self.block_idx, attn_start, local_end_idx)
+                if use_local_cache_rope:
+                    attn_k = self._apply_rope_with_cache_range(
+                        attn_k,
+                        freqs,
+                        h,
+                        w,
+                        1,
+                        0,
+                        attn_start,
+                        local_end_idx,
+                        cache_ref_tokens,
+                        cache_per_frame,
+                    )
                 attn_out = kv_cache.sp_kvcache_attn_head_shard(
                     q=q_rope,
                     k_cache=attn_k,
@@ -404,6 +424,8 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
             attn_v = kv_cache.v_cache(self.block_idx, attn_start, local_end_idx)
             if self.config.get("ar_config", {}).get("kv_quant", {}).get("calibrate", False):
                 kv_cache.capture_attn(self.block_idx, attn_start, local_end_idx)
+            rope_global_end = None if use_local_cache_rope else current_end
+            rope_sink_tokens = 0 if use_local_cache_rope else sink_tokens
             q = self._apply_rope_with_cache_range(
                 q,
                 freqs,
@@ -415,8 +437,8 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
                 local_end_idx,
                 local_ref_tokens,
                 local_per_frame,
-                global_end=current_end,
-                sink_tokens=sink_tokens,
+                global_end=rope_global_end,
+                sink_tokens=rope_sink_tokens,
             )
             attn_k = self._apply_rope_with_cache_range(
                 attn_k,
@@ -429,8 +451,8 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
                 local_end_idx,
                 local_ref_tokens,
                 local_per_frame,
-                global_end=current_end,
-                sink_tokens=sink_tokens,
+                global_end=rope_global_end,
+                sink_tokens=rope_sink_tokens,
             )
             if isinstance(attn_k, tuple):
                 k_lens = torch.empty_like(seq_lens).fill_(attn_k[0].size(0))

From 6badce8bb86e8edb074eb751f792d1d5b68bf5ed Mon Sep 17 00:00:00 2001
From: huochaitiantang <1432249204@qq.com>
Date: Tue, 14 Jul 2026 15:42:52 +0800
Subject: [PATCH 8/9] Support silence prompt

---
 .../models/runners/wan/wan_audio_runner.py    | 38 +++++++++++++++++--
 1 file changed, 35 insertions(+), 3 deletions(-)

diff --git a/lightx2v/models/runners/wan/wan_audio_runner.py b/lightx2v/models/runners/wan/wan_audio_runner.py
index 53ab16bdd..590c8f86e 100755
--- a/lightx2v/models/runners/wan/wan_audio_runner.py
+++ b/lightx2v/models/runners/wan/wan_audio_runner.py
@@ -1004,6 +1004,10 @@ def __init__(self, config):
         super().__init__(config)
         self.prompt_travel_segments = None
         self.prompt_travel_text_encoder_outputs = None
+        self._voice_text_encoder_output = None
+        self._silence_text_encoder_output = None
+        self._voice_prompt_text = None
+        self._silence_prompt_text = None
         self.audio_sliding_processor = CausalAudioSlidingProcessor(
             audio_window=self.config.get("audio_window", 1.0),
             look_ahead=self.config.get("look_ahead", 0.0),
@@ -1198,6 +1202,35 @@ def _apply_prompt_travel_for_latent(self, latent_idx):
             f"start_latent={self.prompt_travel_segments[prompt_idx].start_latent}"
         )
 
+    def _prepare_silence_voice_prompts(self, voice_text_encoder_output):
+        """Optional stream mode: voice uses input_info.prompt, silence uses config silence_prompt."""
+        self._voice_text_encoder_output = None
+        self._silence_text_encoder_output = None
+        self._voice_prompt_text = None
+        self._silence_prompt_text = None
+        silence_prompt = self.config.get("silence_prompt")
+        if not isinstance(silence_prompt, str) or not silence_prompt.strip():
+            return
+        self._voice_text_encoder_output = voice_text_encoder_output
+        self._voice_prompt_text = self.input_info.prompt
+        orig_prompt = self.input_info.prompt
+        self.input_info.prompt = silence_prompt
+        self._silence_text_encoder_output = self.run_text_encoder(self.input_info)
+        self.input_info.prompt = orig_prompt
+        self._silence_prompt_text = silence_prompt
+        logger.info(f"prompt: {orig_prompt}, silence_prompt: {silence_prompt}")
+
+    def _apply_silence_voice_prompt(self, is_silence):
+        if self._silence_text_encoder_output is None or self._voice_text_encoder_output is None:
+            return False
+        if is_silence:
+            self.inputs["text_encoder_output"] = self._silence_text_encoder_output
+            self.input_info.prompt = self._silence_prompt_text
+        else:
+            self.inputs["text_encoder_output"] = self._voice_text_encoder_output
+            self.input_info.prompt = self._voice_prompt_text
+        return True
+
     def _validate_prompt_travel_schedule(self, allow_open_ended=False):
         if not self.prompt_travel_segments:
             return
@@ -1234,6 +1267,7 @@ def _run_input_encoder_local_s2v(self):
             self.input_info.prompt = self.prompt_travel_segments[0].text
         else:
             text_encoder_output = self.run_text_encoder(self.input_info)
+            self._prepare_silence_voice_prompts(text_encoder_output)
 
         torch.cuda.empty_cache()
         gc.collect()
@@ -1539,7 +1573,6 @@ def run_main(self):
             self.input_info.latent_shape = latent_shape
             self._enable_ar_chunk_noise(self.input_info.seed, latent_shape)
             self.model.scheduler.prepare(seed=self._ar_chunk_seed(0), latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
-            self._validate_prompt_travel_schedule(allow_open_ended=True)
             self.init_kv_cache_manager()
             self.prefill_reference_kv()
 
@@ -1569,8 +1602,7 @@ def run_main(self):
                         # reset pause signal
                         self.pause_signal = False
                         self.can_pause = valid_duration <= 1e-5
-                        chunk_size = int(self.model.scheduler.chunk_size)
-                        self._apply_prompt_travel_for_latent(segment_idx * chunk_size)
+                        self._apply_silence_voice_prompt(self.can_pause)
                         self.init_run_segment(segment_idx, origin_audio, latent_audio)
                         self.model.scheduler.prepare(seed=self._ar_chunk_seed(segment_idx), latent_shape=latent_shape, infer_steps=self.config.get("infer_steps"))
                         self.check_stop()

From b7940b913790159cbbde7b70d2b78c3abda732fb Mon Sep 17 00:00:00 2001
From: gushiqiao <975033167@qq.com>
Date: Tue, 14 Jul 2026 09:01:22 +0000
Subject: [PATCH 9/9] fix prompt travel

---
 .../wan/infer/audio/transformer_infer.py      | 82 +++++++++++++++++++
 1 file changed, 82 insertions(+)

diff --git a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
index 415007c83..56559d3e8 100755
--- a/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
+++ b/lightx2v/models/networks/wan/infer/audio/transformer_infer.py
@@ -479,3 +479,85 @@ def infer_self_attn_with_kvcache(self, phase, grid_sizes, x, seq_lens, freqs, sh
                 next_prefetch=None,
             )
         return y
+
+    def infer_cross_attn_with_kvcache(self, phase, x, context, y_out, gate_msa):
+        num_frames = gate_msa.shape[0]
+        frame_seqlen = x.shape[0] // num_frames
+        seg_index = self.scheduler.seg_index
+
+        x.add_((y_out.unflatten(dim=0, sizes=(num_frames, frame_seqlen)) * gate_msa).flatten(0, 1))
+
+        norm3_out = phase.norm3.apply(x)
+
+        if self.task in ["i2v", "flf2v", "animate", "s2v", "rs2v"] and self.config.get("use_image_encoder", True):
+            context_img = context[:257]
+            context = context[257:]
+        else:
+            context_img = None
+
+        if self.sensitive_layer_dtype != self.infer_dtype:
+            context = context.to(self.infer_dtype)
+            if context_img is not None:
+                context_img = context_img.to(self.infer_dtype)
+
+        n, d = self.num_heads, self.head_dim
+        q = phase.cross_attn_norm_q.apply(phase.cross_attn_q.apply(norm3_out)).view(-1, n, d)
+
+        if self.config["ar_config"].get("use_cross_attn_kv_cache", False):
+            cross_kv_cache = self.kv_cache_manager.cross_attn_kv_cache
+            if seg_index == 0:
+                k = phase.cross_attn_norm_k.apply(phase.cross_attn_k.apply(context)).view(-1, n, d)
+                v = phase.cross_attn_v.apply(context).view(-1, n, d)
+                cross_kv_cache.store_kv(k, v, self.block_idx)
+                self._cross_kv_len = k.size(0)
+            else:
+                L = self._cross_kv_len
+                k = cross_kv_cache.k_cache(self.block_idx)[:L]
+                v = cross_kv_cache.v_cache(self.block_idx)[:L]
+        else:
+            k = phase.cross_attn_norm_k.apply(phase.cross_attn_k.apply(context)).view(-1, n, d)
+            v = phase.cross_attn_v.apply(context).view(-1, n, d)
+
+        cu_seqlens_q, cu_seqlens_k = self._calculate_q_k_len(
+            q,
+            k_lens=torch.tensor([k.size(0)], dtype=torch.int32),
+        )
+        attn_out = phase.cross_attn_1.apply(
+            q=q,
+            k=k,
+            v=v,
+            cu_seqlens_q=cu_seqlens_q,
+            cu_seqlens_kv=cu_seqlens_k,
+            max_seqlen_q=q.size(0),
+            max_seqlen_kv=k.size(0),
+        )
+
+        if context_img is not None:
+            k_img = phase.cross_attn_norm_k_img.apply(phase.cross_attn_k_img.apply(context_img)).view(-1, n, d)
+            v_img = phase.cross_attn_v_img.apply(context_img).view(-1, n, d)
+            cu_seqlens_q, cu_seqlens_k = self._calculate_q_k_len(
+                q,
+                k_lens=torch.tensor([k_img.size(0)], dtype=torch.int32),
+            )
+            attn_out.add_(
+                phase.cross_attn_2.apply(
+                    q=q,
+                    k=k_img,
+                    v=v_img,
+                    cu_seqlens_q=cu_seqlens_q,
+                    cu_seqlens_kv=cu_seqlens_k,
+                    max_seqlen_q=q.size(0),
+                    max_seqlen_kv=k_img.size(0),
+                )
+            )
+
+            if self.clean_cuda_cache:
+                del k_img, v_img
+                torch_device_module.empty_cache()
+
+        attn_out = phase.cross_attn_o.apply(attn_out)
+
+        if self.clean_cuda_cache:
+            del q, k, v, norm3_out, context, context_img
+            torch_device_module.empty_cache()
+        return x, attn_out