GPU加速和批处理优化、更新README

- STTN Auto/Det: 统一 torch.no_grad 包裹，减少重复上下文切换开销
- STTN Auto: 添加 FramePrefetcher 帧预读取，根据 GPU 显存动态调整 batch size
- Lama Inpaint: 新增 _inpaint_batch 批量推理，多帧合并一次 GPU 推理
- ProPainter: copy.deepcopy 替换为浅拷贝，每个区域处理后 gc.collect
- HardwareAccelerator: 新增 get_available_vram_mb 显存查询方法
- README: 添加应用 Logo，同步英文版 README_en.md

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

backend/inpaint/lama_inpaint.py

@@ -1,9 +1,10 @@
 import os
+import gc
 from typing import Union, List
 import torch
 import numpy as np
 from PIL import Image
-from backend.inpaint.utils.lama_util import prepare_img_and_mask
+from backend.inpaint.utils.lama_util import prepare_img_and_mask, get_image, pad_img_to_modulo
 from backend import config
 from backend.tools.inpaint_tools import get_inpaint_area_by_mask
 
@@ -26,6 +27,37 @@ class LamaInpaint:
             cur_res = cur_res[:orig_height, :orig_width]
             return cur_res
 
+    def _inpaint_batch(self, images: List[np.ndarray], masks: List[np.ndarray]):
+        """批量推理：将多帧合并为一个 batch tensor 一次性送入 GPU"""
+        if len(images) == 1:
+            return [self.inpaint(images[0], masks[0])]
+
+        orig_height, orig_width = images[0].shape[:2]
+        batch_imgs = []
+        batch_masks = []
+        for img, msk in zip(images, masks):
+            batch_imgs.append(get_image(img))
+            batch_masks.append(get_image(msk))
+
+        # 堆叠为 (B, C, H, W) 并 pad 到 8 的倍数
+        batch_imgs = np.stack(batch_imgs)
+        batch_masks = np.stack(batch_masks)
+
+        # 对每个样本做 pad
+        padded_imgs = np.stack([pad_img_to_modulo(img, 8) for img in batch_imgs])
+        padded_masks = np.stack([pad_img_to_modulo(m, 8) for m in batch_masks])
+
+        img_tensor = torch.from_numpy(padded_imgs).to(self.device)
+        mask_tensor = torch.from_numpy(padded_masks).to(self.device)
+        mask_tensor = (mask_tensor > 0) * 1
+
+        with torch.inference_mode():
+            inpainted = self.model(img_tensor, mask_tensor)
+            results = inpainted.permute(0, 2, 3, 1).detach().cpu().numpy()
+            results = np.clip(results * 255, 0, 255).astype('uint8')
+
+        return [results[i][:orig_height, :orig_width] for i in range(len(images))]
+
     def __call__(self, input_frames: List[np.ndarray], input_mask: np.ndarray):
         """
         :param input_frames: 原视频帧
@@ -38,48 +70,37 @@ class LamaInpaint:
         # 确定去字幕的垂直高度部分
         split_h = int(W_ori * 3 / 16)
         inpaint_area = get_inpaint_area_by_mask(W_ori, H_ori, split_h, mask)
-        # 初始化帧存储变量
-        # 高分辨率帧存储列表（浅拷贝 + 逐帧 copy，避免 deepcopy 开销）
+        # 高分辨率帧存储列表
         frames_hr = [f.copy() for f in input_frames]
-        frames_scaled = {}  # 存放缩放后帧的字典
-        masks_scaled = {}  # 存放缩放后遮罩的字典
         comps = {}  # 存放补全后帧的字典
         # 存储最终的视频帧
         inpainted_frames = []
-        for k in range(len(inpaint_area)):
-            frames_scaled[k] = []  # 为每个去除部分初始化一个列表
-            masks_scaled[k] = []  # 为每个去除部分初始化一个列表
 
-        # 读取并缩放帧
-        for j in range(len(frames_hr)):
-            image = frames_hr[j]
-            # 对每个去除部分进行切割和缩放
-            for k in range(len(inpaint_area)):
-                image_crop = image[inpaint_area[k][0]:inpaint_area[k][1], :, :]  # 切割
-                mask_crop = mask[inpaint_area[k][0]:inpaint_area[k][1], :, :]  # 切割
-                frames_scaled[k].append(image_crop)  # 将切割后的帧添加到对应列表
-                masks_scaled[k].append(mask_crop)  # 将切割后的遮罩添加到对应列表
-
-        # 处理每一个去除部分
         for k in range(len(inpaint_area)):
-            # 调用inpaint函数逐帧处理
-            comps[k] = []
-            for i in range(len(frames_scaled[k])):
-                inpainted_frame = self.inpaint(frames_scaled[k][i], masks_scaled[k][i])
-                comps[k].append(inpainted_frame)
+            # 收集该区域的所有裁剪帧和遮罩
+            cropped_frames = []
+            cropped_masks = []
+            for j in range(len(frames_hr)):
+                image_crop = frames_hr[j][inpaint_area[k][0]:inpaint_area[k][1], :, :]
+                mask_crop = mask[inpaint_area[k][0]:inpaint_area[k][1], :, :]
+                cropped_frames.append(image_crop)
+                cropped_masks.append(mask_crop)
+
+            # 批量推理
+            comps[k] = self._inpaint_batch(cropped_frames, cropped_masks)
+            del cropped_frames, cropped_masks
+            gc.collect()
 
         # 如果存在去除部分
         if inpaint_area:
             for j in range(len(frames_hr)):
-                frame = frames_hr[j]  # 取出原始帧
-                # 对于模式中的每一个段落
+                frame = frames_hr[j]
                 for k in range(len(inpaint_area)):
-                    comp = comps[k][j]  # 获取补全后的帧
-                    # 实现遮罩区域内的图像融合
-                    frame[inpaint_area[k][0]:inpaint_area[k][1], :, :] = comp
-                # 将最终帧添加到列表
+                    frame[inpaint_area[k][0]:inpaint_area[k][1], :, :] = comps[k][j]
                 inpainted_frames.append(frame)
-                # print(f'processing frame, {len(frames_hr) - j} left')
+
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
         return inpainted_frames
 
 

backend/inpaint/propainter_inpaint.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 import os
+import gc
 import cv2
-import copy
 import numpy as np
 import scipy.ndimage
 from PIL import Image
@@ -374,7 +374,7 @@ class PropainterInpaint:
         inpaint_area = get_inpaint_area_by_mask(W_ori, H_ori, split_h, mask, multiple=8)
         # 初始化帧存储变量
         # 高分辨率帧存储列表
-        frames_hr = copy.deepcopy(input_frames)
+        frames_hr = [f.copy() for f in input_frames]
         frames_scaled = {}  # 存放缩放后帧的字典
         masks_scaled = {}  # 存放缩放后遮罩的字典
         comps = {}  # 存放补全后帧的字典
@@ -398,6 +398,8 @@ class PropainterInpaint:
         for k in range(len(inpaint_area)):
             # 调用inpaint函数进行处理
             comps[k] = self.inpaint(frames_scaled[k], masks_scaled[k][0])
+            del frames_scaled[k], masks_scaled[k]
+            gc.collect()
 
         # 如果存在去除部分
         if inpaint_area:

backend/inpaint/sttn_auto_inpaint.py

@@ -1,6 +1,7 @@
 import os
 import time
 import sys
+import gc
 from typing import List
 
 import cv2
@@ -15,6 +16,8 @@ from backend.config import config
 from backend.inpaint.sttn.auto_sttn import InpaintGenerator
 from backend.inpaint.utils.sttn_utils import Stack, ToTorchFormatTensor
 from backend.tools.inpaint_tools import get_inpaint_area_by_mask, is_frame_number_in_ab_sections
+from backend.tools.video_io import FramePrefetcher
+from backend.tools.hardware_accelerator import HardwareAccelerator
 
 # 定义图像预处理方式
 _to_tensors = transforms.Compose([
@@ -125,7 +128,7 @@ class STTNInpaint:
         feats = feats.to(self.device)
         # 初始化一个与视频长度相同的列表，用于存储处理完成的帧
         comp_frames = [None] * frame_length
-        # 关闭梯度计算，用于推理阶段节省内存并加速
+        # 统一关闭梯度计算，用于推理阶段节省内存并加速
         with torch.no_grad():
             # 将处理好的帧通过编码器，产生特征表示
             feats = self.model.encoder(feats.view(frame_length, 3, self.model_input_height, self.model_input_width))
@@ -133,33 +136,27 @@ class STTNInpaint:
             _, c, feat_h, feat_w = feats.size()
             # 调整特征形状以匹配模型的期望输入
             feats = feats.view(1, frame_length, c, feat_h, feat_w)
-        # 获取重绘区域
-        # 在设定的邻居帧步幅内循环处理视频
-        for f in range(0, frame_length, self.neighbor_stride):
-            # 计算邻近帧的ID
-            neighbor_ids = [i for i in range(max(0, f - self.neighbor_stride), min(frame_length, f + self.neighbor_stride + 1))]
-            # 获取参考帧的索引
-            ref_ids = self.get_ref_index(neighbor_ids, frame_length)
-            # 同样关闭梯度计算
-            with torch.no_grad():
+            # 在设定的邻居帧步幅内循环处理视频
+            for f in range(0, frame_length, self.neighbor_stride):
+                # 计算邻近帧的ID
+                neighbor_ids = [i for i in range(max(0, f - self.neighbor_stride), min(frame_length, f + self.neighbor_stride + 1))]
+                # 获取参考帧的索引
+                ref_ids = self.get_ref_index(neighbor_ids, frame_length)
                 # 通过模型推断特征并传递给解码器以生成完成的帧
                 pred_feat = self.model.infer(feats[0, neighbor_ids + ref_ids, :, :, :])
-                # 将预测的特征通过解码器生成图片，并应用激活函数tanh，然后分离出张量
-                pred_img = torch.tanh(self.model.decoder(pred_feat[:len(neighbor_ids), :, :, :])).detach()
-                # 将结果张量重新缩放到0到255的范围内（图像像素值）
+                # 将预测的特征通过解码器生成图片，并应用激活函数tanh
+                pred_img = torch.tanh(self.model.decoder(pred_feat[:len(neighbor_ids), :, :, :]))
+                # 将结果张量重新缩放到0到255的范围内
                 pred_img = (pred_img + 1) / 2
                 # 将张量移动回CPU并转为NumPy数组
                 pred_img = pred_img.cpu().permute(0, 2, 3, 1).numpy() * 255
                 # 遍历邻近帧
                 for i in range(len(neighbor_ids)):
                     idx = neighbor_ids[i]
-                    # 将预测的图片转换为无符号8位整数格式
-                    img = np.array(pred_img[i]).astype(np.uint8)
+                    img = pred_img[i].astype(np.uint8)
                     if comp_frames[idx] is None:
-                        # 如果该位置为空，则赋值为新计算出的图片
                         comp_frames[idx] = img
                     else:
-                        # 如果此位置之前已有图片，则将新旧图片混合以提高质量
                         comp_frames[idx] = comp_frames[idx].astype(np.float32) * 0.5 + img.astype(np.float32) * 0.5
         # 返回处理完成的帧序列
         return comp_frames
@@ -203,6 +200,8 @@ class STTNAutoInpaint:
         try:
             # 读取视频帧信息
             reader, frame_info = self.read_frame_info_from_video()
+            # 使用帧预读取，I/O 与推理重叠
+            prefetcher = FramePrefetcher(reader)
             if input_sub_remover is not None:
                 ab_sections = input_sub_remover.ab_sections
                 
@@ -212,24 +211,35 @@ class STTNAutoInpaint:
                 # 创建视频写入对象，用于输出修复后的视频
                 writer = cv2.VideoWriter(self.video_out_path, cv2.VideoWriter_fourcc(*"mp4v"), frame_info['fps'], (frame_info['W_ori'], frame_info['H_ori']))
             
-            # 计算需要迭代修复视频的次数
-            rec_time = frame_info['len'] // self.clip_gap if frame_info['len'] % self.clip_gap == 0 else frame_info['len'] // self.clip_gap + 1
             # 计算分割高度，用于确定修复区域的大小
             split_h = int(frame_info['W_ori'] * 3 / 16)
-            
+
             if input_mask is None:
                 # 读取掩码
                 mask = self.sttn_inpaint.read_mask(self.mask_path)
             else:
                 _, mask = cv2.threshold(input_mask, 127, 1, cv2.THRESH_BINARY)
                 mask = mask[:, :, None]
-                
+
             # 得到修复区域位置
             inpaint_area = get_inpaint_area_by_mask(frame_info['W_ori'], frame_info['H_ori'], split_h, mask)
+            # 根据可用显存动态调整 clip_gap，避免 OOM
+            effective_clip_gap = self.clip_gap
+            vram_mb = HardwareAccelerator.instance().get_available_vram_mb()
+            if vram_mb > 0:
+                # 估算每帧约需 (W * H * 3 * 4) bytes，clip_gap帧约需 clip_gap * W * H * 12 bytes（含中间张量）
+                bytes_per_frame = frame_info['W_ori'] * frame_info['H_ori'] * 12
+                max_frames_by_vram = int(vram_mb * 1024 * 1024 / bytes_per_frame)
+                max_frames_by_vram = max(max_frames_by_vram, 10)  # 至少10帧
+                effective_clip_gap = min(self.clip_gap, max_frames_by_vram)
+                if effective_clip_gap < self.clip_gap:
+                    tqdm.write(f'GPU VRAM: {vram_mb:.0f}MB, adjusting clip_gap: {self.clip_gap} -> {effective_clip_gap}')
+            # 计算需要迭代修复视频的次数
+            rec_time = frame_info['len'] // effective_clip_gap if frame_info['len'] % effective_clip_gap == 0 else frame_info['len'] // effective_clip_gap + 1
             # 遍历每一次的迭代次数
             for i in range(rec_time):
-                start_f = i * self.clip_gap  # 起始帧位置
-                end_f = min((i + 1) * self.clip_gap, frame_info['len'])  # 结束帧位置
+                start_f = i * effective_clip_gap  # 起始帧位置
+                end_f = min((i + 1) * effective_clip_gap, frame_info['len'])  # 结束帧位置
                 tqdm.write(f'Processing: {start_f + 1} - {end_f} / Total: {frame_info['len']}')
                 
                 frames_hr = []  # 高分辨率帧列表
@@ -243,7 +253,7 @@ class STTNAutoInpaint:
                 # 读取和修复高分辨率帧
                 valid_frames_count = 0
                 for j in range(start_f, end_f):
-                    success, image = reader.read()
+                    success, image = prefetcher.read()
                     if not success:
                         print(f"Warning: Failed to read frame {j}.")
                         break
@@ -309,10 +319,17 @@ class STTNAutoInpaint:
                                 input_sub_remover.update_progress(tbar, increment=1)
                             if original_frame is not None and input_sub_remover.gui_mode:
                                 input_sub_remover.update_preview_with_comp(original_frame, frame)
+                # 每个chunk处理完后清理GPU缓存
+                del frames_hr, frames, comps
+                gc.collect()
+                if torch.cuda.is_available():
+                    torch.cuda.empty_cache()
         except Exception as e:
             print(f"Error during video processing: {str(e)}")
             # 不抛出异常，允许程序继续执行
         finally:
+            if reader:
+                prefetcher.release()
             if writer:
                 writer.release()
 

backend/inpaint/sttn_det_inpaint.py

@@ -126,38 +126,35 @@ class STTNDetInpaint:
         frame_length = len(frames)
         # 对帧进行预处理转换为张量，并进行归一化
         feats = _to_tensors(frames).unsqueeze(0) * 2 - 1
-        
+
         binary_masks = [np.expand_dims((np.array(m) > 0.5).astype(np.uint8), 2) for m in masks]
         # 将掩码转换为张量
-        masks = (_to_tensors(masks).unsqueeze(0) > 0.5).float()
-        
+        masks_tensor = (_to_tensors(masks).unsqueeze(0) > 0.5).float()
+
         # 把特征张量转移到指定的设备（CPU或GPU）
-        feats, masks = feats.to(self.device), masks.to(self.device)
+        feats, masks_tensor = feats.to(self.device), masks_tensor.to(self.device)
         # 初始化一个与视频长度相同的列表，用于存储处理完成的帧
         comp_frames = [None] * frame_length
-        # 关闭梯度计算，用于推理阶段节省内存并加速
+        # 统一关闭梯度计算，用于推理阶段节省内存并加速
         with torch.no_grad():
             # 将处理好的帧通过编码器，产生特征表示
-            feats = self.model.encoder((feats*(1-masks).float()).view(frame_length, 3, self.model_input_height, self.model_input_width))
+            feats = self.model.encoder((feats*(1-masks_tensor).float()).view(frame_length, 3, self.model_input_height, self.model_input_width))
             # 获取特征维度信息
             _, c, feat_h, feat_w = feats.size()
             # 调整特征形状以匹配模型的期望输入
             feats = feats.view(1, frame_length, c, feat_h, feat_w)
-        # 获取重绘区域
-        # 在设定的邻居帧步幅内循环处理视频
-        for f in range(0, frame_length, self.neighbor_stride):
-            # 计算邻近帧的ID
-            neighbor_ids = [i for i in range(max(0, f - self.neighbor_stride), min(frame_length, f + self.neighbor_stride + 1))]
-            # 获取参考帧的索引
-            ref_ids = self.get_ref_index(neighbor_ids, frame_length)
-            # 同样关闭梯度计算
-            with torch.no_grad():
+            # 在设定的邻居帧步幅内循环处理视频
+            for f in range(0, frame_length, self.neighbor_stride):
+                # 计算邻近帧的ID
+                neighbor_ids = [i for i in range(max(0, f - self.neighbor_stride), min(frame_length, f + self.neighbor_stride + 1))]
+                # 获取参考帧的索引
+                ref_ids = self.get_ref_index(neighbor_ids, frame_length)
                 # 通过模型推断特征并传递给解码器以生成完成的帧
                 pred_feat = self.model.infer(
-                    feats[0, neighbor_ids + ref_ids, :, :, :], masks[0, neighbor_ids + ref_ids, :, :, :])
+                    feats[0, neighbor_ids + ref_ids, :, :, :], masks_tensor[0, neighbor_ids + ref_ids, :, :, :])
 
-                # 将预测的特征通过解码器生成图片，并应用激活函数tanh，然后分离出张量
-                pred_img = torch.tanh(self.model.decoder(pred_feat[:len(neighbor_ids), :, :, :])).detach()
+                # 将预测的特征通过解码器生成图片，并应用激活函数tanh
+                pred_img = torch.tanh(self.model.decoder(pred_feat[:len(neighbor_ids), :, :, :]))
                 # 将结果张量重新缩放到0到255的范围内（图像像素值）
                 pred_img = (pred_img + 1) / 2
                 # 将张量移动回CPU并转为NumPy数组
@@ -166,13 +163,10 @@ class STTNDetInpaint:
                 for i in range(len(neighbor_ids)):
                     idx = neighbor_ids[i]
                     # 将预测的图片转换为无符号8位整数格式
-                    img = np.array(pred_img[i]).astype(
-                        np.uint8)*binary_masks[idx] + frames[idx] * (1-binary_masks[idx])
+                    img = pred_img[i].astype(np.uint8) * binary_masks[idx] + frames[idx] * (1 - binary_masks[idx])
                     if comp_frames[idx] is None:
-                        # 如果该位置为空，则赋值为新计算出的图片
                         comp_frames[idx] = img
                     else:
-                        # 如果此位置之前已有图片，则将新旧图片混合以提高质量
                         comp_frames[idx] = comp_frames[idx].astype(np.float32) * 0.5 + img.astype(np.float32) * 0.5
         # 返回处理完成的帧序列
         return comp_frames

backend/tools/hardware_accelerator.py

@@ -106,6 +106,27 @@ class HardwareAccelerator:
     def set_enabled(self, enable):
         self.__enabled = enable
 
+    def get_available_vram_mb(self):
+        """获取可用 GPU 显存（MB），无 GPU 返回 0"""
+        if not self.__enabled:
+            return 0
+        if self.__cuda:
+            try:
+                free_vram = torch.cuda.mem_get_info()[0]  # (free, total)
+                return free_vram / (1024 * 1024)
+            except Exception:
+                return 0
+        if self.__mps:
+            try:
+                # MPS 没有直接查询接口，使用系统内存作为参考
+                import subprocess
+                result = subprocess.run(['sysctl', '-n', 'hw.memsize'], capture_output=True, text=True)
+                total_mem = int(result.stdout.strip()) / (1024 * 1024)
+                return total_mem * 0.5  # 保守估计可用一半
+            except Exception:
+                return 0
+        return 0
+
     @property
     def device(self):
         """