Refactor svs_inference and related functions; Bug fixes and code cleanup

server.py

@@ -86,12 +86,9 @@ async def process_audio(file: UploadFile = File(...)):
         f.write(output)
 
     wav_info = svs_inference(
-        config.model_path,
-        svs_model,
         output,
-        lang=config.lang,
-        random_gen=True,
-        fs=44100
+        svs_model,
+        config,
     )
     sf.write("tmp/response.wav", wav_info, samplerate=44100)
 

svs_utils.py

@@ -1,54 +1,13 @@
-from util import (
-    preprocess_input,
-    postprocess_phn,
-    get_tokenizer,
-    get_pinyin,
-)
-from espnet_model_zoo.downloader import ModelDownloader
-from espnet2.bin.svs_inference import SingingGenerate
+import json
+import random
+
 import librosa
-import torch
 import numpy as np
-import random
-import json
+import torch
+from espnet2.bin.svs_inference import SingingGenerate
+from espnet_model_zoo.downloader import ModelDownloader
 
-import argparse
-import soundfile as sf
-
-# the code below should be in app.py than svs_utils.py
-# espnet_model_dict = {
-#     "Model①(Chinese)-zh": "espnet/aceopencpop_svs_visinger2_40singer_pretrain",
-#     "Model②(Multilingual)-zh": "espnet/mixdata_svs_visinger2_spkembed_lang_pretrained",
-#     "Model②(Multilingual)-jp": "espnet/mixdata_svs_visinger2_spkembed_lang_pretrained",
-# }
-
-
-singer_embeddings = {
-    "espnet/aceopencpop_svs_visinger2_40singer_pretrain": {
-        "singer1 (male)": 1,
-        "singer2 (female)": 12,
-        "singer3 (male)": 23,
-        "singer4 (female)": 29,
-        "singer5 (male)": 18,
-        "singer6 (female)": 8,
-        "singer7 (male)": 25,
-        "singer8 (female)": 5,
-        "singer9 (male)": 10,
-        "singer10 (female)": 15,
-    },
-    "espnet/mixdata_svs_visinger2_spkembed_lang_pretrained": {
-        "singer1 (male)": "resource/singer/singer_embedding_ace-1.npy",
-        "singer2 (female)": "resource/singer/singer_embedding_ace-2.npy",
-        "singer3 (male)": "resource/singer/singer_embedding_ace-3.npy",
-        "singer4 (female)": "resource/singer/singer_embedding_ace-8.npy",
-        "singer5 (male)": "resource/singer/singer_embedding_ace-7.npy",
-        "singer6 (female)": "resource/singer/singer_embedding_itako.npy",
-        "singer7 (male)": "resource/singer/singer_embedding_ofuton.npy",
-        "singer8 (female)": "resource/singer/singer_embedding_kising_orange.npy",
-        "singer9 (male)": "resource/singer/singer_embedding_m4singer_Tenor-1.npy",
-        "singer10 (female)": "resource/singer/singer_embedding_m4singer_Alto-4.npy",
-    },
-}
+from util import get_pinyin, get_tokenizer, postprocess_phn, preprocess_input
 
 
 def svs_warmup(config):
@@ -86,7 +45,7 @@ def svs_text_preprocessor(model_path, texts, lang):
     fs = 44100
 
     if texts is None:
-        return (fs, np.array([0.0])), "Error: No Text provided!"
+        raise ValueError("texts is None when calling svs_text_preprocessor")
 
     # preprocess
     if lang == "zh":
@@ -129,7 +88,7 @@ def svs_text_preprocessor(model_path, texts, lang):
     return lyric_ls, sybs, labels
 
 
-def svs_get_batch(model_path, answer_text, lang, random_gen=True):
+def create_batch_with_randomized_melody(lyric_ls, sybs, labels, config):
     """
     Input:
         - answer_text (str), in Chinese character or Japanese character
@@ -144,72 +103,55 @@ def svs_get_batch(model_path, answer_text, lang, random_gen=True):
      'text': 'n@zh i@zh k@zh e@zh m@zh ei@zh'}
     """
     tempo = 120
-    lyric_ls, sybs, labels = svs_text_preprocessor(model_path, answer_text, lang)
     len_note = len(lyric_ls)
     notes = []
-    if random_gen:
-        # midi_range = (57,69)
-        st = 0
-        for id_lyric in range(len_note):
-            pitch = random.randint(57, 69)
-            period = round(random.uniform(0.1, 0.5), 4)
-            ed = st + period
-            note = [st, ed, lyric_ls[id_lyric], pitch, sybs[id_lyric]]
-            st = ed
-            notes.append(note)
-
-        phns_str = " ".join(labels)
-        batch = {
-            "score": (
-                int(tempo),
-                notes,
-            ),
-            "text": phns_str,
-        }
-
-    # print(batch)
+    # midi_range = (57,69)
+    st = 0
+    for id_lyric in range(len_note):
+        pitch = random.randint(57, 69)
+        period = round(random.uniform(0.1, 0.5), 4)
+        ed = st + period
+        note = [st, ed, lyric_ls[id_lyric], pitch, sybs[id_lyric]]
+        st = ed
+        notes.append(note)
+    phns_str = " ".join(labels)
+    batch = {
+        "score": (
+            int(tempo),
+            notes,
+        ),
+        "text": phns_str,
+    }
     return batch
 
 
-langs = {
-    "zh": 2,
-    "jp": 1,
-    "en": 2,
-}
-
-exist_model = "Null"
-svs = None
-
-
-def svs_inference(model_name, model_svs, answer_text, lang, random_gen=True, fs=44100):
-    batch = svs_get_batch(model_name, answer_text, lang, random_gen=random_gen)
-
-    # Infer
-    spk = "singer1 (male)"
-    global exist_model
-    global svs
-    svs = model_svs
-    exist_model = model_name
-    # if exist_model == "Null" or exist_model != model_name:
-    #     # device = "cpu"
-    #     device = "cuda" if torch.cuda.is_available() else "cpu"
-    #     d = ModelDownloader(cachedir="./cache")
-    #     pretrain_downloaded = d.download_and_unpack(model_name)
-    #     svs = SingingGenerate(
-    #         train_config = pretrain_downloaded["train_config"],
-    #         model_file = pretrain_downloaded["model_file"],
-    #         device = device
-    #     )
-    #     exist_model = model_name
-    if model_name == "Model①(Chinese)-zh":
-        sid = np.array([singer_embeddings[model_name][spk]])
-        output_dict = svs(batch, sids=sid)
+def svs_inference(answer_text, svs_model, config, **kwargs):
+    lyric_ls, sybs, labels = svs_text_preprocessor(
+        config.model_path, answer_text, config.lang
+    )
+    if config.melody_source.startswith("random_generate"):
+        batch = create_batch_with_randomized_melody(lyric_ls, sybs, labels, config)
+    elif config.melody_source.startswith("random_select"):
+        segment_iterator = song_segment_iterator(kwargs["song_db"], kwargs["metadata"])
+        batch = align_score_and_text(segment_iterator, lyric_ls, sybs, labels, config)
     else:
-        lid = np.array([langs[lang]])
-        spk_embed = np.load("resource/singer/singer_embedding_ace-2.npy")
-        output_dict = svs(batch, lids=lid, spembs=spk_embed)
-    wav_info = output_dict["wav"].cpu().numpy()
+        raise NotImplementedError(f"melody source {config.melody_source} not supported")
 
+    if config.model_path == "espnet/aceopencpop_svs_visinger2_40singer_pretrain":
+        sid = np.array([config.speaker])
+        output_dict = svs_model(batch, sids=sid)
+    elif config.model_path == "espnet/mixdata_svs_visinger2_spkembed_lang_pretrained":
+        langs = {
+            "zh": 2,
+            "jp": 1,
+            "en": 2,
+        }
+        lid = np.array([langs[config.lang]])
+        spk_embed = np.load(config.speaker)
+        output_dict = svs_model(batch, lids=lid, spembs=spk_embed)
+    else:
+        raise NotImplementedError(f"Model {config.model_path} not supported")
+    wav_info = output_dict["wav"].cpu().numpy()
     return wav_info
 
 
@@ -230,8 +172,6 @@ def singmos_evaluation(predictor, wav_info, fs):
 
 def estimate_sentence_length(query, config, song2note_lengths):
     if config.melody_source.startswith("random_select"):
-        # random select a song from database, and return its value in the phrase_length column
-        # return phrase_length column and song name
         song_name = random.choice(list(song2note_lengths.keys()))
         phrase_length = song2note_lengths[song_name]
         metadata = {"song_name": song_name}
@@ -263,7 +203,10 @@ def align_score_and_text(segment_iterator, lyric_ls, sybs, labels, config):
                     ]
                 )
                 text.append(reference_note_lyric.strip("<>"))
-            elif reference_note_lyric in ["-", "——"] and config.melody_source == "random_select.take_lyric_continuation":
+            elif (
+                reference_note_lyric in ["-", "——"]
+                and config.melody_source == "random_select.take_lyric_continuation"
+            ):
                 notes_info.append(
                     [
                         note_start_time,
@@ -311,6 +254,8 @@ def song_segment_iterator(song_db, metadata):
 
 
 def load_song_database(config):
+    from datasets import load_dataset
+
     song_db = load_dataset(
         "jhansss/kising_score_segments", cache_dir="cache", split="train"
     ).to_pandas()
@@ -325,6 +270,8 @@ def load_song_database(config):
 
 
 if __name__ == "__main__":
+    import argparse
+    import soundfile as sf
 
     # -------- demo code for generate audio from randomly selected song ---------#
     config = argparse.Namespace(
@@ -333,6 +280,7 @@ if __name__ == "__main__":
         device="cuda", # "cpu"
         melody_source="random_generate", # "random_select.take_lyric_continuation"
         lang="zh",
+        speaker="resource/singer/singer_embedding_ace-2.npy",
     )
 
     # load model
@@ -344,28 +292,22 @@ if __name__ == "__main__":
 
     if config.melody_source.startswith("random_select"):
         # load song database: jhansss/kising_score_segments
-        from datasets import load_dataset
         song2note_lengths, song_db = load_song_database(config)
 
         # get song_name and phrase_length
+        phrase_length, metadata = estimate_sentence_length(
+            None, config, song2note_lengths
+        )
         phrase_length, metadata = estimate_sentence_length(None, config, song2note_lengths)
 
         # then, phrase_length info should be added to llm prompt, and get the answer lyrics from llm
         # e.g. answer_text = "天气真好\n空气清新"
-        lyric_ls, sybs, labels = svs_text_preprocessor(
-            config.model_path, answer_text, config.lang
-        )
-        segment_iterator = song_segment_iterator(song_db, metadata)
-        batch = align_score_and_text(segment_iterator, lyric_ls, sybs, labels, config)
-        singer_embedding = np.load(singer_embeddings[config.model_path]["singer2 (female)"])
-        lid = np.array([langs[config.lang]])
-        output_dict = model(batch, lids=lid, spembs=singer_embedding)
-        wav_info = output_dict["wav"].cpu().numpy()
+        additional_kwargs = {"song_db": song_db, "metadata": metadata}
+    else:
+        additional_kwargs = {}
 
-    
-    elif config.melody_source.startswith("random_generate"):
-        wav_info = svs_inference(config.model_path, model, answer_text, lang=config.lang, random_gen=True, fs=sample_rate)
+    wav_info = svs_inference(answer_text, model, config, **additional_kwargs)
 
     # write wav to output_retrieved.wav
-    save_name = config.melody_source.split('.')[0]
+    save_name = config.melody_source
     sf.write(f"{save_name}.wav", wav_info, samplerate=sample_rate)

util.py

@@ -21,6 +21,7 @@ def postprocess_phn(phns, model_name, lang):
 
 
 def pyopenjtalk_g2p(text) -> List[str]:
+    import pyopenjtalk
     with warnings.catch_warnings(record=True) as w:
         warnings.simplefilter("always")
         # phones is a str object separated by space
@@ -53,20 +54,25 @@ def split_pinyin_py(pinyin: str) -> tuple[str]:
 
 
 def get_tokenizer(model, lang):
-    if lang == "zh":
-        if "Chinese" in model:
-            print("hello")
+    if model == "espnet/aceopencpop_svs_visinger2_40singer_pretrain":
+        if lang == "zh":
             return lambda text: split_pinyin_py(text)
         else:
+            raise ValueError(f"Only support Chinese language for {model}")
+    elif model == "espnet/mixdata_svs_visinger2_spkembed_lang_pretrained":
+        if lang == "zh":
             with open(os.path.join("resource/all_plans.json"), "r") as f:
                 all_plan_dict = json.load(f)
             for plan in all_plan_dict["plans"]:
                 if plan["language"] == "zh":
                     zh_plan = plan
             return lambda text: split_pinyin_ace(text, zh_plan)
-    elif lang == "jp":
-        import pyopenjtalk
-        return pyopenjtalk_g2p
+        elif lang == "jp":
+            return pyopenjtalk_g2p
+        else:
+            raise ValueError(f"Only support Chinese and Japanese language for {model}")
+    else:
+        raise ValueError(f"Only support espnet/aceopencpop_svs_visinger2_40singer_pretrain and espnet/mixdata_svs_visinger2_spkembed_lang_pretrained for now")
 
 
 def get_pinyin(texts):