Replaced Encodec with Vocos

app.py

@@ -1,4 +1,3 @@
-import argparse
 import logging
 import os
 import pathlib
@@ -19,7 +18,6 @@ langid.set_languages(['en', 'zh', 'ja'])
 
 import torch
 import torchaudio
-import random
 
 import numpy as np
 
@@ -35,7 +33,8 @@ from macros import *
 from examples import *
 
 import gradio as gr
-import whisper
+from vocos import Vocos
+from transformers import WhisperProcessor, WhisperForConditionalGeneration
 
 torch._C._jit_set_profiling_executor(False)
 torch._C._jit_set_profiling_mode(False)
@@ -72,8 +71,13 @@ model.eval()
 # Encodec model
 audio_tokenizer = AudioTokenizer(device)
 
+# Vocos decoder
+vocos = Vocos.from_pretrained('charactr/vocos-encodec-24khz').to(device)
+
 # ASR
-whisper_model = whisper.load_model("medium").to(device)
+whisper_processor = WhisperProcessor.from_pretrained("openai/whisper-medium")
+whisper = WhisperForConditionalGeneration.from_pretrained("openai/whisper-medium").to(device)
+whisper.config.forced_decoder_ids = None
 
 # Voice Presets
 preset_list = os.walk("./presets/").__next__()[2]
@@ -89,34 +93,33 @@ def clear_prompts():
                 endfiletime = time.time() - 60
                 if endfiletime > lastmodifytime:
                     os.remove(filename)
+        del path, filename, lastmodifytime, endfiletime
+        gc.collect()
     except:
         return
 
-def transcribe_one(model, audio_path):
-    # load audio and pad/trim it to fit 30 seconds
-    audio = whisper.load_audio(audio_path)
-    audio = whisper.pad_or_trim(audio)
+def transcribe_one(wav, sr):
+    if sr != 16000:
+        wav4trans = torchaudio.transforms.Resample(sr, 16000)(wav)
+    else:
+        wav4trans = wav
 
-    # make log-Mel spectrogram and move to the same device as the model
-    mel = whisper.log_mel_spectrogram(audio).to(model.device)
+    input_features = whisper_processor(wav4trans.squeeze(0), sampling_rate=16000, return_tensors="pt").input_features
 
-    # detect the spoken language
-    _, probs = model.detect_language(mel)
-    print(f"Detected language: {max(probs, key=probs.get)}")
-    lang = max(probs, key=probs.get)
-    # decode the audio
-    options = whisper.DecodingOptions(temperature=1.0, best_of=5, fp16=False if device == torch.device("cpu") else True, sample_len=150)
-    result = whisper.decode(model, mel, options)
+    # generate token ids
+    predicted_ids = whisper.generate(input_features.to(device))
+    lang = whisper_processor.batch_decode(predicted_ids[:, 1])[0].strip("<|>")
+    # decode token ids to text
+    text_pr = whisper_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
 
     # print the recognized text
-    print(result.text)
+    print(text_pr)
 
-    text_pr = result.text
     if text_pr.strip(" ")[-1] not in "?!.,。，？！。、":
         text_pr += "."
 
     # delete all variables
-    del audio, mel, probs, result
+    del wav4trans, input_features, predicted_ids
     gc.collect()
     return lang, text_pr
 
@@ -137,7 +140,7 @@ def make_npz_prompt(name, uploaded_audio, recorded_audio, transcript_content):
     assert wav_pr.ndim and wav_pr.size(0) == 1
 
     if transcript_content == "":
-        text_pr, lang_pr = make_prompt(name, wav_pr, sr, save=False)
+        lang_pr, text_pr = transcribe_one(wav_pr, sr)
     else:
         lang_pr = langid.classify(str(transcript_content))[0]
         lang_token = lang2token[lang_pr]
@@ -147,6 +150,8 @@ def make_npz_prompt(name, uploaded_audio, recorded_audio, transcript_content):
     audio_tokens = encoded_frames[0][0].transpose(2, 1).cpu().numpy()
 
     # tokenize text
+    lang_token = lang2token[lang_pr]
+    text_pr = lang_token + text_pr + lang_token
     phonemes, _ = text_tokenizer.tokenize(text=f"{text_pr}".strip())
     text_tokens, enroll_x_lens = text_collater(
         [
@@ -155,6 +160,8 @@ def make_npz_prompt(name, uploaded_audio, recorded_audio, transcript_content):
     )
 
     message = f"Detected language: {lang_pr}\n Detected text {text_pr}\n"
+    if lang_pr not in ['ja', 'zh', 'en']:
+        return f"Prompt can only made with one of model-supported languages, got {lang_pr} instead", None
 
     # save as npz file
     np.savez(os.path.join(tempfile.gettempdir(), f"{name}.npz"),
@@ -166,30 +173,6 @@ def make_npz_prompt(name, uploaded_audio, recorded_audio, transcript_content):
     return message, os.path.join(tempfile.gettempdir(), f"{name}.npz")
 
 
-def make_prompt(name, wav, sr, save=True):
-    if not isinstance(wav, torch.FloatTensor):
-        wav = torch.tensor(wav)
-    if wav.abs().max() > 1:
-        wav /= wav.abs().max()
-    if wav.size(-1) == 2:
-        wav = wav.mean(-1, keepdim=False)
-    if wav.ndim == 1:
-        wav = wav.unsqueeze(0)
-    assert wav.ndim and wav.size(0) == 1
-    torchaudio.save(f"./prompts/{name}.wav", wav, sr)
-    lang, text = transcribe_one(whisper_model, f"./prompts/{name}.wav")
-    lang_token = lang2token[lang]
-    text = lang_token + text + lang_token
-    with open(f"./prompts/{name}.txt", 'w') as f:
-        f.write(text)
-    if not save:
-        os.remove(f"./prompts/{name}.wav")
-        os.remove(f"./prompts/{name}.txt")
-    # delete all variables
-    del lang_token, wav, sr
-    gc.collect()
-    return text, lang
-
 @torch.no_grad()
 def infer_from_audio(text, language, accent, audio_prompt, record_audio_prompt, transcript_content):
     if len(text) > 150:
@@ -209,7 +192,7 @@ def infer_from_audio(text, language, accent, audio_prompt, record_audio_prompt,
     assert wav_pr.ndim and wav_pr.size(0) == 1
 
     if transcript_content == "":
-        text_pr, lang_pr = make_prompt('dummy', wav_pr, sr, save=False)
+        lang_pr, text_pr = transcribe_one(wav_pr, sr)
     else:
         lang_pr = langid.classify(str(transcript_content))[0]
         lang_token = lang2token[lang_pr]
@@ -222,6 +205,9 @@ def infer_from_audio(text, language, accent, audio_prompt, record_audio_prompt,
     lang = token2lang[lang_token]
     text = lang_token + text + lang_token
 
+    if lang_pr not in ['ja', 'zh', 'en']:
+        return f"Reference audio must be a speech of one of model-supported languages, got {lang_pr} instead", None
+
     # tokenize audio
     encoded_frames = tokenize_audio(audio_tokenizer, (wav_pr, sr))
     audio_prompts = encoded_frames[0][0].transpose(2, 1).to(device)
@@ -237,6 +223,8 @@ def infer_from_audio(text, language, accent, audio_prompt, record_audio_prompt,
 
     enroll_x_lens = None
     if text_pr:
+        lang_token = lang2token[lang_pr]
+        text_pr = lang_token + text_pr + lang_token
         text_prompts, _ = text_tokenizer.tokenize(text=f"{text_pr}".strip())
         text_prompts, enroll_x_lens = text_collater(
             [
@@ -256,15 +244,16 @@ def infer_from_audio(text, language, accent, audio_prompt, record_audio_prompt,
         prompt_language=lang_pr,
         text_language=langs if accent == "no-accent" else lang,
     )
-    samples = audio_tokenizer.decode(
-        [(encoded_frames.transpose(2, 1), None)]
-    )
+    # Decode with Vocos
+    frames = encoded_frames.permute(2,0,1)
+    features = vocos.codes_to_features(frames)
+    samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
 
     message = f"text prompt: {text_pr}\nsythesized text: {text}"
     # delete all variables
     del audio_prompts, text_tokens, text_prompts, phone_tokens, encoded_frames, wav_pr, sr, audio_prompt, record_audio_prompt, transcript_content
     gc.collect()
-    return message, (24000, samples[0][0].cpu().numpy())
+    return message, (24000, samples.squeeze(0).cpu().numpy())
 
 @torch.no_grad()
 def infer_from_prompt(text, language, accent, preset_prompt, prompt_file):
@@ -315,16 +304,17 @@ def infer_from_prompt(text, language, accent, preset_prompt, prompt_file):
         prompt_language=lang_pr,
         text_language=langs if accent == "no-accent" else lang,
     )
-    samples = audio_tokenizer.decode(
-        [(encoded_frames.transpose(2, 1), None)]
-    )
+    # Decode with Vocos
+    frames = encoded_frames.permute(2,0,1)
+    features = vocos.codes_to_features(frames)
+    samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
 
     message = f"sythesized text: {text}"
 
     # delete all variables
     del audio_prompts, text_tokens, text_prompts, phone_tokens, encoded_frames, prompt_file, preset_prompt
     gc.collect()
-    return message, (24000, samples[0][0].cpu().numpy())
+    return message, (24000, samples.squeeze(0).cpu().numpy())
 
 
 from utils.sentence_cutter import split_text_into_sentences
@@ -407,11 +397,13 @@ def infer_long_text(text, preset_prompt, prompt=None, language='auto', accent='n
                 text_language=langs if accent == "no-accent" else lang,
             )
             complete_tokens = torch.cat([complete_tokens, encoded_frames.transpose(2, 1)], dim=-1)
-        samples = audio_tokenizer.decode(
-            [(complete_tokens, None)]
-        )
+        # Decode with Vocos
+        frames = encoded_frames.permute(2, 0, 1)
+        features = vocos.codes_to_features(frames)
+        samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
+
         message = f"Cut into {len(sentences)} sentences"
-        return message, (24000, samples[0][0].cpu().numpy())
+        return message, (24000, samples.squeeze(0).cpu().numpy())
     elif mode == "sliding-window":
         complete_tokens = torch.zeros([1, NUM_QUANTIZERS, 0]).type(torch.LongTensor).to(device)
         original_audio_prompts = audio_prompts
@@ -453,12 +445,14 @@ def infer_long_text(text, preset_prompt, prompt=None, language='auto', accent='n
             else:
                 audio_prompts = original_audio_prompts
                 text_prompts = original_text_prompts
-        samples = audio_tokenizer.decode(
-            [(complete_tokens, None)]
-        )
+        # Decode with Vocos
+        frames = encoded_frames.permute(2, 0, 1)
+        features = vocos.codes_to_features(frames)
+        samples = vocos.decode(features, bandwidth_id=torch.tensor([2], device=device))
+
         message = f"Cut into {len(sentences)} sentences"
 
-        return message, (24000, samples[0][0].cpu().numpy())
+        return message, (24000, samples.squeeze(0).cpu().numpy())
     else:
         raise ValueError(f"No such mode {mode}")
 

requirements.txt

@@ -5,6 +5,7 @@ torchvision==0.15.2
 torchaudio
 tokenizers
 encodec
+vocos
 langid
 unidecode
 pyopenjtalk