Commit 054abe33 authored by myhloli's avatar myhloli

refactor(magic_pdf): implement model singleton pattern for custom models

Introduce a Singleton pattern to manage custom models in the magic_pdf module.
This change improves the efficiency by ensuring that a single instance of the
custom model is created and reused, thereby reducing the overhead of multiple
instantiate calls for the same model configuration.
parent b59d511d
...@@ -48,7 +48,23 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200) -> list: ...@@ -48,7 +48,23 @@ def load_images_from_pdf(pdf_bytes: bytes, dpi=200) -> list:
return images return images
def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False): class ModelSingleton:
_instance = None
_models = {}
def __new__(cls, *args, **kwargs):
if cls._instance is None:
cls._instance = super().__new__(cls)
return cls._instance
def get_model(self, ocr: bool, show_log: bool):
key = (ocr, show_log)
if key not in self._models:
self._models[key] = custom_model_init(ocr=ocr, show_log=show_log)
return self._models[key]
def custom_model_init(ocr: bool = False, show_log: bool = False):
model = None model = None
if model_config.__model_mode__ == "lite": if model_config.__model_mode__ == "lite":
...@@ -76,6 +92,14 @@ def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False): ...@@ -76,6 +92,14 @@ def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
logger.error("use_inside_model is False, not allow to use inside model") logger.error("use_inside_model is False, not allow to use inside model")
exit(1) exit(1)
return custom_model
def doc_analyze(pdf_bytes: bytes, ocr: bool = False, show_log: bool = False):
model_manager = ModelSingleton()
custom_model = model_manager.get_model(ocr, show_log)
images = load_images_from_pdf(pdf_bytes) images = load_images_from_pdf(pdf_bytes)
model_json = [] model_json = []
......
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