feat(vr): add file-based video samplers

This commit is contained in:
gaclove
2026-09-01 15:09:01 +08:00
parent f061716fc2
commit ef7d6df232
4 changed files with 377 additions and 2 deletions
+10 -2
View File
@@ -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",
+88
View File
@@ -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
+167
View File
@@ -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."""
+112
View File
@@ -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,)