Spaces:
Running
Running
| import io | |
| from typing import Iterator, List, Optional | |
| import torch | |
| from torch import Tensor | |
| from torio.io._streaming_media_decoder import _get_afilter_desc, StreamingMediaDecoder as StreamReader | |
| from torio.io._streaming_media_encoder import CodecConfig, StreamingMediaEncoder as StreamWriter | |
| class _StreamingIOBuffer: | |
| """Streaming Bytes IO buffer. Data are dropped when read.""" | |
| def __init__(self): | |
| self._buffer: List(bytes) = [] | |
| def write(self, b: bytes): | |
| if b: | |
| self._buffer.append(b) | |
| return len(b) | |
| def pop(self, n): | |
| """Pop the oldest byte string. It does not necessary return the requested amount""" | |
| if not self._buffer: | |
| return b"" | |
| if len(self._buffer[0]) <= n: | |
| return self._buffer.pop(0) | |
| ret = self._buffer[0][:n] | |
| self._buffer[0] = self._buffer[0][n:] | |
| return ret | |
| def _get_sample_fmt(dtype: torch.dtype): | |
| types = { | |
| torch.uint8: "u8", | |
| torch.int16: "s16", | |
| torch.int32: "s32", | |
| torch.float32: "flt", | |
| torch.float64: "dbl", | |
| } | |
| if dtype not in types: | |
| raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}") | |
| return types[dtype] | |
| class _AudioStreamingEncoder: | |
| """Given a waveform, encode on-demand and return bytes""" | |
| def __init__( | |
| self, | |
| src: Tensor, | |
| sample_rate: int, | |
| effect: str, | |
| muxer: str, | |
| encoder: Optional[str], | |
| codec_config: Optional[CodecConfig], | |
| frames_per_chunk: int, | |
| ): | |
| self.src = src | |
| self.buffer = _StreamingIOBuffer() | |
| self.writer = StreamWriter(self.buffer, format=muxer) | |
| self.writer.add_audio_stream( | |
| num_channels=src.size(1), | |
| sample_rate=sample_rate, | |
| format=_get_sample_fmt(src.dtype), | |
| encoder=encoder, | |
| filter_desc=effect, | |
| codec_config=codec_config, | |
| ) | |
| self.writer.open() | |
| self.fpc = frames_per_chunk | |
| # index on the input tensor (along time-axis) | |
| # we use -1 to indicate that we finished iterating the tensor and | |
| # the writer is closed. | |
| self.i_iter = 0 | |
| def read(self, n): | |
| while not self.buffer._buffer and self.i_iter >= 0: | |
| self.writer.write_audio_chunk(0, self.src[self.i_iter : self.i_iter + self.fpc]) | |
| self.i_iter += self.fpc | |
| if self.i_iter >= self.src.size(0): | |
| self.writer.flush() | |
| self.writer.close() | |
| self.i_iter = -1 | |
| return self.buffer.pop(n) | |
| def _encode( | |
| src: Tensor, | |
| sample_rate: int, | |
| effect: str, | |
| muxer: str, | |
| encoder: Optional[str], | |
| codec_config: Optional[CodecConfig], | |
| ): | |
| buffer = io.BytesIO() | |
| writer = StreamWriter(buffer, format=muxer) | |
| writer.add_audio_stream( | |
| num_channels=src.size(1), | |
| sample_rate=sample_rate, | |
| format=_get_sample_fmt(src.dtype), | |
| encoder=encoder, | |
| filter_desc=effect, | |
| codec_config=codec_config, | |
| ) | |
| with writer.open(): | |
| writer.write_audio_chunk(0, src) | |
| buffer.seek(0) | |
| return buffer | |
| def _get_muxer(dtype: torch.dtype): | |
| # TODO: check if this works in Windows. | |
| types = { | |
| torch.uint8: "u8", | |
| torch.int16: "s16le", | |
| torch.int32: "s32le", | |
| torch.float32: "f32le", | |
| torch.float64: "f64le", | |
| } | |
| if dtype not in types: | |
| raise ValueError(f"Unsupported dtype is provided {dtype}. Supported dtypes are: {types.keys()}") | |
| return types[dtype] | |
| class AudioEffector: | |
| """Apply various filters and/or codecs to waveforms. | |
| .. versionadded:: 2.1 | |
| Args: | |
| effect (str or None, optional): Filter expressions or ``None`` to apply no filter. | |
| See https://ffmpeg.org/ffmpeg-filters.html#Audio-Filters for the | |
| details of filter syntax. | |
| format (str or None, optional): When provided, encode the audio into the | |
| corresponding format. Default: ``None``. | |
| encoder (str or None, optional): When provided, override the encoder used | |
| by the ``format``. Default: ``None``. | |
| codec_config (CodecConfig or None, optional): When provided, configure the encoding codec. | |
| Should be provided in conjunction with ``format`` option. | |
| pad_end (bool, optional): When enabled, and if the waveform becomes shorter after applying | |
| effects/codec, then pad the end with silence. | |
| Example - Basic usage | |
| To use ``AudioEffector``, first instantiate it with a set of | |
| ``effect`` and ``format``. | |
| >>> # instantiate the effector | |
| >>> effector = AudioEffector(effect=..., format=...) | |
| Then, use :py:meth:`~AudioEffector.apply` or :py:meth:`~AudioEffector.stream` | |
| method to apply them. | |
| >>> # Apply the effect to the whole waveform | |
| >>> applied = effector.apply(waveform, sample_rate) | |
| >>> # Apply the effect chunk-by-chunk | |
| >>> for chunk in effector.stream(waveform, sample_rate): | |
| >>> ... | |
| Example - Applying effects | |
| Please refer to | |
| https://ffmpeg.org/ffmpeg-filters.html#Filtergraph-description | |
| for the overview of filter description, and | |
| https://ffmpeg.org/ffmpeg-filters.html#toc-Audio-Filters | |
| for the list of available filters. | |
| Tempo - https://ffmpeg.org/ffmpeg-filters.html#atempo | |
| >>> AudioEffector(effect="atempo=1.5") | |
| Echo - https://ffmpeg.org/ffmpeg-filters.html#aecho | |
| >>> AudioEffector(effect="aecho=0.8:0.88:60:0.4") | |
| Flanger - https://ffmpeg.org/ffmpeg-filters.html#flanger | |
| >>> AudioEffector(effect="aflanger") | |
| Vibrato - https://ffmpeg.org/ffmpeg-filters.html#vibrato | |
| >>> AudioEffector(effect="vibrato") | |
| Tremolo - https://ffmpeg.org/ffmpeg-filters.html#tremolo | |
| >>> AudioEffector(effect="vibrato") | |
| You can also apply multiple effects at once. | |
| >>> AudioEffector(effect="") | |
| Example - Applying codec | |
| One can apply codec using ``format`` argument. ``format`` can be | |
| audio format or container format. If the container format supports | |
| multiple encoders, you can specify it with ``encoder`` argument. | |
| Wav format | |
| (no compression is applied but samples are converted to | |
| 16-bit signed integer) | |
| >>> AudioEffector(format="wav") | |
| Ogg format with default encoder | |
| >>> AudioEffector(format="ogg") | |
| Ogg format with vorbis | |
| >>> AudioEffector(format="ogg", encoder="vorbis") | |
| Ogg format with opus | |
| >>> AudioEffector(format="ogg", encoder="opus") | |
| Webm format with opus | |
| >>> AudioEffector(format="webm", encoder="opus") | |
| Example - Applying codec with configuration | |
| Reference: https://trac.ffmpeg.org/wiki/Encode/MP3 | |
| MP3 with default config | |
| >>> AudioEffector(format="mp3") | |
| MP3 with variable bitrate | |
| >>> AudioEffector(format="mp3", codec_config=CodecConfig(qscale=5)) | |
| MP3 with constant bitrate | |
| >>> AudioEffector(format="mp3", codec_config=CodecConfig(bit_rate=32_000)) | |
| """ | |
| def __init__( | |
| self, | |
| effect: Optional[str] = None, | |
| format: Optional[str] = None, | |
| *, | |
| encoder: Optional[str] = None, | |
| codec_config: Optional[CodecConfig] = None, | |
| pad_end: bool = True, | |
| ): | |
| if format is None: | |
| if encoder is not None or codec_config is not None: | |
| raise ValueError("`encoder` and/or `condec_config` opions are provided without `format` option.") | |
| self.effect = effect | |
| self.format = format | |
| self.encoder = encoder | |
| self.codec_config = codec_config | |
| self.pad_end = pad_end | |
| def _get_reader(self, waveform, sample_rate, output_sample_rate, frames_per_chunk=None): | |
| num_frames, num_channels = waveform.shape | |
| if self.format is not None: | |
| muxer = self.format | |
| encoder = self.encoder | |
| option = {} | |
| # Some formats are headerless, so need to provide these infomation. | |
| if self.format == "mulaw": | |
| option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"} | |
| else: # PCM | |
| muxer = _get_muxer(waveform.dtype) | |
| encoder = None | |
| option = {"sample_rate": f"{sample_rate}", "channels": f"{num_channels}"} | |
| if frames_per_chunk is None: | |
| src = _encode(waveform, sample_rate, self.effect, muxer, encoder, self.codec_config) | |
| else: | |
| src = _AudioStreamingEncoder( | |
| waveform, sample_rate, self.effect, muxer, encoder, self.codec_config, frames_per_chunk | |
| ) | |
| output_sr = sample_rate if output_sample_rate is None else output_sample_rate | |
| filter_desc = _get_afilter_desc(output_sr, _get_sample_fmt(waveform.dtype), num_channels) | |
| if self.pad_end: | |
| filter_desc = f"{filter_desc},apad=whole_len={num_frames}" | |
| reader = StreamReader(src, format=muxer, option=option) | |
| reader.add_audio_stream(frames_per_chunk or -1, -1, filter_desc=filter_desc) | |
| return reader | |
| def apply(self, waveform: Tensor, sample_rate: int, output_sample_rate: Optional[int] = None) -> Tensor: | |
| """Apply the effect and/or codecs to the whole tensor. | |
| Args: | |
| waveform (Tensor): The input waveform. Shape: ``(time, channel)`` | |
| sample_rate (int): Sample rate of the input waveform. | |
| output_sample_rate (int or None, optional): Output sample rate. | |
| If provided, override the output sample rate. | |
| Otherwise, the resulting tensor is resampled to have | |
| the same sample rate as the input. | |
| Default: ``None``. | |
| Returns: | |
| Tensor: | |
| Resulting Tensor. Shape: ``(time, channel)``. The number of frames | |
| could be different from that of the input. | |
| """ | |
| if waveform.ndim != 2: | |
| raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}") | |
| if waveform.numel() == 0: | |
| return waveform | |
| reader = self._get_reader(waveform, sample_rate, output_sample_rate) | |
| reader.process_all_packets() | |
| (applied,) = reader.pop_chunks() | |
| return Tensor(applied) | |
| def stream( | |
| self, waveform: Tensor, sample_rate: int, frames_per_chunk: int, output_sample_rate: Optional[int] = None | |
| ) -> Iterator[Tensor]: | |
| """Apply the effect and/or codecs to the given tensor chunk by chunk. | |
| Args: | |
| waveform (Tensor): The input waveform. Shape: ``(time, channel)`` | |
| sample_rate (int): Sample rate of the waveform. | |
| frames_per_chunk (int): The number of frames to return at a time. | |
| output_sample_rate (int or None, optional): Output sample rate. | |
| If provided, override the output sample rate. | |
| Otherwise, the resulting tensor is resampled to have | |
| the same sample rate as the input. | |
| Default: ``None``. | |
| Returns: | |
| Iterator[Tensor]: | |
| Series of processed chunks. Shape: ``(time, channel)``, where the | |
| the number of frames matches ``frames_per_chunk`` except the | |
| last chunk, which could be shorter. | |
| """ | |
| if waveform.ndim != 2: | |
| raise ValueError(f"Expected the input waveform to be 2D. Found: {waveform.ndim}") | |
| if waveform.numel() == 0: | |
| return waveform | |
| reader = self._get_reader(waveform, sample_rate, output_sample_rate, frames_per_chunk) | |
| for (applied,) in reader.stream(): | |
| yield Tensor(applied) | |