update post process

lightx2v/models/runners/default_runner.py

@@ -522,7 +522,17 @@ def post_prompt_enhancer(self):
                     return enhanced_prompt
 
     def process_images_after_vae_decoder(self):
-        self.gen_video_final = wan_vae_to_comfy(self.gen_video_final)
+        return_result_tensor = self.input_info.return_result_tensor
+        save_result = self.input_info.save_result_path is not None
+        main_process = not dist.is_initialized() or dist.get_rank() == 0
+
+        should_process = return_result_tensor or (save_result and main_process)
+        if not should_process:
+            self.gen_video_final = None
+            return {"video": None}
+
+        with ProfilingContext4DebugL2("wan_vae_to_comfy"):
+            self.gen_video_final = wan_vae_to_comfy(self.gen_video_final)
 
         if "video_frame_interpolation" in self.config:
             assert self.vfi_model is not None and self.config["video_frame_interpolation"].get("target_fps", None) is not None
@@ -534,36 +544,38 @@ def process_images_after_vae_decoder(self):
                 target_fps=target_fps,
             )
 
-        if self.input_info.return_result_tensor:
+        if return_result_tensor:
+            self.gen_video_final = self.gen_video_final.cpu()
             return {"video": self.gen_video_final}
-        elif self.input_info.save_result_path is not None:
-            if "video_frame_interpolation" in self.config and self.config["video_frame_interpolation"].get("target_fps"):
-                fps = self.config["video_frame_interpolation"]["target_fps"]
-            else:
-                fps = self.config.get("fps", 16)
-
-            if not dist.is_initialized() or dist.get_rank() == 0:
-                out_path = self.input_info.save_result_path
-                img_in = (getattr(self.input_info, "image_path", None) or "").strip()
-                vid_in = (getattr(self.input_info, "video_path", None) or "").strip()
-                sr_from_image_only = self.config.get("task") == "sr" and bool(img_in) and not bool(vid_in)
-
-                if sr_from_image_only:
-                    logger.info("🖼 Start to save SR image (image_path input, no video_path) 🖼")
-                    save_to_image(self.gen_video_final, out_path)
-                    logger.info(f"✅ Image saved successfully to: {out_path} ✅")
-                else:
-                    logger.info(f"🎬 Start to save video 🎬")
-
-                    save_to_video(self.gen_video_final, out_path, fps=fps, method="ffmpeg")
-                    if self.config.get("task") in ("sr", "animate"):
-                        input_video_path = getattr(self.input_info, "video_path", "")
-                        if input_video_path:
-                            muxed_path = mux_audio_from_video(input_video_path, out_path)
-                            if muxed_path:
-                                logger.info(f"Audio muxed from input video: {input_video_path}")
-                    logger.info(f"✅ Video saved successfully to: {out_path} ✅")
-            return {"video": None}
+
+        # Reaching here means should_process was true because this is the main
+        # process and a save path was provided.
+        if "video_frame_interpolation" in self.config and self.config["video_frame_interpolation"].get("target_fps"):
+            fps = self.config["video_frame_interpolation"]["target_fps"]
+        else:
+            fps = self.config.get("fps", 16)
+
+        out_path = self.input_info.save_result_path
+        img_in = (getattr(self.input_info, "image_path", None) or "").strip()
+        vid_in = (getattr(self.input_info, "video_path", None) or "").strip()
+        sr_from_image_only = self.config.get("task") == "sr" and bool(img_in) and not bool(vid_in)
+
+        if sr_from_image_only:
+            logger.info("🖼 Start to save SR image (image_path input, no video_path) 🖼")
+            save_to_image(self.gen_video_final, out_path)
+            logger.info(f"✅ Image saved successfully to: {out_path} ✅")
+        else:
+            logger.info(f"🎬 Start to save video 🎬")
+
+            save_to_video(self.gen_video_final, out_path, fps=fps, method="ffmpeg")
+            if self.config.get("task") in ("sr", "animate"):
+                input_video_path = getattr(self.input_info, "video_path", "")
+                if input_video_path:
+                    muxed_path = mux_audio_from_video(input_video_path, out_path)
+                    if muxed_path:
+                        logger.info(f"Audio muxed from input video: {input_video_path}")
+            logger.info(f"✅ Video saved successfully to: {out_path} ✅")
+        return {"video": None}
 
     @ProfilingContext4DebugL1("RUN pipeline", recorder_mode=GET_RECORDER_MODE(), metrics_func=monitor_cli.lightx2v_worker_request_duration, metrics_labels=["DefaultRunner"])
     def run_pipeline(self, input_info):

lightx2v/models/runners/wan/wan_infinitetalk_runner.py

@@ -827,23 +827,35 @@ def process_images_after_vae_decoder(self):
                 logger.info(f"Video saved to {self.input_info.save_result_path}")
             return {"video": None}
 
+        return_result_tensor = self.input_info.return_result_tensor
+        save_result = self.input_info.save_result_path is not None
+        main_process = is_main_process()
+
+        should_process = return_result_tensor or (save_result and main_process)
+        if not should_process:
+            self.gen_video_final = None
+            return {"video": None}
+
         self.gen_video_final = wan_vae_to_comfy(self.gen_video)
-        if self.input_info.return_result_tensor:
+        if return_result_tensor:
+            self.gen_video_final = self.gen_video_final.cpu()
             return {"video": self.gen_video_final}
-        if self.input_info.save_result_path is not None and is_main_process():
-            out_path = self.input_info.save_result_path
-            os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
-            save_to_video(self.gen_video_final, out_path, fps=self.target_fps, method="ffmpeg")
-            # Prefer original input audio (full quality) over the 16kHz resampled wav
-            original_audio = getattr(self.input_info, "audio_path", None) or self.config.get("audio_path", "")
-            original_audio = original_audio.split(",")[0].strip() if original_audio else None
-            mux_audio = original_audio if (original_audio and os.path.isfile(original_audio)) else self.video_audio_path
-            if mux_audio and os.path.isfile(mux_audio):
-                try:
-                    self._mux_audio(out_path, mux_audio)
-                finally:
-                    self._remove_video_audio_path()
-            logger.info(f"Video saved to {out_path}")
+
+        # Reaching here means should_process was true because this is the main
+        # process and a save path was provided.
+        out_path = self.input_info.save_result_path
+        os.makedirs(os.path.dirname(out_path) or ".", exist_ok=True)
+        save_to_video(self.gen_video_final, out_path, fps=self.target_fps, method="ffmpeg")
+        # Prefer original input audio (full quality) over the 16kHz resampled wav
+        original_audio = getattr(self.input_info, "audio_path", None) or self.config.get("audio_path", "")
+        original_audio = original_audio.split(",")[0].strip() if original_audio else None
+        mux_audio = original_audio if (original_audio and os.path.isfile(original_audio)) else self.video_audio_path
+        if mux_audio and os.path.isfile(mux_audio):
+            try:
+                self._mux_audio(out_path, mux_audio)
+            finally:
+                self._remove_video_audio_path()
+        logger.info(f"Video saved to {out_path}")
         return {"video": None}
 
     @staticmethod

lightx2v/utils/utils.py

@@ -56,7 +56,7 @@ def save_videos_grid(videos: torch.Tensor, path: str, rescale=False, n_rows=1, f
         if rescale:
             x = (x + 1.0) / 2.0  # -1,1 -> 0,1
         x = torch.clamp(x, 0, 1)
-        x = (x * 255).numpy().astype(np.uint8)
+        x = (x * 255).to(torch.uint8).cpu().numpy()
         outputs.append(x)
 
     os.makedirs(os.path.dirname(path), exist_ok=True)
@@ -122,6 +122,7 @@ def vae_to_comfyui_image(vae_output: torch.Tensor) -> torch.Tensor:
     Returns:
         ComfyUI Image tensor in range [0, 1]
         Shape: [B, H, W, C] for single frame or [B*T, H, W, C] for video
+        Note: The returned tensor remains on the input device.
     """
     # Handle video tensor (5D) vs image tensor (4D)
     if vae_output.dim() == 5:
@@ -137,7 +138,7 @@ def vae_to_comfyui_image(vae_output: torch.Tensor) -> torch.Tensor:
     images = torch.clamp(images, 0, 1)
 
     # Convert from [B, C, H, W] to [B, H, W, C]
-    images = images.permute(0, 2, 3, 1).cpu()
+    images = images.permute(0, 2, 3, 1)
 
     return images
 
@@ -154,7 +155,7 @@ def vae_to_comfyui_image_inplace(vae_output: torch.Tensor) -> torch.Tensor:
     Returns:
         ComfyUI Image tensor in range [0, 1]
         Shape: [B, H, W, C] for single frame or [B*T, H, W, C] for video
-        Note: The returned tensor is the same object as input (modified in-place)
+        Note: The returned tensor remains on the input device.
     """
     # Handle video tensor (5D) vs image tensor (4D)
     if vae_output.dim() == 5:
@@ -169,8 +170,8 @@ def vae_to_comfyui_image_inplace(vae_output: torch.Tensor) -> torch.Tensor:
     # Clamp values to [0, 1] (inplace)
     vae_output.clamp_(0, 1)
 
-    # Convert from [B, C, H, W] to [B, H, W, C] and move to CPU
-    vae_output = vae_output.permute(0, 2, 3, 1).cpu()
+    # Convert from [B, C, H, W] to [B, H, W, C]
+    vae_output = vae_output.permute(0, 2, 3, 1)
 
     return vae_output
 
@@ -196,10 +197,10 @@ def wan_vae_to_comfy(vae_output: torch.Tensor) -> torch.Tensor:
         # Video: [B, C, T, H, W] -> [B, T, H, W, C]
         vae_output = vae_output.permute(0, 2, 3, 4, 1)
         # -> [B*T, H, W, C]
-        return vae_output.cpu().flatten(0, 1)
+        return vae_output.flatten(0, 1)
     else:
         # Image: [B, C, H, W] -> [B, H, W, C]
-        return vae_output.permute(0, 2, 3, 1).cpu()
+        return vae_output.permute(0, 2, 3, 1)
 
 
 def diffusers_vae_to_comfy(vae_output: torch.Tensor) -> torch.Tensor:
@@ -239,13 +240,11 @@ def save_to_video(
 
     if method == "imageio":
         # Convert to uint8
-        # frames = (images * 255).cpu().numpy().astype(np.uint8)
         frames = (images * 255).to(torch.uint8).cpu().numpy()
         imageio.mimsave(output_path, frames, fps=fps)  # type: ignore
 
     elif method == "ffmpeg":
         # Convert to numpy and scale to [0, 255]
-        # frames = (images * 255).cpu().numpy().clip(0, 255).astype(np.uint8)
         frames = (images * 255).clamp(0, 255).to(torch.uint8).cpu().numpy()
 
         # Convert RGB to BGR for OpenCV/FFmpeg