Эх сурвалжийг харах

perf(model): add threading lock for OCR model initialization

- Introduce a lock to synchronize access to OCR model initialization- This change improves thread safety when multiple threads access the OCR model concurrently
- The lock ensures that the OCR model is initialized only once, even in multi-threaded scenarios
myhloli 11 сар өмнө
parent
commit
044780951a

+ 10 - 2
magic_pdf/model/sub_modules/model_init.py

@@ -82,9 +82,12 @@ def ocr_model_init(show_log: bool = False,
     return model
 
 
+from threading import Lock
+
 class AtomModelSingleton:
     _instance = None
     _models = {}
+    _lock = Lock()
 
     def __new__(cls, *args, **kwargs):
         if cls._instance is None:
@@ -95,8 +98,13 @@ class AtomModelSingleton:
         lang = kwargs.get('lang', None)
         layout_model_name = kwargs.get('layout_model_name', None)
         key = (atom_model_name, layout_model_name, lang)
-        if key not in self._models:
-            self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
+        if atom_model_name == AtomicModel.OCR:
+            with self._lock:
+                if key not in self._models:
+                    self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
+        else:
+            if key not in self._models:
+                self._models[key] = atom_model_init(model_name=atom_model_name, **kwargs)
         return self._models[key]