mirror of
https://github.com/immich-app/immich.git
synced 2024-11-28 09:33:27 +02:00
2b1b43a7e4
* modularize model classes * various fixes * expose port * change response * round coordinates * simplify preload * update server * simplify interface simplify * update tests * composable endpoint * cleanup fixes remove unnecessary interface support text input, cleanup * ew camelcase * update server server fixes fix typing * ml fixes update locustfile fixes * cleaner response * better repo response * update tests formatting and typing rename * undo compose change * linting fix type actually fix typing * stricter typing fix detection-only response no need for defaultdict * update spec file update api linting * update e2e * unnecessary dimension * remove commented code * remove duplicate code * remove unused imports * add batch dim
41 lines
1.7 KiB
Python
41 lines
1.7 KiB
Python
from typing import Any
|
|
|
|
from app.models.base import InferenceModel
|
|
from app.models.clip.textual import MClipTextualEncoder, OpenClipTextualEncoder
|
|
from app.models.clip.visual import OpenClipVisualEncoder
|
|
from app.schemas import ModelSource, ModelTask, ModelType
|
|
|
|
from .constants import get_model_source
|
|
from .facial_recognition.detection import FaceDetector
|
|
from .facial_recognition.recognition import FaceRecognizer
|
|
|
|
|
|
def get_model_class(model_name: str, model_type: ModelType, model_task: ModelTask) -> type[InferenceModel]:
|
|
source = get_model_source(model_name)
|
|
match source, model_type, model_task:
|
|
case ModelSource.OPENCLIP | ModelSource.MCLIP, ModelType.VISUAL, ModelTask.SEARCH:
|
|
return OpenClipVisualEncoder
|
|
|
|
case ModelSource.OPENCLIP, ModelType.TEXTUAL, ModelTask.SEARCH:
|
|
return OpenClipTextualEncoder
|
|
|
|
case ModelSource.MCLIP, ModelType.TEXTUAL, ModelTask.SEARCH:
|
|
return MClipTextualEncoder
|
|
|
|
case ModelSource.INSIGHTFACE, ModelType.DETECTION, ModelTask.FACIAL_RECOGNITION:
|
|
return FaceDetector
|
|
|
|
case ModelSource.INSIGHTFACE, ModelType.RECOGNITION, ModelTask.FACIAL_RECOGNITION:
|
|
return FaceRecognizer
|
|
|
|
case _:
|
|
raise ValueError(f"Unknown model combination: {source}, {model_type}, {model_task}")
|
|
|
|
|
|
def from_model_type(model_name: str, model_type: ModelType, model_task: ModelTask, **kwargs: Any) -> InferenceModel:
|
|
return get_model_class(model_name, model_type, model_task)(model_name, **kwargs)
|
|
|
|
|
|
def get_model_deps(model_name: str, model_type: ModelType, model_task: ModelTask) -> list[tuple[ModelType, ModelTask]]:
|
|
return get_model_class(model_name, model_type, model_task).depends
|