diff --git a/__init__.py b/__init__.py index ff1b1e9..9119c62 100644 --- a/__init__.py +++ b/__init__.py @@ -3,16 +3,19 @@ from .nodes.SaveVideo import VideoBasicSaveVideo from .nodes.VideoUpload import VideoBasicUploadVideo +from .nodes.upscaleNode import VideoUpscaleWithModel NODE_CLASS_MAPPINGS = { "VideoBasicLoadVideo": VideoBasicUploadVideo, "VideoBasicVideoSave": VideoBasicSaveVideo, + "VideoBasicVideoUpscaleWithModel": VideoUpscaleWithModel, } NODE_DISPLAY_NAME_MAPPINGS = { "VideoBasicLoadVideo": "VideoBasic Load Video", "VideoBasicVideoSave": "VideoBasic Save Video", + "VideoBasicVideoUpscaleWithModel": "VideoBasic Video Upscale with Model", } WEB_DIRECTORY = "./web" diff --git a/nodes/upscaleNode.py b/nodes/upscaleNode.py new file mode 100644 index 0000000..f7f9c90 --- /dev/null +++ b/nodes/upscaleNode.py @@ -0,0 +1,137 @@ +from comfy import model_management +import comfy.utils +import torch +import cv2 +import os +import numpy as np +from tqdm import tqdm + + +class VideoUpscaleWithModel: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "upscale_model": ("UPSCALE_MODEL",), + "video_path": ("STRING", {"default": ""}), + "batch_size": ("INT", {"default": 8, "min": 1, "max": 1000, "step": 1}), + "output_path": ("STRING", {"default": "upscaled_video.mp4"}), + }} + RETURN_TYPES = ("STRING",) + RETURN_NAMES = ("output_video_path",) + FUNCTION = "upscale_video" + + CATEGORY = "VideoBasic" + + def upscale_video(self, upscale_model, video_path, batch_size, output_path): + if not os.path.exists(video_path): + raise ValueError(f"Video file not found: {video_path}") + + # 打开视频文件 + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Could not open video file: {video_path}") + + # 获取视频信息 + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + + # 计算上采样后的分辨率 + new_width = int(width * upscale_model.scale) + new_height = int(height * upscale_model.scale) + + # 创建视频写入器 + fourcc = cv2.VideoWriter_fourcc(*'mp4v') + out = cv2.VideoWriter(output_path, fourcc, fps, (new_width, new_height)) + + device = model_management.get_torch_device() + + # 估计所需内存 + memory_required = model_management.module_size(upscale_model.model) + memory_required += (width * height * 3) * 4 * max(upscale_model.scale, 1.0) * 384.0 * batch_size + model_management.free_memory(memory_required, device) + + # 将模型移至设备 + upscale_model.to(device) + + # 设置瓦片参数 + tile = 512 + overlap = 32 + + # 批量处理帧 + frames_processed = 0 + + try: + with tqdm(total=total_frames) as pbar: + while frames_processed < total_frames: + # 读取一批帧 + batch_frames = [] + for _ in range(min(batch_size, total_frames - frames_processed)): + ret, frame = cap.read() + if not ret: + break + # 转换为RGB + frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) + # 归一化到0-1 + frame = frame.astype(np.float32) / 255.0 + batch_frames.append(frame) + + if not batch_frames: + break + + # 转换为PyTorch张量 + batch_tensor = torch.from_numpy(np.stack(batch_frames)).permute(0, 3, 1, 2).to(device) + + # 上采样处理 + upscaled_batch = None + oom = True + current_tile = tile + + while oom: + try: + steps = batch_tensor.shape[0] * comfy.utils.get_tiled_scale_steps( + batch_tensor.shape[3], batch_tensor.shape[2], + tile_x=current_tile, tile_y=current_tile, overlap=overlap + ) + pbar_upscale = comfy.utils.ProgressBar(steps) + + # 逐帧处理批次 + upscaled_frames = [] + for i in range(batch_tensor.shape[0]): + frame = batch_tensor[i:i+1] + upscaled = comfy.utils.tiled_scale( + frame, lambda a: upscale_model(a), + tile_x=current_tile, tile_y=current_tile, + overlap=overlap, upscale_amount=upscale_model.scale, + pbar=pbar_upscale + ) + upscaled_frames.append(upscaled) + + upscaled_batch = torch.cat(upscaled_frames, dim=0) + oom = False + except model_management.OOM_EXCEPTION as e: + current_tile //= 2 + if current_tile < 128: + raise e + + # 转换回numpy并写入视频 + upscaled_batch = torch.clamp(upscaled_batch.permute(0, 2, 3, 1), min=0, max=1.0) + upscaled_batch = upscaled_batch.cpu().numpy() + + for frame in upscaled_batch: + # 转换为BGR并缩放到0-255 + frame_bgr = cv2.cvtColor((frame * 255).astype(np.uint8), cv2.COLOR_RGB2BGR) + out.write(frame_bgr) + + frames_processed += len(batch_frames) + pbar.update(len(batch_frames)) + + finally: + # 清理资源 + upscale_model.to("cpu") + cap.release() + out.release() + + print(f"Video upscaling complete. Output saved to: {output_path}") + return (output_path,) \ No newline at end of file