140 lines
4.1 KiB
Python
140 lines
4.1 KiB
Python
#encoding = utf8
|
|
import heapq
|
|
import logging
|
|
import os
|
|
import shutil
|
|
import time
|
|
from threading import Lock, Thread
|
|
|
|
from eventbus import EventBus
|
|
from utils import save_wav
|
|
from human_handler import AudioHandler
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class TTSAudioHandle(AudioHandler):
|
|
def __init__(self, context, handler):
|
|
super().__init__(context, handler)
|
|
self._sample_rate = 16000
|
|
self._index = -1
|
|
|
|
EventBus().register('stop', self._on_stop)
|
|
|
|
def __del__(self):
|
|
EventBus().unregister('stop', self._on_stop)
|
|
|
|
def _on_stop(self, *args, **kwargs):
|
|
self.stop()
|
|
|
|
@property
|
|
def sample_rate(self):
|
|
return self._sample_rate
|
|
|
|
@sample_rate.setter
|
|
def sample_rate(self, value):
|
|
self._sample_rate = value
|
|
|
|
def get_index(self):
|
|
self._index = self._index + 1
|
|
return self._index
|
|
|
|
def on_handle(self, stream, index):
|
|
pass
|
|
|
|
def stop(self):
|
|
pass
|
|
|
|
def pause_talk(self):
|
|
self._index = -1
|
|
super().pause_talk()
|
|
|
|
|
|
class TTSAudioSplitHandle(TTSAudioHandle):
|
|
def __init__(self, context, handler):
|
|
super().__init__(context, handler)
|
|
self.sample_rate = self._context.sample_rate
|
|
self._chunk = self.sample_rate // self._context.fps
|
|
self._priority_queue = []
|
|
self._lock = Lock()
|
|
self._current = 0
|
|
self._is_running = True
|
|
self._thread = Thread(target=self._process_loop)
|
|
self._thread.start()
|
|
logger.info("TTSAudioSplitHandle init")
|
|
|
|
def _process_loop(self):
|
|
while self._is_running:
|
|
with self._lock:
|
|
if self._priority_queue and self._priority_queue[0][0] == self._current:
|
|
self._current += 1
|
|
chunks, txt = heapq.heappop(self._priority_queue)[1]
|
|
if chunks is not None:
|
|
for chunk in chunks:
|
|
self.on_next_handle((chunk, txt), 0)
|
|
time.sleep(0.001) # Sleep briefly to prevent busy-waiting
|
|
|
|
def on_handle(self, stream, index):
|
|
if not self._is_running:
|
|
logger.info('TTSAudioSplitHandle::on_handle is not running')
|
|
return
|
|
|
|
logger.info(f'TTSAudioSplitHandle::on_handle {index}')
|
|
s, txt = stream
|
|
if s is None:
|
|
heapq.heappush(self._priority_queue, (index, None))
|
|
else:
|
|
stream_len = s.shape[0]
|
|
idx = 0
|
|
chunks = []
|
|
while stream_len >= self._chunk and self._is_running:
|
|
chunks.append(s[idx:idx + self._chunk])
|
|
stream_len -= self._chunk
|
|
idx += self._chunk
|
|
if not self._is_running:
|
|
return
|
|
heapq.heappush(self._priority_queue, (index, (chunks, txt)))
|
|
|
|
def stop(self):
|
|
self._is_running = False
|
|
self._thread.join()
|
|
|
|
def pause_talk(self):
|
|
super().pause_talk()
|
|
with self._lock:
|
|
self._current = 0
|
|
self._priority_queue.clear()
|
|
|
|
|
|
class TTSAudioSaveHandle(TTSAudioHandle):
|
|
def __init__(self, context, handler):
|
|
super().__init__(context, handler)
|
|
self._save_path_dir = '../temp/audio/'
|
|
self._clean()
|
|
|
|
def _clean(self):
|
|
directory = self._save_path_dir
|
|
if not os.path.exists(directory):
|
|
print(f"The directory {directory} does not exist.")
|
|
return
|
|
|
|
for filename in os.listdir(directory):
|
|
file_path = os.path.join(directory, filename)
|
|
|
|
# 如果是文件，删除
|
|
if os.path.isfile(file_path):
|
|
os.remove(file_path)
|
|
print(f"Deleted file: {file_path}")
|
|
# 如果是文件夹，递归删除所有文件夹中的内容
|
|
elif os.path.isdir(file_path):
|
|
shutil.rmtree(file_path)
|
|
print(f"Deleted directory and its contents: {file_path}")
|
|
|
|
def on_handle(self, stream, index):
|
|
file_name = self._save_path_dir + str(index) + '.wav'
|
|
save_wav(stream, file_name, self.sample_rate)
|
|
|
|
def stop(self):
|
|
pass
|
|
|