Files
jarvis-models/src/blackbox/blackbox_factory.py
superobk 68c453b219 rebase
2024-05-24 10:50:45 +08:00

64 lines
2.1 KiB
Python

from blackbox.emotion import Emotion
from .chat import Chat
from .audio_chat import AudioChat
from .sentiment import Sentiment
from .tts import TTS
from .asr import ASR
from .audio_to_text import AudioToText
from .blackbox import Blackbox
from .fastchat import Fastchat
from .g2e import G2E
from .text_and_image import TextAndImage
from .melotts import MeloTTS
from .vlms import VLMS
from .chroma_query import ChromaQuery
from .chroma_upsert import ChromaUpsert
from .chroma_chat import ChromaChat
from injector import inject, singleton
@singleton
class BlackboxFactory:
models = {}
@inject
def __init__(self,
audio_to_text: AudioToText,
asr: ASR,
tts: TTS,
sentiment_engine: Sentiment,
emotion: Emotion,
fastchat: Fastchat,
audio_chat: AudioChat,
g2e: G2E,
text_and_image: TextAndImage,
melotts: MeloTTS,
vlms: VLMS,
chroma_query: ChromaQuery,
chroma_upsert: ChromaUpsert,
chroma_chat: ChromaChat,
chat: Chat) -> None:
self.models["audio_to_text"] = audio_to_text
self.models["asr"] = asr
self.models["tts"] = tts
self.models["sentiment_engine"] = sentiment_engine
self.models["emotion"] = emotion
self.models["fastchat"] = fastchat
self.models["audio_chat"] = audio_chat
self.models["g2e"] = g2e
self.models["text_and_image"] = text_and_image
self.models["chroma_query"] = chroma_query
self.models["chroma_upsert"] = chroma_upsert
self.models["chroma_chat"] = chroma_chat
self.models["melotts"] = melotts
self.models["vlms"] = vlms
self.models["chat"] = chat
def __call__(self, *args, **kwargs):
return self.processing(*args, **kwargs)
def get_blackbox(self, blackbox_name: str) -> Blackbox:
model = self.models.get(blackbox_name)
if model is None:
raise ValueError("Invalid Blackbox Type...")
return model