import json import logging import os import random import re import sys import time import uuid from threading import Thread from typing import Optional import torch import tqdm from torch import nn from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer from transformers.generation import GenerationConfig import torchaudio from vita_audio.data.processor.audio_processor import add_audio_input_contiguous from vita_audio.tokenizer import get_audio_tokenizer logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) torch.manual_seed(1234) device_map = "cuda:0" audio_tokenizer_rank = 0 torch_dtype = torch.bfloat16 # model_name_or_path = sys.argv[1] # audio_tokenizer_path = sys.argv[2] # flow_path = sys.argv[3] if True: # if False: # sensevoice glm4voice tokenizer sys.path.append("third_party/GLM-4-Voice/") sys.path.append("third_party/GLM-4-Voice/cosyvoice/") sys.path.append("third_party/GLM-4-Voice/third_party/Matcha-TTS/") audio_tokenizer_path = "/data/models/THUDM/glm-4-voice-tokenizer" flow_path = "/data/models/THUDM/glm-4-voice-decoder" audio_tokenizer_type = "sensevoice_glm4voice" model_name_or_path = "VITA-MLLM/VITA-Audio-Plus-Vanilla/" # if True: if False: # glm4voice tokenizer sys.path.append("third_party/GLM-4-Voice/") sys.path.append("third_party/GLM-4-Voice/cosyvoice/") sys.path.append("third_party/GLM-4-Voice/third_party/Matcha-TTS/") audio_tokenizer_path = "/data/models/THUDM/glm-4-voice-tokenizer" flow_path = "/data/models/THUDM/glm-4-voice-decoder" audio_tokenizer_type = "glm4voice" # model_name_or_path = "VITA-MLLM/VITA-Audio-Balance" model_name_or_path = "VITA-MLLM/VITA-Audio-Boost" output_dir = "/data/output/LM/inference/" os.makedirs(output_dir, exist_ok=True) class TextAudioIteratorStreamer(TextIteratorStreamer): def __init__( self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs, ): super().__init__(tokenizer, skip_prompt, timeout, **decode_kwargs) # self.audio_offset = tokenizer.convert_tokens_to_ids("<|audio_0|>") self.audio_offset = tokenizer.convert_tokens_to_ids("<|begin_of_audio|>") self.num_decode_tokens = 0 def put(self, value): """ Receives tokens, decodes them, and prints them to stdout as soon as they form entire words. """ if len(value.shape) > 1 and value.shape[0] > 1: raise ValueError("TextStreamer only supports batch size 1") elif len(value.shape) > 1: value = value[0] if self.skip_prompt and self.next_tokens_are_prompt: self.next_tokens_are_prompt = False return self.num_decode_tokens += len(value) # Add the new token to the cache and decodes the entire thing. self.token_cache.extend(value.tolist()) text = self.tokenizer.decode(self.token_cache, **self.decode_kwargs) # After the symbol for a new line, we flush the cache. if text.endswith("\n"): printable_text = text[self.print_len :] self.token_cache = [] self.print_len = 0 # If the last token is a CJK character, we print the characters. elif len(text) > 0 and self._is_chinese_char(ord(text[-1])): printable_text = text[self.print_len :] self.print_len += len(printable_text) elif self.token_cache[-1] >= self.audio_offset: printable_text = text[self.print_len :] self.print_len += len(printable_text) # Otherwise, prints until the last space char (simple heuristic to avoid printing incomplete words, # which may change with the subsequent token -- there are probably smarter ways to do this!) else: printable_text = text[self.print_len : text.rfind(" ") + 1] self.print_len += len(printable_text) self.on_finalized_text(printable_text) while self.text_queue.qsize() > 10: time.sleep(0.01) class BenchmarkIteratorStreamer(TextIteratorStreamer): def __init__( self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, **decode_kwargs, ): super().__init__(tokenizer, skip_prompt, timeout, **decode_kwargs) self.num_decode_tokens = 0 def put(self, value): """ Receives tokens, decodes them, and prints them to stdout as soon as they form entire words. """ if len(value.shape) > 1 and value.shape[0] > 1: raise ValueError("TextStreamer only supports batch size 1") elif len(value.shape) > 1: value = value[0] if self.skip_prompt and self.next_tokens_are_prompt: self.next_tokens_are_prompt = False return self.num_decode_tokens += len(value) printable_text = " ".join([str(x) for x in value.tolist()]) + " " self.on_finalized_text(printable_text) def find_audio_segments_regex(text): """ Find all substrings between <|begin_of_audio|> and <|end_of_audio|> using regex. Args: text (str): The input string to search through Returns: list: A list of all found audio segments (substrings between the delimiters) """ pattern = re.compile(r"<\|begin_of_audio\|>(.*?)<\|end_of_audio\|>", re.DOTALL) segments = pattern.findall(text) return [segment.strip() for segment in segments] def extract_token_ids_as_int(text): pattern = re.compile(r"<\|audio_(\d+)\|>") token_ids = pattern.findall(text) return [int(id) for id in token_ids] def custom_init_weights(module): if isinstance(module, torch.nn.Conv2d) or isinstance(module, torch.nn.Linear): torch.nn.init.xavier_uniform_(module.weight) if module.bias is not None: torch.nn.init.constant_(module.bias, 0) elif isinstance(module, torch.nn.BatchNorm2d) or isinstance(module, torch.nn.BatchNorm1d): torch.nn.init.constant_(module.weight, 1) torch.nn.init.constant_(module.bias, 0) class S2SInference: def __init__( self, model_name_or_path, audio_tokenizer_path, audio_tokenizer_type, flow_path=None ): config = AutoConfig.from_pretrained( model_name_or_path, trust_remote_code=True, ) if "qwen2" in config.model_type.lower(): from evaluation.get_chat_template import qwen2_chat_template as chat_template add_generation_prompt = True default_system_message = [] if "hunyuan" in config.model_type.lower(): from evaluation.get_chat_template import hunyuan_chat_template as chat_template add_generation_prompt = False default_system_message = [ { "role": "system", "content": "You are a helpful AI assistant.", } ] luke_system_message = [ { "role": "system", "content": "Your Name: Luke\nYour Gender: male\n\nRespond in a text-audio interleaved manner.", }, ] tokenizer = AutoTokenizer.from_pretrained( model_name_or_path, trust_remote_code=True, chat_template=chat_template, ) # print(f"{tokenizer=}") print(f"{tokenizer.get_chat_template()=}") model = AutoModelForCausalLM.from_pretrained( model_name_or_path, trust_remote_code=True, device_map=device_map, torch_dtype=torch_dtype, attn_implementation="flash_attention_2", ).eval() # print("model", model) print(f"{model.config.model_type=}") print(f"{model.hf_device_map=}") model.generation_config = GenerationConfig.from_pretrained( model_name_or_path, trust_remote_code=True ) model.generation_config.max_new_tokens = 8192 model.generation_config.chat_format = "chatml" model.generation_config.max_window_size = 8192 model.generation_config.use_cache = True # model.generation_config.use_cache = False model.generation_config.do_sample = False model.generation_config.temperature = 1.0 model.generation_config.top_k = 50 model.generation_config.top_p = 1.0 model.generation_config.num_beams = 1 model.generation_config.pad_token_id = tokenizer.pad_token_id if model.config.model_type == "hunyuan": model.generation_config.eos_token_id = tokenizer.eos_id print(f"{model.generation_config=}") audio_tokenizer = get_audio_tokenizer( audio_tokenizer_path, audio_tokenizer_type, flow_path=flow_path, rank=audio_tokenizer_rank, ) self.model = model self.tokenizer = tokenizer self.audio_tokenizer = audio_tokenizer self.add_generation_prompt = add_generation_prompt self.default_system_message = default_system_message self.luke_system_message = luke_system_message audio_0_id = tokenizer("<|audio_0|>").input_ids[0] print(f"{audio_0_id=}") def benchmark_forward(self, mtp_inference_mode): print("-" * 100) print("benchmark_forward...") print(f"{mtp_inference_mode=}") total_time = 0 past_key_values = None use_cache = True self.model.input_ids = None self.model.inputs_embeds = None self.model.hidden_states = [None] * (self.model.config.num_nextn_predict_layers + 1) self.model.position_ids = None self.model.attention_mask = None self.model.mtp_idx = -1 self.model.num_prefill_tokens = -1 model_max_length = 1024 if mtp_inference_mode is not None: ori_mtp_inference_mode = self.model.generation_config.mtp_inference_mode self.model._prepare_mtp_for_generation(mtp_inference_mode, model_max_length) else: self.model._prepare_mtp_for_generation( self.model.generation_config.mtp_inference_mode, model_max_length ) for i in tqdm.tqdm(range(1, model_max_length + 1)): if use_cache: input_ids = torch.tensor([i - 1], dtype=torch.long).unsqueeze(0).to("cuda") position_ids = torch.tensor([i - 1], dtype=torch.long).unsqueeze(0).to("cuda") else: input_ids = torch.arange(i, dtype=torch.long).unsqueeze(0).to("cuda") position_ids = torch.arange(i, dtype=torch.long).unsqueeze(0).to("cuda") attention_mask = torch.tensor([1] * i, dtype=torch.float).unsqueeze(0).to("cuda") torch.cuda.synchronize() start = time.time() output = self.model( input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, use_cache=use_cache, num_logits_to_keep=1, ) torch.cuda.synchronize() end = time.time() total_time += end - start # print(f"{i=} {total_time=}") past_key_values = output.past_key_values print() print(f"{total_time=}") print(f"second/token {total_time/model_max_length=}") print(f"token/second {model_max_length/total_time=}") if mtp_inference_mode is not None: self.model.mtp_inference_mode = ori_mtp_inference_mode def benchmark_generate(self, mtp_inference_mode): self.model.apply(custom_init_weights) print("-" * 100) print("benchmark_generate...") print(f"{mtp_inference_mode=}") total_time = 0 self.model.generation_config.use_cache = True self.model.generation_config.max_new_tokens = 8192 if mtp_inference_mode is not None: ori_mtp_inference_mode = self.model.generation_config.mtp_inference_mode self.model.generation_config.mtp_inference_mode = mtp_inference_mode input_ids = torch.tensor([0], dtype=torch.long).unsqueeze(0).to("cuda") torch.cuda.synchronize() start = time.time() output = self.model.generate( input_ids, ) # print(f"{output.size()=}") torch.cuda.synchronize() end = time.time() total_time += end - start print() print(f"{total_time=}") print(f"second/token {total_time/output.size(1)=}") print(f"token/second {output.size(1)/total_time=}") if mtp_inference_mode is not None: self.model.generation_config.mtp_inference_mode = ori_mtp_inference_mode def benchmark_generate_stream(self, mtp_inference_mode): print("-" * 100) print("benchmark_generate_stream...") print(f"{mtp_inference_mode=}") self.model.apply(custom_init_weights) total_time = 0 self.model.generation_config.use_cache = True # model_max_length = 8192 model_max_length = 4096 # model_max_length = 2048 # model_max_length = 1024 num_prefill_tokens = 32 self.model.generation_config.max_new_tokens = model_max_length self.model.generation_config.do_sample = False if mtp_inference_mode is not None: ori_mtp_inference_mode = self.model.generation_config.mtp_inference_mode self.model.generation_config.mtp_inference_mode = mtp_inference_mode input_ids = torch.tensor([0] * num_prefill_tokens, dtype=torch.long).unsqueeze(0).to("cuda") streamer = BenchmarkIteratorStreamer(self.tokenizer, skip_prompt=True) generation_kwargs = dict(input_ids=input_ids, streamer=streamer) thread = Thread(target=self.model.generate, kwargs=generation_kwargs) token_decode_time = [] torch.cuda.synchronize() start = time.time() thread.start() generated_text = "" for new_text in tqdm.tqdm(streamer, total=model_max_length): generated_text += new_text end = time.time() token_decode_time.append(end - start) yield new_text # print(f"{len(generated_text)}") torch.cuda.synchronize() end = time.time() total_time += end - start print() print(f"{token_decode_time[-1]=}") print(f"{streamer.num_decode_tokens=}") print(f"second/token {token_decode_time[-1]/streamer.num_decode_tokens=}") print(f"token/second {streamer.num_decode_tokens/token_decode_time[-1]=}") # if mtp_inference_mode is None: # mtp_inference_mode = [] # with open(f'token_decode_time_{str(mtp_inference_mode)}.json', 'w') as f: # json.dump(token_decode_time, f) if mtp_inference_mode is not None: self.model.generation_config.mtp_inference_mode = ori_mtp_inference_mode def run_infer( self, audio_path=None, prompt_audio_path=None, stream_stride=4, max_returned_tokens=4096, sample_rate=16000, request_id="", audio_feats=None, message="", use_past=False, mode="luke", do_sample=False, mtp_inference_mode=None, ): AUD_TAG_TOKEN = "<|audio|>" AUD_CONTEXT_TOKEN = "<|context_of_audio|>" AUD_START_TOKEN = "<|begin_of_audio|>" AUD_END_TOKEN = "<|end_of_audio|>" if prompt_audio_path is not None: system_message = [ { "role": "system", "content": f"Your Voice: <|audio|>\n", }, ] elif mode == "luke": system_message = self.luke_system_message else: system_message = self.default_system_message if prompt_audio_path is not None and self.audio_tokenizer.apply_to_role("user", is_discrete=True): # discrete codec audio_tokens = self.audio_tokenizer.encode(prompt_audio_path) audio_tokens = "".join(f"<|audio_{i}|>" for i in audio_tokens) system_message[-1]["content"] = system_message[-1]["content"].replace( "<|audio|>", f"<|begin_of_audio|>{audio_tokens}<|end_of_audio|>" ) if audio_path is not None: messages = system_message + [ { "role": "user", "content": message + "\n<|audio|>", }, ] else: messages = system_message + [ { "role": "user", "content": message, }, ] if audio_path is not None and self.audio_tokenizer.apply_to_role("user", is_discrete=True): # discrete codec audio_tokens = self.audio_tokenizer.encode(audio_path) audio_tokens = "".join(f"<|audio_{i}|>" for i in audio_tokens) messages[-1]["content"] = messages[-1]["content"].replace( "<|audio|>", f"<|begin_of_audio|>{audio_tokens}<|end_of_audio|>" ) input_ids = self.tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=self.add_generation_prompt, ) if (audio_path is not None or prompt_audio_path is not None) and self.audio_tokenizer.apply_to_role( "user", is_contiguous=True ): # contiguous codec audio_paths = [] if audio_path is not None: audio_paths.append(audio_path) if prompt_audio_path is not None: audio_paths.append(prompt_audio_path) input_ids, audios, audio_indices = add_audio_input_contiguous( input_ids, audio_paths, self.tokenizer, self.audio_tokenizer ) else: audios = None audio_indices = None input_ids = torch.tensor([input_ids], dtype=torch.long).to("cuda") print("input", self.tokenizer.decode(input_ids[0], skip_special_tokens=False), flush=True) self.model.generation_config.do_sample = do_sample if mtp_inference_mode is not None: ori_mtp_inference_mode = self.model.generation_config.mtp_inference_mode self.model.generation_config.mtp_inference_mode = mtp_inference_mode outputs = self.model.generate( input_ids, audios=audios, audio_indices=audio_indices, ) output = self.tokenizer.decode(outputs[0], skip_special_tokens=False) print(f"{output=}", flush=True) audio_offset = self.tokenizer.convert_tokens_to_ids("<|audio_0|>") audio_tokens = [] for token_id in outputs[0]: if token_id >= audio_offset: audio_tokens.append(token_id - audio_offset) if len(audio_tokens) > 0: tts_speech = self.audio_tokenizer.decode( audio_tokens, source_speech_16k=prompt_audio_path ) else: tts_speech = None if mtp_inference_mode is not None: self.model.generation_config.mtp_inference_mode = ori_mtp_inference_mode return output, tts_speech def run_infer_stream( self, audio_path=None, prompt_audio_path=None, stream_stride=4, max_returned_tokens=4096, sample_rate=16000, request_id="", audio_feats=None, message="", use_past=False, mode="luke", do_sample=False, mtp_inference_mode=None, ): if prompt_audio_path is not None: system_message = [ { "role": "system", "content": f"Your Voice: <|audio|>\n", }, ] elif mode == "luke": system_message = self.luke_system_message else: system_message = self.default_system_message if prompt_audio_path is not None and self.audio_tokenizer.apply_to_role("user", is_discrete=True): # discrete codec audio_tokens = self.audio_tokenizer.encode(prompt_audio_path) audio_tokens = "".join(f"<|audio_{i}|>" for i in audio_tokens) system_message[-1]["content"] = system_message[-1]["content"].replace( "<|audio|>", f"<|begin_of_audio|>{audio_tokens}<|end_of_audio|>" ) if audio_path is not None: messages = system_message + [ { "role": "user", "content": message + "\n<|audio|>", }, ] else: messages = system_message + [ { "role": "user", "content": message, }, ] if audio_path is not None and self.audio_tokenizer.apply_to_role("user", is_discrete=True): # discrete codec audio_tokens = self.audio_tokenizer.encode(audio_path) audio_tokens = "".join(f"<|audio_{i}|>" for i in audio_tokens) messages[-1]["content"] = messages[-1]["content"].replace( "<|audio|>", f"<|begin_of_audio|>{audio_tokens}<|end_of_audio|>" ) input_ids = self.tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=self.add_generation_prompt, ) if (audio_path is not None or prompt_audio_path is not None) and self.audio_tokenizer.apply_to_role( "user", is_contiguous=True ): # contiguous codec audio_paths = [] if audio_path is not None: audio_paths.append(audio_path) if prompt_audio_path is not None: audio_paths.append(prompt_audio_path) input_ids, audios, audio_indices = add_audio_input_contiguous( input_ids, audio_paths, self.tokenizer, self.audio_tokenizer ) else: audios = None audio_indices = None input_ids = torch.tensor([input_ids], dtype=torch.long).to("cuda") print("input", self.tokenizer.decode(input_ids[0], skip_special_tokens=False), flush=True) self.model.generation_config.do_sample = do_sample if mtp_inference_mode is not None: ori_mtp_inference_mode = self.model.generation_config.mtp_inference_mode self.model.generation_config.mtp_inference_mode = mtp_inference_mode streamer = TextAudioIteratorStreamer(self.tokenizer, skip_prompt=True) generation_kwargs = dict( input_ids=input_ids, audios=audios, audio_indices=audio_indices, streamer=streamer, ) thread = Thread(target=self.model.generate, kwargs=generation_kwargs) thread.start() # generated_text = "" for new_text in streamer: # generated_text += new_text yield new_text # torch.cuda.synchronize() if mtp_inference_mode is not None: self.model.generation_config.mtp_inference_mode = ori_mtp_inference_mode def benchmark_llm(): for mtp_inference_mode, tag in zip( [ [8192, 0], [1, 4, 3, 8, 4, 10], [1, 10, 4, 10], [1, 10], ], [ "Vanilla", "Balance", "Boost", "Turbo", ], ): print("=" * 100) print("benchmark_llm") print(f"{tag}") s2s_inference.benchmark_forward(mtp_inference_mode) s2s_inference.benchmark_generate(mtp_inference_mode) generated_text = "" for new_text in s2s_inference.benchmark_generate_stream( mtp_inference_mode=mtp_inference_mode ): generated_text += new_text # print(new_text, end="", flush=True) def benchmark_sts(): audio_paths = [ "asset/介绍一下上海.wav", "asset/发表一个悲伤的演讲.wav", "asset/发表一个振奋人心的演讲.wav", ] for _ in range(10): print("=" * 100) print("benchmark_sts") audio_path = random.choice(audio_paths) print(f"{audio_path}") start = time.time() audio_idx = 0 generated_text = "" all_tts_speech = [] past_tts_speech_len = 0 for new_text in s2s_inference.run_infer_stream(audio_path=audio_path): # print(new_text, end="", flush=True) generated_text += new_text if new_text == "<|end_of_audio|>": audio_tokens = extract_token_ids_as_int(generated_text) tts_speech = s2s_inference.audio_tokenizer.decode(audio_tokens, option_steps=1) tts_speech = tts_speech[past_tts_speech_len:] past_tts_speech_len += len(tts_speech) all_tts_speech.append(tts_speech) end = time.time() if audio_idx == 0: print(audio_tokens) print(f"{audio_idx} audio chunk {end - start}") wav_path = os.path.join(output_dir, audio_path[:-4] + f"_{audio_idx}.wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") audio_idx += 1 start = time.time() wav_path = os.path.join(output_dir, audio_path[:-4] + ".wav") tts_speech = torch.cat(all_tts_speech, dim=0) os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") # ============================================================== # Text def text_task(): for text in [ "How many helicopters can a human eat in one sitting?", "你叫什么名字？", "写一首诗", "介绍一下上海", ]: print("=" * 100) print("text_task") print(f"{text=}") output, _ = s2s_inference.run_infer( message=text, mode=None, # do_sample=True, mtp_inference_mode=[8192, 0], ) print(f"{output=}", flush=True) # ============================================================== # Text stream def text_stream_task(): for text in [ "你叫什么名字？", ]: print("=" * 100) print("text_stream_task") print(f"{text=}") generated_text = "" for new_text in s2s_inference.run_infer_stream( message=text, mode=None, # do_sample=True, mtp_inference_mode=[8192, 0], ): generated_text += new_text print(new_text, end="") print("") # ============================================================== # S2S def sts_task(): for audio_path in [ "asset/介绍一下上海.wav", "asset/发表一个悲伤的演讲.wav", "asset/发表一个振奋人心的演讲.wav", "asset/piano.mp3", ]: print("=" * 100) print("sts_task") print(f"{audio_path=}") output, tts_speech = s2s_inference.run_infer( audio_path=audio_path, ) wav_path = os.path.join(output_dir, audio_path[:-4] + ".wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") # ============================================================== # S2S stream def sts_stream_task(): for audio_path in [ "asset/介绍一下上海.wav", ]: print("=" * 100) print("sts_stream_task") print(f"{audio_path=}") generated_text = "" for new_text in s2s_inference.run_infer_stream(audio_path=audio_path): generated_text += new_text print(new_text, end="") print("") audio_decode_time = [] audio_segments = find_audio_segments_regex(generated_text) for audio_idx, audio_segment in enumerate(audio_segments): start = time.time() audio_tokens = extract_token_ids_as_int(audio_segment) # print(audio_tokens) tts_speech = s2s_inference.audio_tokenizer.decode(audio_tokens) end = time.time() audio_decode_time.append(end - start) wav_path = os.path.join(output_dir, audio_path[:-4] + f"_{audio_idx}.wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") # print(f"{audio_decode_time=}") # ============================================================== # ASR def asr_task(): for audio_path in [ "/data/data/wenet-e2e/wenetspeech/data/cuts_TEST_NET.00000000/TES/TEST_NET_Y0000000020_5XD21BihDd8_S00395.wav", "/data/data/wenet-e2e/wenetspeech/data/cuts_TEST_NET.00000000/TES/TEST_NET_Y0000000000_-KTKHdZ2fb8_S00424.wav", "/data/data/wenet-e2e/wenetspeech/data/cuts_TEST_NET.00000000/TES/TEST_NET_Y0000000050_LOLTeK1BNMo_S00045.wav", "/data/data/fixie-ai/librispeech_asr/test.clean/2830-3980-0034.wav", "/data/data/fixie-ai/librispeech_asr/test.clean/237-134500-0040.wav", ]: print("=" * 100) print("asr_task") print(f"{audio_path=}") output, tts_speech = s2s_inference.run_infer( audio_path=audio_path, # message="Translate the speech to text.", message="Convert the speech to text.", mode=None, ) print(f"{output=}", flush=True) # ============================================================== # TTS def tts_task(): TTS_texts = [ "我们将为全球城市的可持续发展贡献力量。", "通天河 灵感大王", "他本是我莲花池里养大的金鱼，每日浮头听经，修成手段。那一柄九瓣铜锤，乃是一枝未开的菡萏，被他运炼成兵。不知是那一日，海潮泛涨，走到此间。我今早扶栏看花，却不见这厮出拜，掐指巡纹，算着他在此成精，害你师父，故此未及梳妆，运神功，织个竹篮儿擒他。", "一二三四五六七八九十", "One Two Tree Four Five Six Seven Eight Night Ten", "1 2 3 4 5 6 7 8 9 10", "12345678910", "两个黄鹂鸣翠柳，一行白鹭上青天。窗含西岭千秋雪，门泊东吴万里船。", "坡上立着一只鹅，坡下就是一条河。宽宽的河，肥肥的鹅，鹅要过河，河要渡鹅不知是鹅过河，还是河渡鹅?", "扁担长，板凳宽，扁担没有板凳宽，板凳没有扁担长。扁担绑在板凳上，板凳不让扁担绑在板凳上。", "化肥会挥发，黑化肥发灰，灰化肥发黑。黑化肥发灰会挥发；灰化肥挥发会发黑。黑化肥挥发发灰会花飞；灰化肥挥发发黑会飞花，黑灰化肥会挥发发灰黑讳为花飞；灰黑化肥会挥发发黑灰为讳飞花。", "圆桌儿、方桌儿没有腿儿，墨水瓶儿里没有水儿，花瓶里有花儿没有叶儿，练习本儿上写字儿没有准儿，甘蔗好吃净是节儿。西瓜挺大没有味儿，坛儿里的小米儿长了虫儿，鸡毛掸子成了棍儿，水缸沿儿上系围裙儿，耗子打更猫打盹儿，新买的小褂儿没钉扣儿，奶奶想说没有劲儿。", "起床歌：小宝宝，起得早，睁开眼，眯眯笑，咿呀呀，学说话，伸伸手，要人抱。穿衣歌小胳膊，穿袖子，穿上衣，扣扣子，小脚丫，穿裤子，穿上袜子穿鞋子。小镜子-小镜子，圆又圆，看宝宝，露笑脸。闭上眼，做个梦，变月亮，挂上天。小铃铛叮铃铃，叮铃铃，一会远，一会近。小宝宝，耳朵灵，听铃声，找到铃。学画画小宝宝，学画画，大蜡笔，手中拿，画小鸭，叫嘎嘎，画小马，骑回家。大鞋子大鞋子，像只船，爸爸穿，我也穿，一二一，向前走，走呀走，翻了船。逛公园逛公园，宝宝笑，东看看，西瞧瞧，花儿香，鸟儿叫，小草绿，小树摇。看画报小娃娃，看画报，睁大眼，仔细瞧，布娃娃，哈哈笑，伸伸手，要你抱。搭积木大积木，红黄兰，小宝宝，最爱玩，搭火车，钻山洞，盖高楼，连着天。小汽车小汽车，嘀嘀嘀，开过来，开过去，小宝宝，当司机，送妈妈，上班去。藏猫猫儿歌：躲猫猫，躲猫猫， 猫猫、猫猫在哪里？喵……猫咪在这里。", ] for text in TTS_texts: print("=" * 100) print("tts_task") print(f"{text=}") output, tts_speech = s2s_inference.run_infer( message="Convert the text to speech.\n" + text, mode=None, do_sample=True, ) wav_path = os.path.join(output_dir, text[:16] + ".wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") # ============================================================== # Clone TTS for text in TTS_texts: for prompt_audio_path in [ "asset/2631296891109983590.wav", "asset/379838640-d5ff0815-74f8-4738-b0f1-477cfc8dcc2d.wav", "asset/4202818730519913143.wav", ]: print("=" * 100) print("tts_task") print(f"{text=} {prompt_audio_path=}") output, tts_speech = s2s_inference.run_infer( prompt_audio_path=prompt_audio_path, # message="Translate the text to speech.\n" + text, message="Convert the text to speech.\n" + text, mode=None, do_sample=True, ) wav_path = os.path.join(output_dir, prompt_audio_path[:16] + "_" + text[:16] + ".wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") # ============================================================== # TTS stream def tts_stream_task(): TTS_texts = [ "他本是我莲花池里养大的金鱼，每日浮头听经，修成手段。那一柄九瓣铜锤，乃是一枝未开的菡萏，被他运炼成兵。不知是那一日，海潮泛涨，走到此间。我今早扶栏看花，却不见这厮出拜，掐指巡纹，算着他在此成精，害你师父，故此未及梳妆，运神功，织个竹篮儿擒他。", ] for text in TTS_texts: print("=" * 100) print("tts_stream_task") print(f"{text=}") generated_text = "" for new_text in s2s_inference.run_infer_stream( message="Convert the text to speech.\n" + text, mode=None, do_sample=True, ): generated_text += new_text print(new_text, end="") print("") audio_segments = find_audio_segments_regex(generated_text) for audio_idx, audio_segment in enumerate(audio_segments): audio_tokens = extract_token_ids_as_int(audio_segment) # print(audio_tokens) tts_speech = s2s_inference.audio_tokenizer.decode(audio_tokens) wav_path = os.path.join(output_dir, text[:16] + f"_{audio_idx}.wav") os.makedirs(os.path.dirname(wav_path), exist_ok=True) torchaudio.save(wav_path, tts_speech.unsqueeze(0), 22050, format="wav") s2s_inference = S2SInference( model_name_or_path, audio_tokenizer_path, audio_tokenizer_type, flow_path=flow_path ) text_task() text_stream_task() sts_task() sts_stream_task() asr_task() tts_task() tts_stream_task() benchmark_sts() benchmark_llm()