update post process#1263
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes video processing after the VAE decoder in both default_runner.py and wan_infinitetalk_runner.py by skipping processing on non-main processes or when results are not needed. It also refactors utility functions in utils.py to keep tensors on their input device rather than moving them to the CPU prematurely. The reviewer recommends clearing internal references to self.gen_video and self.gen_video_final across all exit paths in both runners to prevent GPU/CPU memory leaks, and suggests checking dist.is_available() before dist.is_initialized() for better robustness.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if not should_process: | ||
| self.gen_video_final = None | ||
| return {"video": None} |
There was a problem hiding this comment.
When should_process is False, only self.gen_video_final is cleared. However, self.gen_video still holds a reference to the large decoded video tensor, which will leak GPU memory until the next request. We should clear both self.gen_video and self.gen_video_final to free up GPU memory immediately.
| if not should_process: | |
| self.gen_video_final = None | |
| return {"video": None} | |
| if not should_process: | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": None} |
| if return_result_tensor: | ||
| self.gen_video_final = self.gen_video_final.cpu() | ||
| return {"video": self.gen_video_final} |
There was a problem hiding this comment.
To prevent the runner from holding onto the large video tensors after returning them to the caller, we should clear the internal references self.gen_video and self.gen_video_final before returning.
| if return_result_tensor: | |
| self.gen_video_final = self.gen_video_final.cpu() | |
| return {"video": self.gen_video_final} | |
| if return_result_tensor: | |
| self.gen_video_final = self.gen_video_final.cpu() | |
| video_out = self.gen_video_final | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": video_out} |
| 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} |
There was a problem hiding this comment.
After successfully saving the video, the runner still holds references to self.gen_video and self.gen_video_final in its state. This causes a GPU/CPU memory leak until the next request starts. We should clear these references before returning.
| return {"video": None} | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": None} |
| if not should_process: | ||
| self.gen_video_final = None | ||
| return {"video": None} |
There was a problem hiding this comment.
When should_process is False, self.gen_video still holds a reference to the large decoded video tensor, causing a memory leak. We should clear both self.gen_video and self.gen_video_final to free up memory immediately.
| if not should_process: | |
| self.gen_video_final = None | |
| return {"video": None} | |
| if not should_process: | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": None} |
| if return_result_tensor: | ||
| self.gen_video_final = self.gen_video_final.cpu() | ||
| return {"video": self.gen_video_final} |
There was a problem hiding this comment.
To prevent the runner from holding onto the large video tensors after returning them to the caller, we should clear the internal references self.gen_video and self.gen_video_final before returning.
| if return_result_tensor: | |
| self.gen_video_final = self.gen_video_final.cpu() | |
| return {"video": self.gen_video_final} | |
| if return_result_tensor: | |
| self.gen_video_final = self.gen_video_final.cpu() | |
| video_out = self.gen_video_final | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": video_out} |
| finally: | ||
| self._remove_video_audio_path() | ||
| logger.info(f"Video saved to {out_path}") | ||
| return {"video": None} |
There was a problem hiding this comment.
After successfully saving the video, the runner still holds references to self.gen_video and self.gen_video_final in its state. This causes a memory leak until the next request starts. We should clear these references before returning.
| return {"video": None} | |
| self.gen_video = None | |
| self.gen_video_final = None | |
| return {"video": None} |
| 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 |
There was a problem hiding this comment.
Using dist.is_initialized() directly might fail if the distributed package is not available on the platform. It is safer and more robust to check dist.is_available() first, consistent with the helper function is_main_process().
| main_process = not dist.is_initialized() or dist.get_rank() == 0 | |
| main_process = not dist.is_available() or not dist.is_initialized() or dist.get_rank() == 0 |
2 participants
No description provided.