| from queue import Queue | |
| import threading | |
| import time | |
| from typing import Optional, Tuple, List | |
| import numpy as np | |
| from pathlib import Path | |
| import logging | |
| from datetime import datetime | |
| from .audio_io import save_audio_file | |
| logging.getLogger("phonemizer").setLevel(logging.ERROR) | |
| logging.getLogger("speechbrain.utils.quirks").setLevel(logging.ERROR) | |
| logging.basicConfig(format="%(message)s", level=logging.INFO) | |
| class AudioGenerationQueue: | |
| """ | |
| A queue system for managing asynchronous audio generation from text input. | |
| This class implements a threaded queue system that handles text-to-audio generation | |
| in a background thread. It provides functionality for adding sentences to be processed, | |
| retrieving generated audio, and monitoring the generation process. | |
| Attributes: | |
| generator: Audio generator instance used for text-to-speech conversion | |
| speed (float): Speed multiplier for audio generation | |
| output_dir (Path): Directory where generated audio files are saved | |
| sentences_processed (int): Count of processed sentences | |
| audio_generated (int): Count of successfully generated audio files | |
| failed_sentences (list): List of tuples containing failed sentences and error messages | |
| """ | |
| def __init__( | |
| self, generator, speed: float = 1.0, output_dir: Optional[Path] = None | |
| ): | |
| """ | |
| Initialize the audio generation queue system. | |
| Args: | |
| generator: Audio generator instance for text-to-speech conversion | |
| speed: Speed multiplier for audio generation (default: 1.0) | |
| output_dir: Directory path for saving generated audio files (default: "generated_audio") | |
| """ | |
| self.generator = generator | |
| self.speed = speed | |
| self.lock = threading.Lock() | |
| self.output_dir = output_dir or Path("generated_audio") | |
| self.output_dir.mkdir(exist_ok=True) | |
| self.sentence_queue = Queue() | |
| self.audio_queue = Queue() | |
| self.is_running = False | |
| self.generation_thread = None | |
| self.sentences_processed = 0 | |
| self.audio_generated = 0 | |
| self.failed_sentences = [] | |
| def start(self): | |
| """ | |
| Start the audio generation thread if not already running. | |
| The thread will process sentences from the queue until stopped. | |
| """ | |
| if not self.is_running: | |
| self.is_running = True | |
| self.generation_thread = threading.Thread(target=self._generation_worker) | |
| self.generation_thread.daemon = True | |
| self.generation_thread.start() | |
| def stop(self): | |
| """ | |
| Stop the audio generation thread gracefully. | |
| Waits for the current queue to be processed before stopping. | |
| Outputs final processing statistics. | |
| """ | |
| if self.generation_thread: | |
| while not self.sentence_queue.empty(): | |
| time.sleep(0.1) | |
| time.sleep(0.5) | |
| self.is_running = False | |
| self.generation_thread.join() | |
| self.generation_thread = None | |
| logging.info( | |
| f"\nAudio Generation Complete - Processed: {self.sentences_processed}, Generated: {self.audio_generated}, Failed: {len(self.failed_sentences)}" | |
| ) | |
| def add_sentences(self, sentences: List[str]): | |
| """ | |
| Add a list of sentences to the generation queue. | |
| Args: | |
| sentences: List of text strings to be converted to audio | |
| """ | |
| added_count = 0 | |
| for sentence in sentences: | |
| sentence = sentence.strip() | |
| if sentence: | |
| self.sentence_queue.put(sentence) | |
| added_count += 1 | |
| if not self.is_running: | |
| self.start() | |
| def get_next_audio(self) -> Tuple[Optional[np.ndarray], Optional[Path]]: | |
| """ | |
| Retrieve the next generated audio segment from the queue. | |
| Returns: | |
| Tuple containing: | |
| - numpy array of audio data (or None if queue is empty) | |
| - Path object for the saved audio file (or None if queue is empty) | |
| """ | |
| try: | |
| audio_data, output_path = self.audio_queue.get_nowait() | |
| return audio_data, output_path | |
| except: | |
| return None, None | |
| def clear_queues(self): | |
| """ | |
| Clear both sentence and audio queues, removing all pending items. | |
| Returns immediately without waiting for queue processing. | |
| """ | |
| sentences_cleared = 0 | |
| audio_cleared = 0 | |
| while not self.sentence_queue.empty(): | |
| try: | |
| self.sentence_queue.get_nowait() | |
| sentences_cleared += 1 | |
| except: | |
| pass | |
| while not self.audio_queue.empty(): | |
| try: | |
| self.audio_queue.get_nowait() | |
| audio_cleared += 1 | |
| except: | |
| pass | |
| def _generation_worker(self): | |
| """ | |
| Internal worker method that runs in a separate thread. | |
| Continuously processes sentences from the queue, generating audio | |
| and handling any errors that occur during generation. | |
| """ | |
| while self.is_running or not self.sentence_queue.empty(): | |
| try: | |
| try: | |
| sentence = self.sentence_queue.get_nowait() | |
| self.sentences_processed += 1 | |
| except: | |
| if not self.is_running and self.sentence_queue.empty(): | |
| break | |
| time.sleep(0.01) | |
| continue | |
| try: | |
| audio_data, phonemes = self.generator.generate( | |
| sentence, speed=self.speed | |
| ) | |
| if audio_data is None or len(audio_data) == 0: | |
| raise ValueError("Generated audio data is empty") | |
| output_path = save_audio_file(audio_data, self.output_dir) | |
| self.audio_generated += 1 | |
| self.audio_queue.put((audio_data, output_path)) | |
| except Exception as e: | |
| error_msg = str(e) | |
| self.failed_sentences.append((sentence, error_msg)) | |
| continue | |
| except Exception as e: | |
| if not self.is_running and self.sentence_queue.empty(): | |
| break | |
| time.sleep(0.1) | |