import argparse import datetime as dt import warnings from pathlib import Path import ffmpeg import gradio as gr import IPython.display as ipd import joblib as jl import numpy as np import soundfile as sf import torch from tqdm.auto import tqdm from diff_ttsg.hifigan.config import v1 from diff_ttsg.hifigan.denoiser import Denoiser from diff_ttsg.hifigan.env import AttrDict from diff_ttsg.hifigan.models import Generator as HiFiGAN from diff_ttsg.models.diff_ttsg import Diff_TTSG from diff_ttsg.text import cmudict, sequence_to_text, text_to_sequence from diff_ttsg.text.symbols import symbols from diff_ttsg.utils.model import denormalize from diff_ttsg.utils.utils import intersperse, plot_tensor from pymo.preprocessing import MocapParameterizer from pymo.viz_tools import render_mp4 from pymo.writers import BVHWriter device = torch.device("cuda" if torch.cuda.is_available() else "cpu") DIFF_TTSG_CHECKPOINT = "diff_ttsg_checkpoint.ckpt" HIFIGAN_CHECKPOINT = "g_02500000" MOTION_PIPELINE = "diff_ttsg/resources/data_pipe.expmap_86.1328125fps.sav" CMU_DICT_PATH = "diff_ttsg/resources/cmu_dictionary" OUTPUT_FOLDER = "synth_output" # Model loading tools def load_model(checkpoint_path): model = Diff_TTSG.load_from_checkpoint(checkpoint_path, map_location=device) model.eval() return model # Vocoder loading tools def load_vocoder(checkpoint_path): 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 # Setup text preprocessing cmu = cmudict.CMUDict(CMU_DICT_PATH) def process_text(text: str): x = torch.LongTensor(intersperse(text_to_sequence(text, dictionary=cmu), len(symbols))).to(device)[None] x_lengths = torch.LongTensor([x.shape[-1]]).to(device) x_phones = sequence_to_text(x.squeeze(0).tolist()) return { 'x_orig': text, 'x': x, 'x_lengths': x_lengths, 'x_phones': x_phones } # Setup motion visualisation motion_pipeline = jl.load(MOTION_PIPELINE) bvh_writer = BVHWriter() mocap_params = MocapParameterizer("position") ## Load models model = load_model(DIFF_TTSG_CHECKPOINT) vocoder = load_vocoder(HIFIGAN_CHECKPOINT) denoiser = Denoiser(vocoder, mode='zeros') # Synthesis functions @torch.inference_mode() def synthesise(text, mel_timestep, motion_timestep, length_scale, mel_temp, motion_temp): ## Number of timesteps to run the reverse denoising process n_timesteps = { 'mel': mel_timestep, 'motion': motion_timestep, } ## Sampling temperature temperature = { 'mel': mel_temp, 'motion': motion_temp } text_processed = process_text(text) t = dt.datetime.now() output = model.synthesise( text_processed['x'], text_processed['x_lengths'], n_timesteps=n_timesteps, temperature=temperature, stoc=False, spk=None, length_scale=length_scale ) t = (dt.datetime.now() - t).total_seconds() print(f'RTF: {t * 22050 / (output["mel"].shape[-1] * 256)}') output.update(text_processed) # merge everything to one dict return output @torch.inference_mode() def to_waveform(mel, vocoder): audio = vocoder(mel).clamp(-1, 1) audio = denoiser(audio.squeeze(0)).cpu().squeeze() return audio def to_bvh(motion): with warnings.catch_warnings(): warnings.simplefilter("ignore") return motion_pipeline.inverse_transform([motion.cpu().squeeze(0).T]) def save_to_folder(filename: str, output: dict, folder: str): folder = Path(folder) folder.mkdir(exist_ok=True, parents=True) np.save(folder / f'{filename}', output['mel'].cpu().numpy()) sf.write(folder / f'{filename}.wav', output['waveform'], 22050, 'PCM_24') with open(folder / f'{filename}.bvh', 'w') as f: bvh_writer.write(output['bvh'], f) def to_stick_video(filename, bvh, folder): folder = Path(folder) folder.mkdir(exist_ok=True, parents=True) with warnings.catch_warnings(): warnings.simplefilter("ignore") X_pos = mocap_params.fit_transform([bvh]) print(f"rendering {filename} ...") render_mp4(X_pos[0], folder / f'{filename}.mp4', axis_scale=200) def combine_audio_video(filename: str, folder: str): print("Combining audio and video") folder = Path(folder) folder.mkdir(exist_ok=True, parents=True) input_video = ffmpeg.input(str(folder / f'{filename}.mp4')) input_audio = ffmpeg.input(str(folder / f'{filename}.wav')) output_filename = folder / f'{filename}_audio.mp4' ffmpeg.concat(input_video, input_audio, v=1, a=1).output(str(output_filename)).run(overwrite_output=True) print(f"Final output with audio: {output_filename}") def run(text, output, mel_timestep, motion_timestep, length_scale, mel_temp, motion_temp): print("Running synthesis") output = synthesise(text, mel_timestep, motion_timestep, length_scale, mel_temp, motion_temp) output['waveform'] = to_waveform(output['mel'], vocoder) output['bvh'] = to_bvh(output['motion'])[0] save_to_folder('temp', output, OUTPUT_FOLDER) return ( output, output['x_phones'], plot_tensor(output['mel'].squeeze().cpu().numpy()), plot_tensor(output['motion'].squeeze().cpu().numpy()), str(Path(OUTPUT_FOLDER) / f'temp.wav'), gr.update(interactive=True) ) def visualize_it(output): to_stick_video('temp', output['bvh'], OUTPUT_FOLDER) combine_audio_video('temp', OUTPUT_FOLDER) return str(Path(OUTPUT_FOLDER) / 'temp_audio.mp4') with gr.Blocks() as demo: output = gr.State(value=None) with gr.Box(): with gr.Row(): gr.Markdown("# Diff-TTSG: Denoising probabilistic integrated speech and gesture synthesis") with gr.Row(): gr.Markdown("### Read more about it at: [https://shivammehta25.github.io/Diff-TTSG/](https://shivammehta25.github.io/Diff-TTSG/)") with gr.Row(): gr.Markdown("# Text Input") with gr.Row(): gr.Markdown("Enter , to insert pause and ; for breathing pause.") with gr.Row(): gr.Markdown("It is recommended to give spaces between punctuations and words.") with gr.Row(): text = gr.Textbox(label="Text Input") with gr.Row(): examples = gr.Examples(examples=[ "Hello world ! This is a demo of Diff T T S G .", "And the train stopped, The door opened. I got out first, then Jack Kane got out, Ronan got out, Louise got out.", ], inputs=[text]) with gr.Box(): with gr.Row(): gr.Markdown("### Hyper parameters") with gr.Row(): mel_timestep = gr.Slider(label="Number of timesteps (mel)", minimum=0, maximum=1000, step=1, value=50, interactive=True) motion_timestep = gr.Slider(label="Number of timesteps (motion)", minimum=0, maximum=1000, step=1, value=500, interactive=True) length_scale = gr.Slider(label="Length scale (Speaking rate)", minimum=0.01, maximum=3.0, step=0.05, value=1.15, interactive=True) mel_temp = gr.Slider(label="Sampling temperature (mel)", minimum=0.01, maximum=5.0, step=0.05, value=1.3, interactive=True) motion_temp = gr.Slider(label="Sampling temperature (motion)", minimum=0.01, maximum=5.0, step=0.05, value=1.5, interactive=True) synth_btn = gr.Button("Synthesise") with gr.Box(): with gr.Row(): gr.Markdown("### Phonetised text") with gr.Row(): phonetised_text = gr.Textbox(label="Phonetised text", interactive=False) with gr.Box(): with gr.Row(): mel_spectrogram = gr.Image(interactive=False, label="Mel spectrogram") motion_representation = gr.Image(interactive=False, label="Motion representation") with gr.Row(): audio = gr.Audio(interactive=False, label="Audio") with gr.Box(): with gr.Row(): gr.Markdown("### Generate stick figure visualisation") with gr.Row(): gr.Markdown("(This will take a while)") with gr.Row(): visualize = gr.Button("Visualize", interactive=False) with gr.Row(): video = gr.Video(label="Video", interactive=False) synth_btn.click( fn=run, inputs=[ text, output, mel_timestep, motion_timestep, length_scale, mel_temp, motion_temp ], outputs=[ output, phonetised_text, mel_spectrogram, motion_representation, audio, # video, visualize ], api_name="diff_ttsg") visualize.click( fn=visualize_it, inputs=[output], outputs=[video], ) demo.queue(1) demo.launch()