1# !pip install --upgrade onnxruntime==1.22.1 huggingface_hub==0.34.4 transformers==4.46.3 numpy==2.2.6 tqdm==4.67.1 librosa==0.11.0 soundfile==0.13.1 resemble-perth==1.0.1
2
3import onnxruntime
4
5from huggingface_hub import hf_hub_download
6from transformers import AutoTokenizer
7
8import numpy as np
9from tqdm import tqdm
10import librosa
11import soundfile as sf
12
13S3GEN_SR = 24000
14START_SPEECH_TOKEN = 6561
15STOP_SPEECH_TOKEN = 6562
16
17
18class RepetitionPenaltyLogitsProcessor:
19 def __init__(self, penalty: float):
20 if not isinstance(penalty, float) or not (penalty > 0):
21 raise ValueError(f"`penalty` must be a strictly positive float, but is {penalty}")
22 self.penalty = penalty
23
24 def __call__(self, input_ids: np.ndarray, scores: np.ndarray) -> np.ndarray:
25 score = np.take_along_axis(scores, input_ids, axis=1)
26 score = np.where(score < 0, score * self.penalty, score / self.penalty)
27 scores_processed = scores.copy()
28 np.put_along_axis(scores_processed, input_ids, score, axis=1)
29 return scores_processed
30
31
32def run_inference(
33 text="The Lord of the Rings is the greatest work of literature.",
34 target_voice_path=None,
35 max_new_tokens = 256,
36 exaggeration=0.5,
37 output_dir="converted",
38 output_file_name="output.wav",
39 apply_watermark=True,
40):
41
42 model_id = "onnx-community/chatterbox-onnx"
43 if not target_voice_path:
44 target_voice_path = hf_hub_download(repo_id=model_id, filename="default_voice.wav", local_dir=output_dir)
45
46 ## Load model
47 speech_encoder_path = hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx", local_dir=output_dir, subfolder='onnx')
48 hf_hub_download(repo_id=model_id, filename="speech_encoder.onnx_data", local_dir=output_dir, subfolder='onnx')
49 embed_tokens_path = hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx", local_dir=output_dir, subfolder='onnx')
50 hf_hub_download(repo_id=model_id, filename="embed_tokens.onnx_data", local_dir=output_dir, subfolder='onnx')
51 conditional_decoder_path = hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx", local_dir=output_dir, subfolder='onnx')
52 hf_hub_download(repo_id=model_id, filename="conditional_decoder.onnx_data", local_dir=output_dir, subfolder='onnx')
53 language_model_path = hf_hub_download(repo_id=model_id, filename="language_model.onnx", local_dir=output_dir, subfolder='onnx')
54 hf_hub_download(repo_id=model_id, filename="language_model.onnx_data", local_dir=output_dir, subfolder='onnx')
55
56 # # Start inferense sessions
57 speech_encoder_session = onnxruntime.InferenceSession(speech_encoder_path)
58 embed_tokens_session = onnxruntime.InferenceSession(embed_tokens_path)
59 llama_with_past_session = onnxruntime.InferenceSession(language_model_path)
60 cond_decoder_session = onnxruntime.InferenceSession(conditional_decoder_path)
61
62 def execute_text_to_audio_inference(text):
63 print("Start inference script...")
64
65 audio_values, _ = librosa.load(target_voice_path, sr=S3GEN_SR)
66 audio_values = audio_values[np.newaxis, :].astype(np.float32)
67
68 ## Prepare input
69 tokenizer = AutoTokenizer.from_pretrained(model_id)
70 input_ids = tokenizer(text, return_tensors="np")["input_ids"].astype(np.int64)
71
72 position_ids = np.where(
73 input_ids >= START_SPEECH_TOKEN,
74 0,
75 np.arange(input_ids.shape[1])[np.newaxis, :] - 1
76 )
77
78 ort_embed_tokens_inputs = {
79 "input_ids": input_ids,
80 "position_ids": position_ids,
81 "exaggeration": np.array([exaggeration], dtype=np.float32)
82 }
83
84 ## Instantiate the logits processors.
85 repetition_penalty = 1.2
86 repetition_penalty_processor = RepetitionPenaltyLogitsProcessor(penalty=repetition_penalty)
87
88 num_hidden_layers = 30
89 num_key_value_heads = 16
90 head_dim = 64
91
92 generate_tokens = np.array([[START_SPEECH_TOKEN]], dtype=np.long)
93
94 # ---- Generation Loop using kv_cache ----
95 for i in tqdm(range(max_new_tokens), desc="Sampling", dynamic_ncols=True):
96
97 inputs_embeds = embed_tokens_session.run(None, ort_embed_tokens_inputs)[0]
98 if i == 0:
99 ort_speech_encoder_input = {
100 "audio_values": audio_values,
101 }
102 cond_emb, prompt_token, ref_x_vector, prompt_feat = speech_encoder_session.run(None, ort_speech_encoder_input)
103 inputs_embeds = np.concatenate((cond_emb, inputs_embeds), axis=1)
104
105 ## Prepare llm inputs
106 batch_size, seq_len, _ = inputs_embeds.shape
107 past_key_values = {
108 f"past_key_values.{layer}.{kv}": np.zeros([batch_size, num_key_value_heads, 0, head_dim], dtype=np.float32)
109 for layer in range(num_hidden_layers)
110 for kv in ("key", "value")
111 }
112 attention_mask = np.ones((batch_size, seq_len), dtype=np.int64)
113
114 logits, *present_key_values = llama_with_past_session.run(None, dict(
115 inputs_embeds=inputs_embeds,
116 attention_mask=attention_mask,
117 **past_key_values,
118 ))
119
120 logits = logits[:, -1, :]
121 next_token_logits = repetition_penalty_processor(generate_tokens, logits)
122
123 next_token = np.argmax(next_token_logits, axis=-1, keepdims=True).astype(np.int64)
124 generate_tokens = np.concatenate((generate_tokens, next_token), axis=-1)
125 if (next_token.flatten() == STOP_SPEECH_TOKEN).all():
126 break
127
128 # Get embedding for the new token.
129 position_ids = np.full(
130 (input_ids.shape[0], 1),
131 i + 1,
132 dtype=np.int64,
133 )
134 ort_embed_tokens_inputs["input_ids"] = next_token
135 ort_embed_tokens_inputs["position_ids"] = position_ids
136
137 ## Update values for next generation loop
138 attention_mask = np.concatenate([attention_mask, np.ones((batch_size, 1), dtype=np.int64)], axis=1)
139 for j, key in enumerate(past_key_values):
140 past_key_values[key] = present_key_values[j]
141
142 speech_tokens = generate_tokens[:, 1:-1]
143 speech_tokens = np.concatenate([prompt_token, speech_tokens], axis=1)
144 return speech_tokens, ref_x_vector, prompt_feat
145
146 speech_tokens, speaker_embeddings, speaker_features = execute_text_to_audio_inference(text)
147 cond_incoder_input = {
148 "speech_tokens": speech_tokens,
149 "speaker_embeddings": speaker_embeddings,
150 "speaker_features": speaker_features,
151 }
152 wav = cond_decoder_session.run(None, cond_incoder_input)[0]
153 wav = np.squeeze(wav, axis=0)
154
155 # Optional: Apply watermark
156 if apply_watermark:
157 import perth
158 watermarker = perth.PerthImplicitWatermarker()
159 wav = watermarker.apply_watermark(wav, sample_rate=S3GEN_SR)
160
161 sf.write(output_file_name, wav, S3GEN_SR)
162 print(f"{output_file_name} was successfully saved")
163
164if __name__ == "__main__":
165 run_inference(
166 text="Ezreal and Jinx teamed up with Ahri, Yasuo, and Teemo to take down the enemy's Nexus in an epic late-game pentakill.",
167 exaggeration=0.5,
168 output_file_name="output.wav",
169 apply_watermark=False,
170 )
