app.py

import os
import torchaudio
import torch
import gradio as gr


from pflow.models.pflow_tts import pflowTTS
from pflow.text import  text_to_sequence, sequence_to_text, cleaned_text_to_sequence
from pflow.utils.utils import intersperse
from pflow.data.text_mel_datamodule import mel_spectrogram
from pflow.utils.model import normalize
from vocos import Vocos


PFLOW_MODEL_PATH = 'checkpoints/checkpoint_epoch=100.ckpt'
VOCODER22_MODEL_PATH = 'BSC-LT/vocos-mel-22khz'
VOCODER44_MODEL_PATH = 'patriotyk/vocos-mel-hifigan-compat-44100khz'


volnorm = torchaudio.transforms.Vol(gain=-15, gain_type="db")
prompts_dir = 'prompts'
prompts_list = sorted(os.listdir(prompts_dir), key=lambda x: x.split('.')[0])

def process_text(text: str, device: torch.device, ipa=False):
    if ipa:
        seq = cleaned_text_to_sequence(text)
    else:
        seq = text_to_sequence(text, ["ukr_cleaners"])
    
    x = torch.tensor(
        intersperse(seq, 0),
        dtype=torch.long,
        device=device,
    )[None]
    x_lengths = torch.tensor([x.shape[-1]], dtype=torch.long, device=device)
    x_phones = sequence_to_text(x.squeeze(0).tolist())
    return {"x_orig": text, "x": x, "x_lengths": x_lengths, 'x_phones':x_phones}



def load_vocos(checkpoint_path, config_path, device):
    model = Vocos.from_hparams(config_path).to(device)

    raw_model = torch.load(checkpoint_path, map_location=device)
    raw_model = raw_model if 'state_dict' not in raw_model else raw_model['state_dict']
    model.load_state_dict(raw_model, strict=False)
    model.eval()
    return model



def get_device():
    if torch.cuda.is_available():
        print("[+] GPU Available! Using GPU")
        device = torch.device("cuda")
    else:
        print("[-] GPU not available or forced CPU run! Using CPU")
        device = torch.device("cpu")
    return device


device = get_device()
model = pflowTTS.load_from_checkpoint(PFLOW_MODEL_PATH, map_location=device)
_ = model.eval()

vocos_44100 = Vocos.from_pretrained(VOCODER44_MODEL_PATH).to(device)
vocos_22050 = Vocos.from_pretrained(VOCODER22_MODEL_PATH).to(device)


@torch.inference_mode()
def synthesise(text, ipa, prompt_selection, audio_prompt, temperature, speed):
    print(text, prompt_selection, temperature, speed)
    if len(text) > 1000:
        raise gr.Error("Текст повинен бути коротшим за 1000 символів.")

    if audio_prompt:
        wav, sr = torchaudio.load(audio_prompt)
        wav = torchaudio.functional.resample(wav, orig_freq=sr, new_freq=44100)
    else:
        prompt_audio_path = os.path.join(prompts_dir, prompt_selection)
        wav, _ = torchaudio.load(prompt_audio_path)

    if ipa:
        text_processed = process_text(ipa, device, ipa=True)
    else:
        text_processed = process_text(text.strip(), device, ipa=False)

    prompt = mel_spectrogram(volnorm(wav), 2048, 80, 44100, 512, 2048, 0, 8000, center=False)[:,:,:264]

    
    output = model.synthesise(
        text_processed["x"].to(device),
        text_processed["x_lengths"].to(device),
        n_timesteps=40,
        temperature=temperature,
        length_scale=1/speed,
        prompt=normalize(prompt, model.mel_mean, model.mel_std).to(device),

        guidance_scale=1.8

    )
    
    waveform_vocos_44100 = vocos_44100.decode(output["mel"]).cpu().squeeze()
    waveform_vocos = vocos_22050.decode(output["mel"]).cpu().squeeze()
    return text_processed['x_phones'][1::2], (44100, waveform_vocos_44100.numpy()), (22050, waveform_vocos.numpy())


description = f'''
Модель натренована на приватному датасеті з аудіо книжок створненому за допомогою програми
[narizaka](https://github.com/patriotyk/narizaka).
Програма може не коректно визначати деякі наголоси і не дуже добре перетворює цифри, акроніми і різні скорочення в словесну форму.
На даний момент, відкритого рішення для української мови для цих проблем нема, тому якщо у вас є запитання,
чи ви хочете допомогти їх вирішити приєднуйтесь до нашого чату в [телеграм](https://t.me/speech_synthesis_uk) або [discord](https://discord.gg/yVAjkBgmt4) 
'''


if __name__ == "__main__":
    i = gr.Interface(
        fn=synthesise,
        description=description,
        inputs=[
            gr.Text(label='Текст для синтезу:', lines=5, max_lines=10),
            gr.Text(label='Aбо IPA:', lines=5, max_lines=10),
            gr.Dropdown(label="Виберіть промт", choices=prompts_list, value=prompts_list[0]),
            gr.Audio(label="Або завантажте свій:", interactive=True, type='filepath', max_length=300, waveform_options={'waveform_progress_color': '#3C82F6'}),
            gr.Slider(minimum=0.0, maximum=1.0, label="Шум", value=0.7),
            gr.Slider(minimum=0.6, maximum=2.0, label="Швидкість", value=1.1)
        ],
        outputs=[
            gr.Text(label='Фонемізований текст:', lines=5),
            gr.Audio(
                        label="Vocos 44100 аудіо:",
                        autoplay=False,
                        streaming=False,
                        type="numpy",
                    ),
            gr.Audio(
                        label="Vocos 22050 аудіо:",
                        autoplay=False,
                        streaming=False,
                        type="numpy",
                    ),
            
        ],
        allow_flagging ='manual',
        cache_examples=True,
        title='Генерація української мови за допомогою pflowtts.',
        examples=[ 
            ['Мені тринадцятий минало. Я пас ягнята за селом. Чи то так сонечко сіяло, Чи так мені чого було? Мені так любо, любо стало, Неначе в бога. Уже прокликали до паю, А я собі у бур\'яні Молюся богу І не знаю, Чого маленькому мені Тоді так приязно молилось, Чого так весело було?', "meˈnʲi trɪˈnad͡zʲt͡sʲɐtɪi̯ mɪˈnaɫɔ. jɐ pɐs jɐɦˈnʲatɐ zɐ seˈɫɔm. t͡ʃɪ tɔ tɐk ˈsɔnet͡ʃkɔ sʲiˈjɐɫɔ, t͡ʃɪ tɐk meˈnʲi t͡ʃɔˈɦɔ bʊˈɫɔ? meˈnʲi tɐk ˈlʲubɔ, ˈlʲubɔ ˈstaɫɔ, neˈnat͡ʃe ʋ ˈbɔɦɐ.  ʊˈʒɛ prɔkɫɪkɐɫɪ dɔ ˈpajʊ, ɐ jɐ soˈbʲi ʊ bur-jɐˈnʲi moˈlʲusʲɐ ˈbɔɦʊ i ne ˈznajʊ, t͡ʃɔˈɦɔ mɐˈɫɛnʲkɔmʊ meˈnʲi toˈdʲi tɐk ˈprɪjɐznɔ mɔˈɫɪɫɔsʲ, t͡ʃɔˈɦɔ tɐk ˈʋɛseɫɔ bʊˈɫɔ?", '', 'prompts/speaker_22.wav', 0.6, 1.1],
            ['Ти, малий, скажи малому, хай малий малому скаже, хай малий теля прив\'яже.', '',  '', 'prompts/speaker_11.wav', 0.4, 1.1 ],
            ['По мірі розвитку клубу зростатиме і кількість його членів, а отже, команда менеджменту теж буде пропорційно збільшуватись. Яка ж команда потрібна клубу, що налічує, скажімо, сто осіб, і які компетенції повинна мати?', '',  '', 'prompts/speaker_20.wav', 0.7, 1.1],
            ['Да ти дєтка гоніш! один рік? І що? Як ви задрали нити, рік вона не може, в когось діти мруть в день народження, викидні, а вона, бляха, рік не може, купи собі рожеве поні і реви побільше, дурепа.', 'dɐ tɪ dʲetkɐ ɦɔnʲiʃ! ɔˈdɪn rʲik? i ʃt͡ʃɔ? jɐk ʋɪ zɐˈdraɫɪ ˈnɪtɪ, rʲik wɔˈna ne ˈmɔʒe, ʋ kɔɦɔsʲ ˈdʲitɪ mrʊtʲ ʋ denʲ nɐˈrɔd͡ʒenʲːɐ, ˈʋɪkɪdʲnʲi, ɐ wɔˈna, ˈblʲaxɐ, rʲik ne ˈmɔʒe, kʊpɪ soˈbʲi rɔˈʒɛʋe ˈpɔnʲi i reʋɪ poˈbʲilʲʃe, dʊˈrɛpɐ.', '', 'prompts/speaker_5.wav', 0.7, 1.2]
        ],
    )
    i.queue(max_size=20, default_concurrency_limit=4)
    i.launch(share=False, server_name="0.0.0.0")