diff --git a/examples/40_Text_to_Speech_Generation.ipynb b/examples/40_Text_to_Speech_Generation.ipynb index 6b294da1..070d39d1 100644 --- a/examples/40_Text_to_Speech_Generation.ipynb +++ b/examples/40_Text_to_Speech_Generation.ipynb @@ -128,7 +128,7 @@ "\n", "# Print waveplot\n", "plt.figure(figsize=(15, 5))\n", - "plot = librosa.display.waveshow(speech, sr=rate)" + "plot = librosa.display.waveshow(speech[0], sr=speech[1])" ] }, { @@ -179,7 +179,7 @@ "\n", "def play(speech):\n", " # Convert to MP3 to save space\n", - " sf.write(\"speech.wav\", speech, 22050)\n", + " sf.write(\"speech.wav\", speech[0], speech[1])\n", " !ffmpeg -i speech.wav -y -b:a 64 speech.mp3 2> /dev/null\n", "\n", " # Play speech\n", @@ -259,7 +259,7 @@ "cell_type": "code", "source": [ "text = \"This is streaming speech generation. It's designed to take output tokens from a streaming LLM. It returns snippets of speech.\".split()\n", - "for speech in tts(text, stream=True):\n", + "for speech, _ in tts(text, stream=True):\n", " print(speech.shape)" ], "metadata": { diff --git a/src/python/txtai/models/pooling/factory.py b/src/python/txtai/models/pooling/factory.py index 3f114838..131367d0 100644 --- a/src/python/txtai/models/pooling/factory.py +++ b/src/python/txtai/models/pooling/factory.py @@ -5,8 +5,7 @@ Factory module import json import os -from huggingface_hub import hf_hub_download -from huggingface_hub.utils import EntryNotFoundError +from transformers.utils import cached_file from .base import Pooling from .cls import ClsPooling @@ -77,7 +76,7 @@ class PoolingFactory: method = "clspooling" # Ignore this error - except EntryNotFoundError: + except OSError: pass return method @@ -103,7 +102,7 @@ class PoolingFactory: maxlength = config.get("max_seq_length") # Ignore this error - except EntryNotFoundError: + except OSError: pass return maxlength @@ -122,6 +121,6 @@ class PoolingFactory: """ # Download file and parse JSON - path = hf_hub_download(repo_id=path, filename=name) + path = cached_file(path_or_repo_id=path, filename=name) with open(path, encoding="utf-8") as f: return json.load(f) diff --git a/src/python/txtai/pipeline/audio/__init__.py b/src/python/txtai/pipeline/audio/__init__.py index 5e7ce561..e3a1d8e3 100644 --- a/src/python/txtai/pipeline/audio/__init__.py +++ b/src/python/txtai/pipeline/audio/__init__.py @@ -2,7 +2,10 @@ Audio imports """ +from .audiomixer import AudioMixer from .audiostream import AudioStream from .microphone import Microphone +from .signal import Signal +from .texttoaudio import TextToAudio from .texttospeech import TextToSpeech from .transcription import Transcription diff --git a/src/python/txtai/pipeline/audio/audiomixer.py b/src/python/txtai/pipeline/audio/audiomixer.py new file mode 100644 index 00000000..31eb1f6a --- /dev/null +++ b/src/python/txtai/pipeline/audio/audiomixer.py @@ -0,0 +1,58 @@ +""" +AudioMixer module +""" + +from ..base import Pipeline +from .signal import Signal, SCIPY + + +class AudioMixer(Pipeline): + """ + Mixes multiple audio streams into a single stream. + """ + + def __init__(self, rate=None): + """ + Creates an AudioMixer pipeline. + + Args: + rate: optional target sample rate, otherwise uses input target rate with each audio segment + """ + + if not SCIPY: + raise ImportError('AudioMixer pipeline is not available - install "pipeline" extra to enable.') + + # Target sample rate + self.rate = rate + + def __call__(self, segment, scale1=1, scale2=1): + """ + Mixes multiple audio streams into a single stream. + + Args: + segment: ((audio1, sample rate), (audio2, sample rate))|list + scale1: optional scaling factor for segment1 + scale2: optional scaling factor for segment2 + + Returns: + list of (audio, sample rate) + """ + + # Convert single element to list + segments = [segment] if not isinstance(segment, list) else segment + + results = [] + for segment1, segment2 in segments: + audio1, rate1 = segment1 + audio2, rate2 = segment2 + + # Resample audio, as necessary + target = self.rate if self.rate else rate1 + audio1 = Signal.resample(audio1, rate1, target) + audio2 = Signal.resample(audio2, rate2, target) + + # Mix audio into single segment + results.append((Signal.mix(audio1, audio2, scale1, scale2), target)) + + # Return single element if single element passed in + return results[0] if not isinstance(segment, list) else results diff --git a/src/python/txtai/pipeline/audio/audiostream.py b/src/python/txtai/pipeline/audio/audiostream.py index 1c564149..04d3e56a 100644 --- a/src/python/txtai/pipeline/audio/audiostream.py +++ b/src/python/txtai/pipeline/audio/audiostream.py @@ -5,15 +5,15 @@ AudioStream module from queue import Queue from threading import Thread -import numpy as np - # Conditional import try: import sounddevice as sd - SOUNDDEVICE = True + from .signal import Signal, SCIPY + + AUDIOSTREAM = SCIPY except (ImportError, OSError): - SOUNDDEVICE = False + AUDIOSTREAM = False from ..base import Pipeline @@ -24,20 +24,26 @@ class AudioStream(Pipeline): to run on local machines given that it requires access to write to an output device. """ - COMPLETE = 1 + # End of stream message + COMPLETE = (1, None) - def __init__(self, rate=22050): + def __init__(self, rate=None): """ Creates an AudioStream pipeline. Args: - rate: sample rate of audio segments + rate: optional target sample rate, otherwise uses input target rate with each audio segment """ - if not SOUNDDEVICE: - raise ImportError("SoundDevice library not installed or portaudio library not found") + if not AUDIOSTREAM: + raise ImportError( + ( + 'AudioStream pipeline is not available - install "pipeline" extra to enable. ' + "Also check that the portaudio system library is available." + ) + ) - # Sampler rate + # Target sample rate self.rate = rate self.queue = Queue() @@ -45,14 +51,24 @@ class AudioStream(Pipeline): self.thread.start() def __call__(self, segment): + """ + Queues audio segments for the audio player. + + Args: + segment: (audio, sample rate)|list + + Returns: + segment + """ + # Convert single element to list - segments = [segment] if isinstance(segment, np.ndarray) else segment + segments = [segment] if not isinstance(segment, list) else segment for x in segments: self.queue.put(x) # Return single element if single element passed in - return segments[0] if isinstance(segment, np.ndarray) else segments + return segments[0] if not isinstance(segment, list) else segments def wait(self): """ @@ -66,7 +82,13 @@ class AudioStream(Pipeline): Reads audio segments from queue. This method runs in a separate non-blocking thread. """ - audio = self.queue.get() - while not isinstance(audio, int) or audio != AudioStream.COMPLETE: - sd.play(audio, self.rate, blocking=True) - audio = self.queue.get() + audio, rate = self.queue.get() + while not isinstance(audio, int) or (audio, rate) != AudioStream.COMPLETE: + # Resample to target sample rate, if necessary + audio, rate = (Signal.resample(audio, rate, self.rate), self.rate) if self.rate else (audio, rate) + + # Play audio segment + sd.play(audio, rate, blocking=True) + + # Get next segment + audio, rate = self.queue.get() diff --git a/src/python/txtai/pipeline/audio/microphone.py b/src/python/txtai/pipeline/audio/microphone.py index e37c9e1d..f9eba6c8 100644 --- a/src/python/txtai/pipeline/audio/microphone.py +++ b/src/python/txtai/pipeline/audio/microphone.py @@ -14,9 +14,11 @@ try: from scipy.fft import rfft, rfftfreq from scipy.signal import butter, sosfilt - SOUNDDEVICE = True + from .signal import Signal, SCIPY + + MICROPHONE = SCIPY except (ImportError, OSError): - SOUNDDEVICE = False + MICROPHONE = False from ..base import Pipeline @@ -45,8 +47,13 @@ class Microphone(Pipeline): pause: number of non-speech chunks to keep before considering speech complete, defaults to 8 """ - if not SOUNDDEVICE: - raise ImportError("SoundDevice library not installed or portaudio library not found") + if not MICROPHONE: + raise ImportError( + ( + 'Microphone pipeline is not available - install "pipeline" extra to enable. ' + "Also check that the portaudio system library is available." + ) + ) # Sample rate self.rate = rate @@ -65,6 +72,16 @@ class Microphone(Pipeline): self.pause = pause def __call__(self, device=None): + """ + Reads audio from an input device. + + Args: + device: optional input device id, otherwise uses system default + + Returns: + list of (audio, sample rate) + """ + # Listen for audio audio = self.listen(device[0] if isinstance(device, list) else device) @@ -124,7 +141,7 @@ class Microphone(Pipeline): # Convert to float32 and return audio = np.frombuffer(b"".join(chunks), np.int16) - return self.float32(audio) + return Signal.float32(audio) def isspeech(self, chunks): """ @@ -185,7 +202,7 @@ class Microphone(Pipeline): """ # Upsample to float32 - audio = self.float32(audio) + audio = Signal.float32(audio) # Human voice frequency range low = self.voicestart / (0.5 * self.rate) @@ -196,7 +213,7 @@ class Microphone(Pipeline): audio = sosfilt(sos, audio) # Scale back to int16 - audio = self.int16(audio) + audio = Signal.int16(audio) # Pass filtered signal to WebRTC VAD return self.detect(audio.tobytes()) @@ -238,35 +255,3 @@ class Microphone(Pipeline): ratio = speechenergy / sum(energyfreq.values()) logger.debug("SPEECH %.4f", ratio) return ratio >= self.vadthreshold - - def float32(self, audio): - """ - Converts an input NumPy array with 16-bit ints to 32-bit floats. - - Args: - audio: input audio array as 16-bit ints - - Returns: - audio array as 32-bit floats - """ - - i = np.iinfo(audio.dtype) - abs_max = 2 ** (i.bits - 1) - offset = i.min + abs_max - return (audio.astype(np.float32) - offset) / abs_max - - def int16(self, audio): - """ - Converts an input NumPy array with 32-bit floats to 16-bit ints. - - Args: - audio: input audio array as 32-bit floats - - Returns: - audio array as 16-bit ints - """ - - i = np.iinfo(np.int16) - absmax = 2 ** (i.bits - 1) - offset = i.min + absmax - return (audio * absmax + offset).clip(i.min, i.max).astype(np.int16) diff --git a/src/python/txtai/pipeline/audio/signal.py b/src/python/txtai/pipeline/audio/signal.py new file mode 100644 index 00000000..43f76ca6 --- /dev/null +++ b/src/python/txtai/pipeline/audio/signal.py @@ -0,0 +1,118 @@ +""" +Signal module +""" + +import numpy as np + +# Conditional import +try: + from scipy import signal + + SCIPY = True +except ImportError: + SCIPY = False + + +class Signal: + """ + Utility methods for audio signal processing. + """ + + @staticmethod + def mono(audio): + """ + Convert stereo to mono audio. + + Args: + audio: audio data + + Returns: + audio data with a single channel + """ + + return audio.mean(axis=1) if len(audio.shape) > 1 else audio + + @staticmethod + def resample(audio, rate, target): + """ + Resample audio if the sample rate doesn't match the target sample rate. + + Args: + audio: audio data + rate: current sample rate + target: target sample rate + + Returns: + audio resampled if necessary or original audio + """ + + if rate != target: + # Transpose audio + audio = audio.T + + # Resample audio and tranpose back + samples = round(len(audio) * float(target) / rate) + audio = signal.resample(audio, samples).T + + return audio + + @staticmethod + def float32(audio): + """ + Converts an input NumPy array with 16-bit ints to 32-bit floats. + + Args: + audio: input audio array as 16-bit ints + + Returns: + audio array as 32-bit floats + """ + + i = np.iinfo(audio.dtype) + abs_max = 2 ** (i.bits - 1) + offset = i.min + abs_max + return (audio.astype(np.float32) - offset) / abs_max + + @staticmethod + def int16(audio): + """ + Converts an input NumPy array with 32-bit floats to 16-bit ints. + + Args: + audio: input audio array as 32-bit floats + + Returns: + audio array as 16-bit ints + """ + + i = np.iinfo(np.int16) + absmax = 2 ** (i.bits - 1) + offset = i.min + absmax + return (audio * absmax + offset).clip(i.min, i.max).astype(np.int16) + + @staticmethod + def mix(audio1, audio2, scale1=1, scale2=1): + """ + Mixes audio1 and audio 2 into a single output audio segment. + + Args: + audio1: audio segment 1 + audio2: audio segment 2 + scale1: scale factor for audio segment 1 + scale2: scale factor for audio segment 2 + """ + + # Reshape audio, as necessary + audio1 = audio1.reshape(1, -1) if len(audio1.shape) <= 1 else audio1 + audio2 = audio2.reshape(1, -1) if len(audio2.shape) <= 1 else audio2 + + # Scale audio + audio1 = audio1 * scale1 + audio2 = audio2 * scale2 + + # Make audio files the same length + large, small = (audio1, audio2) if audio1.shape[1] > audio2.shape[1] else (audio2, audio1) + small = np.tile(small, (large.shape[1] // small.shape[1]) + 1).take(axis=1, indices=range(0, large.shape[1])) + + # Mix audio together + return small + large diff --git a/src/python/txtai/pipeline/audio/texttoaudio.py b/src/python/txtai/pipeline/audio/texttoaudio.py new file mode 100644 index 00000000..cfb700b9 --- /dev/null +++ b/src/python/txtai/pipeline/audio/texttoaudio.py @@ -0,0 +1,60 @@ +""" +TextToAudio module +""" + +from ..hfpipeline import HFPipeline +from .signal import Signal, SCIPY + + +class TextToAudio(HFPipeline): + """ + Generates audio from text. + """ + + def __init__(self, path=None, quantize=False, gpu=True, model=None, rate=None, **kwargs): + if not SCIPY: + raise ImportError('TextToAudio pipeline is not available - install "pipeline" extra to enable.') + + # Call parent constructor + super().__init__("text-to-audio", path, quantize, gpu, model, **kwargs) + + # Target sample rate, defaults to model sample rate + self.rate = rate + + def __call__(self, text, maxlength=512): + """ + Generates audio from text. + + This method supports text as a string or a list. If the input is a string, + the return type is a single audio output. If text is a list, the return type is a list. + + Args: + text: text|list + maxlength: maximum audio length to generate + + Returns: + list of (audio, sample rate) + """ + + # Format inputs + texts = [text] if isinstance(text, str) else text + + # Run pipeline + results = [self.convert(x) for x in self.pipeline(texts, forward_params={"max_new_tokens": maxlength})] + + # Extract results + return results[0] if isinstance(text, str) else results + + def convert(self, result): + """ + Converts audio result to target sample rate for this pipeline, if set. + + Args: + result: dict with audio samples and sample rate + + Returns: + (audio, sample rate) + """ + + audio, rate = result["audio"].squeeze(), result["sampling_rate"] + return (Signal.resample(audio, rate, self.rate), self.rate) if self.rate else (audio, rate) diff --git a/src/python/txtai/pipeline/audio/texttospeech.py b/src/python/txtai/pipeline/audio/texttospeech.py index 9a7f3cf6..e5ab6acc 100644 --- a/src/python/txtai/pipeline/audio/texttospeech.py +++ b/src/python/txtai/pipeline/audio/texttospeech.py @@ -10,7 +10,9 @@ try: from ttstokenizer import TTSTokenizer - TTS = True + from .signal import Signal, SCIPY + + TTS = SCIPY except ImportError: TTS = False @@ -19,7 +21,8 @@ import yaml import numpy as np -from huggingface_hub import hf_hub_download +from transformers import SpeechT5Processor +from transformers.utils import cached_file from ..base import Pipeline @@ -32,13 +35,14 @@ class TextToSpeech(Pipeline): Generates speech from text. """ - def __init__(self, path=None, maxtokens=512): + def __init__(self, path=None, maxtokens=512, rate=22050): """ Creates a new TextToSpeech pipeline. Args: - path: optional Hugging Face model hub id + path: optional model path maxtokens: maximum number of tokens model can process, defaults to 512 + rate: target sample rate, defaults to 22050 """ if not TTS: @@ -47,37 +51,19 @@ class TextToSpeech(Pipeline): # Default path path = path if path else "neuml/ljspeech-jets-onnx" - # Get path to model and config - config = hf_hub_download(path, filename="config.yaml") - model = hf_hub_download(path, filename="model.onnx") + # Target sample rate + self.rate = rate - # Read yaml config - with open(config, "r", encoding="utf-8") as f: - config = yaml.safe_load(f) - - # Create tokenizer - tokens = config.get("token", {}).get("list") - self.tokenizer = TTSTokenizer(tokens) - - # Create ONNX Session - self.model = ort.InferenceSession(model, ort.SessionOptions(), self.providers()) - - # Max number of input tokens model can handle - self.maxtokens = maxtokens - - # Get model input name, typically "text" - self.input = self.model.get_inputs()[0].name - - # Get parameter names - self.params = set(x.name for x in self.model.get_inputs()) + # Load target tts pipeline + self.pipeline = ESPnet(path, maxtokens, self.providers()) if self.hasfile(path, "model.onnx") else SpeechT5(path, maxtokens, self.providers()) def __call__(self, text, stream=False, speaker=1): """ Generates speech from text. Text longer than maxtokens will be batched and returned as a single waveform per text input. - This method supports files as a string or a list. If the input is a string, - the return type is string. If text is a list, the return type is a list. + This method supports text as a string or a list. If the input is a string, + the return type is audio. If text is a list, the return type is a list. Args: text: text|list @@ -85,7 +71,7 @@ class TextToSpeech(Pipeline): speaker: speaker id, defaults to 1 Returns: - list of speech as NumPy array waveforms + list of (audio, sample rate) """ # Convert results to a list if necessary @@ -117,6 +103,26 @@ class TextToSpeech(Pipeline): # Default when CUDA provider isn't available return ["CPUExecutionProvider"] + def hasfile(self, path, name): + """ + Tests if a file exists in a local or remote repo. + + Args: + path: model path + name: file name + + Returns: + True if name exists in path, False otherwise + """ + + try: + # Check if file exists + cached_file(path_or_repo_id=path, filename=name) + except OSError: + return False + + return True + def stream(self, texts, speaker): """ Iterates over texts, splits into segments and yields snippets of audio. @@ -134,7 +140,7 @@ class TextToSpeech(Pipeline): for x in texts: buffer.append(x) - if x == "\n" or x.strip().endswith("."): + if x == "\n" or (x.strip().endswith(".") and len([y for y in buffer if y]) > 2): data, buffer = "".join(buffer), [] yield self.execute(data, speaker) @@ -152,7 +158,65 @@ class TextToSpeech(Pipeline): speaker: speaker id Returns: - waveform as NumPy array + (audio, sample rate) + """ + + # Run pipeline model + audio, rate = self.pipeline(text, speaker) + + # Resample, if necessary and return + return (Signal.resample(audio, rate, self.rate), self.rate) if self.rate else (audio, rate) + + +class ESPnet(Pipeline): + """ + Text to Speech pipeline with an ESPnet ONNX model. + """ + + def __init__(self, path, maxtokens, providers): + """ + Creates a new ESPnet pipeline. + + Args: + path: model path + maxtokens: maximum number of tokens model can process + providers: list of supported ONNX providers + """ + + # Get path to model and config + config = cached_file(path_or_repo_id=path, filename="config.yaml") + model = cached_file(path_or_repo_id=path, filename="model.onnx") + + # Read yaml config + with open(config, "r", encoding="utf-8") as f: + config = yaml.safe_load(f) + + # Create tokenizer + tokens = config.get("token", {}).get("list") + self.tokenizer = TTSTokenizer(tokens) + + # Create ONNX Session + self.model = ort.InferenceSession(model, ort.SessionOptions(), providers) + + # Max number of input tokens model can handle + self.maxtokens = maxtokens + + # Get model input name, typically "text" + self.input = self.model.get_inputs()[0].name + + # Get parameter names + self.params = set(x.name for x in self.model.get_inputs()) + + def __call__(self, text, speaker): + """ + Executes a model run. This method will build batches of tokens when len(tokens) > maxtokens. + + Args: + text: text to tokenize and pass to model + speaker: speaker id + + Returns: + (audio, sample rate) """ # Debug logging for input text @@ -173,4 +237,162 @@ class TextToSpeech(Pipeline): results.append(output[0]) # Concatenate results and return - return np.concatenate(results) + return (np.concatenate(results), 22050) + + +class SpeechT5(Pipeline): + """ + Text to Speech pipeline with a SpeechT5 ONNX model. + """ + + def __init__(self, path, maxtokens, providers): + """ + Creates a new SpeechT5 pipeline. + + Args: + path: model path + maxtokens: maximum number of tokens model can process + providers: list of supported ONNX providers + """ + + self.encoder = ort.InferenceSession(cached_file(path_or_repo_id=path, filename="encoder_model.onnx"), providers=providers) + self.decoder = ort.InferenceSession(cached_file(path_or_repo_id=path, filename="decoder_model_merged.onnx"), providers=providers) + self.vocoder = ort.InferenceSession(cached_file(path_or_repo_id=path, filename="decoder_postnet_and_vocoder.onnx"), providers=providers) + + self.processor = SpeechT5Processor.from_pretrained(path) + self.defaultspeaker = np.load(cached_file(path_or_repo_id=path, filename="speaker.npy")) + + # Max number of input tokens model can handle + self.maxtokens = maxtokens + + def __call__(self, text, speaker): + """ + Executes a model run. This method will build batches of tokens when len(tokens) > maxtokens. + + Args: + text: text to tokenize and pass to model + speaker: speaker embeddings + + Returns: + (audio, sample rate) + """ + + # Debug logging for input text + logger.debug("%s", text) + + # Tokenize text + inputs = self.processor(text=text, return_tensors="np", normalize=True) + + # Split into batches and process + results = [] + for x in self.batch(inputs["input_ids"][0], self.maxtokens): + # Run text through TTS model and save waveform + results.append(self.process(np.array([x]), speaker)) + + # Concatenate results and return + return (np.concatenate(results), 16000) + + def process(self, inputs, speaker): + """ + Runs model inference. + + Args: + inputs: input token ids + speaker: speaker embeddings + + Returns: + waveform as NumPy array + """ + + # Run through encoder model + outputs = self.encoder.run(None, {"input_ids": inputs}) + outputs = {key.name: outputs[x] for x, key in enumerate(self.encoder.get_outputs())} + + # Encoder outputs and parameters + hiddenstate, attentionmask = outputs["encoder_outputs"], outputs["encoder_attention_mask"] + minlenratio, maxlenratio = 0.0, 20.0 + reduction, threshold, melbins = 2, 0.5, 80 + + maxlen = int(hiddenstate.shape[1] * maxlenratio / reduction) + minlen = int(hiddenstate.shape[1] * minlenratio / reduction) + + # Main processing loop + spectrogram, index, crossattention, branch, outputs = [], 0, None, False, {} + while True: + index += 1 + + inputs = { + "use_cache_branch": np.array([branch]), + "encoder_attention_mask": attentionmask, + "speaker_embeddings": speaker if speaker is not None and isinstance(speaker, np.ndarray) else self.defaultspeaker, + } + + if index == 1: + inputs = self.placeholders(inputs) + inputs["output_sequence"] = np.zeros((1, 1, melbins)).astype(np.float32) + inputs["encoder_hidden_states"] = hiddenstate + branch = True + else: + inputs = self.inputs(inputs, outputs, crossattention) + inputs["output_sequence"] = outputs["output_sequence_out"] + inputs["encoder_hidden_states"] = np.zeros((1, 0, 768)).astype(np.float32) + + # Run inputs through decoder + outputs = self.decoder.run(None, inputs) + outputs = {key.name: outputs[x] for x, key in enumerate(self.decoder.get_outputs())} + + # Get cross attention with 1st pass + if index == 1: + crossattention = {key: val for key, val in outputs.items() if ("encoder" in key and "present" in key)} + + # Decoder outputs + prob = outputs["prob"] + spectrum = outputs["spectrum"] + spectrogram.append(spectrum) + + # Done when stop token or maximum length is reached. + if index >= minlen and (int(sum(prob >= threshold)) > 0 or index >= maxlen): + spectrogram = np.concatenate(spectrogram) + return self.vocoder.run(None, {"spectrogram": spectrogram})[0] + + def placeholders(self, inputs): + """ + Creates decoder model inputs for initial inference pass. + + Args: + inputs: current decoder inputs + + Returns: + updated decoder inputs + """ + + length = inputs["encoder_attention_mask"].shape[1] + + for x in range(6): + inputs[f"past_key_values.{x}.encoder.key"] = np.zeros((1, 12, length, 64)).astype(np.float32) + inputs[f"past_key_values.{x}.encoder.value"] = np.zeros((1, 12, length, 64)).astype(np.float32) + inputs[f"past_key_values.{x}.decoder.key"] = np.zeros((1, 12, 1, 64)).astype(np.float32) + inputs[f"past_key_values.{x}.decoder.value"] = np.zeros((1, 12, 1, 64)).astype(np.float32) + + return inputs + + def inputs(self, inputs, previous, crossattention): + """ + Creates decoder model inputs for follow-on inference passes. + + Args: + inputs: current decoder inputs + previous: previous decoder outputs + crossattention: crossattention parameters + + Returns: + updated decoder inputs + """ + + for x in range(6): + inputs[f"past_key_values.{x}.encoder.key"] = crossattention[f"present.{x}.encoder.key"] + inputs[f"past_key_values.{x}.encoder.value"] = crossattention[f"present.{x}.encoder.value"] + inputs[f"past_key_values.{x}.decoder.key"] = previous[f"present.{x}.decoder.key"] + inputs[f"past_key_values.{x}.decoder.value"] = previous[f"present.{x}.decoder.value"] + + return inputs diff --git a/src/python/txtai/pipeline/audio/transcription.py b/src/python/txtai/pipeline/audio/transcription.py index 19a47d27..f121470b 100644 --- a/src/python/txtai/pipeline/audio/transcription.py +++ b/src/python/txtai/pipeline/audio/transcription.py @@ -5,11 +5,11 @@ Transcription module try: import soundfile as sf - from scipy import signal + from .signal import Signal, SCIPY - SOUNDFILE = True + TRANSCRIPTION = SCIPY except (ImportError, OSError): - SOUNDFILE = False + TRANSCRIPTION = False from ..hfpipeline import HFPipeline @@ -20,8 +20,10 @@ class Transcription(HFPipeline): """ def __init__(self, path=None, quantize=False, gpu=True, model=None, **kwargs): - if not SOUNDFILE: - raise ImportError("SoundFile library not installed or libsndfile not found") + if not TRANSCRIPTION: + raise ImportError( + 'Transcription pipeline is not available - install "pipeline" extra to enable. Also check that libsndfile is available.' + ) # Call parent constructor super().__init__("automatic-speech-recognition", path, quantize, gpu, model, **kwargs) @@ -67,7 +69,7 @@ class Transcription(HFPipeline): rate: optional sample rate Returns: - List of (audio data, sample rate) + list of (audio data, sample rate) """ speech = [] @@ -163,50 +165,18 @@ class Transcription(HFPipeline): Args: raw: raw audio data + rate: target sample rate Returns: audio data ready for pipeline model """ # Convert stereo to mono, if necessary - raw = self.mono(raw) + raw = Signal.mono(raw) - # Resample to model sample rate - raw, rate = self.resample(raw, rate) - - return {"raw": raw, "sampling_rate": rate} - - def mono(self, raw): - """ - Convert stereo to mono audio. - - Args: - raw: raw audio data - - Returns: - audio data with a single channel - """ - - return raw.mean(axis=1) if len(raw.shape) > 1 else raw - - def resample(self, raw, rate): - """ - Resample raw audio if the sample rate doesn't match the sample rate required for this model. - - Args: - raw: raw audio data - rate: sample rate - - Returns: - raw audio resampled if necessary or original raw audio - """ - - targetrate = self.pipeline.feature_extractor.sampling_rate - if rate != targetrate: - samples = round(len(raw) * float(targetrate) / rate) - raw = signal.resample(raw, samples) - - return raw, targetrate + # Resample to target sample rate + target = self.pipeline.feature_extractor.sampling_rate + return {"raw": Signal.resample(raw, rate, target), "sampling_rate": target} def clean(self, text): """ diff --git a/test/python/testoptional.py b/test/python/testoptional.py index 9394e28e..fcae4d4b 100644 --- a/test/python/testoptional.py +++ b/test/python/testoptional.py @@ -43,8 +43,9 @@ class TestOptional(unittest.TestCase): "pgvector", "PIL", "rich", - "sklearn.decomposition", + "scipy", "sentence_transformers", + "sklearn.decomposition", "sounddevice", "soundfile", "sqlalchemy", @@ -179,6 +180,7 @@ class TestOptional(unittest.TestCase): """ from txtai.pipeline import ( + AudioMixer, AudioStream, Caption, HFOnnx, @@ -192,11 +194,15 @@ class TestOptional(unittest.TestCase): Segmentation, Tabular, Textractor, + TextToAudio, TextToSpeech, Transcription, Translation, ) + with self.assertRaises(ImportError): + AudioMixer() + with self.assertRaises(ImportError): AudioStream() @@ -236,6 +242,9 @@ class TestOptional(unittest.TestCase): with self.assertRaises(ImportError): Textractor() + with self.assertRaises(ImportError): + TextToAudio() + with self.assertRaises(ImportError): TextToSpeech() diff --git a/test/python/testpipeline/testaudiomixer.py b/test/python/testpipeline/testaudiomixer.py new file mode 100644 index 00000000..8cd32898 --- /dev/null +++ b/test/python/testpipeline/testaudiomixer.py @@ -0,0 +1,29 @@ +""" +AudioMixer module tests +""" + +import unittest + +import numpy as np + +from txtai.pipeline import AudioMixer + + +class TestAudioStream(unittest.TestCase): + """ + AudioStream tests. + """ + + def testAudioStream(self): + """ + Test mixing audio streams + """ + + audio1 = np.random.rand(2, 5000), 100 + audio2 = np.random.rand(2, 5000), 100 + + mixer = AudioMixer() + audio, rate = mixer((audio1, audio2)) + + self.assertEqual(audio.shape, (2, 5000)) + self.assertEqual(rate, 100) diff --git a/test/python/testpipeline/testaudiostream.py b/test/python/testpipeline/testaudiostream.py index 4a9529e4..4705e505 100644 --- a/test/python/testpipeline/testaudiostream.py +++ b/test/python/testpipeline/testaudiostream.py @@ -28,10 +28,10 @@ class TestAudioStream(unittest.TestCase): play.return_value = True # Read audio data - raw, samplerate = sf.read(Utils.PATH + "/Make_huge_profits.wav") + audio, rate = sf.read(Utils.PATH + "/Make_huge_profits.wav") - audio = AudioStream(samplerate) - self.assertIsNotNone(audio([raw, AudioStream.COMPLETE])) + stream = AudioStream() + self.assertIsNotNone(stream([(audio, rate), AudioStream.COMPLETE])) # Wait for completion - audio.wait() + stream.wait() diff --git a/test/python/testpipeline/testtexttoaudio.py b/test/python/testpipeline/testtexttoaudio.py new file mode 100644 index 00000000..2369fa6e --- /dev/null +++ b/test/python/testpipeline/testtexttoaudio.py @@ -0,0 +1,26 @@ +""" +TextToAudio module tests +""" + +import unittest + +from txtai.pipeline import TextToAudio + + +class TestTextToAudio(unittest.TestCase): + """ + TextToAudio tests. + """ + + def testTextToAudio(self): + """ + Test generating audio for text + """ + + tta = TextToAudio("hf-internal-testing/tiny-random-MusicgenForConditionalGeneration") + + # Check that data is generated + audio, rate = tta("This is a test") + + self.assertGreater(len(audio), 0) + self.assertEqual(rate, 24000) diff --git a/test/python/testpipeline/testtexttospeech.py b/test/python/testpipeline/testtexttospeech.py index 3729485b..2b7c27a4 100644 --- a/test/python/testpipeline/testtexttospeech.py +++ b/test/python/testpipeline/testtexttospeech.py @@ -14,15 +14,18 @@ class TestTextToSpeech(unittest.TestCase): TextToSpeech tests. """ - def testTextToSpeech(self): + def testESPnet(self): """ - Test generating speech for text + Test generating speech for text with an ESPnet model """ tts = TextToSpeech() # Check that data is generated - self.assertGreater(len(tts("This is a test")), 0) + speech, rate = tts("This is a test") + + self.assertGreater(len(speech), 0) + self.assertEqual(rate, 22050) @patch("onnxruntime.get_available_providers") @patch("torch.cuda.is_available") @@ -38,6 +41,19 @@ class TestTextToSpeech(unittest.TestCase): tts = TextToSpeech() self.assertEqual(tts.providers()[0][0], "CUDAExecutionProvider") + def testSpeechT5(self): + """ + Test generating speech for text with a SpeechT5 model + """ + + tts = TextToSpeech("NeuML/txtai-speecht5-onnx") + + # Check that data is generated + speech, rate = tts("This is a test") + + self.assertGreater(len(speech), 0) + self.assertEqual(rate, 22050) + def testStreaming(self): """ Test streaming speech generation @@ -46,4 +62,8 @@ class TestTextToSpeech(unittest.TestCase): tts = TextToSpeech() # Check that data is generated - self.assertGreater(len(list(tts("This is a test. And another".split(), stream=True))), 0) + speech, rate = list(tts("This is a test. And another".split(), stream=True))[0] + + # Check that data is generated + self.assertGreater(len(speech), 0) + self.assertEqual(rate, 22050)