154 lines
6.0 KiB
Python
154 lines
6.0 KiB
Python
# !/usr/bin/env python
|
||
# -*- coding: UTF-8 -*-
|
||
|
||
import numpy as np
|
||
import torch
|
||
import os
|
||
import math
|
||
import weakref
|
||
import folder_paths
|
||
from comfy_api.latest import io, ui, Types
|
||
from comfy_api.input_impl import VideoFromFile
|
||
|
||
MAX_SEED = np.iinfo(np.int32).max
|
||
node_cr_path = os.path.dirname(os.path.abspath(__file__))
|
||
|
||
device = torch.device(
|
||
"cuda:0") if torch.cuda.is_available() else torch.device(
|
||
"mps") if torch.backends.mps.is_available() else torch.device(
|
||
"cpu")
|
||
|
||
weigths_FlashVSR_current_path = os.path.join(folder_paths.models_dir, "FlashVSR")
|
||
if not os.path.exists(weigths_FlashVSR_current_path):
|
||
os.makedirs(weigths_FlashVSR_current_path)
|
||
|
||
folder_paths.add_model_folder_path("FlashVSR", weigths_FlashVSR_current_path) # FlashVSR dir
|
||
|
||
|
||
# ---- 切分预览缓存：同一 trimmed 视频对象复用同一段临时预览文件，避免重复转码 ----
|
||
_trim_preview_results: "weakref.WeakKeyDictionary" = weakref.WeakKeyDictionary()
|
||
|
||
|
||
def _save_trim_preview(video) -> ui.PreviewVideo:
|
||
"""把迭代切分出的视频段落写入临时目录并返回预览（参考官方 VideoTrim 的 save_video_preview）"""
|
||
cached = _trim_preview_results.get(video)
|
||
if cached is not None and os.path.isfile(cached[0]):
|
||
return ui.PreviewVideo([cached[1]])
|
||
|
||
full_output_folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
||
"FlashVSR_temp_video", folder_paths.get_temp_directory(), 0, 0
|
||
)
|
||
file = f"{filename}_{counter:05}_.{Types.VideoContainer.get_extension(Types.VideoContainer.MP4)}"
|
||
full_path = os.path.join(full_output_folder, file)
|
||
video.save_to(
|
||
full_path,
|
||
format=Types.VideoContainer.MP4,
|
||
codec="auto",
|
||
preset="ultrafast",
|
||
)
|
||
result = ui.SavedResult(file, subfolder, io.FolderType.temp)
|
||
_trim_preview_results[video] = (full_path, result)
|
||
return ui.PreviewVideo([result])
|
||
|
||
|
||
class FlashVSR_SM_VideoTrimLoop(io.ComfyNode):
|
||
"""
|
||
可视化视频循环切分节点（参考官方 LoadVideo + VideoTrim）。
|
||
|
||
直接在节点上选择一个完整视频文件（无需预先切分），
|
||
每次迭代按 split_length（秒）输出下一段视频，输出到最后一段后从头循环。
|
||
"""
|
||
@classmethod
|
||
def __init__(cls):
|
||
cls.counters = {}
|
||
|
||
@classmethod
|
||
def define_schema(cls):
|
||
# 与官方 LoadVideo 一致：列出 input 目录下可播放的视频文件，并支持上传按钮
|
||
input_dir = folder_paths.get_input_directory()
|
||
files = [f for f in os.listdir(input_dir) if os.path.isfile(os.path.join(input_dir, f))]
|
||
files = folder_paths.filter_files_content_types(files, ["video"])
|
||
return io.Schema(
|
||
node_id="FlashVSR_SM_VideoTrimLoop",
|
||
display_name="FlashVSR_SM_VideoTrimLoop",
|
||
category="FlashVSR",
|
||
has_intermediate_output=True,
|
||
inputs=[
|
||
io.Combo.Input("video_file", options=sorted(files), upload=io.UploadType.video,
|
||
tooltip="选择完整视频文件，无需预先切分"),
|
||
io.Float.Input("split_length", default=5.0, min=0.1, max=3600.0, step=0.1,
|
||
tooltip="每次迭代输出的切分时长（秒）"),
|
||
io.Int.Input("seed", default=0, min=0, max=MAX_SEED),
|
||
io.Int.Input("reset_bool", default=0, optional=True,
|
||
tooltip="输入值 >= 1 时，迭代计数器归零（从头开始）"),
|
||
],
|
||
outputs=[
|
||
io.Video.Output(),
|
||
io.Custom("NUMBER").Output(display_name="number"),
|
||
io.Int.Output(display_name="seed"),
|
||
io.String.Output(display_name="filename"),
|
||
io.Float.Output(display_name="start_time"),
|
||
],
|
||
)
|
||
|
||
@classmethod
|
||
def execute(cls, video_file, split_length, seed, reset_bool=0) -> io.NodeOutput:
|
||
video_path = folder_paths.get_annotated_filepath(video_file)
|
||
if not os.path.isfile(video_path):
|
||
raise FileNotFoundError(f"Invalid video file: {video_file}")
|
||
|
||
# ---- 迭代计数器：以文件绝对路径为 key ----
|
||
UNIQUE_ID = os.path.normpath(video_path)
|
||
start = 0
|
||
step = 1
|
||
counter = start
|
||
if cls.counters.__contains__(UNIQUE_ID):
|
||
counter = cls.counters[UNIQUE_ID]
|
||
if round(reset_bool) >= 1:
|
||
counter = start
|
||
counter += step
|
||
cls.counters[UNIQUE_ID] = counter
|
||
result = int(counter)
|
||
|
||
source = VideoFromFile(video_path)
|
||
total_duration = float(source.get_duration())
|
||
split_length = max(float(split_length), 0.1)
|
||
num_segments = max(1, math.ceil(total_duration / split_length))
|
||
|
||
# ---- 本次迭代对应的切分段：超出总段数后回到第 1 段（循环） ----
|
||
segment_index = (result - 1) % num_segments
|
||
start_time = segment_index * split_length
|
||
seg_duration = min(split_length, max(total_duration - start_time, 0.1))
|
||
|
||
trimmed = source.as_trimmed(start_time, seg_duration, strict_duration=False)
|
||
if trimmed is None:
|
||
raise ValueError(
|
||
f"Failed to trim video:\nSource duration: {total_duration}\n"
|
||
f"Start time: {start_time}\nTarget duration: {seg_duration}"
|
||
)
|
||
|
||
print(f"[FlashVSR_SM_VideoTrimLoop] segment {segment_index + 1}/{num_segments}: "
|
||
f"start={start_time:.3f}s duration={seg_duration:.3f}s ({video_file})")
|
||
|
||
filename = os.path.basename(video_path)
|
||
return io.NodeOutput(
|
||
trimmed, result, seed, filename, float(start_time),
|
||
ui=_save_trim_preview(trimmed),
|
||
)
|
||
|
||
@classmethod
|
||
def validate_inputs(cls, video_file, **kwargs):
|
||
if not folder_paths.exists_annotated_filepath(video_file):
|
||
return "Invalid video file: {}".format(video_file)
|
||
return True
|
||
|
||
@classmethod
|
||
def fingerprint_inputs(cls, **kwargs):
|
||
# 每次入队都重新执行，保证迭代计数器推进
|
||
return ""
|
||
|
||
|
||
|
||
|
||
|