Browse Source

Relocate `magic_pdf imports` to `setup` to Ensure Proper CUDA Initialization

Hui Kang 10 tháng trước cách đây
mục cha
commit
0160e5d265
1 tập tin đã thay đổi với 8 bổ sung4 xóa
  1. 8 4
      projects/multi_gpu/server.py

+ 8 - 4
projects/multi_gpu/server.py

@@ -10,8 +10,6 @@ import filetype
 import litserve as ls
 from pathlib import Path
 from fastapi import HTTPException
-from magic_pdf.tools.cli import do_parse, convert_file_to_pdf
-from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
 
 
 class MinerUAPI(ls.LitAPI):
@@ -24,6 +22,12 @@ class MinerUAPI(ls.LitAPI):
             if torch.cuda.device_count() > 1:
                 raise RuntimeError("Remove any CUDA actions before setting 'CUDA_VISIBLE_DEVICES'.")
 
+        from magic_pdf.tools.cli import do_parse, convert_file_to_pdf
+        from magic_pdf.model.doc_analyze_by_custom_model import ModelSingleton
+
+        self.do_parse = do_parse
+        self.convert_file_to_pdf = convert_file_to_pdf
+
         model_manager = ModelSingleton()
         model_manager.get_model(True, False)
         model_manager.get_model(False, False)
@@ -41,7 +45,7 @@ class MinerUAPI(ls.LitAPI):
         try:
             pdf_name = str(uuid.uuid4())
             output_dir = self.output_dir.joinpath(pdf_name)
-            do_parse(self.output_dir, pdf_name, inputs[0], [], **inputs[1])
+            self.do_parse(self.output_dir, pdf_name, inputs[0], [], **inputs[1])
             return output_dir
         except Exception as e:
             shutil.rmtree(output_dir, ignore_errors=True)
@@ -73,7 +77,7 @@ class MinerUAPI(ls.LitAPI):
                         return f.convert_to_pdf()
                 else:
                     temp_file.write_bytes(file_bytes)
-                    convert_file_to_pdf(temp_file, temp_dir)
+                    self.convert_file_to_pdf(temp_file, temp_dir)
                     return temp_file.with_suffix('.pdf').read_bytes()
             else:
                 raise Exception('Unsupported file format')