|
import gc |
|
import os |
|
from contextlib import contextmanager |
|
from time import time |
|
from typing import Optional |
|
|
|
import streamlit as st |
|
|
|
from tortoise.api import TextToSpeech |
|
from tortoise.utils.audio import load_voices |
|
|
|
|
|
@contextmanager |
|
def timeit(desc=""): |
|
start = time() |
|
yield |
|
print(f"{desc} took {time() - start:.2f} seconds") |
|
|
|
|
|
@st.cache_resource(max_entries=1) |
|
def load_model( |
|
model_dir, |
|
high_vram, |
|
kv_cache, |
|
ar_checkpoint, |
|
diff_checkpoint, |
|
): |
|
gc.collect() |
|
return TextToSpeech( |
|
models_dir=model_dir, |
|
high_vram=high_vram, |
|
kv_cache=kv_cache, |
|
ar_checkpoint=ar_checkpoint, |
|
diff_checkpoint=diff_checkpoint, |
|
) |
|
|
|
|
|
@st.cache_data |
|
def list_voices(extra_voices_dir: Optional[str]): |
|
voices = ["random"] |
|
if extra_voices_dir and os.path.isdir(extra_voices_dir): |
|
voices.extend(os.listdir(extra_voices_dir)) |
|
extra_voices_ls = [extra_voices_dir] |
|
else: |
|
extra_voices_ls = [] |
|
voices.extend( |
|
[v for v in os.listdir("tortoise/voices") if v != "cond_latent_example"] |
|
) |
|
|
|
return voices, extra_voices_ls |
|
|
|
|
|
@st.cache_resource(max_entries=1) |
|
def load_voice_conditionings(voice, extra_voices_ls): |
|
gc.collect() |
|
voice_samples, conditioning_latents = load_voices(voice, extra_voices_ls) |
|
return voice_samples, conditioning_latents |
|
|
