Spaces:
Running
Running
import os | |
# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" | |
# os.environ["CUDA_VISIBLE_DEVICES"] = "5" # these are only used if developping locally | |
import gradio as gr | |
import torch | |
import torchaudio | |
from data.tokenizer import ( | |
AudioTokenizer, | |
TextTokenizer, | |
) | |
from models import voicecraft | |
import io | |
import numpy as np | |
import random | |
import spaces | |
whisper_model, voicecraft_model = None, None | |
def seed_everything(seed): | |
if seed != -1: | |
os.environ['PYTHONHASHSEED'] = str(seed) | |
random.seed(seed) | |
np.random.seed(seed) | |
torch.manual_seed(seed) | |
torch.cuda.manual_seed(seed) | |
torch.backends.cudnn.benchmark = False | |
torch.backends.cudnn.deterministic = True | |
def load_models(whisper_model_choice, voicecraft_model_choice): | |
global whisper_model, voicecraft_model | |
if whisper_model_choice is not None: | |
import whisper | |
from whisper.tokenizer import get_tokenizer | |
whisper_model = { | |
"model": whisper.load_model(whisper_model_choice), | |
"tokenizer": get_tokenizer(multilingual=False) | |
} | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
voicecraft_name = f"{voicecraft_model_choice}.pth" | |
ckpt_fn = f"./pretrained_models/{voicecraft_name}" | |
encodec_fn = "./pretrained_models/encodec_4cb2048_giga.th" | |
if not os.path.exists(ckpt_fn): | |
os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/{voicecraft_name}\?download\=true") | |
os.system(f"mv {voicecraft_name}\?download\=true ./pretrained_models/{voicecraft_name}") | |
if not os.path.exists(encodec_fn): | |
os.system(f"wget https://huggingface.co/pyp1/VoiceCraft/resolve/main/encodec_4cb2048_giga.th") | |
os.system(f"mv encodec_4cb2048_giga.th ./pretrained_models/encodec_4cb2048_giga.th") | |
ckpt = torch.load(ckpt_fn, map_location="cpu") | |
model = voicecraft.VoiceCraft(ckpt["config"]) | |
model.load_state_dict(ckpt["model"]) | |
model.to(device) | |
model.eval() | |
voicecraft_model = { | |
"ckpt": ckpt, | |
"model": model, | |
"text_tokenizer": TextTokenizer(backend="espeak"), | |
"audio_tokenizer": AudioTokenizer(signature=encodec_fn) | |
} | |
return gr.Accordion() | |
def transcribe(seed, audio_path): | |
if whisper_model is None: | |
raise gr.Error("Whisper model not loaded") | |
seed_everything(seed) | |
number_tokens = [ | |
i | |
for i in range(whisper_model["tokenizer"].eot) | |
if all(c in "0123456789" for c in whisper_model["tokenizer"].decode([i]).removeprefix(" ")) | |
] | |
result = whisper_model["model"].transcribe(audio_path, suppress_tokens=[-1] + number_tokens, word_timestamps=True) | |
words = [word_info for segment in result["segments"] for word_info in segment["words"]] | |
transcript = result["text"] | |
transcript_with_start_time = " ".join([f"{word['start']} {word['word']}" for word in words]) | |
transcript_with_end_time = " ".join([f"{word['word']} {word['end']}" for word in words]) | |
choices = [f"{word['start']} {word['word']} {word['end']}" for word in words] | |
return [ | |
transcript, transcript_with_start_time, transcript_with_end_time, | |
gr.Dropdown(value=choices[-1], choices=choices, interactive=True), # prompt_to_word | |
gr.Dropdown(value=choices[0], choices=choices, interactive=True), # edit_from_word | |
gr.Dropdown(value=choices[-1], choices=choices, interactive=True), # edit_to_word | |
words | |
] | |
def get_output_audio(audio_tensors, codec_audio_sr): | |
result = torch.cat(audio_tensors, 1) | |
buffer = io.BytesIO() | |
torchaudio.save(buffer, result, int(codec_audio_sr), format="wav") | |
buffer.seek(0) | |
return buffer.read() | |
def run(seed, left_margin, right_margin, codec_audio_sr, codec_sr, top_k, top_p, temperature, | |
stop_repetition, sample_batch_size, kvcache, silence_tokens, | |
audio_path, word_info, transcript, smart_transcript, | |
mode, prompt_end_time, edit_start_time, edit_end_time, | |
split_text, selected_sentence, previous_audio_tensors): | |
if voicecraft_model is None: | |
raise gr.Error("VoiceCraft model not loaded") | |
if smart_transcript and (word_info is None): | |
raise gr.Error("Can't use smart transcript: whisper transcript not found") | |
seed_everything(seed) | |
if mode == "Long TTS": | |
if split_text == "Newline": | |
sentences = transcript.split('\n') | |
else: | |
from nltk.tokenize import sent_tokenize | |
sentences = sent_tokenize(transcript.replace("\n", " ")) | |
elif mode == "Rerun": | |
colon_position = selected_sentence.find(':') | |
selected_sentence_idx = int(selected_sentence[:colon_position]) | |
sentences = [selected_sentence[colon_position + 1:]] | |
else: | |
sentences = [transcript.replace("\n", " ")] | |
device = "cuda" if torch.cuda.is_available() else "cpu" | |
info = torchaudio.info(audio_path) | |
audio_dur = info.num_frames / info.sample_rate | |
audio_tensors = [] | |
inference_transcript = "" | |
for sentence in sentences: | |
decode_config = {"top_k": top_k, "top_p": top_p, "temperature": temperature, "stop_repetition": stop_repetition, | |
"kvcache": kvcache, "codec_audio_sr": codec_audio_sr, "codec_sr": codec_sr, | |
"silence_tokens": silence_tokens, "sample_batch_size": sample_batch_size} | |
if mode != "Edit": | |
from inference_tts_scale import inference_one_sample | |
if smart_transcript: | |
target_transcript = "" | |
for word in word_info: | |
if word["end"] < prompt_end_time: | |
target_transcript += word["word"] | |
elif (word["start"] + word["end"]) / 2 < prompt_end_time: | |
# include part of the word it it's big, but adjust prompt_end_time | |
target_transcript += word["word"] | |
prompt_end_time = word["end"] | |
break | |
else: | |
break | |
target_transcript += f" {sentence}" | |
else: | |
target_transcript = sentence | |
inference_transcript += target_transcript + "\n" | |
prompt_end_frame = int(min(audio_dur, prompt_end_time) * info.sample_rate) | |
_, gen_audio = inference_one_sample(voicecraft_model["model"], | |
voicecraft_model["ckpt"]["config"], | |
voicecraft_model["ckpt"]["phn2num"], | |
voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"], | |
audio_path, target_transcript, device, decode_config, | |
prompt_end_frame) | |
else: | |
from inference_speech_editing_scale import inference_one_sample | |
if smart_transcript: | |
target_transcript = "" | |
for word in word_info: | |
if word["start"] < edit_start_time: | |
target_transcript += word["word"] | |
else: | |
break | |
target_transcript += f" {sentence}" | |
for word in word_info: | |
if word["end"] > edit_end_time: | |
target_transcript += word["word"] | |
else: | |
target_transcript = sentence | |
inference_transcript += target_transcript + "\n" | |
morphed_span = (max(edit_start_time - left_margin, 1 / codec_sr), min(edit_end_time + right_margin, audio_dur)) | |
mask_interval = [[round(morphed_span[0]*codec_sr), round(morphed_span[1]*codec_sr)]] | |
mask_interval = torch.LongTensor(mask_interval) | |
_, gen_audio = inference_one_sample(voicecraft_model["model"], | |
voicecraft_model["ckpt"]["config"], | |
voicecraft_model["ckpt"]["phn2num"], | |
voicecraft_model["text_tokenizer"], voicecraft_model["audio_tokenizer"], | |
audio_path, target_transcript, mask_interval, device, decode_config) | |
gen_audio = gen_audio[0].cpu() | |
audio_tensors.append(gen_audio) | |
if mode != "Rerun": | |
output_audio = get_output_audio(audio_tensors, codec_audio_sr) | |
sentences = [f"{idx}: {text}" for idx, text in enumerate(sentences)] | |
component = gr.Dropdown(choices=sentences, value=sentences[0]) | |
return output_audio, inference_transcript, component, audio_tensors | |
else: | |
previous_audio_tensors[selected_sentence_idx] = audio_tensors[0] | |
output_audio = get_output_audio(previous_audio_tensors, codec_audio_sr) | |
sentence_audio = get_output_audio(audio_tensors, codec_audio_sr) | |
return output_audio, inference_transcript, sentence_audio, previous_audio_tensors | |
def update_input_audio(audio_path): | |
if audio_path is None: | |
return 0, 0, 0 | |
info = torchaudio.info(audio_path) | |
max_time = round(info.num_frames / info.sample_rate, 2) | |
return [ | |
gr.Slider(maximum=max_time, value=max_time), | |
gr.Slider(maximum=max_time, value=0), | |
gr.Slider(maximum=max_time, value=max_time), | |
] | |
def change_mode(mode): | |
tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor | |
return [ | |
gr.Group(visible=mode != "Edit"), | |
gr.Group(visible=mode == "Edit"), | |
gr.Radio(visible=mode == "Edit"), | |
gr.Radio(visible=mode == "Long TTS"), | |
gr.Group(visible=mode == "Long TTS"), | |
] | |
def load_sentence(selected_sentence, codec_audio_sr, audio_tensors): | |
if selected_sentence is None: | |
return None | |
colon_position = selected_sentence.find(':') | |
selected_sentence_idx = int(selected_sentence[:colon_position]) | |
return get_output_audio([audio_tensors[selected_sentence_idx]], codec_audio_sr) | |
def update_bound_word(is_first_word, selected_word, edit_word_mode): | |
if selected_word is None: | |
return None | |
word_start_time = float(selected_word.split(' ')[0]) | |
word_end_time = float(selected_word.split(' ')[-1]) | |
if edit_word_mode == "Replace half": | |
bound_time = (word_start_time + word_end_time) / 2 | |
elif is_first_word: | |
bound_time = word_start_time | |
else: | |
bound_time = word_end_time | |
return bound_time | |
def update_bound_words(from_selected_word, to_selected_word, edit_word_mode): | |
return [ | |
update_bound_word(True, from_selected_word, edit_word_mode), | |
update_bound_word(False, to_selected_word, edit_word_mode), | |
] | |
smart_transcript_info = """ | |
If enabled, the target transcript will be constructed for you:</br> | |
- In TTS and Long TTS mode just write the text you want to synthesize.</br> | |
- In Edit mode just write the text to replace selected editing segment.</br> | |
If disabled, you should write the target transcript yourself:</br> | |
- In TTS mode write prompt transcript followed by generation transcript.</br> | |
- In Long TTS select split by newline (<b>SENTENCE SPLIT WON'T WORK</b>) and start each line with a prompt transcript.</br> | |
- In Edit mode write full prompt</br> | |
""" | |
demo_original_transcript = " But when I had approached so near to them, the common object, which the sense deceives, lost not by distance any of its marks." | |
demo_text = { | |
"TTS": { | |
"smart": "I cannot believe that the same model can also do text to speech synthesis as well!", | |
"regular": "But when I had approached so near to them, the common I cannot believe that the same model can also do text to speech synthesis as well!" | |
}, | |
"Edit": { | |
"smart": "saw the mirage of the lake in the distance,", | |
"regular": "But when I saw the mirage of the lake in the distance, which the sense deceives, Lost not by distance any of its marks," | |
}, | |
"Long TTS": { | |
"smart": "You can run generation on a big text!\n" | |
"Just write it line-by-line. Or sentence-by-sentence.\n" | |
"If some sentences sound odd, just rerun generation on them, no need to generate the whole text again!", | |
"regular": "But when I had approached so near to them, the common You can run generation on a big text!\n" | |
"But when I had approached so near to them, the common Just write it line-by-line. Or sentence-by-sentence.\n" | |
"But when I had approached so near to them, the common If some sentences sound odd, just rerun generation on them, no need to generate the whole text again!" | |
} | |
} | |
all_demo_texts = {vv for k, v in demo_text.items() for kk, vv in v.items()} | |
demo_words = [ | |
"0.03 but 0.18", | |
"0.18 when 0.32", | |
"0.32 i 0.48", | |
"0.48 had 0.64", | |
"0.64 approached 1.19", | |
"1.22 so 1.58", | |
"1.58 near 1.91", | |
"1.91 to 2.07", | |
"2.07 them 2.42", | |
"2.53 the 2.61", | |
"2.61 common 3.01", | |
"3.05 object 3.62", | |
"3.68 which 3.93", | |
"3.93 the 4.02", | |
"4.02 sense 4.34", | |
"4.34 deceives 4.97", | |
"5.04 lost 5.54", | |
"5.54 not 6.00", | |
"6.00 by 6.14", | |
"6.14 distance 6.67", | |
"6.79 any 7.05", | |
"7.05 of 7.18", | |
"7.18 its 7.34", | |
"7.34 marks 7.87" | |
] | |
demo_word_info = [ | |
{"word": "but", "start": 0.03, "end": 0.18}, | |
{"word": "when", "start": 0.18, "end": 0.32}, | |
{"word": "i", "start": 0.32, "end": 0.48}, | |
{"word": "had", "start": 0.48, "end": 0.64}, | |
{"word": "approached", "start": 0.64, "end": 1.19}, | |
{"word": "so", "start": 1.22, "end": 1.58}, | |
{"word": "near", "start": 1.58, "end": 1.91}, | |
{"word": "to", "start": 1.91, "end": 2.07}, | |
{"word": "them", "start": 2.07, "end": 2.42}, | |
{"word": "the", "start": 2.53, "end": 2.61}, | |
{"word": "common", "start": 2.61, "end": 3.01}, | |
{"word": "object", "start": 3.05, "end": 3.62}, | |
{"word": "which", "start": 3.68, "end": 3.93}, | |
{"word": "the", "start": 3.93, "end": 4.02}, | |
{"word": "sense", "start": 4.02, "end": 4.34}, | |
{"word": "deceives", "start": 4.34, "end": 4.97}, | |
{"word": "lost", "start": 5.04, "end": 5.54}, | |
{"word": "not", "start": 5.54, "end": 6.0}, | |
{"word": "by", "start": 6.0, "end": 6.14}, | |
{"word": "distance", "start": 6.14, "end": 6.67}, | |
{"word": "any", "start": 6.79, "end": 7.05}, | |
{"word": "of", "start": 7.05, "end": 7.18}, | |
{"word": "its", "start": 7.18, "end": 7.34}, | |
{"word": "marks", "start": 7.34, "end": 7.87} | |
] | |
def update_demo(mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word): | |
if transcript not in all_demo_texts: | |
return transcript, edit_from_word, edit_to_word | |
replace_half = edit_word_mode == "Replace half" | |
change_edit_from_word = edit_from_word == demo_words[2] or edit_from_word == demo_words[3] | |
change_edit_to_word = edit_to_word == demo_words[11] or edit_to_word == demo_words[12] | |
demo_edit_from_word_value = demo_words[2] if replace_half else demo_words[3] | |
demo_edit_to_word_value = demo_words[12] if replace_half else demo_words[11] | |
return [ | |
demo_text[mode]["smart" if smart_transcript else "regular"], | |
demo_edit_from_word_value if change_edit_from_word else edit_from_word, | |
demo_edit_to_word_value if change_edit_to_word else edit_to_word, | |
] | |
with gr.Blocks() as app: | |
with gr.Row(): | |
with gr.Column(scale=2): | |
load_models_btn = gr.Button(value="Load models") | |
with gr.Column(scale=5): | |
with gr.Accordion("Select models", open=False) as models_selector: | |
with gr.Row(): | |
voicecraft_model_choice = gr.Radio(label="VoiceCraft model", value="giga830M", choices=["giga330M", "giga830M"]) | |
whisper_model_choice = gr.Radio(label="Whisper model", value="base.en", | |
choices=[None, "tiny.en", "base.en", "small.en", "medium.en", "large"]) | |
with gr.Row(): | |
with gr.Column(scale=2): | |
input_audio = gr.Audio(value="./demo/84_121550_000074_000000.wav", label="Input Audio", type="filepath") | |
with gr.Group(): | |
original_transcript = gr.Textbox(label="Original transcript", lines=5, value=demo_original_transcript, interactive=False, | |
info="Use whisper model to get the transcript. Fix it if necessary.") | |
with gr.Accordion("Word start time", open=False): | |
transcript_with_start_time = gr.Textbox(label="Start time", lines=5, interactive=False, info="Start time before each word") | |
with gr.Accordion("Word end time", open=False): | |
transcript_with_end_time = gr.Textbox(label="End time", lines=5, interactive=False, info="End time after each word") | |
transcribe_btn = gr.Button(value="Transcribe") | |
with gr.Column(scale=3): | |
with gr.Group(): | |
transcript = gr.Textbox(label="Text", lines=7, value=demo_text["TTS"]["smart"]) | |
with gr.Row(): | |
smart_transcript = gr.Checkbox(label="Smart transcript", value=True) | |
with gr.Accordion(label="?", open=False): | |
info = gr.Markdown(value=smart_transcript_info) | |
with gr.Row(): | |
mode = gr.Radio(label="Mode", choices=["TTS", "Edit", "Long TTS"], value="TTS") | |
split_text = gr.Radio(label="Split text", choices=["Newline", "Sentence"], value="Newline", | |
info="Split text into parts and run TTS for each part.", visible=False) | |
edit_word_mode = gr.Radio(label="Edit word mode", choices=["Replace half", "Replace all"], value="Replace half", | |
info="What to do with first and last word", visible=False) | |
with gr.Group() as tts_mode_controls: | |
prompt_to_word = gr.Dropdown(label="Last word in prompt", choices=demo_words, value=demo_words[10], interactive=True) | |
prompt_end_time = gr.Slider(label="Prompt end time", minimum=0, maximum=7.93, step=0.01, value=3.01) | |
with gr.Group(visible=False) as edit_mode_controls: | |
with gr.Row(): | |
edit_from_word = gr.Dropdown(label="First word to edit", choices=demo_words, value=demo_words[2], interactive=True) | |
edit_to_word = gr.Dropdown(label="Last word to edit", choices=demo_words, value=demo_words[12], interactive=True) | |
with gr.Row(): | |
edit_start_time = gr.Slider(label="Edit from time", minimum=0, maximum=7.93, step=0.01, value=0.35) | |
edit_end_time = gr.Slider(label="Edit to time", minimum=0, maximum=7.93, step=0.01, value=3.75) | |
run_btn = gr.Button(value="Run") | |
with gr.Column(scale=2): | |
output_audio = gr.Audio(label="Output Audio") | |
with gr.Accordion("Inference transcript", open=False): | |
inference_transcript = gr.Textbox(label="Inference transcript", lines=5, interactive=False, | |
info="Inference was performed on this transcript.") | |
with gr.Group(visible=False) as long_tts_sentence_editor: | |
sentence_selector = gr.Dropdown(label="Sentence", value=None, | |
info="Select sentence you want to regenerate") | |
sentence_audio = gr.Audio(label="Sentence Audio", scale=2) | |
rerun_btn = gr.Button(value="Rerun") | |
with gr.Row(): | |
with gr.Accordion("VoiceCraft config", open=False): | |
seed = gr.Number(label="seed", value=-1, precision=0) | |
left_margin = gr.Number(label="left_margin", value=0.08) | |
right_margin = gr.Number(label="right_margin", value=0.08) | |
codec_audio_sr = gr.Number(label="codec_audio_sr", value=16000) | |
codec_sr = gr.Number(label="codec_sr", value=50) | |
top_k = gr.Number(label="top_k", value=0) | |
top_p = gr.Number(label="top_p", value=0.8) | |
temperature = gr.Number(label="temperature", value=1) | |
stop_repetition = gr.Radio(label="stop_repetition", choices=[-1, 1, 2, 3], value=3, | |
info="if there are long silence in the generated audio, reduce the stop_repetition to 3, 2 or even 1, -1 = disabled") | |
sample_batch_size = gr.Number(label="sample_batch_size", value=4, precision=0, | |
info="generate this many samples and choose the shortest one") | |
kvcache = gr.Radio(label="kvcache", choices=[0, 1], value=1, | |
info="set to 0 to use less VRAM, but with slower inference") | |
silence_tokens = gr.Textbox(label="silence tokens", value="[1388,1898,131]") | |
audio_tensors = gr.State() | |
word_info = gr.State(value=demo_word_info) | |
mode.change(fn=update_demo, | |
inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word], | |
outputs=[transcript, edit_from_word, edit_to_word]) | |
edit_word_mode.change(fn=update_demo, | |
inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word], | |
outputs=[transcript, edit_from_word, edit_to_word]) | |
smart_transcript.change(fn=update_demo, | |
inputs=[mode, smart_transcript, edit_word_mode, transcript, edit_from_word, edit_to_word], | |
outputs=[transcript, edit_from_word, edit_to_word]) | |
load_models_btn.click(fn=load_models, | |
inputs=[whisper_model_choice, voicecraft_model_choice], | |
outputs=[models_selector]) | |
input_audio.upload(fn=update_input_audio, | |
inputs=[input_audio], | |
outputs=[prompt_end_time, edit_start_time, edit_end_time]) | |
transcribe_btn.click(fn=transcribe, | |
inputs=[seed, input_audio], | |
outputs=[original_transcript, transcript_with_start_time, transcript_with_end_time, | |
prompt_to_word, edit_from_word, edit_to_word, word_info]) | |
mode.change(fn=change_mode, | |
inputs=[mode], | |
outputs=[tts_mode_controls, edit_mode_controls, edit_word_mode, split_text, long_tts_sentence_editor]) | |
run_btn.click(fn=run, | |
inputs=[ | |
seed, left_margin, right_margin, | |
codec_audio_sr, codec_sr, | |
top_k, top_p, temperature, | |
stop_repetition, sample_batch_size, | |
kvcache, silence_tokens, | |
input_audio, word_info, transcript, smart_transcript, | |
mode, prompt_end_time, edit_start_time, edit_end_time, | |
split_text, sentence_selector, audio_tensors | |
], | |
outputs=[output_audio, inference_transcript, sentence_selector, audio_tensors]) | |
sentence_selector.change(fn=load_sentence, | |
inputs=[sentence_selector, codec_audio_sr, audio_tensors], | |
outputs=[sentence_audio]) | |
rerun_btn.click(fn=run, | |
inputs=[ | |
seed, left_margin, right_margin, | |
codec_audio_sr, codec_sr, | |
top_k, top_p, temperature, | |
stop_repetition, sample_batch_size, | |
kvcache, silence_tokens, | |
input_audio, word_info, transcript, smart_transcript, | |
gr.State(value="Rerun"), prompt_end_time, edit_start_time, edit_end_time, | |
split_text, sentence_selector, audio_tensors | |
], | |
outputs=[output_audio, inference_transcript, sentence_audio, audio_tensors]) | |
prompt_to_word.change(fn=update_bound_word, | |
inputs=[gr.State(False), prompt_to_word, gr.State("Replace all")], | |
outputs=[prompt_end_time]) | |
edit_from_word.change(fn=update_bound_word, | |
inputs=[gr.State(True), edit_from_word, edit_word_mode], | |
outputs=[edit_start_time]) | |
edit_to_word.change(fn=update_bound_word, | |
inputs=[gr.State(False), edit_to_word, edit_word_mode], | |
outputs=[edit_end_time]) | |
edit_word_mode.change(fn=update_bound_words, | |
inputs=[edit_from_word, edit_to_word, edit_word_mode], | |
outputs=[edit_start_time, edit_end_time]) | |
if __name__ == "__main__": | |
app.launch() |