mirror of
https://github.com/agent0ai/agent-zero.git
synced 2026-08-30 10:44:01 +00:00
5a96eda078
Cache the underlying SentenceTransformer by its effective constructor options so startup preload and first-chat initialization share one instance while wrappers keep their own runtime configuration. Pass the runtime embedding configuration through preload and cover reuse, concurrency, option changes, and failure handling.
64 lines
2.3 KiB
Python
64 lines
2.3 KiB
Python
import asyncio
|
|
from helpers import runtime
|
|
from helpers.print_style import PrintStyle
|
|
import models
|
|
from plugins._kokoro_tts.helpers import runtime as kokoro_tts_runtime
|
|
from plugins._whisper_stt.helpers import runtime as whisper_stt_runtime
|
|
|
|
|
|
async def preload():
|
|
try:
|
|
# preload whisper model
|
|
async def preload_whisper():
|
|
if not whisper_stt_runtime.is_globally_enabled():
|
|
return None
|
|
try:
|
|
config = whisper_stt_runtime.get_config()
|
|
return await whisper_stt_runtime.preload(str(config["model_size"]))
|
|
except Exception as e:
|
|
PrintStyle().error(f"Error in preload_whisper: {e}")
|
|
|
|
# preload embedding model
|
|
async def preload_embedding():
|
|
try:
|
|
from plugins._model_config.helpers.model_config import get_embedding_model_config_object
|
|
emb_cfg = get_embedding_model_config_object()
|
|
if emb_cfg.provider.lower() == "huggingface":
|
|
emb_mod = models.get_embedding_model(
|
|
emb_cfg.provider,
|
|
emb_cfg.name,
|
|
model_config=emb_cfg,
|
|
**emb_cfg.build_kwargs(),
|
|
)
|
|
emb_txt = await emb_mod.aembed_query("test")
|
|
return emb_txt
|
|
except Exception as e:
|
|
PrintStyle().error(f"Error in preload_embedding: {e}")
|
|
|
|
# preload kokoro tts model if enabled
|
|
async def preload_kokoro():
|
|
if not kokoro_tts_runtime.is_globally_enabled():
|
|
return None
|
|
try:
|
|
return await kokoro_tts_runtime.preload()
|
|
except Exception as e:
|
|
PrintStyle().error(f"Error in preload_kokoro: {e}")
|
|
|
|
# async tasks to preload
|
|
tasks = [
|
|
preload_embedding(),
|
|
# preload_whisper(),
|
|
# preload_kokoro()
|
|
]
|
|
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
PrintStyle().print("Preload completed.")
|
|
except Exception as e:
|
|
PrintStyle().error(f"Error in preload: {e}")
|
|
|
|
|
|
# preload transcription model
|
|
if __name__ == "__main__":
|
|
PrintStyle().print("Running preload...")
|
|
runtime.initialize()
|
|
asyncio.run(preload())
|