from pathlib import Path import torchaudio import gradio as gr import numpy as np import torch import json from hifigan.config import v1 from hifigan.denoiser import Denoiser from hifigan.env import AttrDict from hifigan.models import Generator as HiFiGAN from pflow.models.pflow_tts import pflowTTS from pflow.text import text_to_sequence, sequence_to_text 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=649.ckpt' #PFLOW_MODEL_PATH = 'checkpoint_m_epoch=054.ckpt' VOCODER_MODEL_PATH = 'checkpoints/pytorch_model.bin' HIFIGAN_MODEL_PATH = 'checkpoints/g_00120000' transform = torchaudio.transforms.Vol(gain=-32, gain_type="db") wav, sr = torchaudio.load('prompt.wav') prompt = mel_spectrogram( wav, 1024, 80, 22050, 256, 1024, 0, 8000, center=False, )[:,:,:264] def process_text(text: str, device: torch.device): x = torch.tensor( intersperse(text_to_sequence(text, ["ukr_cleaners"]), 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_hifigan(checkpoint_path, device): h = AttrDict(v1) hifigan = HiFiGAN(h).to(device) hifigan.load_state_dict(torch.load(checkpoint_path, map_location=device)["generator"]) _ = hifigan.eval() hifigan.remove_weight_norm() return hifigan def load_vocos(checkpoint_path, config_path, device): model = Vocos.from_hparams(config_path) raw_model = torch.load(checkpoint_path, map_location=torch.device('cpu')) 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 to_waveform(mel, vocoder, denoiser=None): return vocoder.decode(mel).cpu().squeeze() # audio = vocoder(mel).clamp(-1, 1) # if denoiser is not None: # audio = denoiser(audio.squeeze(), strength=0.00025).cpu().squeeze() # return audio.cpu().squeeze() 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() #hifigan = load_hifigan(HIFIGAN_MODEL_PATH, device) vocos = load_vocos(VOCODER_MODEL_PATH, 'config.yaml', device) #vocos_44100 = load_vocos('checkpoints/vocos_checkpoint_epoch=4_step=93440_val_loss=5.2596_44100_10.ckpt', 'vocos.yaml', device) denoiser = None#Denoiser(vocoder, mode="zeros") @torch.inference_mode() def synthesise(text, speed): if len(text) > 1000: raise gr.Error("Текст повинен бути коротшим за 1000 символів.") text_processed = process_text(text.strip(), device) output = model.synthesise( text_processed["x"].to(device), text_processed["x_lengths"].to(device), n_timesteps=40, temperature=0.0, length_scale=1/speed, prompt=normalize(prompt, model.mel_mean, model.mel_std).to(device), guidance_scale=1.5 ) waveform_vocos = vocos.decode(output["mel"]).cpu().squeeze() #waveform_vocos_44100 = vocos_44100.decode(output["mel"]).cpu().squeeze() #waveform_hifigan = hifigan(output["mel"]).clamp(-1, 1).cpu().squeeze() #transform = torchaudio.transforms.Vol(gain=-18, gain_type="db") return text_processed['x_phones'][1::2], (22050, waveform_vocos.numpy()) description = f''' # Експериментальна апка для генерації аудіо з тексту. pflow checkpoint {PFLOW_MODEL_PATH} vocoder: Vocos - {VOCODER_MODEL_PATH} ''' if __name__ == "__main__": i = gr.Interface( fn=synthesise, description=description, inputs=[ gr.Text(label='Текст для синтезу:', lines=5, max_lines=10), gr.Slider(minimum=0.6, maximum=2.0, label="Швидкість", value=1.0) ], outputs=[ gr.Text(label='Фонемізований текст:', lines=5), # gr.Audio( # label="Vocos 44100 аудіо:", # autoplay=False, # streaming=False, # type="numpy", # ), gr.Audio( label="Vocos аудіо:", autoplay=False, streaming=False, type="numpy", ), # gr.Audio( # label="HIFIGAN аудіо:", # autoplay=False, # streaming=False, # type="numpy", # ) ], allow_flagging ='manual', flagging_options=[("Якщо дуже погоне аудіо, тисни цю кнопку.", "negative")], cache_examples=True, title='', # description=description, # article=article, # examples=examples, ) i.queue(max_size=20, default_concurrency_limit=4) i.launch(share=False, server_name="0.0.0.0")