stream_save_video for infinitetalk (#1204)

Author: gushiqiao

configs/infinitetalk/5090/infinitetalk_single_distilled_8gpus.json

@@ -0,0 +1,51 @@
+{
+    "infer_steps": 4,
+    "target_video_length": 81,
+    "motion_frame": 9,
+    "target_fps": 25,
+    "video_duration": 300,
+    "infinitetalk_mode": "single",
+    "infinitetalk_size": "infinitetalk-720",
+    "dit_quantized_ckpt": "/path/to/InfiniteTalk/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors",
+    "dit_quantized": true,
+    "dit_quant_scheme": "fp8-sgl",
+    "adapter_model_path": "/path/to/InfiniteTalk/single/single/infinitetalk-fp8.safetensors",
+    "adapter_quantized": true,
+    "adapter_quant_scheme": "fp8-sgl",
+    "audio_encoder_path": "/path/to/InfiniteTalk/TencentGameMate/chinese-wav2vec2-base",
+    "clip_quantized": true,
+    "clip_quant_scheme": "fp8-sgl",
+    "t5_quantized": true,
+    "t5_quant_scheme": "fp8-sgl",
+    "audio_sample_rate": 16000,
+    "sample_shift": 7,
+    "sample_text_guide_scale": 1.0,
+    "sample_audio_guide_scale": 1.0,
+    "enable_cfg": false,
+    "enable_text_cfg": false,
+    "use_image_encoder": true,
+    "feature_caching": "NoCaching",
+    "cpu_offload": true,
+    "offload_granularity": "block",
+    "t5_cpu_offload": false,
+    "vae_cpu_offload": false,
+    "clip_cpu_offload": false,
+    "self_attn_1_type": "sage_attn2",
+    "cross_attn_1_type": "sage_attn2",
+    "cross_attn_2_type": "sage_attn2",
+    "audio_cross_attn_type": "sage_attn2",
+    "text_len": 512,
+    "target_height": 720,
+    "target_width": 1280,
+    "audio_window": 5,
+    "infinitetalk_vae_scale": 4,
+    "infinitetalk_context_tokens": 32,
+    "infinitetalk_audio_output_dim": 768,
+    "norm_output_audio": true,
+    "mxfp8_fuse_enable": false,
+    "use_timestep_transform": true,
+    "parallel": {
+        "seq_p_size": 8,
+        "seq_p_attn_type": "ulysses-4090"
+    }
+}

configs/infinitetalk/h100/infinitetalk_single_distilled.json

@@ -6,13 +6,13 @@
     "video_duration": 300,
     "infinitetalk_mode": "single",
     "infinitetalk_size": "infinitetalk-720",
-    "dit_quantized_ckpt": "/path/to/InfiniteTalk/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors",
+    "dit_quantized_ckpt": "/data/nvme5/gushiqiao/models/InfiniteTalk/seko/InfiniteTalk-4StepDistill-Mean-w-ITAudioAdaptorV6.1-fp8.safetensors",
     "dit_quantized": true,
     "dit_quant_scheme": "fp8-sgl",
-    "adapter_model_path": "/path/to/InfiniteTalk/single/single/infinitetalk-fp8.safetensors",
+    "adapter_model_path": "/data/nvme5/gushiqiao/models/InfiniteTalk/single/single/infinitetalk-fp8.safetensors",
     "adapter_quantized": true,
     "adapter_quant_scheme": "fp8-sgl",
-    "audio_encoder_path": "/path/to/InfiniteTalk/TencentGameMate/chinese-wav2vec2-base",
+    "audio_encoder_path": "/data/nvme5/gushiqiao/models/InfiniteTalk/TencentGameMate/chinese-wav2vec2-base",
     "clip_quantized": true,
     "clip_quant_scheme": "fp8-sgl",
     "t5_quantized": true,

lightx2v/models/runners/wan/wan_infinitetalk_runner.py

@@ -23,6 +23,7 @@
 from lightx2v.utils.profiler import ProfilingContext4DebugL1, ProfilingContext4DebugL2
 from lightx2v.utils.registry_factory import RUNNER_REGISTER
 from lightx2v.utils.utils import is_main_process, save_to_video, wan_vae_to_comfy
+from lightx2v.utils.va_controller import VAController
 from lightx2v_platform.base.global_var import AI_DEVICE
 
 torch_device_module = getattr(torch, AI_DEVICE)
@@ -106,8 +107,11 @@ def __init__(self, config):
         self.audio_sample_rate = int(self.config.get("audio_sample_rate", 16000))
         self.target_fps = int(self.config.get("target_fps", 25))
         self.video_audio_path = None
+        self.video_audio_array = None
         self.cond_video_temp_path = None
         self.cond_video_duration = None
+        self.va_controller = None
+        self.stream_save_video = False
 
     def init_scheduler(self):
         self.scheduler = InfiniteTalkScheduler(self.config)
@@ -289,6 +293,7 @@ def _audio_prepare_multi(self, left_path, right_path, audio_type):
         return new_speech1, new_speech2, new_speech1 + new_speech2
 
     def _write_sum_audio(self, input_data, audio_arrays):
+        self.video_audio_array = np.asarray(audio_arrays, dtype=np.float32)
         if sf is not None:
             fd, audio_path = tempfile.mkstemp(prefix="infinitetalk_sum_", suffix=".wav")
             os.close(fd)
@@ -506,6 +511,7 @@ def _run_input_encoder_local_s2v(self):
         self.cond_video_duration = self._get_cond_video_duration(self.cond_file_path)
         first_image = self._extract_specific_frame(self.cond_file_path, 0)
         self.src_h, self.src_w, self.target_h, self.target_w = self._select_target_size(first_image)
+        self.input_info.target_shape = [self.target_h, self.target_w]
 
         full_audio_embs = self._prepare_audio_embeddings(input_data)
         if any(audio_emb.shape[0] <= 0 for audio_emb in full_audio_embs):
@@ -623,7 +629,7 @@ def init_run(self):
 
         self.cond_image = self._prepare_cond_image(0)
         self.cond_frame = None
-        self.gen_video_list = []
+        self.gen_video_list = None if self.stream_save_video else []
 
     def get_video_segment_num(self):
         if self.expected_frames <= self.frame_num:
@@ -683,6 +689,47 @@ def run_segment(self, segment_idx=0):
         self._run_dit_clip(self.dit_inputs)
         return self.scheduler.latents
 
+    def _should_stream_save_video(self):
+        return bool(self.config.get("stream_save_video", True) and not self.input_info.return_result_tensor and getattr(self.input_info, "save_result_path", None))
+
+    def _init_stream_video_controller(self):
+        if not self.stream_save_video:
+            return
+        self.va_controller = VAController(self)
+        logger.info(f"init va_recorder: {self.va_controller.recorder} and va_reader: {self.va_controller.reader}")
+
+    def _get_audio_segment(self, start_frame, frame_count):
+        audio_sample_start = int(round(start_frame * self.audio_sample_rate / self.target_fps))
+        audio_sample_end = int(round((start_frame + frame_count) * self.audio_sample_rate / self.target_fps))
+        audio_sample_count = max(audio_sample_end - audio_sample_start, 0)
+        if audio_sample_count == 0:
+            return torch.zeros(0, dtype=torch.float32)
+
+        if self.video_audio_array is None:
+            return torch.zeros(audio_sample_count, dtype=torch.float32)
+
+        audio = self.video_audio_array.reshape(-1)
+        audio_seg = audio[audio_sample_start : min(audio_sample_end, audio.shape[0])]
+        if audio_seg.shape[0] < audio_sample_count:
+            audio_seg = np.pad(audio_seg, (0, audio_sample_count - audio_seg.shape[0]))
+        return torch.from_numpy(audio_seg.astype(np.float32, copy=False))
+
+    def _publish_video_segment(self, videos, start_frame):
+        if self.va_controller is None or self.va_controller.recorder is None:
+            return
+        frame_count = videos.shape[2]
+        if frame_count <= 0:
+            return
+        video_seg = videos[:, :, :frame_count].to(torch.float32)
+        comfy_video = wan_vae_to_comfy(video_seg.cpu())
+        audio_seg = self._get_audio_segment(start_frame, frame_count)
+        self.va_controller.pub_livestream(
+            comfy_video,
+            audio_seg,
+            video_seg.cpu(),
+            valid_duration=frame_count / self.target_fps,
+        )
+
     @ProfilingContext4DebugL1(
         "End run segment",
         recorder_mode=GET_RECORDER_MODE(),
@@ -692,9 +739,19 @@ def run_segment(self, segment_idx=0):
     def end_run_segment(self, segment_idx, latents):
         videos = self.run_vae_decoder(latents).cpu()
         if self.is_first_segment:
-            self.gen_video_list.append(videos)
+            output_videos = videos
+            output_start_frame = 0
         else:
-            self.gen_video_list.append(videos[:, :, self.current_motion_frames_num :])
+            output_videos = videos[:, :, self.current_motion_frames_num :]
+            output_start_frame = self.audio_start_idx + self.current_motion_frames_num
+
+        valid_frames = min(output_videos.shape[2], max(self.expected_frames - output_start_frame, 0))
+        if valid_frames > 0:
+            output_videos = output_videos[:, :, :valid_frames]
+            if self.stream_save_video:
+                self._publish_video_segment(output_videos, output_start_frame)
+            else:
+                self.gen_video_list.append(output_videos)
 
         if segment_idx < self.video_segment_num - 1:
             self.cond_frame = videos[:, :, -self.motion_frame :].to(torch.float32).to(AI_DEVICE)
@@ -706,8 +763,10 @@ def end_run_segment(self, segment_idx, latents):
 
     @ProfilingContext4DebugL2("Run DiT + decode")
     def run_main(self):
+        self.stream_save_video = self._should_stream_save_video()
         self.init_run()
         self.get_video_segment_num()
+        self._init_stream_video_controller()
 
         for segment_idx in range(self.video_segment_num):
             logger.info(f"start InfiniteTalk segment {segment_idx + 1}/{self.video_segment_num}")
@@ -716,11 +775,19 @@ def run_main(self):
                 latents = self.run_segment(segment_idx)
                 self.end_run_segment(segment_idx, latents)
 
+        if self.stream_save_video:
+            return self.process_images_after_vae_decoder()
+
         self.gen_video = torch.cat(self.gen_video_list, dim=2)[:, :, : self.expected_frames].to(torch.float32)
         return self.process_images_after_vae_decoder()
 
     @ProfilingContext4DebugL1("Process after vae decoder")
     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():
+                logger.info(f"Video saved to {self.input_info.save_result_path}")
+            return {"video": None}
+
         self.gen_video_final = wan_vae_to_comfy(self.gen_video)
         if self.input_info.return_result_tensor:
             return {"video": self.gen_video_final}
@@ -771,8 +838,13 @@ def _mux_audio(video_path, audio_path):
                 os.remove(tmp_path)
 
     def end_run(self):
+        if self.va_controller is not None:
+            self.va_controller.clear()
+            self.va_controller = None
         self._remove_video_audio_path()
         self._remove_cond_video_temp_path()
+        self.video_audio_array = None
+        self.stream_save_video = False
         if hasattr(self, "inputs"):
             del self.inputs
         torch.cuda.empty_cache()

scripts/infinitetalk/run_infinitetalk_single.sh

@@ -1,10 +1,10 @@
 #!/bin/bash
 
 # set path firstly
-lightx2v_path=/path/to/LightX2V
-model_path=/path/to/InfiniteTalk
+lightx2v_path=/data/nvme4/gushiqiao/new/debug/LightX2V
+model_path=/data/nvme5/gushiqiao/models/InfiniteTalk
 
-export CUDA_VISIBLE_DEVICES=0
+export CUDA_VISIBLE_DEVICES=7
 
 
 # set environment variables
@@ -14,7 +14,7 @@ python -m lightx2v.infer \
 --model_cls infinitetalk \
 --task s2v \
 --model_path $model_path \
---config_json ${lightx2v_path}/configs/infinitetalk/fp8/infinitetalk_single_distilled.json \
+--config_json ${lightx2v_path}/configs/infinitetalk/h100/infinitetalk_single_distilled.json \
 --prompt  "让角色根据音频内容自然说话" \
 --negative_prompt 色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走 \
 --image_path /data/nvme5/gushiqiao/cases/wecom-temp-3950334-bfa56035a08485356431b5a1c5c28a82.png \

scripts/infinitetalk/run_infinitetalk_single_dist.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+
+# set path firstly
+lightx2v_path=/data/nvme4/gushiqiao/new/debug/LightX2V
+model_path=/data/nvme5/gushiqiao/models/InfiniteTalk
+
+export CUDA_VISIBLE_DEVICES=7
+
+
+# set environment variables
+source ${lightx2v_path}/scripts/base/base.sh
+
+python -m lightx2v.infer \
+--model_cls infinitetalk \
+--task s2v \
+--model_path $model_path \
+--config_json ${lightx2v_path}/configs/infinitetalk/5090/infinitetalk_single_distilled_8gpus.json \
+--prompt  "让角色根据音频内容自然说话" \
+--negative_prompt 色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走 \
+--image_path /data/nvme5/gushiqiao/cases/wecom-temp-3950334-bfa56035a08485356431b5a1c5c28a82.png \
+--audio_path ${lightx2v_path}/assets/inputs/audio/seko_input.mp3 \
+--save_result_path ${lightx2v_path}/save_results/infinitetalk_single_720p_dist_8gpus.mp4 \
+--seed 42