feat(vr): add file-based video samplers
This commit is contained in:
+10
-2
@@ -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",
|
||||
|
||||
@@ -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
@@ -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."""
|
||||
|
||||
|
||||
@@ -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,)
|
||||
|
||||
Reference in New Issue
Block a user