From 1da61ed0757fb99d67d1b37e9895900f465a0ec9 Mon Sep 17 00:00:00 2001 From: gaclove Date: Mon, 8 Jun 2026 01:13:41 +0800 Subject: [PATCH] fix(seedvr2): improve runtime cleanup and preview compatibility --- lightx2v | 2 +- nodes/seedvr.py | 43 +++++++++++++++++++++++++++++++++---------- 2 files changed, 34 insertions(+), 11 deletions(-) diff --git a/lightx2v b/lightx2v index f32c5db..5fc9f69 160000 --- a/lightx2v +++ b/lightx2v @@ -1 +1 @@ -Subproject commit f32c5dbc83daf8b7ba4fdbc00d7a8b40ca82e4eb +Subproject commit 5fc9f69495afa9831164b6a2c5892ba0bab3674a diff --git a/nodes/seedvr.py b/nodes/seedvr.py index c5d6142..945bf2f 100644 --- a/nodes/seedvr.py +++ b/nodes/seedvr.py @@ -15,10 +15,12 @@ import argparse import gc import logging import math +import shutil import subprocess import tempfile import types import wave +from collections.abc import Mapping from pathlib import Path import folder_paths @@ -103,8 +105,8 @@ def _output_video_full_path(filename, validate_exists=True): def _audio_to_wav(audio, wav_path): - if not isinstance(audio, dict): - logger.info("[LightX2VOutputVideoPreview] skip audio mux: optional AUDIO input is empty") + if not isinstance(audio, Mapping): + logger.info("[LightX2VOutputVideoPreview] skip audio mux: unsupported AUDIO input type=%s", type(audio).__name__) return False if audio.get("waveform") is None or audio.get("sample_rate") is None: logger.info("[LightX2VOutputVideoPreview] skip audio mux: AUDIO has no waveform/sample_rate") @@ -136,6 +138,11 @@ def _audio_to_wav(audio, wav_path): return True +def _audio_mux_path(video_path): + video_path = Path(video_path) + return video_path.with_name(f"{video_path.stem}-audio{video_path.suffix}") + + def _mux_audio_into_video(video_path, audio): from imageio_ffmpeg import get_ffmpeg_exe @@ -145,9 +152,10 @@ def _mux_audio_into_video(video_path, audio): with tempfile.TemporaryDirectory(prefix=".lightx2v_audio_mux.", dir=str(video_path.parent)) as tmp_dir: wav_path = Path(tmp_dir) / "audio.wav" - muxed_path = Path(tmp_dir) / "muxed.mp4" + muxed_tmp_path = Path(tmp_dir) / "muxed.mp4" + muxed_path = _audio_mux_path(video_path) if not _audio_to_wav(audio, wav_path): - return False + return None command = [ get_ffmpeg_exe(), @@ -167,14 +175,16 @@ def _mux_audio_into_video(video_path, audio): "-b:a", "192k", "-shortest", - str(muxed_path), + str(muxed_tmp_path), ] process = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, check=False) if process.returncode != 0: stderr = process.stderr.decode("utf-8", errors="replace").strip() raise RuntimeError(f"FFmpeg audio mux failed: {stderr}") - muxed_path.replace(video_path) - return True + shutil.copy2(muxed_tmp_path, muxed_path) + muxed_tmp_path.replace(video_path) + logger.info("[LightX2VOutputVideoPreview] muxed audio into %s and updated %s", muxed_path, video_path) + return muxed_path def _install_tensor_input_shim(runner, frames_u8, fps): @@ -182,6 +192,9 @@ def _install_tensor_input_shim(runner, frames_u8, fps): frames_u8: torch.uint8 [T, C, H, W] on CPU (same format as torchvision.io.read_video output). """ + if not hasattr(runner, "_lightx2v_original_run_input_encoder_local_sr"): + runner._lightx2v_original_run_input_encoder_local_sr = runner._run_input_encoder_local_sr.__func__ + runner._tensor_input = frames_u8 runner._tensor_input_fps = float(fps) @@ -196,7 +209,7 @@ def _install_tensor_input_shim(runner, frames_u8, fps): return torch.empty(0, 3, 0, 0, dtype=torch.uint8) return seg - original_encoder = runner._run_input_encoder_local_sr.__func__ + original_encoder = runner._lightx2v_original_run_input_encoder_local_sr def _run_input_encoder_local_sr(self): if getattr(self, "_sr_segment", None) is None: @@ -213,6 +226,12 @@ def _install_tensor_input_shim(runner, frames_u8, fps): runner.run_input_encoder = runner._run_input_encoder_local_sr +def _clear_tensor_input_shim(runner): + for attr in ("_tensor_input", "_tensor_input_fps"): + if hasattr(runner, attr): + delattr(runner, attr) + + class LightX2VSeedVR2Loader: """Load a SeedVR2 DiT checkpoint from models/lightx2v/seedvr2/.""" @@ -503,6 +522,7 @@ class LightX2VSeedVR2Sampler: try: result = runner.run_pipeline(input_info) finally: + _clear_tensor_input_shim(runner) torch.cuda.empty_cache() gc.collect() @@ -561,8 +581,11 @@ class LightX2VOutputVideoPreview: CATEGORY = "LightX2V/Output" def preview(self, filename, validate_exists, mux_audio, audio=None): + preview_filename = filename if bool(mux_audio) and audio is not None: video_path = _output_video_full_path(filename, bool(validate_exists)) - _mux_audio_into_video(video_path, audio) - file_info, relative_name = _output_video_file_info(filename, bool(validate_exists)) + muxed_path = _mux_audio_into_video(video_path, audio) + if muxed_path is not None: + preview_filename = str(muxed_path) + file_info, relative_name = _output_video_file_info(preview_filename, bool(validate_exists)) return {"ui": {"images": [file_info], "animated": (True,)}, "result": (relative_name,)}