import json import numpy as np import re import argparse import os import onnxruntime as ort from optimum.onnxruntime import ORTModelForCausalLM from transformers import AutoTokenizer, PretrainedConfig, GenerationConfig import soundfile as sf GENDER_MAP = { "female": 0, "male": 1, } LEVELS_MAP = { "very_low": 0, "low": 1, "moderate": 2, "high": 3, "very_high": 4, } TASK_TOKEN_MAP = { "vc": "<|task_vc|>", "tts": "<|task_tts|>", "asr": "<|task_asr|>", "s2s": "<|task_s2s|>", "t2s": "<|task_t2s|>", "understand": "<|task_understand|>", "caption": "<|task_cap|>", "controllable_tts": "<|task_controllable_tts|>", "prompt_tts": "<|task_prompt_tts|>", "speech_edit": "<|task_edit|>", } def process_prompt( text: str, prompt_speech_path, audio_tokenizer, prompt_text: str = None, ): global_token_ids, semantic_token_ids = audio_tokenizer.tokenize( prompt_speech_path ) global_tokens = "".join( [f"<|bicodec_global_{i}|>" for i in global_token_ids.squeeze()] ) # Prepare the input tokens for the model if prompt_text is not None: semantic_tokens = "".join( [f"<|bicodec_semantic_{i}|>" for i in semantic_token_ids.squeeze()] ) inputs = [ TASK_TOKEN_MAP["tts"], "<|start_content|>", prompt_text, text, "<|end_content|>", "<|start_global_token|>", global_tokens, "<|end_global_token|>", "<|start_semantic_token|>", semantic_tokens, ] else: inputs = [ TASK_TOKEN_MAP["tts"], "<|start_content|>", text, "<|end_content|>", "<|start_global_token|>", global_tokens, "<|end_global_token|>", ] inputs = "".join(inputs) return inputs, global_token_ids def process_prompt_control(gender, pitch, speed, text): gender_id = GENDER_MAP[gender] pitch_level_id = LEVELS_MAP[pitch] speed_level_id = LEVELS_MAP[speed] pitch_label_tokens = f"<|pitch_label_{pitch_level_id}|>" speed_label_tokens = f"<|speed_label_{speed_level_id}|>" gender_tokens = f"<|gender_{gender_id}|>" attribte_tokens = "".join( [gender_tokens, pitch_label_tokens, speed_label_tokens] ) control_tts_inputs = [ TASK_TOKEN_MAP["controllable_tts"], "<|start_content|>", text, "<|end_content|>", "<|start_style_label|>", attribte_tokens, "<|end_style_label|>", ] return "".join(control_tts_inputs) def parse_arguments(): parser = argparse.ArgumentParser(description="Spark TTS inference script") parser.add_argument("--text", type=str, required=True, help="Text for TTS generation") parser.add_argument("--prompt", type=str, help="Transcript of prompt audio") parser.add_argument("--gender_voice", type=str, default="male", help="Voice gender") parser.add_argument("--clone_voice", type=str, default=None, help="Path to voice clone file") parser.add_argument("--model_dir", type=str, required=True, help="Path to the model directory") parser.add_argument("--flavor", type=str, default="q4", help="Model flavor: FP32, FP6, or quantized.") parser.add_argument("--num_gpus", type=int, default=1, help="Number of GPUs to use.") parser.add_argument("--pitch", type=str, default="moderate", help="Voice pitch.") parser.add_argument("--speed", type=str, default="moderate", help="Voice pitch.") return parser.parse_args() def main(): args = parse_arguments() use_gpu = False # "CUDAExecutionProvider" in ort.get_available_providers() providers = [ ("CUDAExecutionProvider", {"device_id": args.num_gpus - 1}) ] if use_gpu else [] providers.append("CPUExecutionProvider") work_dir = os.path.join(args.model_dir, "LLM") if os.path.exists(work_dir): config = PretrainedConfig.from_pretrained(work_dir) gen_config = GenerationConfig.from_pretrained(work_dir) suffix = "" if ( args.flavor is None or args.flavor == "" ) else f"_{args.flavor}" model_path = os.path.join(work_dir, "onnx", f"model{suffix}.onnx") sess_options = ort.SessionOptions() ort_model = ort.InferenceSession(model_path, sess_options, providers=providers) llm_model = ORTModelForCausalLM( session=ort_model, config=config, generation_config=gen_config, use_io_binding=True, use_cache=True, ) tokenizer = AutoTokenizer.from_pretrained(work_dir) else: raise ValueError(f"{model_path} does not exist.") sess_options = ort.SessionOptions() audio_detokenizer = ort.InferenceSession( os.path.join(args.model_dir, "bicodec.onnx"), sess_options, providers=providers ) # Process the prompt and clone_voice text = args.text text = "Some keyword arguments were passed to the ORTModelForCausalLM constructor that are not part of its signature: use_cache. These arguments will be ignored in the current version and will raise an error in the next version." if args.clone_voice: raise NotImplementedError() print(f"Using voice clone: {args.clone_voice}") prompt, global_tokens = process_prompt(args.text, args.clone_voice, args.prompt) else: print(f"Using {args.gender_voice} voice ") prompt = process_prompt_control(args.gender_voice, args.pitch, args.speed, text) print(f"Using prompt: {prompt}") inputs = tokenizer([prompt], return_tensors="pt") generated_ids = llm_model.generate( **inputs, max_new_tokens=3000, do_sample=True, top_k=50, top_p=0.95, temperature=0.8, ) # Trim the output tokens to remove the input tokens generated_ids = [ output_ids[len(input_ids) :] for input_ids, output_ids in zip(inputs.input_ids, generated_ids) ] # Decode the generated tokens into text predicts = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0] semantic_tokens = np.array( [int(token) for token in re.findall(r"bicodec_semantic_(\d+)", predicts)] )[None] if args.clone_voice is None: global_token = np.array( [int(token) for token in re.findall(r"bicodec_global_(\d+)", predicts)] )[None, None] print(semantic_tokens.shape) wav = audio_detokenizer.run( ["audio"], {"semantic_tokens": semantic_tokens, "global_tokens": global_token} )[0] sf.write("test_spark.wav", wav[0, 0], 16000) return # run: python test_spark_tts.py --text "Your text to synthesize" --model_dir "/path/to/model" if __name__ == "__main__": main()