diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..2fbed63 --- /dev/null +++ b/.gitignore @@ -0,0 +1,62 @@ +# Python +__pycache__/ +*.py[cod] +*.so +*.pyd +*.pyo +*.pdb +*.egg-info/ +*.egg +*.log + +# Jupyter +.ipynb_checkpoints/ + +# VSCode +.vscode/ + +# 环境/依赖 +.env +.venv +venv/ +ENV/ + +# 临时文件 +*.tmp +*.temp +*.bak +*.swp +*.swo +*.DS_Store +Thumbs.db + +# 数据/输出 +output/ +*.ckpt +*.pth +*.h5 +*.pt +*.onnx +*.npz +*.npy +*.wav +*.mp4 +*.avi +*.mov +*.mkv +*.webm +*.gif +*.png +*.jpg +*.jpeg +*.bmp +*.tiff +*.tif + +# 云配置 +cloud_config.json + +# 其他 +*.sqlite3 +*.db +*.lock diff --git a/node/cloud_node.py b/node/cloud_node.py index e9cae9f..fec8622 100644 --- a/node/cloud_node.py +++ b/node/cloud_node.py @@ -189,6 +189,7 @@ class CloudImagesToVideoAndUpload: "folder": ("STRING", {"default": "video"}), "key_prefix": ("STRING", {"default": "comfyui_"}), "ext": (["mp4", "mov", "avi", "mkv"], {"default": "mp4"}), + "audio": ("AUDIO", {"default": None}), # 改为AUDIO类型 } } @@ -198,7 +199,8 @@ class CloudImagesToVideoAndUpload: CATEGORY = "云服务" OUTPUT_NODE = True - def images_to_video_and_upload(self, images, fps, cloud_type, access_key, secret_key, bucket_name, domain, folder, key_prefix, ext): + def images_to_video_and_upload(self, images, fps, cloud_type, access_key, secret_key, bucket_name, domain, folder, key_prefix, ext, audio=None): + import torchaudio config = load_cloud_config()[1] access_key = access_key or config.get("access_key", "") secret_key = secret_key or config.get("secret_key", "") @@ -212,6 +214,8 @@ class CloudImagesToVideoAndUpload: pix_fmt = 'rgb24' with tempfile.NamedTemporaryFile(suffix=f'.{ext}', delete=False) as tmpfile: tmp_path = tmpfile.name + # 先生成无音频视频，临时文件名用 _noaudio 结尾但扩展名标准 + tmp_video_path = tmp_path.replace(f'.{ext}', f'_noaudio.{ext}') cmd = [ imageio_ffmpeg.get_ffmpeg_exe(), '-y', @@ -224,19 +228,45 @@ class CloudImagesToVideoAndUpload: '-an', '-vcodec', 'libx264', '-pix_fmt', 'yuv420p', - tmp_path + tmp_video_path ] proc = subprocess.Popen(cmd, stdin=subprocess.PIPE) for im in img_list: proc.stdin.write(im.convert('RGB').tobytes()) proc.stdin.close() proc.wait() + # 如果有AUDIO，保存为wav临时文件再合成 + if audio is not None and isinstance(audio, dict) and "waveform" in audio and "sample_rate" in audio: + with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_audio: + audio_path = tmp_audio.name + waveform = audio["waveform"] + sample_rate = audio["sample_rate"] + # waveform shape: (1, channels, samples) or (channels, samples) + if waveform.dim() == 3: + waveform = waveform.squeeze(0) + torchaudio.save(audio_path, waveform, sample_rate) + # 合成音视频到tmp_path + merge_cmd = [ + imageio_ffmpeg.get_ffmpeg_exe(), + '-y', + '-i', tmp_video_path, + '-i', audio_path, + '-c:v', 'copy', + '-c:a', 'aac', + '-shortest', + tmp_path + ] + subprocess.run(merge_cmd, check=True) + os.remove(tmp_video_path) + os.remove(audio_path) + else: + # 无音频直接重命名 + os.rename(tmp_video_path, tmp_path) # 3. 上传到云存储 if cloud_type == "qiniu": uploader = QiniuUploader(access_key, secret_key, bucket_name, domain) elif cloud_type == "jdcloud": # uploader = JDCloudUploader(access_key, secret_key, bucket_name, domain) - # else: raise NotImplementedError(f"暂不支持的云类型: {cloud_type}") random_name = uuid.uuid4().hex folder_path = folder.strip().strip('/') diff --git a/node/video_split_node.py b/node/video_split_node.py index b5163b9..4918ec9 100644 --- a/node/video_split_node.py +++ b/node/video_split_node.py @@ -1,10 +1,14 @@ import cv2 import torch import numpy as np +import torchaudio +import tempfile +import subprocess +import os class SplitVideoByFrames: """ - 传入视频路径，根据帧率将视频拆分为多段，每段帧数上限可控，输出为张量序列列表 + 用OpenCV拆分视频为多个片段，每段帧数不超过max_frames_per_clip，音频输出为ComfyUI官方格式 """ @classmethod def INPUT_TYPES(cls): @@ -15,17 +19,32 @@ class SplitVideoByFrames: } } - RETURN_TYPES = ("LIST", "INT") - RETURN_NAMES = ("clips", "num_clips") + RETURN_TYPES = ("INT", "LIST", "AUDIO") + RETURN_NAMES = ("num_clips", "clips", "audio") FUNCTION = "split_video" CATEGORY = "云服务" def split_video(self, video_path, max_frames_per_clip): - clips = [] + # 1. 提取音频为wav临时文件 + with tempfile.NamedTemporaryFile(suffix='.wav', delete=False) as tmp_audio: + audio_path = tmp_audio.name + ffmpeg_bin = "ffmpeg" # 假设已在环境变量 + cmd = [ffmpeg_bin, '-y', '-i', video_path, '-vn', '-acodec', 'pcm_s16le', '-ar', '44100', '-ac', '2', audio_path] + subprocess.run(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + # 2. 读取音频为tensor + waveform, sample_rate = torchaudio.load(audio_path) + os.remove(audio_path) + # 保证shape为(1,channels,samples) + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) + audio_dict = {"waveform": waveform, "sample_rate": sample_rate} + + # 3. 视频分帧（不做resize，假设视频分辨率一致） cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise FileNotFoundError(f"无法打开视频文件: {video_path}") frames = [] + clips = [] while True: ret, frame = cap.read() if not ret: @@ -34,78 +53,195 @@ class SplitVideoByFrames: np_image = frame.astype(np.float32) / 255.0 frames.append(torch.from_numpy(np_image)) if len(frames) == max_frames_per_clip: - clips.append(torch.stack(frames, dim=0)) + clips.append({"frames": torch.stack(frames, dim=0), "audio": audio_dict}) frames = [] if frames: - clips.append(torch.stack(frames, dim=0)) + clips.append({"frames": torch.stack(frames, dim=0), "audio": audio_dict}) cap.release() - return (clips, len(clips)) + return (len(clips), clips, audio_dict) class GetVideoClipByIndex: """ - 传入 clips（张量序列列表）和索引，返回该片段合成的视频（mp4二进制） + 输入clips和索引，输出该片段的图片数组（IMAGE）、音频（官方格式）和帧数（num_frames） """ @classmethod def INPUT_TYPES(cls): return { "required": { - "clips": ("LIST", {}), "index": ("INT", {"default": 0, "min": 0}), - "fps": ("INT", {"default": 24, "min": 1, "max": 120}), - "ext": (["mp4", "avi", "mov"], {"default": "mp4"}), + "clips": ("LIST", {}), } } - RETURN_TYPES = ("BYTES",) - RETURN_NAMES = ("video_bytes",) - FUNCTION = "get_clip_video" + RETURN_TYPES = ("IMAGE", "AUDIO", "INT") + RETURN_NAMES = ("images", "audio", "num_frames") + FUNCTION = "get_clip" CATEGORY = "云服务" - def get_clip_video(self, clips, index, fps, ext): - import tempfile - import imageio_ffmpeg - import os - import io + def get_clip(self, index, clips): if not clips or index < 0 or index >= len(clips): raise IndexError("索引超出clips范围") clip = clips[index] - arr = clip.cpu().numpy() if hasattr(clip, 'cpu') else clip - img_list = [np.clip(255. * img, 0, 255).astype(np.uint8) for img in arr] - height, width = img_list[0].shape[0], img_list[0].shape[1] - pix_fmt = 'rgb24' - with tempfile.NamedTemporaryFile(suffix=f'.{ext}', delete=False) as tmpfile: - tmp_path = tmpfile.name - cmd = [ - imageio_ffmpeg.get_ffmpeg_exe(), - '-y', - '-f', 'rawvideo', - '-vcodec', 'rawvideo', - '-s', f'{width}x{height}', - '-pix_fmt', pix_fmt, - '-r', str(fps), - '-i', '-', - '-an', - '-vcodec', 'libx264', - '-pix_fmt', 'yuv420p', - tmp_path - ] - import subprocess - proc = subprocess.Popen(cmd, stdin=subprocess.PIPE) - for im in img_list: - proc.stdin.write(im.tobytes()) - proc.stdin.close() - proc.wait() - with open(tmp_path, "rb") as f: - video_bytes = f.read() - os.remove(tmp_path) - return (video_bytes,) + frames = clip["frames"] + audio = clip["audio"] + print(f"获取片段 {index}，帧数: {frames.shape[0]}, 音频采样率: {audio['sample_rate']}") + num_frames = frames.shape[0] + return (frames, audio, num_frames) + +class CreateEmptyImageBatch: + """ + 创建一个空的图像batch，若传入images则追加到空batch后返回。 + """ + @classmethod + def INPUT_TYPES(cls): + return { + "optional": { + "images": ("IMAGE", {}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_batch",) + FUNCTION = "create" + CATEGORY = "云服务/集合工具" + + def create(self, images=None): + import torch + def to_3ch(img): + if img is None: + return img + if img.dim() == 3: + if img.shape[-1] == 1: + img = img.repeat(1, 1, 3) + elif img.shape[-1] == 4: + img = img[..., :3] + elif img.dim() == 4: + if img.shape[-1] == 1: + img = img.repeat(1, 1, 1, 3) + elif img.shape[-1] == 4: + img = img[..., :3] + return img + batch = torch.empty((0, 0, 0, 0), dtype=torch.float32) + if images is not None: + images = to_3ch(images) + if images.dim() == 3: + images = images.unsqueeze(0) + # 直接用AppendImagesToBatch的逻辑拼接 + if batch.numel() == 0 or batch.shape[0] == 0: + return (images,) + else: + new_batch = torch.cat([batch, images], dim=0) + return (new_batch,) + return (batch,) + +class AppendImagesToBatch: + """ + 向已有图像batch追加图片，支持单张或多张图片，输出合并后的IMAGE batch。 + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_batch": ("IMAGE", {}), + "images_to_add": ("IMAGE", {}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_batch",) + FUNCTION = "append" + CATEGORY = "云服务/集合工具" + + def append(self, image_batch, images_to_add): + import torch + # 如果 image_batch 是空的（shape为(0, 0, 0, 0)），直接返回 images_to_add + if image_batch.numel() == 0 or image_batch.shape[0] == 0: + if images_to_add.dim() == 3: + images_to_add = images_to_add.unsqueeze(0) + return (images_to_add,) + if image_batch.dim() == 3: + image_batch = image_batch.unsqueeze(0) + if images_to_add.dim() == 3: + images_to_add = images_to_add.unsqueeze(0) + new_batch = torch.cat([image_batch, images_to_add], dim=0) + return (new_batch,) + +class GetFirstImageFromBatch: + """ + 获取图片batch的首帧或末帧，输出单帧IMAGE batch（shape为(1, H, W, C)）及宽高。 + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_batch": ("IMAGE", {}), + "mode": (["first", "last"], {"default": "first"}), + } + } + + RETURN_TYPES = ("IMAGE", "INT", "INT") + RETURN_NAMES = ("image", "width", "height") + FUNCTION = "get" + CATEGORY = "云服务/集合工具" + + def get(self, image_batch, mode="first"): + import torch + if image_batch.dim() == 3: + img = image_batch.unsqueeze(0) + else: + if image_batch.shape[0] == 0: + raise ValueError("image_batch is empty!") + if mode == "first": + img = image_batch[0:1] + else: + img = image_batch[-1:] + # img shape: (1, H, W, C) + _, h, w, _ = img.shape + return (img, w, h) + +class RemoveFirstOrLastImageFromBatch: + """ + 删除图片batch的首帧或末帧，输出剩余IMAGE batch。 + """ + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image_batch": ("IMAGE", {}), + "mode": (["first", "last"], {"default": "first"}), + } + } + + RETURN_TYPES = ("IMAGE",) + RETURN_NAMES = ("image_batch",) + FUNCTION = "remove" + CATEGORY = "云服务/集合工具" + + def remove(self, image_batch, mode="first"): + import torch + if image_batch.dim() == 3: + # 只有一张，删除后为空batch + return (torch.empty((0, *image_batch.shape), dtype=image_batch.dtype, device=image_batch.device),) + if image_batch.shape[0] == 0: + return (image_batch,) + if mode == "first": + return (image_batch[1:],) + else: + return (image_batch[:-1],) NODE_CLASS_MAPPINGS = { "SplitVideoByFrames": SplitVideoByFrames, - "GetVideoClipByIndex": GetVideoClipByIndex + "GetVideoClipByIndex": GetVideoClipByIndex, + "CreateEmptyImageBatch": CreateEmptyImageBatch, + "AppendImagesToBatch": AppendImagesToBatch, + "GetFirstImageFromBatch": GetFirstImageFromBatch, + "RemoveFirstOrLastImageFromBatch": RemoveFirstOrLastImageFromBatch } - NODE_DISPLAY_NAME_MAPPINGS = { "SplitVideoByFrames": "🎬 视频分段拆帧", - "GetVideoClipByIndex": "🎞️ 获取片段视频" -} + "GetVideoClipByIndex": "🖼️ 获取片段图片和音频", + "CreateEmptyImageBatch": "📂 创建空图像Batch", + "AppendImagesToBatch": "➕ 追加图片到Batch", + "GetFirstImageFromBatch": "🔍 获取首/末帧图片", + "RemoveFirstOrLastImageFromBatch": "❌ 删除首/末帧图片" +} \ No newline at end of file