Merge pull request #1517 from randydl/fix-multi-gpu

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')