truss 0.10.0rc1__py3-none-any.whl → 0.60.0__py3-none-any.whl

truss/templates/trtllm-audio/packages/whisper_trt/utils.py

@@ -0,0 +1,245 @@
+# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import json
+import os
+from collections import OrderedDict
+from functools import lru_cache
+from pathlib import Path
+from subprocess import CalledProcessError, run
+from typing import Optional, Union
+
+import numpy as np
+import soundfile
+import torch
+import torch.nn.functional as F
+from anyio import to_thread
+
+Pathlike = Union[str, Path]
+
+SAMPLE_RATE = 16000
+N_FFT = 400
+HOP_LENGTH = 160
+CHUNK_LENGTH = 30
+N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
+
+
+def load_audio(file: str, sr: int = SAMPLE_RATE):
+    """
+    Open an audio file and read as mono waveform, resampling as necessary
+
+    Parameters
+    ----------
+    file: str
+        The audio file to open
+
+    sr: int
+        The sample rate to resample the audio if necessary
+
+    Returns
+    -------
+    A NumPy array containing the audio waveform, in float32 dtype.
+    """
+
+    # This launches a subprocess to decode audio while down-mixing
+    # and resampling as necessary.  Requires the ffmpeg CLI in PATH.
+    # fmt: off
+    cmd = [
+        "ffmpeg", "-nostdin", "-threads", "0", "-i", file, "-f", "s16le", "-ac",
+        "1", "-acodec", "pcm_s16le", "-ar",
+        str(sr), "-"
+    ]
+    # fmt: on
+    try:
+        out = run(cmd, capture_output=True, check=True).stdout
+    except CalledProcessError as e:
+        raise RuntimeError(f"Failed to load audio: {e.stderr.decode()}") from e
+
+    return np.frombuffer(out, np.int16).flatten().astype(np.float32) / 32768.0
+
+
+def load_audio_wav_format(wav_path):
+    # make sure audio in .wav format
+    assert wav_path.endswith(".wav"), f"Only support .wav format, but got {wav_path}"
+    waveform, sample_rate = soundfile.read(wav_path)
+    assert sample_rate == 16000, f"Only support 16k sample rate, but got {sample_rate}"
+    return waveform, sample_rate
+
+
+def pad_or_trim(array, length: int = N_SAMPLES, *, axis: int = -1):
+    """
+    Pad or trim the audio array to N_SAMPLES, as expected by the encoder.
+    """
+    if torch.is_tensor(array):
+        if array.shape[axis] > length:
+            array = array.index_select(
+                dim=axis, index=torch.arange(length, device=array.device)
+            )
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = F.pad(array, [pad for sizes in pad_widths[::-1] for pad in sizes])
+    else:
+        if array.shape[axis] > length:
+            array = array.take(indices=range(length), axis=axis)
+
+        if array.shape[axis] < length:
+            pad_widths = [(0, 0)] * array.ndim
+            pad_widths[axis] = (0, length - array.shape[axis])
+            array = np.pad(array, pad_widths)
+
+    return array
+
+
+@lru_cache(maxsize=None)
+def mel_filters(
+    device, n_mels: int, mel_filters_dir: Optional[str] = None
+) -> torch.Tensor:
+    """
+    load the mel filterbank matrix for projecting STFT into a Mel spectrogram.
+    Allows decoupling librosa dependency; saved using:
+
+        np.savez_compressed(
+            "mel_filters.npz",
+            mel_80=librosa.filters.mel(sr=16000, n_fft=400, n_mels=80),
+        )
+    """
+    assert n_mels in {80, 128}, f"Unsupported n_mels: {n_mels}"
+    if mel_filters_dir is None:
+        mel_filters_path = os.path.join(
+            os.path.dirname(__file__), "assets", "mel_filters.npz"
+        )
+    else:
+        mel_filters_path = os.path.join(mel_filters_dir, "mel_filters.npz")
+    with np.load(mel_filters_path) as f:
+        return torch.from_numpy(f[f"mel_{n_mels}"]).to(device)
+
+
+def _log_mel_spectrogram(
+    audio: Union[str, np.ndarray, torch.Tensor],
+    n_mels: int,
+    padding: int = 0,
+    device: Optional[Union[str, torch.device]] = None,
+    mel_filters_dir: Optional[str] = None,
+):
+    """
+    Compute the log-Mel spectrogram of
+
+    Parameters
+    ----------
+    audio: Union[str, np.ndarray, torch.Tensor], shape = (*)
+        The path to audio or either a NumPy array or Tensor containing the audio waveform in 16 kHz
+
+    n_mels: int
+        The number of Mel-frequency filters, only 80 and 128 are supported
+
+    padding: int
+        Number of zero samples to pad to the right
+
+    device: Optional[Union[str, torch.device]]
+        If given, the audio tensor is moved to this device before STFT
+
+    Returns
+    -------
+    torch.Tensor, shape = (80 or 128, n_frames)
+        A Tensor that contains the Mel spectrogram
+    """
+    if not torch.is_tensor(audio):
+        if isinstance(audio, str):
+            if audio.endswith(".wav"):
+                audio, _ = load_audio_wav_format(audio)
+            else:
+                audio = load_audio(audio)
+        assert isinstance(audio, np.ndarray), f"Unsupported audio type: {type(audio)}"
+        audio = pad_or_trim(audio, N_SAMPLES)
+        audio = audio.astype(np.float32)
+        audio = torch.from_numpy(audio)
+
+    if device is not None:
+        audio = audio.to(device)  # type: ignore
+    if padding > 0:
+        audio = F.pad(audio, (0, padding))
+    window = torch.hann_window(N_FFT).to(audio.device)  # type: ignore
+    stft = torch.stft(audio, N_FFT, HOP_LENGTH, window=window, return_complex=True)
+    magnitudes = stft[..., :-1].abs() ** 2
+
+    filters = mel_filters(audio.device, n_mels, mel_filters_dir)  # type: ignore
+    mel_spec = filters @ magnitudes
+
+    log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+    log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+    log_spec = (log_spec + 4.0) / 4.0
+    return log_spec
+
+
+async def log_mel_spectrogram(
+    audio: Union[str, np.ndarray, torch.Tensor],
+    n_mels: int,
+    padding: int = 0,
+    device: Optional[Union[str, torch.device]] = None,
+    mel_filters_dir: Optional[str] = None,
+):
+    log_spec = await to_thread.run_sync(
+        lambda: _log_mel_spectrogram(audio, n_mels, padding, device, mel_filters_dir)
+    )
+    return log_spec
+
+
+def remove_tensor_padding(input_tensor, input_tensor_lengths=None, pad_value=0):
+    if input_tensor.dim() == 2:
+        # Text tensor case: batch, seq_len
+        assert torch.all(input_tensor[:, 0] != pad_value), (
+            "First token in each sequence should not be pad_value"
+        )
+        assert input_tensor_lengths is None
+
+        # Create a mask for all non-pad tokens
+        mask = input_tensor != pad_value
+
+        # Apply the mask to input_tensor to remove pad tokens
+        output_tensor = input_tensor[mask].view(1, -1)
+
+    elif input_tensor.dim() == 3:
+        # Audio tensor case: batch, seq_len, feature_len
+        assert input_tensor_lengths is not None, (
+            "input_tensor_lengths must be provided for 3D input_tensor"
+        )
+        batch_size, seq_len, feature_len = input_tensor.shape
+
+        # Initialize a list to collect valid sequences
+        valid_sequences = []
+
+        for i in range(batch_size):
+            valid_length = input_tensor_lengths[i]
+            valid_sequences.append(input_tensor[i, :valid_length, :])
+
+        # Concatenate all valid sequences along the batch dimension
+        output_tensor = torch.cat(valid_sequences, dim=0)
+
+    else:
+        raise ValueError("Input tensor must have 2 or 3 dimensions")
+
+    return output_tensor
+
+
+def read_config(component, engine_dir):
+    config_path = engine_dir / component / "config.json"
+    with open(config_path, "r") as f:
+        config = json.load(f)
+    model_config = OrderedDict()
+    model_config.update(config["pretrained_config"])
+    model_config.update(config["build_config"])
+    return model_config

truss/templates/trtllm-briton/src/extension.py

@@ -0,0 +1,64 @@
+from briton.spec_dec_truss_model import Model as SpecDecModel
+from briton.trtllm_config import TRTLLMConfiguration
+from briton.truss_model import Model
+
+# TODO(pankaj) Define an ABC base class for this. That baseclass should live in
+# a new, smaller truss sub-library, perhaps called `truss-runtime`` for inclusion
+# in Truss runtime. Once we have that sub-library, we should define the Extension
+# base class there and derive Extension class below from it.
+#
+# That base class would look like:
+# class TrussExtension(ABC):
+#     @abstracemethod
+#     def model_override(self):
+#         pass
+
+#     @abstractmethod
+#     def model_args(self) -> dict:
+#         pass
+
+
+#     @abstractmethod
+#     def load(self) -> dict:
+#         pass
+class Extension:
+    """
+    trt_llm truss extension.
+
+    Provides model_args to supply to model class, which contain the trtllm
+    engine that corresponds to provided config.
+
+    This extension also provides a full replacement of the model class, which is
+    to be used if user doesn't supply it. This may be desired behavior in many
+    cases where users want to just go by config and don't want to do any pre or
+    post-processing.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self._config = kwargs["config"]
+        trt_llm_config = self._config.get("trt_llm")
+        config = TRTLLMConfiguration(**trt_llm_config)
+        if config.build.speculator:
+            self._model = SpecDecModel(*args, **kwargs)
+        else:
+            self._model = Model(*args, **kwargs)
+
+    def model_override(self):
+        """Return a model object.
+
+        This is used if model.py is omitted, which is allowed when using trt_llm.
+        """
+        return self._model
+
+    def model_args(self) -> dict:
+        """Return args to supply as input to Model class' __init__ method.
+
+        Model class can use this to invoke the trt_llm engine.
+
+        Returned engine is a typical Truss model class that provides a predict
+        function. The predict function is async and returns an async generator.
+        """
+        return {"engine": self._model}
+
+    def load(self):
+        self._model.load()