Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3import soundfile as sf
4
5llasa_3b ='HKUSTAudio/Llasa-3B'
6
7tokenizer = AutoTokenizer.from_pretrained(llasa_3b)
8model = AutoModelForCausalLM.from_pretrained(llasa_3b)
9model.eval()
10model.to('cuda')
11
12from xcodec2.modeling_xcodec2 import XCodec2Model
13
14model_path = "HKUST-Audio/xcodec2"
15
16Codec_model = XCodec2Model.from_pretrained(model_path)
17Codec_model.eval().cuda()
18
19input_text = 'Dealing with family secrets is never easy. Yet, sometimes, omission is a form of protection, intending to safeguard some from the harsh truths. One day, I hope you understand the reasons behind my actions. Until then, Anna, please, bear with me.'
20# input_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道："我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？"'
21def ids_to_speech_tokens(speech_ids):
22
23 speech_tokens_str = []
24 for speech_id in speech_ids:
25 speech_tokens_str.append(f"<|s_{speech_id}|>")
26 return speech_tokens_str
27
28def extract_speech_ids(speech_tokens_str):
29
30 speech_ids = []
31 for token_str in speech_tokens_str:
32 if token_str.startswith('<|s_') and token_str.endswith('|>'):
33 num_str = token_str[4:-2]
34
35 num = int(num_str)
36 speech_ids.append(num)
37 else:
38 print(f"Unexpected token: {token_str}")
39 return speech_ids
40
41#TTS start!
42with torch.no_grad():
43
44 formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
45
46 # Tokenize the text
47 chat = [
48 {"role": "user", "content": "Convert the text to speech:" + formatted_text},
49 {"role": "assistant", "content": "<|SPEECH_GENERATION_START|>"}
50 ]
51
52 input_ids = tokenizer.apply_chat_template(
53 chat,
54 tokenize=True,
55 return_tensors='pt',
56 continue_final_message=True
57 )
58 input_ids = input_ids.to('cuda')
59 speech_end_id = tokenizer.convert_tokens_to_ids('<|SPEECH_GENERATION_END|>')
60
61 # Generate the speech autoregressively
62 outputs = model.generate(
63 input_ids,
64 max_length=2048, # We trained our model with a max length of 2048
65 eos_token_id= speech_end_id ,
66 do_sample=True,
67 top_p=1, # Adjusts the diversity of generated content
68 temperature=0.8, # Controls randomness in output
69 )
70 # Extract the speech tokens
71 generated_ids = outputs[0][input_ids.shape[1]:-1]
72
73 speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
74
75 # Convert token <|s_23456|> to int 23456
76 speech_tokens = extract_speech_ids(speech_tokens)
77
78 speech_tokens = torch.tensor(speech_tokens).cuda().unsqueeze(0).unsqueeze(0)
79
80 # Decode the speech tokens to speech waveform
81 gen_wav = Codec_model.decode_code(speech_tokens)
82
83
84sf.write("gen.wav", gen_wav[0, 0, :].cpu().numpy(), 16000)1from transformers import AutoTokenizer, AutoModelForCausalLM
2import torch
3import soundfile as sf
4
5llasa_3b ='HKUST-Audio/Llasa-3B'
6
7tokenizer = AutoTokenizer.from_pretrained(llasa_3b)
8model = AutoModelForCausalLM.from_pretrained(llasa_3b)
9model.eval()
10model.to('cuda')
11
12from xcodec2.modeling_xcodec2 import XCodec2Model
13
14model_path = "HKUST-Audio/xcodec2"
15
16Codec_model = XCodec2Model.from_pretrained(model_path)
17Codec_model.eval().cuda()
18# only 16khz speech support!
19prompt_wav, sr = sf.read("太乙真人.wav") # you can find wav in Files
20#prompt_wav, sr = sf.read("Anna.wav") # English prompt
21prompt_wav = torch.from_numpy(prompt_wav).float().unsqueeze(0)
22
23prompt_text ="对，这就是我万人敬仰的太乙真人，虽然有点婴儿肥，但也掩不住我逼人的帅气。"
24#promt_text = "A chance to leave him alone, but... No. She just wanted to see him again. Anna, you don't know how it feels to lose a sister. Anna, I'm sorry, but your father asked me not to tell you anything."
25target_text = '突然，身边一阵笑声。我看着他们，意气风发地挺直了胸膛，甩了甩那稍显肉感的双臂，轻笑道："我身上的肉，是为了掩饰我爆棚的魅力，否则，岂不吓坏了你们呢？"'
26#target_text = "Dealing with family secrets is never easy. Yet, sometimes, omission is a form of protection, intending to safeguard some from the harsh truths. One day, I hope you understand the reasons behind my actions. Until then, Anna, please, bear with me."
27input_text = prompt_text + target_text
28
29def ids_to_speech_tokens(speech_ids):
30
31 speech_tokens_str = []
32 for speech_id in speech_ids:
33 speech_tokens_str.append(f"<|s_{speech_id}|>")
34 return speech_tokens_str
35
36def extract_speech_ids(speech_tokens_str):
37
38 speech_ids = []
39 for token_str in speech_tokens_str:
40 if token_str.startswith('<|s_') and token_str.endswith('|>'):
41 num_str = token_str[4:-2]
42
43 num = int(num_str)
44 speech_ids.append(num)
45 else:
46 print(f"Unexpected token: {token_str}")
47 return speech_ids
48
49#TTS start!
50with torch.no_grad():
51 # Encode the prompt wav
52 vq_code_prompt = Codec_model.encode_code(input_waveform=prompt_wav)
53 print("Prompt Vq Code Shape:", vq_code_prompt.shape )
54
55 vq_code_prompt = vq_code_prompt[0,0,:]
56 # Convert int 12345 to token <|s_12345|>
57 speech_ids_prefix = ids_to_speech_tokens(vq_code_prompt)
58
59 formatted_text = f"<|TEXT_UNDERSTANDING_START|>{input_text}<|TEXT_UNDERSTANDING_END|>"
60
61 # Tokenize the text and the speech prefix
62 chat = [
63 {"role": "user", "content": "Convert the text to speech:" + formatted_text},
64 {"role": "assistant", "content": "<|SPEECH_GENERATION_START|>" + ''.join(speech_ids_prefix)}
65 ]
66
67 input_ids = tokenizer.apply_chat_template(
68 chat,
69 tokenize=True,
70 return_tensors='pt',
71 continue_final_message=True
72 )
73 input_ids = input_ids.to('cuda')
74 speech_end_id = tokenizer.convert_tokens_to_ids('<|SPEECH_GENERATION_END|>')
75
76 # Generate the speech autoregressively
77 outputs = model.generate(
78 input_ids,
79 max_length=2048, # We trained our model with a max length of 2048
80 eos_token_id= speech_end_id ,
81 do_sample=True,
82 top_p=1,
83 temperature=0.8,
84 )
85 # Extract the speech tokens
86 generated_ids = outputs[0][input_ids.shape[1]-len(speech_ids_prefix):-1]
87
88 speech_tokens = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)
89
90 # Convert token <|s_23456|> to int 23456
91 speech_tokens = extract_speech_ids(speech_tokens)
92
93 speech_tokens = torch.tensor(speech_tokens).cuda().unsqueeze(0).unsqueeze(0)
94
95 # Decode the speech tokens to speech waveform
96 gen_wav = Codec_model.decode_code(speech_tokens)
97
98 # if only need the generated part
99 # gen_wav = gen_wav[:,:,prompt_wav.shape[1]:]
100
101sf.write("gen.wav", gen_wav[0, 0, :].cpu().numpy(), 16000)