From ef7d6df232bcbb01a5f59b4d589ef12cffb85e5e Mon Sep 17 00:00:00 2001 From: gaclove Date: Tue, 1 Sep 2026 15:09:01 +0800 Subject: [PATCH] feat(vr): add file-based video samplers --- nodes/__init__.py | 12 +++- nodes/file_input.py | 88 +++++++++++++++++++++++ nodes/seedvr.py | 167 ++++++++++++++++++++++++++++++++++++++++++++ nodes/swiftvr.py | 112 +++++++++++++++++++++++++++++ 4 files changed, 377 insertions(+), 2 deletions(-) create mode 100644 nodes/file_input.py diff --git a/nodes/__init__.py b/nodes/__init__.py index 5ca501b..47c6d64 100644 --- a/nodes/__init__.py +++ b/nodes/__init__.py @@ -6,6 +6,7 @@ Each submodule groups a category of nodes: - ``talk`` : talk-object input/combiner nodes - ``combiner`` : config combiners (V1/V2/V3) that aggregate the above - ``inference`` : the modular inference runner +- ``file_input``: validated, upload-backed media paths - ``seedvr`` : SeedVR2 super-resolution runner - ``swiftvr`` : SwiftVR restoration runner """ @@ -20,10 +21,11 @@ from .config import ( LightX2VQuantization, LightX2VTeaCache, ) +from .file_input import LightX2VInputVideoPath from .inference import LightX2VModularInferenceV2 from .lora import LightX2VLoRALoader -from .seedvr import LightX2VOutputVideoPreview, LightX2VSeedVR2Loader, LightX2VSeedVR2Sampler -from .swiftvr import LightX2VSwiftVRLoader, LightX2VSwiftVRSampler +from .seedvr import LightX2VOutputVideoPreview, LightX2VSeedVR2FileSampler, LightX2VSeedVR2Loader, LightX2VSeedVR2Sampler +from .swiftvr import LightX2VSwiftVRFileSampler, LightX2VSwiftVRLoader, LightX2VSwiftVRSampler from .talk import ( TalkObjectInput, TalkObjectsCombiner, @@ -40,10 +42,13 @@ NODE_CLASS_MAPPINGS = { "LightX2VConfigCombinerV2": LightX2VConfigCombinerV2, "LightX2VConfigCombinerV3": LightX2VConfigCombinerV3, "LightX2VModularInferenceV2": LightX2VModularInferenceV2, + "LightX2VInputVideoPath": LightX2VInputVideoPath, "LightX2VSeedVR2Loader": LightX2VSeedVR2Loader, "LightX2VSeedVR2Sampler": LightX2VSeedVR2Sampler, + "LightX2VSeedVR2FileSampler": LightX2VSeedVR2FileSampler, "LightX2VSwiftVRLoader": LightX2VSwiftVRLoader, "LightX2VSwiftVRSampler": LightX2VSwiftVRSampler, + "LightX2VSwiftVRFileSampler": LightX2VSwiftVRFileSampler, "LightX2VOutputVideoPreview": LightX2VOutputVideoPreview, "LightX2VTalkObjectInput": TalkObjectInput, "LightX2VTalkObjectsCombiner": TalkObjectsCombiner, @@ -60,10 +65,13 @@ NODE_DISPLAY_NAME_MAPPINGS = { "LightX2VConfigCombinerV2": "LightX2V Config Combiner V2", "LightX2VConfigCombinerV3": "LightX2V Config Combiner V3", "LightX2VModularInferenceV2": "LightX2V Modular Inference V2", + "LightX2VInputVideoPath": "LightX2V Input Video Path (Upload)", "LightX2VSeedVR2Loader": "LightX2V SeedVR2 Loader", "LightX2VSeedVR2Sampler": "LightX2V SeedVR2 Sampler", + "LightX2VSeedVR2FileSampler": "LightX2V SeedVR2 File Sampler", "LightX2VSwiftVRLoader": "LightX2V SwiftVR Loader", "LightX2VSwiftVRSampler": "LightX2V SwiftVR Sampler", + "LightX2VSwiftVRFileSampler": "LightX2V SwiftVR File Sampler", "LightX2VOutputVideoPreview": "LightX2V Output Video Preview", "LightX2VTalkObjectInput": "LightX2V Talk Object Input (Single)", "LightX2VTalkObjectsCombiner": "LightX2V Talk Objects Combiner", diff --git a/nodes/file_input.py b/nodes/file_input.py new file mode 100644 index 0000000..ff6768e --- /dev/null +++ b/nodes/file_input.py @@ -0,0 +1,88 @@ +"""Validated file-path inputs for native LightX2V media runners.""" + +from pathlib import Path + +import folder_paths + + +def _input_video_files(): + input_dir = folder_paths.get_input_directory() + files, _ = folder_paths.recursive_search(input_dir) + return sorted(folder_paths.filter_files_content_types(files, ["video"])) + + +def resolve_input_video_path(filename) -> Path: + """Resolve a ComfyUI input filename without allowing directory escape.""" + + raw = str(filename or "").strip() + if not raw: + raise ValueError("video is required") + + input_dir = Path(folder_paths.get_input_directory()).resolve() + candidate = Path(folder_paths.get_annotated_filepath(raw)).resolve() + try: + candidate.relative_to(input_dir) + except ValueError as exc: + raise ValueError(f"Expected a video under ComfyUI input, got: {filename}") from exc + + if not candidate.is_file(): + raise FileNotFoundError(f"Input video does not exist: {candidate}") + if not folder_paths.filter_files_content_types([candidate.name], ["video"]): + raise ValueError(f"Input file is not recognized as video: {candidate}") + return candidate + + +def probe_video_file(video_path: Path): + """Read only video metadata and the first frame dimensions via decord.""" + + from decord import VideoReader + + reader = VideoReader(str(video_path)) + if len(reader) < 1: + raise ValueError(f"Input video contains no frames: {video_path}") + first_frame = reader[0] + height, width = int(first_frame.shape[0]), int(first_frame.shape[1]) + fps = float(reader.get_avg_fps() or 0.0) + if height <= 0 or width <= 0: + raise ValueError(f"Invalid input video dimensions: {width}x{height}") + return width, height, fps + + +class LightX2VInputVideoPath: + """Upload/select a video under ComfyUI input and expose its absolute path.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "video": ( + _input_video_files(), + { + "video_upload": True, + "tooltip": "Upload or select a video under ComfyUI input. The absolute path is resolved only while executing.", + }, + ) + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("video_path",) + FUNCTION = "resolve" + CATEGORY = "LightX2V/Input" + + def resolve(self, video): + return (str(resolve_input_video_path(video)),) + + @classmethod + def IS_CHANGED(cls, video): + path = resolve_input_video_path(video) + stat = path.stat() + return f"{stat.st_mtime_ns}:{stat.st_size}" + + @classmethod + def VALIDATE_INPUTS(cls, video): + try: + resolve_input_video_path(video) + except (OSError, ValueError) as exc: + return str(exc) + return True diff --git a/nodes/seedvr.py b/nodes/seedvr.py index bf2de7a..89f199b 100644 --- a/nodes/seedvr.py +++ b/nodes/seedvr.py @@ -5,6 +5,8 @@ Split into: models/lightx2v/seedvr2/, load it into VRAM, return a SEEDVR_MODEL handle. - LightX2VSeedVR2Sampler: takes SEEDVR_MODEL + IMAGE + per-call params, returns upscaled IMAGE frames. +- LightX2VSeedVR2FileSampler: takes a validated input video path and streams + the restored result to ComfyUI output while preserving source audio. The sampler installs a small shim on the runner so input frames come from the IMAGE tensor (no temp file, no re-encode); the runner's segmenting logic still @@ -18,6 +20,7 @@ import math import shutil import subprocess import tempfile +import threading import types import wave from collections.abc import Mapping @@ -27,8 +30,12 @@ import folder_paths import torch from comfy.utils import ProgressBar +from .file_input import probe_video_file, resolve_input_video_path + logger = logging.getLogger(__name__) +_SEEDVR_RUN_LOCK = threading.Lock() + def _seedvr2_model_dir() -> Path: return Path(folder_paths.models_dir) / "lightx2v" / "seedvr2" @@ -194,6 +201,8 @@ def _install_tensor_input_shim(runner, frames_u8, fps): """ 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__ + if not hasattr(runner, "_lightx2v_original_run_input_encoder"): + runner._lightx2v_original_run_input_encoder = runner.run_input_encoder runner._tensor_input = frames_u8 runner._tensor_input_fps = float(fps) @@ -230,6 +239,12 @@ def _clear_tensor_input_shim(runner): for attr in ("_tensor_input", "_tensor_input_fps"): if hasattr(runner, attr): delattr(runner, attr) + for attr in ("_probe_video", "_read_video_segment", "_run_input_encoder_local_sr"): + if attr in runner.__dict__: + delattr(runner, attr) + original_run_input_encoder = getattr(runner, "_lightx2v_original_run_input_encoder", None) + if original_run_input_encoder is not None: + runner.run_input_encoder = original_run_input_encoder class LightX2VSeedVR2Loader: @@ -545,6 +560,158 @@ class LightX2VSeedVR2Sampler: return (video, "") +class LightX2VSeedVR2FileSampler: + """Run SeedVR2 on an input video using segmented file I/O.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("SEEDVR_MODEL",), + "video_path": ( + "STRING", + { + "forceInput": True, + "tooltip": "Absolute path produced by LightX2V Input Video Path. The file must remain under ComfyUI input.", + }, + ), + "target_width": ("INT", {"default": 1920, "min": 64, "max": 7680, "step": 8}), + "target_height": ( + "INT", + { + "default": 1080, + "min": 64, + "max": 4320, + "step": 8, + "tooltip": "Target output frame height.", + }, + ), + "infer_steps": ("INT", {"default": 1, "min": 1, "max": 50}), + "segment_length": ( + "INT", + { + "default": 81, + "min": 16, + "max": 512, + "step": 1, + "tooltip": "Frames decoded and restored per segment. Long videos do not materialize as a full IMAGE batch.", + }, + ), + "segment_overlap": ("INT", {"default": 1, "min": 0, "max": 32}), + "seed": ("INT", {"default": 42, "min": 0, "max": 2**32 - 1}), + "filename_prefix": ("STRING", {"default": "lightx2v_seedvr2/SeedVR2"}), + "color_fix": ( + ["gpu", "off", "cpu"], + { + "default": "gpu", + "tooltip": "SeedVR color correction after VAE decode.", + }, + ), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + FUNCTION = "sample" + CATEGORY = "LightX2V/SeedVR" + + def sample( + self, + model, + video_path, + target_height, + target_width, + infer_steps, + segment_length, + segment_overlap, + seed, + filename_prefix, + color_fix, + ): + from ..lightx2v.lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict + + input_path = resolve_input_video_path(video_path) + source_width, source_height, source_fps = probe_video_file(input_path) + effective_fps = source_fps if source_fps > 0 else 16.0 + source_geom = math.sqrt(source_height * source_width) + target_geom = math.sqrt(int(target_height) * int(target_width)) + sr_ratio = max(target_geom / source_geom, 1.0) if source_geom > 0 else 1.0 + if target_geom < source_geom: + logger.warning( + "[SeedVR2FileSampler] target (%sx%s) is smaller than input (%sx%s); SR will run at input scale before final sizing.", + target_width, + target_height, + source_width, + source_height, + ) + + full_path, output_file, output_subfolder = _prepare_output_video(filename_prefix, target_width, target_height) + save_path = str(full_path) + input_info = init_empty_input_info("sr") + update_input_info_from_dict( + input_info, + { + "video_path": str(input_path), + "image_path": "", + "prompt": "", + "negative_prompt": "", + "seed": int(seed), + "sr_ratio": float(sr_ratio), + "save_result_path": save_path, + "return_result_tensor": False, + }, + ) + + runner = model["runner"] + progress = ProgressBar(100) + logger.info( + "[SeedVR2FileSampler] input=%s (%sx%s @ %.3f fps), target=%sx%s, segment=%s/%s", + input_path, + source_width, + source_height, + source_fps, + target_width, + target_height, + segment_length, + segment_overlap, + ) + try: + with _SEEDVR_RUN_LOCK: + _clear_tensor_input_shim(runner) + runner.set_config( + { + "sr_ratio": float(sr_ratio), + "target_height": int(target_height), + "target_width": int(target_width), + "target_video_length": int(segment_length), + "sr_segment_length": int(segment_length), + "sr_overlap": int(segment_overlap), + "stream_save_video": True, + "infer_steps": int(infer_steps), + "seed": int(seed), + "fps": float(effective_fps), + "video_path": str(input_path), + "image_path": "", + "prompt": "", + "negative_prompt": "", + "save_result_path": save_path, + "return_result_tensor": False, + "color_fix": str(color_fix), + } + ) + if hasattr(runner, "set_progress_callback"): + runner.set_progress_callback(lambda current, _total: progress.update_absolute(current)) + runner.run_pipeline(input_info) + finally: + torch.cuda.empty_cache() + gc.collect() + + if not Path(save_path).is_file(): + raise RuntimeError(f"SeedVR2 did not create expected output video: {save_path}") + relative_name = str(Path(output_subfolder) / output_file) if output_subfolder else output_file + return (relative_name,) + + class LightX2VOutputVideoPreview: """Expose an existing ComfyUI output video to the history/view API.""" diff --git a/nodes/swiftvr.py b/nodes/swiftvr.py index a30d56e..bf3afcf 100644 --- a/nodes/swiftvr.py +++ b/nodes/swiftvr.py @@ -12,6 +12,7 @@ import folder_paths import torch from comfy.utils import ProgressBar +from .file_input import probe_video_file, resolve_input_video_path from .seedvr import _prepare_output_video logger = logging.getLogger(__name__) @@ -430,3 +431,114 @@ class LightX2VSwiftVRSampler: return (placeholder, relative_name) return (memory_writer.as_images().clamp_(0.0, 1.0), "") + + +class LightX2VSwiftVRFileSampler: + """Restore an input video from disk without materializing it as IMAGE.""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("SWIFTVR_MODEL",), + "video_path": ( + "STRING", + { + "forceInput": True, + "tooltip": "Absolute path produced by LightX2V Input Video Path. The file must remain under ComfyUI input.", + }, + ), + "target_short_edge": ( + "INT", + { + "default": 1080, + "min": 64, + "max": _MAX_OUTPUT_DIMENSION, + "step": 8, + "tooltip": "Output short edge. SwiftVR preserves aspect ratio.", + }, + ), + "filename_prefix": ("STRING", {"default": "lightx2v_swiftvr/SwiftVR"}), + "video_codec": (["libx265", "libx264"], {"default": "libx265"}), + "quality": ("INT", {"default": 60, "min": 0, "max": 100, "step": 1}), + "ffmpeg_preset": ( + ["ultrafast", "superfast", "veryfast", "faster", "fast", "medium", "slow"], + {"default": "ultrafast"}, + ), + } + } + + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("filename",) + FUNCTION = "sample" + CATEGORY = "LightX2V/SwiftVR" + + def sample(self, model, video_path, target_short_edge, filename_prefix, video_codec, quality, ffmpeg_preset): + from ..lightx2v.lightx2v.utils.input_info import init_empty_input_info, update_input_info_from_dict + + input_path = resolve_input_video_path(video_path) + raw_width, raw_height, source_fps = probe_video_file(input_path) + source_height, source_width = raw_height // 8 * 8, raw_width // 8 * 8 + if source_height <= 0 or source_width <= 0: + raise ValueError(f"SwiftVR input is too small after 8-pixel alignment: {raw_height}x{raw_width}") + + output_height, output_width, sr_ratio = _resolve_output_size( + source_width, + source_height, + target_short_edge, + require_even=True, + ) + full_path, output_file, output_subfolder = _prepare_output_video(filename_prefix, output_width, output_height) + save_path = str(full_path) + target_shape = [output_height, output_width] + + input_info = init_empty_input_info("sr") + update_input_info_from_dict( + input_info, + { + "video_path": str(input_path), + "image_path": "", + "sr_ratio": float(sr_ratio), + "target_shape": target_shape, + "save_result_path": save_path, + "return_result_tensor": False, + }, + ) + + runner = model["runner"] + progress = ProgressBar(100) + logger.info( + "[SwiftVRFileSampler] input=%s (%sx%s @ %.3f fps), output=%sx%s", + input_path, + source_width, + source_height, + source_fps, + output_width, + output_height, + ) + try: + with _SWIFTVR_RUN_LOCK: + runner.set_config( + { + "fps": 0.0, + "video_codec": str(video_codec), + "quality": int(quality), + "ffmpeg_preset": str(ffmpeg_preset), + "video_path": str(input_path), + "image_path": "", + "sr_ratio": float(sr_ratio), + "target_shape": target_shape, + "return_result_tensor": False, + } + ) + if hasattr(runner, "set_progress_callback"): + runner.set_progress_callback(lambda current, _total: progress.update_absolute(current)) + runner.run_pipeline(input_info) + finally: + torch.cuda.empty_cache() + gc.collect() + + if not Path(save_path).is_file(): + raise RuntimeError(f"SwiftVR did not create expected output video: {save_path}") + relative_name = str(Path(output_subfolder) / output_file) if output_subfolder else output_file + return (relative_name,)