File size: 2,334 Bytes
f4b0f7d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 | 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())
|