feat: optimize VRAM management by cleaning up GPU memory during batch processing

This commit is contained in:
myhloli
2026-03-26 10:46:26 +08:00
parent 834dd496a3
commit 185ab841ae

View File

@@ -345,13 +345,16 @@ class BatchAnalyze:
images_mfd_res[image_index], images_formula_list[image_index]
):
formula_res["latex"] = formula_with_latex.get("latex", "")
# 清理显存
clean_vram(self.model.device, vram_threshold=8)
else:
for layout_res in images_layout_res:
# 移除所有的"inline_formula"
layout_res[:] = [res for res in layout_res if res.get("label") != "inline_formula"]
# 清理显存
clean_vram(self.model.device, vram_threshold=8)
ocr_res_list_all_page = []
table_res_list_all_page = []
@@ -691,6 +694,9 @@ class BatchAnalyze:
)
ocr_res_list_dict['layout_res'].extend(ocr_result_list)
# 清理显存
clean_vram(self.model.device, vram_threshold=8)
else:
# 原始单张处理模式
for ocr_res_list_dict in tqdm(ocr_res_list_all_page, desc="OCR-det Predict"):