MertandGitHub
23f1f87dd0
* add immich_model * graph per shape * decouple ocr from ort * prepared graphs * use f16 for accelerators * preload all shapes * warm-up pass after preload * keep precision for fp32 graphs * ensure padding for text recognition * optimize postprocessing for fp16 outputs * handle raw ctc logits * use f16 output for rknn * mmap token embedding * read device metadata
43 lines
1.9 KiB
Python
43 lines
1.9 KiB
Python
from typing import Any
|
|
|
|
from immich_ml.config import log, settings
|
|
from immich_ml.models.base import InferenceModel
|
|
from immich_ml.schemas import ModelFormat, ModelType
|
|
|
|
|
|
# TODO: Remove once everything uses the new model graphs
|
|
class TextModel(InferenceModel):
|
|
def __init__(self, model_name: str, **model_kwargs: Any) -> None:
|
|
if settings.legacy_models:
|
|
model_kwargs["model_format"] = ModelFormat.ONNX # the older exports come in no other format
|
|
super().__init__(model_name, **model_kwargs)
|
|
|
|
def download(self) -> None:
|
|
if not settings.legacy_models or self.cached:
|
|
return super().download()
|
|
# here, as only an install still on the older exports has a use for it
|
|
from rapidocr.inference_engine.base import FileInfo, InferSession
|
|
from rapidocr.utils.download_file import DownloadFile, DownloadFileInput
|
|
from rapidocr.utils.typings import EngineType, LangDet, LangRec, OCRVersion, TaskType
|
|
from rapidocr.utils.typings import ModelType as RapidModelType
|
|
|
|
detects = self.model_type == ModelType.DETECTION
|
|
language = self.model_name.split("__")[0] if "__" in self.model_name else "CH"
|
|
model_info = InferSession.get_model_url(
|
|
FileInfo(
|
|
engine_type=EngineType.ONNXRUNTIME,
|
|
ocr_version=OCRVersion.PPOCRV5,
|
|
task_type=TaskType.DET if detects else TaskType.REC,
|
|
lang_type=LangDet.CH if detects else LangRec[language],
|
|
model_type=RapidModelType.MOBILE if "mobile" in self.model_name else RapidModelType.SERVER,
|
|
)
|
|
)
|
|
log.info(
|
|
f"Downloading {self.model_type} model '{self.model_name}' to {self.model_path}. This may take a while."
|
|
)
|
|
DownloadFile.run(
|
|
DownloadFileInput(
|
|
file_url=model_info["model_dir"], sha256=model_info["SHA256"], save_path=self.model_path, logger=log
|
|
)
|
|
)
|