immich/machine-learning/app/models/__init__.py
Mert 092a23fd7f
feat(server,ml): remove image tagging (#5903)
* remove image tagging

* updated lock

* fixed tests, improved logging

* be nice

* fixed tests
2023-12-20 20:47:56 -05:00

25 lines
911 B
Python

from typing import Any
from app.schemas import ModelType
from .base import InferenceModel
from .clip import MCLIPEncoder, OpenCLIPEncoder
from .constants import is_insightface, is_mclip, is_openclip
from .facial_recognition import FaceRecognizer
def from_model_type(model_type: ModelType, model_name: str, **model_kwargs: Any) -> InferenceModel:
match model_type:
case ModelType.CLIP:
if is_openclip(model_name):
return OpenCLIPEncoder(model_name, **model_kwargs)
elif is_mclip(model_name):
return MCLIPEncoder(model_name, **model_kwargs)
case ModelType.FACIAL_RECOGNITION:
if is_insightface(model_name):
return FaceRecognizer(model_name, **model_kwargs)
case _:
raise ValueError(f"Unknown model type {model_type}")
raise ValueError(f"Unknown ${model_type} model {model_name}")