feat: models lazy loader

src/blackbox/blackbox_factory.py

@@ -1,57 +1,112 @@
-from .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
+from injector import inject, singleton, Injector
+from ..configuration import BlackboxConf
+
+blackboxConf = Injector().get(BlackboxConf)
+
+def model_loader(lazy=blackboxConf.lazyloading):
+    def load(init_fun):
+        model = None
+        if not lazy:
+            model = init_fun()
+        def inner():
+            nonlocal model
+            if model is None:
+                model = init_fun()
+            return model
+        return inner
+    return load
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def text_and_image_loader():
+    from .text_and_image import TextAndImage
+    return Injector().get(TextAndImage)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def audio_chat_loader():
+    from .audio_chat import AudioChat
+    return Injector().get(AudioChat)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def sentiment_loader():
+    from .sentiment import Sentiment
+    return Injector().get(Sentiment)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def g2e_loader():
+    from .g2e import G2E
+    return Injector().get(G2E)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def asr_loader():
+    from .asr import ASR
+    return Injector().get(ASR)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def vlms_loader():
+    from .vlms import VLMS
+    return Injector().get(VLMS)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def melotts_loader():
+    from .melotts import MeloTTS
+    return Injector().get(MeloTTS)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def tts_loader():
+    from .tts import TTS
+    return Injector().get(TTS)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def emotion_loader():
+    from .emotion import Emotion
+    return Injector().get(Emotion)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def fastchat_loader():
+    from .fastchat import Fastchat
+    return Injector().get(Fastchat)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def chat_loader():
+    from .chat import Chat
+    return Injector().get(Chat)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def chroma_query_loader():
+    from .chroma_query import ChromaQuery
+    return Injector().get(ChromaQuery)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def chroma_upsert_loader():
+    from .chroma_upsert import ChromaUpsert
+    return Injector().get(ChromaUpsert)
+
+@model_loader(lazy=blackboxConf.lazyloading)
+def chroma_chat_load():
+    from .chroma_chat import ChromaChat
+    return Injector().get(ChromaChat)
 
 @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 __init__(self,) -> None:
+        self.models["asr"] = asr_loader
+        self.models["tts"] = tts_loader
+        self.models["sentiment_engine"] = sentiment_loader
+        self.models["emotion"] = emotion_loader
+        self.models["fastchat"] = fastchat_loader
+        self.models["audio_chat"] = audio_chat_loader
+        self.models["g2e"] = g2e_loader
+        self.models["text_and_image"] = text_and_image_loader
+        self.models["chroma_query"] = chroma_query_loader
+        self.models["chroma_upsert"] = chroma_upsert_loader
+        self.models["chroma_chat"] = chroma_chat_load
+        self.models["melotts"] = melotts_loader
+        self.models["vlms"] = vlms_loader
+        self.models["chat"] = chat_loader
 
     def __call__(self, *args, **kwargs):
         return self.processing(*args, **kwargs)
@@ -60,4 +115,4 @@ class BlackboxFactory:
         model = self.models.get(blackbox_name)
         if model is None:
             raise ValueError("Invalid Blackbox Type...")
-        return model
+        return model()

src/configuration.py

@@ -1,4 +1,5 @@
 
+from dataclasses import dataclass
 from injector import inject,singleton
 import yaml
 import sys
@@ -102,3 +103,12 @@ class EnvConf():
         self.version = "0.0.1"
         self.host = config.get("env.host", default="0.0.0.0")
         self.port = config.get("env.port", default="8080")
+
+@singleton
+@dataclass
+class BlackboxConf():
+    lazyloading: bool
+
+    @inject
+    def __init__(self, config: Configuration) -> None:
+        self.lazyloading = bool(config.get("blackbox.lazyloading", default=False))