ToluMichael/Chatterbox-Multilingual-TTS
1
1import random
2import numpy as np
3import torch
4from src.chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
5import gradio as gr
6import spaces
7
8DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
9print(f"🚀 Running on device: {DEVICE}")
10
11# --- Global Model Initialization ---
12MODEL = None
13
14LANGUAGE_CONFIG = {
15 "en": {
16 "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/en_f1.flac",
17 "text": "Last month, we reached a new milestone with two billion views on our YouTube channel."
18 },
19 "fr": {
20 "audio": "https://storage.googleapis.com/chatterbox-demo-samples/mtl_prompts/fr_f1.flac",
21 "text": "Le mois dernier, nous avons atteint un nouveau jalon avec deux milliards de vues sur notre chaîne YouTube."
22 },
23 # Add other languages as needed...
24}
25
26# --- UI Helpers ---
27def default_audio_for_ui(lang: str) -> str | None:
28 return LANGUAGE_CONFIG.get(lang, {}).get("audio")
29
30def default_text_for_ui(lang: str) -> str:
31 return LANGUAGE_CONFIG.get(lang, {}).get("text", "")
32
33def get_supported_languages_display() -> str:
34 language_items = []
35 for code, name in sorted(SUPPORTED_LANGUAGES.items()):
36 language_items.append(f"**{name}** (`{code}`)")
37 mid = len(language_items) // 2
38 line1 = " • ".join(language_items[:mid])
39 line2 = " • ".join(language_items[mid:])
40 return f"""
41### 🌍 Supported Languages ({len(SUPPORTED_LANGUAGES)} total)
42{line1}
43
44{line2}
45"""
46
47def get_or_load_model():
48 global MODEL
49 if MODEL is None:
50 print("Model not loaded, initializing...")
51 MODEL = ChatterboxMultilingualTTS.from_pretrained(DEVICE)
52 if hasattr(MODEL, 'to') and str(MODEL.device) != DEVICE:
53 MODEL.to(DEVICE)
54 print(f"Model loaded successfully. Internal device: {getattr(MODEL, 'device', 'N/A')}")
55 return MODEL
56
57try:
58 get_or_load_model()
59except Exception as e:
60 print(f"CRITICAL: Failed to load model. Error: {e}")
61
62def set_seed(seed: int):
63 torch.manual_seed(seed)
64 if DEVICE == "cuda":
65 torch.cuda.manual_seed(seed)
66 torch.cuda.manual_seed_all(seed)
67 random.seed(seed)
68 np.random.seed(seed)
69
70def chunk_text(text: str, max_len: int = 300):
71 """Split text into manageable chunks (~sentences or max_len)."""
72 import re
73 sentences = re.split(r'(?<=[.!?]) +', text)
74 chunks, current = [], ""
75 for sent in sentences:
76 if len(current) + len(sent) <= max_len:
77 current += " " + sent
78 else:
79 chunks.append(current.strip())
80 current = sent
81 if current:
82 chunks.append(current.strip())
83 return chunks
84
85def resolve_audio_prompt(language_id: str, provided_path: str | None) -> str | None:
86 if provided_path and str(provided_path).strip():
87 return provided_path
88 return LANGUAGE_CONFIG.get(language_id, {}).get("audio")
89
90@spaces.GPU
91def generate_tts_audio(
92 text_input: str,
93 language_id: str,
94 audio_prompt_path_input: str = None,
95 exaggeration_input: float = 0.5,
96 temperature_input: float = 0.8,
97 seed_num_input: int = 0,
98 cfgw_input: float = 0.5
99) -> tuple[int, np.ndarray]:
100
101 current_model = get_or_load_model()
102 if current_model is None:
103 raise RuntimeError("TTS model is not loaded.")
104
105 if seed_num_input != 0:
106 set_seed(int(seed_num_input))
107
108 chosen_prompt = audio_prompt_path_input or default_audio_for_ui(language_id)
109 generate_kwargs = {
110 "exaggeration": exaggeration_input,
111 "temperature": temperature_input,
112 "cfg_weight": cfgw_input,
113 }
114 if chosen_prompt:
115 generate_kwargs["audio_prompt_path"] = chosen_prompt
116
117 text_chunks = chunk_text(text_input, max_len=300)
118 print(f"Splitting text into {len(text_chunks)} chunks.")
119
120 audio_pieces = []
121 for idx, chunk in enumerate(text_chunks):
122 print(f"Generating chunk {idx+1}/{len(text_chunks)}: '{chunk[:50]}...' ")
123 wav = current_model.generate(chunk, language_id=language_id, **generate_kwargs)
124 audio_pieces.append(wav.squeeze(0).numpy())
125
126 final_audio = np.concatenate(audio_pieces)
127 print("Audio generation complete.")
128 return (current_model.sr, final_audio)
129
130with gr.Blocks() as demo:
131 gr.Markdown("""
132 # Chatterbox Multilingual Demo
133 Generate high-quality multilingual speech from text with reference audio styling, supporting 23 languages.
134 """)
135
136 gr.Markdown(get_supported_languages_display())
137
138 with gr.Row():
139 with gr.Column():
140 initial_lang = "fr"
141 text = gr.Textbox(
142 value=default_text_for_ui(initial_lang),
143 label="Text to synthesize (no character limit)",
144 max_lines=10
145 )
146
147 language_id = gr.Dropdown(
148 choices=list(ChatterboxMultilingualTTS.get_supported_languages().keys()),
149 value=initial_lang,
150 label="Language"
151 )
152
153 ref_wav = gr.Audio(
154 sources=["upload", "microphone"],
155 type="filepath",
156 label="Reference Audio File (Optional)",
157 value=default_audio_for_ui(initial_lang)
158 )
159
160 exaggeration = gr.Slider(0.25, 2, step=.05, label="Exaggeration", value=.5)
161 cfg_weight = gr.Slider(0.2, 1, step=.05, label="CFG/Pace", value=0.5)
162
163 with gr.Accordion("More options", open=False):
164 seed_num = gr.Number(value=0, label="Random seed (0 for random)")
165 temp = gr.Slider(0.05, 5, step=.05, label="Temperature", value=.8)
166
167 run_btn = gr.Button("Generate", variant="primary")
168
169 with gr.Column():
170 audio_output = gr.Audio(label="Output Audio")
171
172 def on_language_change(lang, current_ref, current_text):
173 return default_audio_for_ui(lang), default_text_for_ui(lang)
174
175 language_id.change(
176 fn=on_language_change,
177 inputs=[language_id, ref_wav, text],
178 outputs=[ref_wav, text],
179 show_progress=False
180 )
181
182 run_btn.click(
183 fn=generate_tts_audio,
184 inputs=[text, language_id, ref_wav, exaggeration, temp, seed_num, cfg_weight],
185 outputs=[audio_output],
186 )
187
188demo.launch(mcp_server=True, share=True)
189