From 9970d3aec0c87585c80f279a1e5ddcd85bd1d56c Mon Sep 17 00:00:00 2001 From: jax Date: Tue, 8 Apr 2025 00:30:26 +0800 Subject: [PATCH] fix --- README.md | 2 ++ nodes/SaveVideo.py | 18 +++++++++--------- nodes/upscaleNode.py | 38 +++++++++++++++++++------------------- 3 files changed, 30 insertions(+), 28 deletions(-) diff --git a/README.md b/README.md index e54f363..8d1ba69 100644 --- a/README.md +++ b/README.md @@ -7,7 +7,9 @@ Used to solve the OOM (Out Of Memory) issue caused by loading all frames of a vi todo: 1. ~~Upscale~~ + 2. Face Swap + 3. Latent Lip Sync diff --git a/nodes/SaveVideo.py b/nodes/SaveVideo.py index 1f449bb..f021cd5 100644 --- a/nodes/SaveVideo.py +++ b/nodes/SaveVideo.py @@ -23,7 +23,7 @@ class VideoBasicSaveVideo: def save(self, file_path): output_dir = os.path.abspath(folder_paths.get_output_directory()) - # 处理相对路径，将其转换为绝对路径 + # Handle relative paths, convert them to absolute paths if not os.path.isabs(file_path): file_path = os.path.join(output_dir, file_path) else: @@ -31,26 +31,26 @@ class VideoBasicSaveVideo: print(f"file_path is {file_path}, output_dir is {output_dir}") - # 判断路径是否在输出目录中且文件存在 + # Check if the path is in the output directory and the file exists if file_path.startswith(output_dir) and os.path.exists(file_path): - # 获取相对于output_dir的路径 + # Get the path relative to output_dir rel_path = os.path.relpath(file_path, output_dir) - # 分离文件名和子文件夹路径 + # Separate filename and subfolder path subfolder = os.path.dirname(rel_path) filename = os.path.basename(file_path) elif os.path.exists(file_path): - # 文件存在但不在output目录，需要复制到output目录 + # File exists but not in the output directory, need to copy to output directory filename = os.path.basename(file_path) - subfolder = "" # 默认保存到output根目录 + subfolder = "" # Default save to output root directory - # 创建目标路径 + # Create destination path dest_path = os.path.join(output_dir, filename) - # 复制文件到output目录 + # Copy file to output directory shutil.copy2(file_path, dest_path) print(f"Copied file from {file_path} to {dest_path}") - # 更新file_path为新的路径 + # Update file_path to the new path file_path = dest_path else: filename = None diff --git a/nodes/upscaleNode.py b/nodes/upscaleNode.py index 3297ff0..32d8f60 100644 --- a/nodes/upscaleNode.py +++ b/nodes/upscaleNode.py @@ -26,75 +26,75 @@ class VideoUpscaleWithModel: if not os.path.exists(video_path): raise ValueError(f"Video file not found: {video_path}") - # 获取临时目录 + # Get temporary directory output_dir = folder_paths.get_temp_directory() os.makedirs(output_dir, exist_ok=True) - # 生成输出文件名 + # Generate output filename video_filename = os.path.basename(video_path) output_filename = f"upscaled_{video_filename}" if not output_filename.lower().endswith('.mp4'): output_filename += '.mp4' output_path = os.path.join(output_dir, output_filename) - # 打开视频文件 + # Open video file cap = cv2.VideoCapture(video_path) if not cap.isOpened(): raise ValueError(f"Could not open video file: {video_path}") - # 获取视频信息 + # Get video information 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)) - # 计算上采样后的分辨率 + # Calculate upsampled resolution new_width = int(width * upscale_model.scale) new_height = int(height * upscale_model.scale) - # 创建视频写入器 + # Create video writer fourcc = cv2.VideoWriter_fourcc(*'mp4v') out = cv2.VideoWriter(output_path, fourcc, fps, (new_width, new_height)) device = model_management.get_torch_device() - # 估计所需内存 + # Estimate required memory 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) - # 将模型移至设备 + # Move model to device upscale_model.to(device) - # 设置瓦片参数 + # Set tile parameters tile = 512 overlap = 32 - # 批量处理帧 + # Process frames in batches frames_processed = 0 try: with tqdm(total=total_frames) as pbar: while frames_processed < total_frames: - # 读取一批帧 + # Read a batch of frames batch_frames = [] for _ in range(min(batch_size, total_frames - frames_processed)): ret, frame = cap.read() if not ret: break - # 转换为RGB + # Convert to RGB frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) - # 归一化到0-1 + # Normalize to 0-1 frame = frame.astype(np.float32) / 255.0 batch_frames.append(frame) if not batch_frames: break - # 转换为PyTorch张量 + # Convert to PyTorch tensor batch_tensor = torch.from_numpy(np.stack(batch_frames)).permute(0, 3, 1, 2).to(device) - # 上采样处理 + # Upsampling process upscaled_batch = None oom = True current_tile = tile @@ -107,7 +107,7 @@ class VideoUpscaleWithModel: ) pbar_upscale = comfy.utils.ProgressBar(steps) - # 逐帧处理批次 + # Process frames in batch one by one upscaled_frames = [] for i in range(batch_tensor.shape[0]): frame = batch_tensor[i:i+1] @@ -126,12 +126,12 @@ class VideoUpscaleWithModel: if current_tile < 128: raise e - # 转换回numpy并写入视频 + # Convert back to numpy and write to video 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 + # Convert to BGR and scale to 0-255 frame_bgr = cv2.cvtColor((frame * 255).astype(np.uint8), cv2.COLOR_RGB2BGR) out.write(frame_bgr) @@ -139,7 +139,7 @@ class VideoUpscaleWithModel: pbar.update(len(batch_frames)) finally: - # 清理资源 + # Clean up resources upscale_model.to("cpu") cap.release() out.release()