Add GPU memory clean up

mineru_single.py

@@ -4,6 +4,8 @@ import uuid
 import json
 import requests
 import logging
+import torch
+import gc
 
 from magic_pdf.data.dataset import PymuDocDataset
 from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
@@ -48,40 +50,57 @@ class Processor:
             logger.error("Failed to initialize Processor: %s", str(e))
             raise
 
+    def cleanup_gpu(self):
+        """
+        Releases GPU memory, use garbage collection to clear PyTorch's CUDA cache.
+        This helps prevent VRAM accumulation.
+        """
+        try:
+            gc.collect()               #garbage collection
+            torch.cuda.empty_cache()   # Clear memory cache on GPU
+            logger.info("GPU memory cleaned up.")
+        except Exception as e:
+            logger.error("Error during GPU cleanup: %s", e)
+
     def process(self, file_url: str, key: str) -> str:
         """
         Process a single PDF, returning final Markdown with irrelevant images removed.
         """
         logger.info("Processing file: %s", file_url)
 
-        response = requests.get(file_url)
-        if response.status_code != 200:
-            logger.error("Failed to download PDF from %s. Status code: %d", file_url, response.status_code)
-            raise Exception(f"Failed to download PDF: {file_url}")
-
-        pdf_bytes = response.content
-        logger.info("Downloaded %d bytes for file_url='%s'", len(pdf_bytes), file_url)
-
-        # Analyze PDF with OCR
-        dataset = PymuDocDataset(pdf_bytes)
-        inference = doc_analyze(
-            dataset,
-            ocr=True,
-            lang=self.language,
-            layout_model=self.layout_mode,
-            formula_enable=self.formula_enable,
-            table_enable=self.table_enable
-        )
-        logger.info("doc_analyze complete for key='%s'. Started to extracting images...", key)
-
-        # Classify images, remove irrelevant
-        image_writer = ImageWriter(self.s3_writer, self.svm_model)
-        pipe_result = inference.pipe_ocr_mode(image_writer, lang=self.language)
-
-        md_content = pipe_result.get_markdown(self.prefix + key + "/")
-        final_markdown = image_writer.remove_redundant_images(md_content)
-        logger.info("Completed PDF process for key='%s'. Final MD length=%d", key, len(final_markdown))
-        return final_markdown
+        try:
+            response = requests.get(file_url)
+            if response.status_code != 200:
+                logger.error("Failed to download PDF from %s. Status code: %d", file_url, response.status_code)
+                raise Exception(f"Failed to download PDF: {file_url}")
+
+            pdf_bytes = response.content
+            logger.info("Downloaded %d bytes for file_url='%s'", len(pdf_bytes), file_url)
+
+            # Analyze PDF with OCR
+            dataset = PymuDocDataset(pdf_bytes)
+            inference = doc_analyze(
+                dataset,
+                ocr=True,
+                lang=self.language,
+                layout_model=self.layout_mode,
+                formula_enable=self.formula_enable,
+                table_enable=self.table_enable
+            )
+            logger.info("doc_analyze complete for key='%s'. Started extracting images...", key)
+
+            # Classify images and remove irrelevant ones
+            image_writer = ImageWriter(self.s3_writer, self.svm_model)
+            pipe_result = inference.pipe_ocr_mode(image_writer, lang=self.language)
+            logger.info("OCR pipeline completed for key='%s'.", key)
+
+            md_content = pipe_result.get_markdown(self.prefix + key + "/")
+            final_markdown = image_writer.remove_redundant_images(md_content)
+            logger.info("Completed PDF process for key='%s'. Final MD length=%d", key, len(final_markdown))
+            return final_markdown
+        finally:
+            # GPU memory is cleaned up after each processing.
+            self.cleanup_gpu()
 
 class ImageWriter(DataWriter):
     """