-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathgeneration.py
More file actions
358 lines (296 loc) · 12 KB
/
Copy pathgeneration.py
File metadata and controls
358 lines (296 loc) · 12 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
import os
import asyncio
import queue
import threading
import json
import itertools
import librosa
from loguru import logger
import torch
import numpy as np
from TTS.tts.configs.xtts_config import XttsConfig
from TTS.tts.models.xtts import Xtts
from audio import FakeAudioStream
from text import split_text_into_speech_segments
cuda_lock = threading.Lock()
class VoiceGenerator:
_model_cache = {}
def __init__(
self,
model_name=None,
voice=None,
cached=False,
voice_speed=None,
leading_silence=None,
trailing_silence=None
):
"""
valid model_name values are listed in tts_models.json
"""
self.model_name = "multispeaker_original" if model_name is None else model_name
self.language = 'en'
self.tts_temperature = 0.7
self.speed = voice_speed if voice_speed is not None else 1.0
self.leading_silence = 0.0 if leading_silence is None else leading_silence
self.trailing_silence = 0.0 if trailing_silence is None else trailing_silence
self.tts_model_params = self.get_model(self.model_name, cached=cached)
self.tts_model = self.tts_model_params["model"]
self.tts_voices = self.tts_model_params["voices"]
default_voice = list(self.tts_voices.keys())[0]
self.set_voice(default_voice if voice is None else voice)
def start(self):
self.running = True
self.text_queue = queue.Queue()
self.thread = threading.Thread(target=self._processing_loop)
self.thread.start()
def set_model(self, model_name, cached=False):
self.model_name = model_name
self.tts_model, self.tts_voices = self.get_model(self.model_name, cached=cached)
def set_voice(self, voice):
if voice not in self.tts_voices:
raise ValueError(f"Voice {voice} not found. Available voices: {self.tts_voices.keys()}")
self.voice = voice
def maybe_set_voice_tone(self, voice_tone):
if ("voice_tone_map" in self.tts_model_params and
voice_tone in self.tts_model_params["voice_tone_map"]):
self.set_voice(self.tts_model_params["voice_tone_map"][voice_tone])
else:
logger.warning(f"Voice tone {voice_tone} not found, ignoring")
@staticmethod
def get_model(model_name, cached=False):
if model_name in VoiceGenerator._model_cache:
if cached: # return cached model
return VoiceGenerator._model_cache[model_name]
else: # re-load the model
VoiceGenerator._model_cache[model_name]["model"].cpu()
del VoiceGenerator._model_cache[model_name]
torch.cuda.empty_cache()
model_file = os.path.join(os.path.dirname(__file__), "tts_models.json")
with open(model_file, "r") as f:
model_config = json.load(f)[model_name]
# TTS lib will silently fail if the tokenizer file does not exist
assert os.path.exists(model_config["tokenizer"]), f"Tokenizer file does not exist: {model_config['tokenizer']}"
config = XttsConfig()
config.load_json(model_config["config"])
model = Xtts.init_from_config(config)
model.load_checkpoint(
config,
checkpoint_path=model_config["checkpoint"],
vocab_path=model_config["tokenizer"],
use_deepspeed=False
)
model.cuda()
voices = torch.load(model_config["voices"])
model_params = {
"model_name": model_name,
"model": model,
"voices": voices,
"avg_text_len": model_config["avg_text_len"]
}
if "voice_tone_map" in model_config:
model_params["voice_tone_map"] = model_config["voice_tone_map"]
VoiceGenerator._model_cache[model_name] = model_params
return model_params
def _sanitize_text(self, text):
return text
def _generate_silence(self, duration, streaming=False):
if streaming:
return itertools.repeat(np.zeros(int(duration * self.sample_rate), dtype=np.float32), 1)
else:
return np.zeros(int(duration * self.sample_rate), dtype=np.float32)
def generate(self, text, streaming=False):
# TODO: split text into chunks
if streaming:
return itertools.chain.from_iterable([
self._generate_silence(self.leading_silence, streaming=True),
self._stream_generator(text),
self._generate_silence(self.trailing_silence, streaming=True)
])
else:
with cuda_lock:
audio = self.tts_model.inference(
text,
language=self.language,
temperature=self.tts_temperature,
gpt_cond_latent=self.tts_voices[self.voice]["gpt_cond_latent"],
speaker_embedding=self.tts_voices[self.voice]["speaker_embedding"],
speed=self.speed
)["wav"]
return np.concatenate([
self._generate_silence(self.leading_silence),
audio,
self._generate_silence(self.trailing_silence)
])
def _stream_generator(
self,
text
):
audio_chunks = self.tts_model.inference_stream(
text=text,
language=self.language,
temperature=self.tts_temperature,
gpt_cond_latent=self.tts_voices[self.voice]["gpt_cond_latent"],
speaker_embedding=self.tts_voices[self.voice]["speaker_embedding"],
speed=self.speed
)
with cuda_lock:
for chunk in audio_chunks:
yield chunk.cpu().numpy()
@property
def sample_rate(self):
return 24000
def stop(self):
self.running = False
self.text_queue.put(None)
self.thread.join(timeout=5)
self.text_queue = None
class MultiVoiceGenerator:
def __init__(
self,
generators
):
assert isinstance(generators, dict), f"generators must be a \{'role': model\} like dict, but got {generators.__class__}"
self.generators = generators
@staticmethod
def from_config(config, cached=False):
generators = {}
for role, params in config.items():
generators[role] = VoiceGenerator(
model_name=params["model"],
voice=params.get("voice"),
leading_silence=params.get("leading_silence"),
trailing_silence=params.get("trailing_silence"),
voice_speed=params.get("speed"),
cached=cached,
)
return MultiVoiceGenerator(generators)
def _get_segments(self, text):
segments = split_text_into_speech_segments(text, avg_text_len=100)
roles = set([s["role"] for s in segments])
generators = set(self.generators.keys())
if roles > generators:
logger.warning(f"There aren't generators for these roles: {roles - generators}")
return segments
def generate(self, text, streaming=False):
segments = self._get_segments(text)
if streaming:
return itertools.chain.from_iterable(
self.generators[segment["role"]].generate(segment["text"], streaming=True)
for segment in segments if segment["role"] in self.generators
)
else:
audio_chunks = []
for segment in segments:
if segment["role"] in self.generators:
chunk = self.generators[segment["role"]].generate(segment["text"], streaming=False)
audio_chunks.append(chunk)
return np.concatenate(audio_chunks)
@property
def sample_rate(self):
return self.generators[list(self.generators.keys())[0]].sample_rate
def maybe_set_voice_tone(self, voice_tone):
assert "character" in self.generators, "Only character voice tone is supported for now, but character voice generator is not found"
self.generators["character"].maybe_set_voice_tone(voice_tone)
class DummyVoiceGenerator:
def __init__(self, *args, **kwargs):
pass
def generate(self, text, streaming=False):
zero_chunk = np.zeros(int(1 * self.sample_rate)) # 1 second of silence
if streaming:
return itertools.repeat(zero_chunk, 1)
else:
return zero_chunk
def maybe_set_voice_tone(self, voice_tone, role="character"):
pass
@property
def sample_rate(self):
return 24000
class AsyncVoiceGenerator:
def __init__(
self,
voice_generator,
generated_audio_cb,
error_cb=None,
event_loop=None
):
"""
generated_audio_cb receives f32le audio chunks
"""
self.voice_generator = voice_generator
self.generated_audio_cb = generated_audio_cb
self.error_cb = error_cb
self.running = False
self.text_queue = None
self._interrupt_flag = False
self._generating = False
self._event_loop = event_loop or asyncio.get_running_loop()
def generate(self, text, id=None):
if not self.running:
raise RuntimeError("AsyncVoiceGenerator is not running")
self.interrupt()
self.text_queue.put({"text": text, "id": id})
def start(self):
self.running = True
self.text_queue = queue.Queue()
self.thread = threading.Thread(target=self._processing_loop)
self.thread.start()
def stop(self):
self.running = False
self.text_queue.put(None)
self.thread.join(timeout=5)
self.text_queue = None
def interrupt(self):
if self._generating:
self._interrupt_flag = True
def _processing_loop(self):
def need_interrupt(text, id):
interrupt_flag = self._interrupt_flag
self._interrupt_flag = False
if interrupt_flag or not self.running:
logger.debug(f"Interrupting generation for text: {text}, id: {id}.")
return True
return False
try:
while self.running:
item = self.text_queue.get()
if item is None:
break
text, id = item["text"], item["id"]
logger.debug(f"Generating audio for text: {text}, id: {id}")
self._generating = True
if text.startswith("file:"):
stream = FakeAudioStream(text[5:], chunk_length=0.1, sr=self.sample_rate)
duration = stream.duration
for chunk, _ in stream:
asyncio.run_coroutine_threadsafe(
self.generated_audio_cb(audio_chunk=chunk, speech_id=id),
self._event_loop
)
if need_interrupt(text, id):
break
else:
duration = 0.0
for chunk in self.voice_generator.generate(text, streaming=True):
asyncio.run_coroutine_threadsafe(
self.generated_audio_cb(audio_chunk=chunk, speech_id=id),
self._event_loop
)
duration += len(chunk) / self.voice_generator.sample_rate
if need_interrupt(text, id):
break
self._generating = False
asyncio.run_coroutine_threadsafe(
self.generated_audio_cb(audio_chunk=None, speech_id=id, duration=duration),
self._event_loop
)
except Exception as e:
self.emit_error(e)
raise e
def emit_error(self, error):
if self.error_cb:
asyncio.run_coroutine_threadsafe(self.error_cb(error), self._event_loop)
@property
def sample_rate(self):
return self.voice_generator.sample_rate
def maybe_set_voice_tone(self, voice_tone):
self.voice_generator.maybe_set_voice_tone(voice_tone)