import sys sys.path.append(".") from modules.tokenizer.tokenizer import get_tokenizer_and_extra_tokens from modules.audio_tokenizer.audio_tokenizer import get_audio_tokenizer from modules.audio_detokenizer.audio_detokenizer import get_audio_detokenizer, detokenize, detokenize_noref import torch import os from glob import glob import base64 import io import torchaudio from transformers import AutoModelForCausalLM, GenerationConfig import librosa from tqdm import tqdm class Model(object): def __init__(self): self.tokenizer, self.extra_tokens = get_tokenizer_and_extra_tokens() self.speech_token_offset = 163840 print(self.extra_tokens) self.assistant_ids = self.tokenizer.encode("assistant") # [110866] self.user_ids = self.tokenizer.encode("user") # [1495] self.audio_ids = self.tokenizer.encode("audio") # [26229] self.spk_0_ids = self.tokenizer.encode("0") # [501] self.spk_1_ids = self.tokenizer.encode("1") # [503] self.msg_end = self.extra_tokens.msg_end # 260 self.user_msg_start = self.extra_tokens.user_msg_start # 261 self.assistant_msg_start = self.extra_tokens.assistant_msg_start # 262 self.name_end = self.extra_tokens.name_end # 272 self.media_begin = self.extra_tokens.media_begin # 273 self.media_content = self.extra_tokens.media_content # 274 self.media_end = self.extra_tokens.media_end # 275 self.audio_tokenizer = get_audio_tokenizer() self.audio_detokenizer = get_audio_detokenizer() model_path = "resources/text2semantic" self.model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda:0", torch_dtype=torch.bfloat16, trust_remote_code=True, force_download=True).to(torch.cuda.current_device()) self.generate_config = GenerationConfig( max_new_tokens=200 * 50, # no more than 200s per turn do_sample=True, top_k=30, top_p=0.8, temperature=0.8, eos_token_id=self.media_end, ) def _clean_text(self, text): # you can add front-end processing here text = text.replace("“", "") text = text.replace("”", "") text = text.replace("...", " ") text = text.replace("…", " ") text = text.replace("*", "") text = text.replace(":", ",") text = text.replace("‘", "'") text = text.replace("’", "'") text = text.strip() return text @torch.inference_mode() def _process_text(self, js): if "role_mapping" in js: for role in js["role_mapping"].keys(): js["role_mapping"][role]["ref_bpe_ids"] = self.tokenizer.encode(self._clean_text(js["role_mapping"][role]["ref_text"])) for turn in js["dialogue"]: turn["bpe_ids"] = self.tokenizer.encode(self._clean_text(turn["text"])) return js def inference(self, js, streaming=False): js = self._process_text(js) if "role_mapping" not in js: return self.infer_without_prompt(js, streaming) else: return self.infer_with_prompt(js, streaming) @torch.inference_mode() def infer_with_prompt(self, js, streaming=False): user_role_0_ids = [self.user_msg_start] + self.user_ids + self.spk_0_ids + [self.name_end] user_role_1_ids = [self.user_msg_start] + self.user_ids + self.spk_1_ids + [self.name_end] assistant_role_0_ids = [self.assistant_msg_start] + self.assistant_ids + self.spk_0_ids + [self.name_end] assistant_role_1_ids = [self.assistant_msg_start] + self.assistant_ids + self.spk_1_ids + [self.name_end] media_start = [self.media_begin] + self.audio_ids + [self.media_content] media_end = [self.media_end] + [self.msg_end] assistant_role_0_ids = torch.LongTensor(assistant_role_0_ids).unsqueeze(0).to(torch.cuda.current_device()) assistant_role_1_ids = torch.LongTensor(assistant_role_1_ids).unsqueeze(0).to(torch.cuda.current_device()) media_start = torch.LongTensor(media_start).unsqueeze(0).to(torch.cuda.current_device()) media_end = torch.LongTensor(media_end).unsqueeze(0).to(torch.cuda.current_device()) prompt = [] cur_role_dict = dict() for role, role_item in js["role_mapping"].items(): waveform_24k = librosa.load(role_item["ref_audio"], sr=24000)[0] waveform_24k = torch.tensor(waveform_24k).unsqueeze(0).to(torch.cuda.current_device()) waveform_16k = librosa.load(role_item["ref_audio"], sr=16000)[0] waveform_16k = torch.tensor(waveform_16k).unsqueeze(0).to(torch.cuda.current_device()) semantic_tokens = self.audio_tokenizer.tokenize(waveform_16k) semantic_tokens = semantic_tokens.to(torch.cuda.current_device()) prompt_ids = semantic_tokens + self.speech_token_offset cur_role_dict[role] = { "ref_bpe_ids": role_item["ref_bpe_ids"], "wav_24k": waveform_24k, "semantic_tokens": semantic_tokens, "prompt_ids": prompt_ids } prompt = prompt + user_role_0_ids + cur_role_dict["0"]["ref_bpe_ids"] + [self.msg_end] prompt = prompt + user_role_1_ids + cur_role_dict["1"]["ref_bpe_ids"] + [self.msg_end] for seg_id, turn in enumerate(js["dialogue"]): role_id = turn["role"] cur_user_ids = user_role_0_ids if role_id == "0" else user_role_1_ids cur_start_ids = cur_user_ids + turn["bpe_ids"] + [self.msg_end] prompt = prompt + cur_start_ids prompt = torch.LongTensor(prompt).unsqueeze(0).to(torch.cuda.current_device()) prompt = torch.cat([prompt, assistant_role_0_ids, media_start, cur_role_dict["0"]["prompt_ids"], media_end], dim=-1) prompt = torch.cat([prompt, assistant_role_1_ids, media_start, cur_role_dict["1"]["prompt_ids"], media_end], dim=-1) generation_config = self.generate_config # you can modify sampling strategy here wav_list = [] for seg_id, turn in tqdm(enumerate(js["dialogue"])): role_id = turn["role"] cur_assistant_ids = assistant_role_0_ids if role_id == "0" else assistant_role_1_ids prompt = torch.cat([prompt, cur_assistant_ids, media_start], dim=-1) len_prompt = prompt.shape[1] generation_config.min_length = len_prompt + 2 # print(generation_config) # todo: add streaming support for generate function outputs = self.model.generate(prompt, generation_config=generation_config) if outputs[0, -1] == self.media_end: outputs = outputs[:, :-1] output_token = outputs[:, len_prompt:] prompt = torch.cat([outputs, media_end], dim=-1) torch_token = output_token - self.speech_token_offset if streaming: # gen_speech_fm = detokenize(self.audio_detokenizer, torch_token, cur_role_dict[role_id]["wav_24k"], cur_role_dict[role_id]["semantic_tokens"]) # yield detokenize(self.audio_detokenizer, torch_token, cur_role_dict[role_id]["wav_24k"], cur_role_dict[role_id]["semantic_tokens"]) for cur_chunk in detokenize(self.audio_detokenizer, torch_token, cur_role_dict[role_id]["wav_24k"], cur_role_dict[role_id]["semantic_tokens"], streaming=True): cur_chunk = cur_chunk.cpu() cur_chunk = cur_chunk / cur_chunk.abs().max() cur_buffer = io.BytesIO() torchaudio.save(cur_buffer, cur_chunk, sample_rate=24000, format="mp3") audio_bytes = cur_buffer.getvalue() audio_b64 = base64.b64encode(audio_bytes).decode("utf-8") yield audio_b64 else: gen_speech_fm = detokenize(self.audio_detokenizer, torch_token, cur_role_dict[role_id]["wav_24k"], cur_role_dict[role_id]["semantic_tokens"]) gen_speech_fm = gen_speech_fm.cpu() gen_speech_fm = gen_speech_fm / gen_speech_fm.abs().max() wav_list.append(gen_speech_fm) del torch_token if not streaming: concat_wav = torch.cat(wav_list, dim=-1).cpu() # print(concat_wav.shape) buffer = io.BytesIO() torchaudio.save(buffer, concat_wav, sample_rate=24000, format="mp3") audio_bytes = buffer.getvalue() audio_b64 = base64.b64encode(audio_bytes).decode("utf-8") return audio_b64 @torch.inference_mode() def infer_without_prompt(self, js, streaming=False): user_role_0_ids = [self.user_msg_start] + self.user_ids + self.spk_0_ids + [self.name_end] user_role_1_ids = [self.user_msg_start] + self.user_ids + self.spk_1_ids + [self.name_end] assistant_role_0_ids = [self.assistant_msg_start] + self.assistant_ids + self.spk_0_ids + [self.name_end] assistant_role_1_ids = [self.assistant_msg_start] + self.assistant_ids + self.spk_1_ids + [self.name_end] media_start = [self.media_begin] + self.audio_ids + [self.media_content] media_end = [self.media_end] + [self.msg_end] assistant_role_0_ids = torch.LongTensor(assistant_role_0_ids).unsqueeze(0).to(torch.cuda.current_device()) assistant_role_1_ids = torch.LongTensor(assistant_role_1_ids).unsqueeze(0).to(torch.cuda.current_device()) media_start = torch.LongTensor(media_start).unsqueeze(0).to(torch.cuda.current_device()) media_end = torch.LongTensor(media_end).unsqueeze(0).to(torch.cuda.current_device()) prompt = [] for seg_id, turn in enumerate(js["dialogue"]): role_id = turn["role"] cur_user_ids = user_role_0_ids if role_id == "0" else user_role_1_ids cur_start_ids = cur_user_ids + turn["bpe_ids"] + [self.msg_end] prompt = prompt + cur_start_ids prompt = torch.LongTensor(prompt).unsqueeze(0).to(torch.cuda.current_device()) generation_config = self.generate_config # you can modify sampling strategy here wav_list = [] for seg_id, turn in tqdm(enumerate(js["dialogue"])): role_id = turn["role"] cur_assistant_ids = assistant_role_0_ids if role_id == "0" else assistant_role_1_ids prompt = torch.cat([prompt, cur_assistant_ids, media_start], dim=-1) len_prompt = prompt.shape[1] generation_config.min_length = len_prompt + 2 # print(generation_config) # todo: add streaming support for generate function outputs = self.model.generate(prompt, generation_config=generation_config) if outputs[0, -1] == self.media_end: outputs = outputs[:, :-1] output_token = outputs[:, len_prompt:] prompt = torch.cat([outputs, media_end], dim=-1) torch_token = output_token - self.speech_token_offset if streaming: for cur_chunk in detokenize_noref(self.audio_detokenizer, torch_token, streaming=True): cur_chunk = cur_chunk.cpu() cur_chunk = cur_chunk / cur_chunk.abs().max() cur_buffer = io.BytesIO() torchaudio.save(cur_buffer, cur_chunk, sample_rate=24000, format="mp3") audio_bytes = cur_buffer.getvalue() audio_b64 = base64.b64encode(audio_bytes).decode("utf-8") yield audio_b64 else: gen_speech_fm = detokenize_noref(self.audio_detokenizer, torch_token) gen_speech_fm = gen_speech_fm.cpu() gen_speech_fm = gen_speech_fm / gen_speech_fm.abs().max() wav_list.append(gen_speech_fm) del torch_token if not streaming: concat_wav = torch.cat(wav_list, dim=-1).cpu() # print(concat_wav.shape) buffer = io.BytesIO() torchaudio.save(buffer, concat_wav, sample_rate=24000, format="mp3") audio_bytes = buffer.getvalue() audio_b64 = base64.b64encode(audio_bytes).decode("utf-8") return audio_b64 if __name__ == "__main__": model = Model() # speaker should be interleaved zh_test_json = { "role_mapping": { "0": { "ref_audio": "./zh_prompt0.wav", "ref_text": "可以每天都骑并且可能会让你爱上骑车，然后通过爱上骑车的你省了很多很多钱。", #asr output }, "1": { "ref_audio": "./zh_prompt1.wav", "ref_text": "他最后就能让同样食材炒出来的菜味道大大提升。" #asr output } }, "dialogue": [ { "role": "0", "text": "我觉得啊，就是经历了这么多年的经验， 就是补剂的作用就是九分的努力， 十分之一的补剂。 嗯，选的话肯定是九分更重要，但是我觉得补剂它能够让你九分的努力更加的有效率，更加的避免徒劳无功。 嗯，就是你，你你得先得真的锻炼，真的努力，真的健康饮食，然后再考虑补剂， 那你再加十十分之一的补剂的话，他可能就是说啊， 一半是心理作用，" }, { "role": "1", "text": "对，其实很多时候心理作用是非常重要的。嗯，然后我每次用补剂的时候，我就会更加努力，就比如说我在健身之前我喝了一勺蛋白粉，我就会督促自己多练，" }, { "role": "0", "text": "其实心理作用只要能实现你的预期目的就可以了。 就比如说给自行车链条加油， 它其实不是必要的，但是它可以让你骑行更顺畅， 然后提高你骑行的频率。" } ] } audio_bytes = model.inference(zh_test_json) file_to_save = open(f"tmp_generated_zh.mp3", "wb") file_to_save.write(base64.b64decode(audio_bytes)) print("zh done") # speaker should be interleaved en_test_json = { "role_mapping": { "0": { "ref_audio": "./en_prompt0.wav", "ref_text": "Yeah, no, this is my backyard. It's never ending So just the way I like it. So social distancing has never been a problem.", #asr output }, "1": { "ref_audio": "./en_prompt1.wav", "ref_text": "I'm doing great And. Look, it couldn't be any better than having you at your set, which is the outdoors." #asr output } }, "dialogue": [ { "role": "0", "text": "In an awesome time, And, we're even gonna do a second episode too So. This is part one part two, coming at some point in the future There. We are.", }, { "role": "1", "text": "I love it. So grateful Thank you So I'm really excited. That's awesome. Yeah." }, { "role": "0", "text": "All I was told, which is good because I don't want to really talk too much more is that you're really really into fitness and nutrition And overall holistic I love it Yes." }, { "role": "1", "text": "Yeah So I started around thirteen Okay But my parents were fitness instructors as well. Awesome So I came from the beginning, and now it's this transition into this wholeness because I had to chart my. Own path and they weren't into nutrition at all So I had to learn that part." } ] } audio_bytes = model.inference(en_test_json) file_to_save = open(f"tmp_generated_en.mp3", "wb") file_to_save.write(base64.b64decode(audio_bytes)) print("en done") # also support inference without prompt # speaker should be interleaved without_prompt_test_json = { "dialogue": [ { "role": "0", "text": "我觉得啊，就是经历了这么多年的经验， 就是补剂的作用就是九分的努力， 十分之一的补剂。 嗯，选的话肯定是九分更重要，但是我觉得补剂它能够让你九分的努力更加的有效率，更加的避免徒劳无功。 嗯，就是你，你你得先得真的锻炼，真的努力，真的健康饮食，然后再考虑补剂， 那你再加十十分之一的补剂的话，他可能就是说啊， 一半是心理作用，" }, { "role": "1", "text": "对，其实很多时候心理作用是非常重要的。嗯，然后我每次用补剂的时候，我就会更加努力，就比如说我在健身之前我喝了一勺蛋白粉，我就会督促自己多练，" }, { "role": "0", "text": "其实心理作用只要能实现你的预期目的就可以了。 就比如说给自行车链条加油， 它其实不是必要的，但是它可以让你骑行更顺畅， 然后提高你骑行的频率。" } ] } audio_bytes = model.inference(without_prompt_test_json) file_to_save = open(f"tmp_generated_woprompt.mp3", "wb") file_to_save.write(base64.b64decode(audio_bytes)) print("without prompt done")