Commit 0160e5d2 authored by Hui Kang's avatar Hui Kang
Browse files

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

parent 989465e8
...@@ -10,8 +10,6 @@ import filetype ...@@ -10,8 +10,6 @@ import filetype
import litserve as ls import litserve as ls
from pathlib import Path from pathlib import Path
from fastapi import HTTPException 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): class MinerUAPI(ls.LitAPI):
...@@ -24,6 +22,12 @@ class MinerUAPI(ls.LitAPI): ...@@ -24,6 +22,12 @@ class MinerUAPI(ls.LitAPI):
if torch.cuda.device_count() > 1: if torch.cuda.device_count() > 1:
raise RuntimeError("Remove any CUDA actions before setting 'CUDA_VISIBLE_DEVICES'.") 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 = ModelSingleton()
model_manager.get_model(True, False) model_manager.get_model(True, False)
model_manager.get_model(False, False) model_manager.get_model(False, False)
...@@ -41,7 +45,7 @@ class MinerUAPI(ls.LitAPI): ...@@ -41,7 +45,7 @@ class MinerUAPI(ls.LitAPI):
try: try:
pdf_name = str(uuid.uuid4()) pdf_name = str(uuid.uuid4())
output_dir = self.output_dir.joinpath(pdf_name) 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 return output_dir
except Exception as e: except Exception as e:
shutil.rmtree(output_dir, ignore_errors=True) shutil.rmtree(output_dir, ignore_errors=True)
...@@ -73,7 +77,7 @@ class MinerUAPI(ls.LitAPI): ...@@ -73,7 +77,7 @@ class MinerUAPI(ls.LitAPI):
return f.convert_to_pdf() return f.convert_to_pdf()
else: else:
temp_file.write_bytes(file_bytes) 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() return temp_file.with_suffix('.pdf').read_bytes()
else: else:
raise Exception('Unsupported file format') raise Exception('Unsupported file format')
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment