Add Text to Audio pipeline, closes #792. Add support for SpeechT5, closes #793. Updated cache methods, closes #794. Add audio signal processing, closes #795.

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": {

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)

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

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

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()

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)

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

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)

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

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):
         """

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()
 

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)

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()

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)

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)