import os import subprocess import folder_paths import re from pathlib import Path class NVVFRNode: """ ComfyUI NVVFR节点 - 使用NVEncC进行视频超分补帧 """ def __init__(self): self.nvenc_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "nvenc", "NVEncC64.exe") self.output_dir = folder_paths.get_output_directory() @classmethod def INPUT_TYPES(cls): return { "required": { "output_prefix": ("STRING", {"default": "ComfyUI"}), "codec": (["h264", "h265"], {"default": "h265"}), "quality": (["high", "medium", "low"], {"default": "high"}), "enable_superres": ("BOOLEAN", {"default": True}), "superres_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1}), "enable_double_frame": ("BOOLEAN", {"default": False}), "width": ("INT", {"default": 1920, "min": 64, "max": 8192}), "height": ("INT", {"default": 1080, "min": 64, "max": 8192}), "depth": (["8bit", "10bit"], {"default": "10bit"}), }, "optional": { "video_path": ("STRING", {"default": ""}), "imagelist": ("STRING", {"default": ""}), "original_frame_rate": ("FLOAT", {"default": 30.0, "min": 1.0, "max": 120.0, "step": 0.1}), } } RETURN_TYPES = ("STRING",) RETURN_NAMES = ("output_path",) FUNCTION = "process_video" CATEGORY = "video/processing" OUTPUT_NODE = True def process_video(self, video_path, output_prefix, codec, quality, enable_superres, enable_double_frame, width=1920, height=1080, depth="10bit", superres_strength=1.0, imagelist="", original_frame_rate=30.0): """ 处理视频文件或图像序列 """ # 检查NVEncC64.exe是否存在 if not os.path.exists(self.nvenc_path): raise FileNotFoundError(f"NVEncC64.exe not found at {self.nvenc_path}") # 检查video_path和imagelist不能同时为空 if not video_path and not imagelist: raise ValueError("Either video_path or imagelist must be provided") # 如果有imagelist，检查original_frame_rate是否为空 if imagelist and original_frame_rate <= 0: raise ValueError("original_frame_rate must be provided when imagelist is used") # 如果有video_path，检查文件是否存在 if video_path and not os.path.exists(video_path): raise FileNotFoundError(f"Input video file not found: {video_path}") avs_file_path = None # 如果有imagelist，生成avs文件 if imagelist: # 解析imagelist（假设是逗号分隔的图像路径） image_paths = [path.strip() for path in imagelist.split(',') if path.strip()] if not image_paths: raise ValueError("imagelist must contain valid image paths") # 检查图像文件是否存在 for img_path in image_paths: if not os.path.exists(img_path): raise FileNotFoundError(f"Image file not found: {img_path}") # 在ComfyUI临时文件夹创建avs文件 temp_dir = folder_paths.get_temp_directory() avs_file_path = os.path.join(temp_dir, f"{output_prefix}_input.avs") # 生成avs文件内容（UTF-8编码） # 使用ImageSource处理图像序列 if len(image_paths) == 1: # 单张图像 avs_content = f"ImageSource(\"{image_paths[0]}\", fps={original_frame_rate})\n" else: # 多张图像序列 # 提取公共路径和文件名模式 first_path = image_paths[0] last_path = image_paths[-1] # 尝试提取文件名模式（假设文件名有数字序列） import glob dir_path = os.path.dirname(first_path) if not dir_path: dir_path = "." # 获取目录中所有匹配的图像文件 pattern = os.path.join(dir_path, "*.png") # 支持PNG image_files = sorted(glob.glob(pattern)) if not image_files: pattern = os.path.join(dir_path, "*.jpg") # 支持JPG image_files = sorted(glob.glob(pattern)) if not image_files: pattern = os.path.join(dir_path, "*.jpeg") # 支持JPEG image_files = sorted(glob.glob(pattern)) if not image_files: pattern = os.path.join(dir_path, "*.bmp") # 支持BMP image_files = sorted(glob.glob(pattern)) if image_files: # 使用ImageSource的序列模式 avs_content = f"ImageSource(\"{image_files[0]}\", end={len(image_files)}, fps={original_frame_rate})\n" else: # 回退到使用Interleave avs_content = "Interleave(\n" for i, img_path in enumerate(image_paths): avs_content += f" ImageSource(\"{img_path}\", fps={original_frame_rate})" if i < len(image_paths) - 1: avs_content += ",\n" else: avs_content += "\n" avs_content += ")\n" # 写入avs文件（UTF-8编码） with open(avs_file_path, 'w', encoding='utf-8') as f: f.write(avs_content) print(f"Generated AVS file: {avs_file_path}") # 检查编码器和深度的兼容性 if codec.lower() == "h264" and depth.lower() == "10bit": print("Warning: H.264 codec does not support 10-bit depth. Forcing to 8-bit depth.") depth = "8bit" max_counter = 0 # Loop through the existing files matcher = re.compile(f"{re.escape(output_prefix)}_(\\d+)\\D*\\..+", re.IGNORECASE) for existing_file in os.listdir(self.output_dir): # Check if the file matches the expected format match = matcher.fullmatch(existing_file) if match: # Extract the numeric portion of the filename file_counter = int(match.group(1)) # Update the maximum counter value if necessary if file_counter > max_counter: max_counter = file_counter counter = max_counter + 1 file_name = f"{output_prefix}_{counter:05}" output_filename = f"{output_prefix}_{counter:05}.mp4" output_path = os.path.join(self.output_dir, output_filename) # 构建命令行参数 cmd = [self.nvenc_path] # 添加编码器参数 if codec.lower() == "h265": cmd.extend(["-c", "hevc"]) else: # h264 cmd.extend(["-c", "h264"]) # 添加质量参数 if quality.lower() == "quality": cmd.extend(["-u", "P7"]) elif quality.lower() == "medium": cmd.extend(["-u", "P4"]) else: cmd.extend(["-u", "P1"]) # 添加输出分辨率 output_resolution = f"{width}x{height}" cmd.extend(["--output-res", output_resolution]) # 添加输出深度 if depth.lower() == "8bit": cmd.extend(["--output-depth", "8"]) else: cmd.extend(["--output-depth", "10"]) # 添加超分辨率参数 if enable_superres: cmd.extend(["--vpp-resize", f"algo=nvvfx-superres,superres-mode=1,superres-strength={superres_strength}"]) # 添加帧率转换参数 if enable_double_frame: cmd.extend(["--vpp-fruc", "double"]) # 添加输入和输出文件 if avs_file_path: cmd.extend(["--avs", avs_file_path]) else: cmd.extend(["-i", video_path]) cmd.extend(["--output", output_path]) try: # 在后台执行命令 print(f"Executing command: {' '.join(cmd)}") process = subprocess.Popen( cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, creationflags=subprocess.CREATE_NO_WINDOW ) # 等待进程完成 stdout, stderr = process.communicate() if process.returncode != 0: raise RuntimeError(f"NVEncC failed with return code {process.returncode}. Error: {stderr.decode('utf-8')}") print(f"Video processing completed successfully. Output saved to: {output_path}") # 返回输出路径给前端用于预览 return {"ui": {"output_path": ((output_path),)}, "result": ((output_path),)} except Exception as e: raise RuntimeError(f"Error processing video: {str(e)}") finally: # 清理临时avs文件 if avs_file_path and os.path.exists(avs_file_path): try: os.remove(avs_file_path) print(f"Cleaned up temporary AVS file: {avs_file_path}") except Exception as e: print(f"Warning: Failed to remove temporary AVS file: {e}") WEB_DIRECTORY = "./js" # 节点注册 NODE_CLASS_MAPPINGS = { "NVVFR": NVVFRNode } NODE_DISPLAY_NAME_MAPPINGS = { "NVVFR": "NVVFR Video Processor" } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]