From 4925a6ad6c03d90f0620d6cf600ee3dec560bb30 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 15:34:32 +0100 Subject: [PATCH 1/6] Decode audio in the Blocks API, as RawAudioSamples PacketDecoder now follows the demuxer it was built from: RawFrames for a VideoDemuxer, RawAudioSamples for an AudioDemuxer. Decoding is the same avcodec_send_packet / avcodec_receive_frame pair either way, so this stays one block rather than two. RawAudioSamples carries the samples in the codec's own sample type (uint8 for u8, int16 for s16, float32 for flt...), always as a contiguous [num_channels, num_samples] tensor. That is a copy, not a view: planar formats put each channel in its own allocation and packed ones interleave them, so neither is that shape as it stands, and an audio frame is only a few kB. Normalizing to [-1, 1] is the converter's job, not the decoder's. Audio is CPU-only, so a non-CPU device raises rather than being silently ignored, and an unspecified one doesn't inherit a non-CPU default device. --- src/torchcodec/_core/AudioCommon.cpp | 38 ++++++ src/torchcodec/_core/AudioCommon.h | 23 ++++ src/torchcodec/_core/PacketDecoder.cpp | 120 +++++++++++++++-- src/torchcodec/_core/PacketDecoder.h | 20 ++- src/torchcodec/_core/StableABICompat.h | 1 + src/torchcodec/_core/_ffmpeg_op_names.py | 1 + src/torchcodec/_core/_ffmpeg_ops.py | 3 + src/torchcodec/_core/custom_ops.cpp | 45 +++++++ src/torchcodec/_core/sources.bzl | 1 + src/torchcodec/decoders/_blocks/__init__.py | 3 +- src/torchcodec/decoders/_blocks/_frame.py | 49 +++++++ .../decoders/_blocks/_packet_decoder.py | 69 ++++++++-- test/test_decoders.py | 125 ++++++++++++++++++ 13 files changed, 472 insertions(+), 26 deletions(-) create mode 100644 src/torchcodec/_core/AudioCommon.cpp create mode 100644 src/torchcodec/_core/AudioCommon.h diff --git a/src/torchcodec/_core/AudioCommon.cpp b/src/torchcodec/_core/AudioCommon.cpp new file mode 100644 index 000000000..64abd9b55 --- /dev/null +++ b/src/torchcodec/_core/AudioCommon.cpp @@ -0,0 +1,38 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#include "AudioCommon.h" + +namespace facebook::torchcodec { + +torch::headeronly::ScalarType sample_format_dtype( + AVSampleFormat sample_format) { + switch (av_get_packed_sample_fmt(sample_format)) { + case AV_SAMPLE_FMT_U8: + return kStableUInt8; + case AV_SAMPLE_FMT_S16: + return kStableInt16; + case AV_SAMPLE_FMT_S32: + return kStableInt32; + case AV_SAMPLE_FMT_S64: + return kStableInt64; + case AV_SAMPLE_FMT_FLT: + return kStableFloat32; + case AV_SAMPLE_FMT_DBL: + return kStableFloat64; + default: + break; + } + const char* name = av_get_sample_fmt_name(sample_format); + STD_TORCH_CHECK( + false, + "Unsupported sample format '", + name == nullptr ? "unknown" : name, + "'."); + return kStableUInt8; +} + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/AudioCommon.h b/src/torchcodec/_core/AudioCommon.h new file mode 100644 index 000000000..482bd86b4 --- /dev/null +++ b/src/torchcodec/_core/AudioCommon.h @@ -0,0 +1,23 @@ +// Copyright (c) Meta Platforms, Inc. and affiliates. +// All rights reserved. +// +// This source code is licensed under the BSD-style license found in the +// LICENSE file in the root directory of this source tree. + +#pragma once + +#include "FFMPEGCommon.h" +#include "StableABICompat.h" + +// Where FFmpeg's audio samples meet torch tensors. FFMPEGCommon deliberately +// knows nothing about tensors, and these helpers are shared by the decode, the +// conversion and the SingleStreamDecoder paths, so they live on their own. + +namespace facebook::torchcodec { + +// The dtype that holds `sample_format`'s samples exactly. Planar and packed +// variants of a format share a sample type, which is why this doesn't care +// which one it is given. +torch::headeronly::ScalarType sample_format_dtype(AVSampleFormat sample_format); + +} // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 1f64165f9..2ae3bd4a5 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -6,7 +6,10 @@ #include "PacketDecoder.h" +#include "AudioCommon.h" + #include +#include namespace facebook::torchcodec { @@ -57,7 +60,13 @@ const AVCodec* find_decoder( PacketDecoder::PacketDecoder( const Demuxer& demuxer, const StableDevice& device, - std::optional ffmpeg_thread_count) { + std::optional ffmpeg_thread_count) + : media_type_(demuxer.media_type()) { + bool is_audio = media_type_ == AVMEDIA_TYPE_AUDIO; + STD_TORCH_CHECK( + !is_audio || device.type() == kStableCPU, + "Audio can only be decoded on the CPU."); + device_interface_ = create_device_interface(device); STD_TORCH_CHECK( device_interface_ != nullptr, @@ -65,18 +74,35 @@ PacketDecoder::PacketDecoder( AVStream* stream = demuxer.active_stream(); time_base_ = stream->time_base; - is_mpeg_ps_ = - std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; - if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { - display_matrix_.emplace(); - std::copy( - matrix, matrix + display_matrix_->size(), display_matrix_->begin()); + + if (is_audio) { + // Audio codecs are hardcoded to a single FFmpeg thread, see + // https://github.com/pytorch/torchcodec/issues/1253. + ffmpeg_thread_count = 1; + } else { + is_mpeg_ps_ = + std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; + if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { + display_matrix_.emplace(); + std::copy( + matrix, matrix + display_matrix_->size(), display_matrix_->begin()); + } } + const AVCodec* av_codec = find_decoder(stream, device_interface_.get()); codec_context_ = create_and_open_codec_context( stream, av_codec, device_interface_.get(), ffmpeg_thread_count); device_interface_->initialize(codec_context_); + if (is_audio) { + // Nothing else to set up: unlike video, we hand out the samples in the + // codec's own format, so no conversion state is needed here. Note we + // deliberately do NOT set request_sample_fmt: what SingleStreamDecoder + // asks for (FLTP) is an optimization for its own conversion, and here it + // would hide what the codec natively produces. + return; + } + const AVPixFmtDescriptor* stream_desc = av_pix_fmt_desc_get(codec_context_->pix_fmt); int stream_bit_depth = stream_desc ? stream_desc->comp[0].depth : 8; @@ -133,10 +159,12 @@ int PacketDecoder::receive_frame(UniqueAVFrame& av_frame) { int status = device_interface_->receive_frame(av_frame); if (status == AVSUCCESS) { device_interface_->make_frame_standalone(av_frame); - // Attach a copy of the display matrix to the frame, so the ColorConverter - // can use it. - set_display_matrix_on_frame( - *av_frame, display_matrix_ ? display_matrix_->data() : nullptr); + if (media_type_ == AVMEDIA_TYPE_VIDEO) { + // Attach a copy of the display matrix to the frame, so the ColorConverter + // can use it. + set_display_matrix_on_frame( + *av_frame, display_matrix_ ? display_matrix_->data() : nullptr); + } } return status; } @@ -240,4 +268,74 @@ std::vector frame_planes( return planes; } +namespace { +// Scatters `num_channels`-interleaved samples into one contiguous row per +// channel. Templated on an integer of the right width rather than the actual +// sample type: we're only moving bytes around, so all that matters is size. +template +void deinterleave( + const uint8_t* src, + uint8_t* dst, + int num_channels, + int num_samples) { + const T* in = reinterpret_cast(src); + T* out = reinterpret_cast(dst); + for (int channel = 0; channel < num_channels; ++channel) { + T* row = out + static_cast(channel) * num_samples; + for (int sample = 0; sample < num_samples; ++sample) { + row[sample] = in[static_cast(sample) * num_channels + channel]; + } + } +} +} // namespace + +torch::stable::Tensor audio_samples(const AVFrame& av_frame) { + auto sample_format = static_cast(av_frame.format); + int num_channels = get_num_channels(av_frame); + int64_t num_samples = av_frame.nb_samples; + + torch::stable::Tensor samples = torch::stable::empty( + {num_channels, num_samples}, sample_format_dtype(sample_format)); + if (num_samples == 0) { + return samples; + } + + int bytes_per_sample = av_get_bytes_per_sample(sample_format); + auto* dst = static_cast(samples.mutable_data_ptr()); + int64_t bytes_per_channel = num_samples * bytes_per_sample; + + if (av_sample_fmt_is_planar(sample_format)) { + for (int channel = 0; channel < num_channels; ++channel) { + // extended_data rather than data: the latter only holds + // AV_NUM_DATA_POINTERS (8) pointers, and we support more channels. + std::memcpy( + dst + channel * bytes_per_channel, + av_frame.extended_data[channel], + bytes_per_channel); + } + return samples; + } + + const uint8_t* src = av_frame.extended_data[0]; + int num_samples_int = static_cast(num_samples); + switch (bytes_per_sample) { + case 1: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 2: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 4: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 8: + deinterleave(src, dst, num_channels, num_samples_int); + break; + default: + STD_TORCH_CHECK( + false, "Unexpected sample width: ", bytes_per_sample, " bytes."); + } + return samples; +} + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h index 6ad481588..b8e3ea95a 100644 --- a/src/torchcodec/_core/PacketDecoder.h +++ b/src/torchcodec/_core/PacketDecoder.h @@ -29,8 +29,9 @@ SharedAVCodecContext create_and_open_codec_context( DeviceInterface* device_interface, std::optional thread_count); -// Decode building block: turns compressed packets into decoded (YUV) frames. -// Configured from a Demuxer's active stream; stateful. Not thread-safe. +// Decode building block: turns compressed packets into decoded frames - (YUV) +// pictures for a video stream, samples in the codec's own format for an audio +// one. Configured from a Demuxer's active stream; stateful. Not thread-safe. class FORCE_PUBLIC_VISIBILITY PacketDecoder { public: explicit PacketDecoder( @@ -67,10 +68,15 @@ class FORCE_PUBLIC_VISIBILITY PacketDecoder { return time_base_; } + AVMediaType media_type() const { + return media_type_; + } + private: std::unique_ptr device_interface_; SharedAVCodecContext codec_context_; AVRational time_base_ = {}; + AVMediaType media_type_ = AVMEDIA_TYPE_VIDEO; // Stamped onto every frame we hand out, so downstream blocks can read the // rotation off the frame itself instead of knowing about the stream. Held by // value: we're only handed the Demuxer at construction and it may well be @@ -111,4 +117,14 @@ FORCE_PUBLIC_VISIBILITY std::vector frame_planes( const StableDevice& device, const torch::stable::Tensor& tensor_handle); +// A decoded audio frame's samples as a contiguous [num_channels, num_samples] +// tensor whose dtype is the frame's own sample type: uint8 for u8, int16 for +// s16, float32 for flt, and so on, planar or not. This is a copy rather than a +// view: planar formats put each channel in its own allocation and packed ones +// interleave them, so neither is a [C, N] tensor as it stands. An audio frame +// is a few kB, so normalizing here buys a uniform layout for the price of a +// memcpy - and it means a converter can treat the result as planar-of-dtype. +FORCE_PUBLIC_VISIBILITY torch::stable::Tensor audio_samples( + const AVFrame& av_frame); + } // namespace facebook::torchcodec diff --git a/src/torchcodec/_core/StableABICompat.h b/src/torchcodec/_core/StableABICompat.h index 514482f77..fe1a4ccc3 100644 --- a/src/torchcodec/_core/StableABICompat.h +++ b/src/torchcodec/_core/StableABICompat.h @@ -63,6 +63,7 @@ constexpr auto kStableXPU = torch::headeronly::DeviceType::XPU; // Scalar type constants constexpr auto kStableUInt8 = torch::headeronly::ScalarType::Byte; constexpr auto kStableUInt16 = torch::headeronly::ScalarType::UInt16; +constexpr auto kStableInt16 = torch::headeronly::ScalarType::Short; constexpr auto kStableInt32 = torch::headeronly::ScalarType::Int; constexpr auto kStableInt64 = torch::headeronly::ScalarType::Long; constexpr auto kStableFloat32 = torch::headeronly::ScalarType::Float; diff --git a/src/torchcodec/_core/_ffmpeg_op_names.py b/src/torchcodec/_core/_ffmpeg_op_names.py index 8523199c9..194a9803a 100644 --- a/src/torchcodec/_core/_ffmpeg_op_names.py +++ b/src/torchcodec/_core/_ffmpeg_op_names.py @@ -41,6 +41,7 @@ "_blocks_packet_decoder_send_eof", "_blocks_packet_decoder_reset", "_blocks_packet_decoder_receive_frame", + "_blocks_audio_packet_decoder_receive_frame", "_blocks_create_color_converter", "_blocks_convert_frame", "_blocks_frame_metadata", diff --git a/src/torchcodec/_core/_ffmpeg_ops.py b/src/torchcodec/_core/_ffmpeg_ops.py index d22a3b8f9..036ed50b7 100644 --- a/src/torchcodec/_core/_ffmpeg_ops.py +++ b/src/torchcodec/_core/_ffmpeg_ops.py @@ -121,6 +121,9 @@ def add_video_stream( _blocks_packet_decoder_receive_frame = ( torch.ops.torchcodec_ns._blocks_packet_decoder_receive_frame.default ) +_blocks_audio_packet_decoder_receive_frame = ( + torch.ops.torchcodec_ns._blocks_audio_packet_decoder_receive_frame.default +) _blocks_create_color_converter = ( torch.ops.torchcodec_ns._blocks_create_color_converter.default ) diff --git a/src/torchcodec/_core/custom_ops.cpp b/src/torchcodec/_core/custom_ops.cpp index 483f0cd9f..5c0ae4f49 100644 --- a/src/torchcodec/_core/custom_ops.cpp +++ b/src/torchcodec/_core/custom_ops.cpp @@ -91,6 +91,8 @@ STABLE_TORCH_LIBRARY_FRAGMENT(torchcodec_ns, m) { m.def("_blocks_packet_decoder_reset(Tensor(a!) decoder) -> ()"); m.def( "_blocks_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, Device, Tensor)"); + m.def( + "_blocks_audio_packet_decoder_receive_frame(Tensor(a!) decoder) -> (Tensor, int, float, float, int, str)"); m.def( "_blocks_create_color_converter(str device=\"cpu\", str output_dtype=\"uint8\") -> Tensor"); m.def( @@ -978,6 +980,46 @@ OpsReceiveFrameOutput _blocks_packet_decoder_receive_frame( storage); } +// (samples, status, pts_seconds, duration_seconds, sample_rate, +// sample_format). `samples` is [num_channels, num_samples] in the frame's own +// sample type; there is no frame handle because, unlike a video frame, nothing +// downstream needs the AVFrame itself. +using OpsReceiveAudioFrameOutput = std:: + tuple; + +OpsReceiveAudioFrameOutput _blocks_audio_packet_decoder_receive_frame( + torch::stable::Tensor& decoder) { + PacketDecoder* decoder_ptr = unwrap_tensor_to_pointer(decoder); + STD_TORCH_CHECK( + decoder_ptr->media_type() == AVMEDIA_TYPE_AUDIO, + "This PacketDecoder decodes video, not audio."); + + UniqueAVFrame av_frame(av_frame_alloc()); + STD_TORCH_CHECK(av_frame != nullptr, "Failed to allocate AVFrame"); + int status = decoder_ptr->receive_frame(av_frame); + if (status != AVSUCCESS) { + return std::make_tuple( + torch::stable::empty({int64_t(0)}, kStableUInt8), + static_cast(status), + 0.0, + 0.0, + static_cast(0), + std::string("")); + } + + AVRational time_base = decoder_ptr->time_base(); + const char* sample_format_name = + av_get_sample_fmt_name(static_cast(av_frame->format)); + return std::make_tuple( + audio_samples(*av_frame), + static_cast(0), + pts_to_seconds(get_pts_or_dts(*av_frame), time_base), + pts_to_seconds(get_duration(*av_frame), time_base), + static_cast(av_frame->sample_rate), + std::string( + sample_format_name == nullptr ? "unknown" : sample_format_name)); +} + torch::stable::Tensor _blocks_create_color_converter( std::string device, std::string output_dtype) { @@ -1613,6 +1655,9 @@ STABLE_TORCH_LIBRARY_IMPL(torchcodec_ns, CPU, m) { m.impl( "_blocks_packet_decoder_receive_frame", TORCH_BOX(&_blocks_packet_decoder_receive_frame)); + m.impl( + "_blocks_audio_packet_decoder_receive_frame", + TORCH_BOX(&_blocks_audio_packet_decoder_receive_frame)); m.impl("_blocks_convert_frame", TORCH_BOX(&_blocks_convert_frame)); m.impl("_blocks_frame_metadata", TORCH_BOX(&_blocks_frame_metadata)); m.impl("_blocks_frame_planes", TORCH_BOX(&_blocks_frame_planes)); diff --git a/src/torchcodec/_core/sources.bzl b/src/torchcodec/_core/sources.bzl index 6e2facadd..f1c9cb903 100644 --- a/src/torchcodec/_core/sources.bzl +++ b/src/torchcodec/_core/sources.bzl @@ -33,6 +33,7 @@ decoder_core_sources = [ "CpuDeviceInterface.cpp", "Demuxer.cpp", "PacketDecoder.cpp", + "AudioCommon.cpp", "ColorConverter.cpp", "SingleStreamDecoder.cpp", "Encoder.cpp", diff --git a/src/torchcodec/decoders/_blocks/__init__.py b/src/torchcodec/decoders/_blocks/__init__.py index 34fd5c45f..c7de94d67 100644 --- a/src/torchcodec/decoders/_blocks/__init__.py +++ b/src/torchcodec/decoders/_blocks/__init__.py @@ -21,7 +21,7 @@ from ._color_converter import ColorConverter from ._demuxer import AudioDemuxer, StreamIndex, VideoDemuxer -from ._frame import Packet, RawFrame +from ._frame import Packet, RawAudioSamples, RawFrame from ._packet_decoder import PacketDecoder __all__ = [ @@ -31,6 +31,7 @@ "ColorConverter", "Packet", "RawFrame", + "RawAudioSamples", "StreamIndex", ] diff --git a/src/torchcodec/decoders/_blocks/_frame.py b/src/torchcodec/decoders/_blocks/_frame.py index 5842dde1e..256ae6e1d 100644 --- a/src/torchcodec/decoders/_blocks/_frame.py +++ b/src/torchcodec/decoders/_blocks/_frame.py @@ -6,6 +6,7 @@ from __future__ import annotations +from dataclasses import dataclass from typing import NamedTuple import torch @@ -162,3 +163,51 @@ def planes(self) -> tuple[torch.Tensor, ...]: # views. self._planes = tuple(p for p in planes if p.numel() > 0) return self._planes + + +# TODO_API_BREAKDOWN DESIGN P1: API design - the class name, and whether +# sample_format is worth carrying now that the layout it describes has been +# normalized away. +@dataclass +class RawAudioSamples: + """One decoded audio frame's samples, as the decoder produced them. + + Produced by :class:`PacketDecoder` for an :class:`AudioDemuxer`'s stream, + consumed by :class:`AudioConverter`. This is the audio counterpart of + :class:`RawFrame`, and like it, nothing has been converted: the samples are + in the codec's own sample type. + + It is not a handle, unlike :class:`RawFrame`. Audio frames are a few kB, so + the samples are copied out of the ``AVFrame`` rather than viewed, which + also lets ``[num_channels, num_samples]`` be the layout for every format: + planar ones store each channel in its own allocation and packed ones + interleave them, so neither is that shape as it stands. + + Attributes: + data (torch.Tensor): ``[num_channels, num_samples]``, in the dtype that + holds the source's samples exactly: ``uint8`` for ``u8``, ``int16`` + for ``s16``, ``int32`` for ``s32``, ``float32`` for ``flt``, + ``float64`` for ``dbl``. Note the integer ones are *not* normalized + to ``[-1, 1]``; :class:`AudioConverter` is what does that. + sample_rate (int): The source's sample rate, in Hz. + sample_format (str): FFmpeg sample-format name, e.g. ``"s16p"`` or + ``"fltp"``. This is the format the samples were decoded in, kept + for provenance: the trailing ``p`` (planar) no longer describes + :attr:`data`, whose layout is always the same. + pts_seconds (float): Presentation timestamp of the first sample. + duration_seconds (float): How long these samples last. + """ + + data: torch.Tensor + sample_rate: int + sample_format: str + pts_seconds: float + duration_seconds: float + + @property + def num_channels(self) -> int: + return self.data.shape[0] + + @property + def num_samples(self) -> int: + return self.data.shape[1] diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index 82916c942..29f12ac7b 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -9,6 +9,7 @@ import torch from torchcodec._core.ops import ( + _blocks_audio_packet_decoder_receive_frame, _blocks_create_packet_decoder, _blocks_packet_decoder_receive_frame, _blocks_packet_decoder_reset, @@ -17,8 +18,8 @@ ) from .._decoder_utils import convert_device_to_str -from ._demuxer import VideoDemuxer -from ._frame import Packet, RawFrame +from ._demuxer import _BaseDemuxer +from ._frame import Packet, RawAudioSamples, RawFrame # TODO_API_BREAKDOWN DOC P1 revisit every single docstring / comments at some point. @@ -26,25 +27,45 @@ class PacketDecoder: """Decode building block: turns compressed :class:`Packet`\\ s into decoded - (YUV) :class:`RawFrame`\\ s. + frames. - Built from a :class:`VideoDemuxer` (for its codec parameters) and stateful (it - holds the codec's reference-frame buffer). Passive and *not* thread-safe: + What comes out follows the demuxer it was built from: (YUV) + :class:`RawFrame`\\ s for a :class:`VideoDemuxer`, :class:`RawAudioSamples` + for an :class:`AudioDemuxer`. Decoding is the same operation either way, + which is why this is one block rather than two. + + Built from a demuxer (for its codec parameters) and stateful: it holds the + codec's buffered state, which is reference frames for video and + overlap-add state for a lossy audio codec. Passive and *not* thread-safe: use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count is kept at 1 for now (not exposed); parallelism comes from composing blocks on your own threads. ``device`` accepts a string or a ``torch.device``. It defaults to ``None``, which means the current default device (see ``torch.set_default_device``). + Audio is decoded on the CPU only: passing a non-CPU ``device`` alongside an + :class:`AudioDemuxer` raises, and leaving it unspecified gives CPU rather + than picking up a non-CPU default. """ - def __init__(self, demuxer: VideoDemuxer, device: str | torch.device | None = None): + def __init__(self, demuxer: _BaseDemuxer, device: str | torch.device | None = None): + self._is_audio = demuxer._media_type == "audio" + if self._is_audio: + if device is not None and torch.device(device).type != "cpu": + raise ValueError( + f"Got device={device}, but audio can only be decoded on the " + "CPU. Leave device unspecified, or pass 'cpu'." + ) + device_str = "cpu" + else: + device_str = convert_device_to_str(device) + self._handle = _blocks_create_packet_decoder( - demuxer._handle, num_threads=1, device=convert_device_to_str(device) + demuxer._handle, num_threads=1, device=device_str ) self._drained = False - def _receive_ready_frames(self) -> list[RawFrame]: + def _receive_ready_video_frames(self) -> list[RawFrame]: frames = [] while True: handle, status, pts_seconds, duration_seconds, device, storage = ( @@ -63,9 +84,33 @@ def _receive_ready_frames(self) -> list[RawFrame]: ) return frames - def decode(self, packet: Packet) -> list[RawFrame]: - """Send one packet and return whatever frames are now ready (possibly - empty, e.g. while the codec buffers B-frames).""" + def _receive_ready_audio_samples(self) -> list[RawAudioSamples]: + samples = [] + while True: + data, status, pts_seconds, duration_seconds, sample_rate, sample_format = ( + _blocks_audio_packet_decoder_receive_frame(self._handle) + ) + if status != 0: # EAGAIN (need more packets) or EOF: nothing ready + break + samples.append( + RawAudioSamples( + data=data, + sample_rate=sample_rate, + sample_format=sample_format, + pts_seconds=pts_seconds, + duration_seconds=duration_seconds, + ) + ) + return samples + + def _receive_ready_frames(self) -> list[RawFrame] | list[RawAudioSamples]: + if self._is_audio: + return self._receive_ready_audio_samples() + return self._receive_ready_video_frames() + + def decode(self, packet: Packet) -> list[RawFrame] | list[RawAudioSamples]: + """Send one packet and return whatever is now ready (possibly empty, + e.g. while the codec buffers B-frames).""" if self._drained: raise RuntimeError( "This PacketDecoder has been drained, and a codec that has been " @@ -77,7 +122,7 @@ def decode(self, packet: Packet) -> list[RawFrame]: raise RuntimeError(f"Failed to send packet to decoder (status {status})") return self._receive_ready_frames() - def drain(self) -> list[RawFrame]: + def drain(self) -> list[RawFrame] | list[RawAudioSamples]: _blocks_packet_decoder_send_eof(self._handle) frames = self._receive_ready_frames() self._drained = True diff --git a/test/test_decoders.py b/test/test_decoders.py index d4605343b..abfc314ff 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -46,6 +46,7 @@ ColorConverter, Packet, PacketDecoder, + RawAudioSamples, RawFrame, VideoDemuxer, ) @@ -5057,6 +5058,130 @@ def test_audio_demuxer_seek(self): assert 0 < num_packets_after_seek < num_packets_from_start + # ===== Audio decoding: RawAudioSamples ===== + + @staticmethod + def _decode_audio(asset, stream_index=None): + demuxer = AudioDemuxer(asset.path, stream_index=stream_index) + decoder = PacketDecoder(demuxer) + chunks = [] + for packet in demuxer: + chunks += decoder.decode(packet) + chunks += decoder.drain() + return chunks + + @pytest.mark.parametrize( + "asset, sample_format, dtype", + ( + (SINE_MONO_U8, "u8", torch.uint8), + (SINE_MONO_S16, "s16", torch.int16), + (SINE_MONO_S32, "s32", torch.int32), + # FFmpeg has no 24-bit sample format, so a 24-bit source is s32. + (SINE_MONO_S24, "s32", torch.int32), + (SINE_MONO_F32, "flt", torch.float32), + (SINE_MONO_F64, "dbl", torch.float64), + (SINE_STEREO_MP2_MPEG_PS, "s16p", torch.int16), + (SINE_16_CHANNEL_S16, "s16", torch.int16), + (NASA_AUDIO_MP3, "fltp", torch.float32), + ), + ) + def test_audio_raw_samples_dtype_and_shape(self, asset, sample_format, dtype): + # The decoder hands out the codec's own sample type, always as + # [num_channels, num_samples]. The packed formats above (no trailing + # 'p') are the ones exercising the de-interleaving path. + chunks = self._decode_audio(asset) + assert len(chunks) > 0 + + for chunk in chunks: + assert isinstance(chunk, RawAudioSamples) + assert chunk.sample_format == sample_format + assert chunk.data.dtype == dtype + assert chunk.data.ndim == 2 + assert chunk.data.is_contiguous() + assert chunk.num_channels == asset.num_channels + assert chunk.sample_rate == asset.sample_rate + assert chunk.duration_seconds >= 0 + + @pytest.mark.parametrize( + "asset", + ( + SINE_MONO_U8, + SINE_MONO_S16, + SINE_MONO_S32, + SINE_MONO_F32, + SINE_MONO_F64, + SINE_STEREO_MP2_MPEG_PS, + SINE_16_CHANNEL_S16, + NASA_AUDIO_MP3, + NASA_AUDIO, + ), + ) + def test_audio_raw_samples_match_audio_decoder(self, asset): + # We hand out the true source samples: normalizing them the way FFmpeg + # does reproduces AudioDecoder's output bit for bit. This is also what + # pins the de-interleaving, most visibly on the 16-channel asset. + raw = torch.cat([chunk.data for chunk in self._decode_audio(asset)], dim=1) + + if raw.dtype == torch.uint8: + got = (raw.to(torch.float32) - 128) / 128 + elif raw.dtype in (torch.int16, torch.int32): + got = raw.to(torch.float32) / -float(torch.iinfo(raw.dtype).min) + else: + got = raw.to(torch.float32) + + expected = AudioDecoder(asset.path).get_all_samples().data + torch.testing.assert_close(got, expected, atol=0, rtol=0) + + def test_audio_raw_samples_pts(self): + chunks = self._decode_audio(SINE_MONO_S16) + pts = [chunk.pts_seconds for chunk in chunks] + assert pts == sorted(pts) + assert pts[0] == pytest.approx(0, abs=1e-6) + + def test_audio_decoder_output_type_follows_the_demuxer(self): + # Same PacketDecoder class, different output type: that's the whole + # reason it isn't split in two. + for demuxer_class, expected_type in ( + (AudioDemuxer, RawAudioSamples), + (VideoDemuxer, RawFrame), + ): + demuxer = demuxer_class(NASA_VIDEO.path) + decoder = PacketDecoder(demuxer) + # A codec needs more than one packet before it outputs anything. + decoded = [] + while not decoded: + decoded = decoder.decode(demuxer.next_packet()) + assert isinstance(decoded[0], expected_type) + + @pytest.mark.parametrize("device", ("cuda", "cuda:0", torch.device("cuda"))) + def test_audio_decoder_non_cpu_device_raises(self, device): + # Not gated on CUDA being available: we reject the request before ever + # touching a device. + demuxer = AudioDemuxer(NASA_AUDIO_MP3.path) + with pytest.raises(ValueError, match="audio can only be decoded on the CPU"): + PacketDecoder(demuxer, device=device) + + @needs_cuda + def test_audio_decoder_ignores_non_cpu_default_device(self): + # An unspecified device means CPU for audio, not whatever + # torch.get_default_device() happens to be. Compare against a plain + # run, so this also catches us silently decoding something else. + expected = self._decode_audio(NASA_AUDIO_MP3)[0].data + + with torch.device("cuda"): + got = self._decode_audio(NASA_AUDIO_MP3)[0].data + assert got.device.type == "cpu" + torch.testing.assert_close(got, expected, atol=0, rtol=0) + + original_device = torch.get_default_device() + try: + torch.set_default_device("cuda") + got = self._decode_audio(NASA_AUDIO_MP3)[0].data + finally: + torch.set_default_device(original_device) + assert got.device.type == "cpu" + torch.testing.assert_close(got, expected, atol=0, rtol=0) + @pytest.mark.parametrize("stream_index", (-1, 6, 1000)) def test_invalid_stream_index_raises(self, stream_index): with pytest.raises(RuntimeError, match="is not a valid stream"): From 191fc572485d9ee1c0246f3b9eb16c1cf6d1f1c1 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 17:13:21 +0100 Subject: [PATCH 2/6] Split PacketDecoder into Video/AudioPacketDecoder, and decode audio The blocks decoder now handles audio, producing RawAudioSamples. Decoding is the same avcodec_send_packet / avcodec_receive_frame pair for both media, so the C++ PacketDecoder stays one class; the split is Python-only, so that each decoder has an exact output type and only the arguments that apply to it. AudioPacketDecoder has no device parameter at all: audio is CPU-only, and that shouldn't change with torch.set_default_device. RawAudioSamples carries the samples in the codec's own sample type (uint8 for u8, int16 for s16, float32 for flt...), always as a contiguous [num_channels, num_samples] tensor. That is a copy, not a view: planar formats put each channel in its own allocation and packed ones interleave them, so neither is that shape as it stands, and an audio frame is only a few kB. Normalizing to [-1, 1] is the converter's job, not the decoder's. is_mpeg_ps_ now applies to audio too. Seeking an MPEG program stream lands on a container-level byte offset, so the parser resumes mid-frame and rebuilds invalid packets until it resyncs - a property of the container, not of the codec. Gating it to video made seeking sine_stereo_mp2.mpg's audio stream fail outright with AVERROR_INVALIDDATA; there's now a test for it. --- benchmarks/bench_blocks.py | 10 +- examples/decoding/blocks.py | 26 ++-- src/torchcodec/_core/PacketDecoder.cpp | 17 ++- src/torchcodec/decoders/_blocks/__init__.py | 14 +- src/torchcodec/decoders/_blocks/_demuxer.py | 12 +- src/torchcodec/decoders/_blocks/_frame.py | 6 +- .../decoders/_blocks/_packet_decoder.py | 143 ++++++++++-------- test/test_decoders.py | 77 ++++++---- 8 files changed, 173 insertions(+), 132 deletions(-) diff --git a/benchmarks/bench_blocks.py b/benchmarks/bench_blocks.py index a11a202b2..01ea29596 100644 --- a/benchmarks/bench_blocks.py +++ b/benchmarks/bench_blocks.py @@ -15,7 +15,7 @@ import torch from torchcodec.decoders import VideoDecoder -from torchcodec.decoders._blocks import ColorConverter, PacketDecoder, VideoDemuxer +from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, VideoPacketDecoder # Kept minimal on purpose; the filename is derived from exactly these. _DURATION_S = 10 @@ -113,7 +113,7 @@ def _consume(frames): def _decode_sequential(path, device="cpu"): demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) _consume(_convert(converter, _decode(decoder, _demux(demuxer)))) @@ -121,7 +121,7 @@ def _decode_sequential(path, device="cpu"): def _decode_prefetch_frames(path, device="cpu"): # [demux + decode] on one thread || [color-convert] on another. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) frames = prefetch(_decode(decoder, _demux(demuxer))) _consume(_convert(converter, frames)) @@ -130,7 +130,7 @@ def _decode_prefetch_frames(path, device="cpu"): def _decode_prefetch_packets(path, device="cpu"): # [demux] on one thread || [decode + color-convert] on another. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) packets = prefetch(_demux(demuxer)) _consume(_convert(converter, _decode(decoder, packets))) @@ -139,7 +139,7 @@ def _decode_prefetch_packets(path, device="cpu"): def _decode_prefetch_packets_and_frames(path, device="cpu"): # [demux] || [decode] || [color-convert], each on its own thread. demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) packets = prefetch(_demux(demuxer)) frames = prefetch(_decode(decoder, packets)) diff --git a/examples/decoding/blocks.py b/examples/decoding/blocks.py index f29701444..5c19a382c 100644 --- a/examples/decoding/blocks.py +++ b/examples/decoding/blocks.py @@ -21,7 +21,7 @@ .. code-block:: - VideoDemuxer -> PacketDecoder -> ColorConverter + VideoDemuxer -> VideoPacketDecoder -> ColorConverter Packet RawFrame RGB Frame The blocks are passive: they never create threads, and they release the GIL. @@ -63,14 +63,14 @@ # it can output a frame, and it buffers a few frames that ``drain()`` returns # at the end. # -# ``PacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``: +# ``VideoPacketDecoder`` and ``ColorConverter`` both accept ``device="cuda"``: # decoding then runs on NVDEC and the color conversion on the GPU, and the # frames never leave the device. Demuxing always happens on the CPU. Left # unspecified, ``device`` is the current default device. -from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, PacketDecoder +from torchcodec.decoders._blocks import ColorConverter, VideoDemuxer, VideoPacketDecoder demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) frames = [] @@ -134,7 +134,7 @@ def drain(): def sequential(): # demux -> decode -> color-convert, all on the calling thread. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) return color_convert(color_converter, decode(packet_decoder, demux(demuxer))) @@ -142,7 +142,7 @@ def sequential(): def convert_on_own_thread(): # [demux + decode] on one thread || [color-convert] on another. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) raw_frames = prefetch(decode(packet_decoder, demux(demuxer))) return color_convert(color_converter, raw_frames) @@ -153,7 +153,7 @@ def demux_on_own_thread(): # natural split on CUDA: demuxing is CPU and I/O work, while decoding and # color conversion both happen on the GPU, so they belong together. demuxer = VideoDemuxer(video_path) - packet_decoder = PacketDecoder(demuxer, device=device) + packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) packets = prefetch(demux(demuxer)) return color_convert(color_converter, decode(packet_decoder, packets)) @@ -179,9 +179,9 @@ def demux_on_own_thread(): # drop them until you reach the timestamp you asked for. # # The seek also invalidates the frames the decoder is holding on to, so the -# ``PacketDecoder`` must be ``reset()``. +# ``VideoPacketDecoder`` must be ``reset()``. demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) seconds = 2.5 @@ -212,7 +212,7 @@ def demux_on_own_thread(): # costs one pass over the file, and it leaves the demuxer back at the start. demuxer = VideoDemuxer(video_path) index = demuxer.scan() -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) print(f"{len(index)} frames at {index.average_fps} fps, " @@ -281,7 +281,7 @@ def demux_on_own_thread(): # Color conversion is optional. A ``RawFrame`` can hand out the decoder's own # planes as tensor views, with no copy and no conversion. demuxer = VideoDemuxer(video_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) raw_frame = next(decode(packet_decoder, demux(demuxer))) Y, U, V = raw_frame.planes @@ -353,7 +353,7 @@ def upsample(plane): ) hdr_demuxer = VideoDemuxer(hdr_video_path) -hdr_packet_decoder = PacketDecoder(hdr_demuxer, device=device) +hdr_packet_decoder = VideoPacketDecoder(hdr_demuxer, device=device) hdr_raw = next(decode(hdr_packet_decoder, demux(hdr_demuxer))) hdr_Y = hdr_raw.planes[0] @@ -413,7 +413,7 @@ def start_live_stream(): # The blocks just stream it, and we stop whenever we want: ffmpeg = start_live_stream() demuxer = VideoDemuxer(fifo_path) -packet_decoder = PacketDecoder(demuxer, device=device) +packet_decoder = VideoPacketDecoder(demuxer, device=device) color_converter = ColorConverter(device=device) frames = [] diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 2ae3bd4a5..652c15937 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -75,18 +75,19 @@ PacketDecoder::PacketDecoder( AVStream* stream = demuxer.active_stream(); time_base_ = stream->time_base; + // This is a property of the container, not of the codec: an audio stream in + // an MPEG program stream resyncs after a seek exactly like a video one does. + is_mpeg_ps_ = + std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; + if (is_audio) { // Audio codecs are hardcoded to a single FFmpeg thread, see // https://github.com/pytorch/torchcodec/issues/1253. ffmpeg_thread_count = 1; - } else { - is_mpeg_ps_ = - std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; - if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { - display_matrix_.emplace(); - std::copy( - matrix, matrix + display_matrix_->size(), display_matrix_->begin()); - } + } else if (const int32_t* matrix = get_display_matrix_from_stream(stream)) { + display_matrix_.emplace(); + std::copy( + matrix, matrix + display_matrix_->size(), display_matrix_->begin()); } const AVCodec* av_codec = find_decoder(stream, device_interface_.get()); diff --git a/src/torchcodec/decoders/_blocks/__init__.py b/src/torchcodec/decoders/_blocks/__init__.py index c7de94d67..3126a9dc7 100644 --- a/src/torchcodec/decoders/_blocks/__init__.py +++ b/src/torchcodec/decoders/_blocks/__init__.py @@ -7,13 +7,14 @@ """Private, experimental building-block decode API. Exposes the three decode stages -- :class:`VideoDemuxer`, -:class:`PacketDecoder`, :class:`ColorConverter` -- as passive, composable, +:class:`VideoPacketDecoder`, :class:`ColorConverter` -- as passive, composable, GIL-releasing units, so a caller can build its own (threaded) decode pipeline and tune how the stages overlap. The blocks do no threading themselves. -:class:`AudioDemuxer` is the audio counterpart of :class:`VideoDemuxer`; -:class:`PacketDecoder` is shared, since decoding is the same operation either -way. +Audio has the same three stages: :class:`AudioDemuxer` and +:class:`AudioPacketDecoder`. The two decoders are a single class in C++, since +decoding is the same operation either way; they are separate here because what +they hand out, and how they are configured, isn't. This is experimental and private; the API may change. See API_breakdown_claude_plan.md for the design and rationale. @@ -22,12 +23,13 @@ from ._color_converter import ColorConverter from ._demuxer import AudioDemuxer, StreamIndex, VideoDemuxer from ._frame import Packet, RawAudioSamples, RawFrame -from ._packet_decoder import PacketDecoder +from ._packet_decoder import AudioPacketDecoder, VideoPacketDecoder __all__ = [ "VideoDemuxer", "AudioDemuxer", - "PacketDecoder", + "VideoPacketDecoder", + "AudioPacketDecoder", "ColorConverter", "Packet", "RawFrame", diff --git a/src/torchcodec/decoders/_blocks/_demuxer.py b/src/torchcodec/decoders/_blocks/_demuxer.py index dff25cfcb..b40323ed7 100644 --- a/src/torchcodec/decoders/_blocks/_demuxer.py +++ b/src/torchcodec/decoders/_blocks/_demuxer.py @@ -182,7 +182,7 @@ def seek(self, seconds: float) -> None: """Move the demuxer to ``seconds``. A seek invalidates whatever the decoder is holding on to, so the - :class:`PacketDecoder` must be ``reset()`` afterwards. + packet decoder must be ``reset()`` afterwards. Where you land, and what comes out first, depends on the medium. For video, a decoder can only start on a :term:`keyframe`, so this lands on @@ -212,8 +212,8 @@ class VideoDemuxer(_BaseDemuxer): of the file, or from wherever :meth:`seek` left it. A :class:`VideoDemuxer` also carries the stream configuration used to build a - :class:`PacketDecoder`, so that is constructed from a demuxer and no extra - container is opened. + :class:`VideoPacketDecoder`, so that is constructed from a demuxer and no + extra container is opened. Args: source (str, ``Pathlib.path``, bytes, ``torch.Tensor`` or file-like object): The source of the video: @@ -246,7 +246,7 @@ def scan(self) -> StreamIndex: The index always covers the entire stream, wherever the demuxer currently is, and the demuxer is left back at the start - so a - :class:`PacketDecoder` built from it must be ``reset()``, as after a + :class:`VideoPacketDecoder` built from it must be ``reset()``, as after a ``seek()``. Nothing is cached: calling this twice scans twice. """ pts, duration, is_key_frame, time_base_num, time_base_den = ( @@ -270,8 +270,8 @@ class AudioDemuxer(_BaseDemuxer): of the file, or from wherever :meth:`seek` left it. An :class:`AudioDemuxer` also carries the stream configuration used to build - a :class:`PacketDecoder`, so that is constructed from a demuxer and no extra - container is opened. + an :class:`AudioPacketDecoder`, so that is constructed from a demuxer and + no extra container is opened. Unlike :class:`VideoDemuxer` there is no ``scan()``: a :class:`StreamIndex` describes keyframes and frame indices, and audio has neither. diff --git a/src/torchcodec/decoders/_blocks/_frame.py b/src/torchcodec/decoders/_blocks/_frame.py index 256ae6e1d..29b3d49cb 100644 --- a/src/torchcodec/decoders/_blocks/_frame.py +++ b/src/torchcodec/decoders/_blocks/_frame.py @@ -29,7 +29,7 @@ class _Metadata(NamedTuple): class Packet: """Opaque, thread-movable handle to a demuxed (compressed) packet. - Produced by :class:`VideoDemuxer`, consumed by :class:`PacketDecoder`. It wraps a raw + Produced by :class:`VideoDemuxer`, consumed by :class:`VideoPacketDecoder`. It wraps a raw pointer, so it is only valid within the process that created it (it cannot cross a process boundary). """ @@ -44,7 +44,7 @@ class RawFrame: """A decoded (YUV) frame, as the decoder produced it: an opaque, thread-movable handle to the frame plus everything describing it. - Produced by :class:`PacketDecoder`, consumed by :class:`ColorConverter`. The + Produced by :class:`VideoPacketDecoder`, consumed by :class:`ColorConverter`. The handle wraps a raw pointer and is process-local. ``pts_seconds`` and ``duration_seconds`` are stamped by the decoder (which knows the stream time base) and carried here so the :class:`ColorConverter` need not be bound to @@ -172,7 +172,7 @@ def planes(self) -> tuple[torch.Tensor, ...]: class RawAudioSamples: """One decoded audio frame's samples, as the decoder produced them. - Produced by :class:`PacketDecoder` for an :class:`AudioDemuxer`'s stream, + Produced by :class:`VideoPacketDecoder` for an :class:`AudioDemuxer`'s stream, consumed by :class:`AudioConverter`. This is the audio counterpart of :class:`RawFrame`, and like it, nothing has been converted: the samples are in the codec's own sample type. diff --git a/src/torchcodec/decoders/_blocks/_packet_decoder.py b/src/torchcodec/decoders/_blocks/_packet_decoder.py index 29f12ac7b..8b7457b86 100644 --- a/src/torchcodec/decoders/_blocks/_packet_decoder.py +++ b/src/torchcodec/decoders/_blocks/_packet_decoder.py @@ -6,6 +6,8 @@ from __future__ import annotations +from typing import Generic, TypeVar + import torch from torchcodec._core.ops import ( @@ -18,54 +20,80 @@ ) from .._decoder_utils import convert_device_to_str -from ._demuxer import _BaseDemuxer +from ._demuxer import AudioDemuxer, VideoDemuxer from ._frame import Packet, RawAudioSamples, RawFrame # TODO_API_BREAKDOWN DOC P1 revisit every single docstring / comments at some point. +_Decoded = TypeVar("_Decoded", RawFrame, RawAudioSamples) -class PacketDecoder: - """Decode building block: turns compressed :class:`Packet`\\ s into decoded - frames. - - What comes out follows the demuxer it was built from: (YUV) - :class:`RawFrame`\\ s for a :class:`VideoDemuxer`, :class:`RawAudioSamples` - for an :class:`AudioDemuxer`. Decoding is the same operation either way, - which is why this is one block rather than two. - Built from a demuxer (for its codec parameters) and stateful: it holds the - codec's buffered state, which is reference frames for video and - overlap-add state for a lossy audio codec. Passive and *not* thread-safe: - use one ``PacketDecoder`` per thread. FFmpeg's internal codec thread count - is kept at 1 for now (not exposed); parallelism comes from composing blocks - on your own threads. +class _BasePacketDecoder(Generic[_Decoded]): + """Shared machinery for :class:`VideoPacketDecoder` and + :class:`AudioPacketDecoder`. - ``device`` accepts a string or a ``torch.device``. It defaults to ``None``, - which means the current default device (see ``torch.set_default_device``). - Audio is decoded on the CPU only: passing a non-CPU ``device`` alongside an - :class:`AudioDemuxer` raises, and leaving it unspecified gives CPU rather - than picking up a non-CPU default. + Decoding is the same ``avcodec_send_packet`` / ``avcodec_receive_frame`` + pair for both, so the C++ side is a single class; what differs is only what + a decoded frame is turned into, which is :meth:`_receive_ready_frames`. """ - def __init__(self, demuxer: _BaseDemuxer, device: str | torch.device | None = None): - self._is_audio = demuxer._media_type == "audio" - if self._is_audio: - if device is not None and torch.device(device).type != "cpu": - raise ValueError( - f"Got device={device}, but audio can only be decoded on the " - "CPU. Leave device unspecified, or pass 'cpu'." - ) - device_str = "cpu" - else: - device_str = convert_device_to_str(device) - + def __init__(self, demuxer, device_str: str): self._handle = _blocks_create_packet_decoder( demuxer._handle, num_threads=1, device=device_str ) self._drained = False - def _receive_ready_video_frames(self) -> list[RawFrame]: + def _receive_ready_frames(self) -> list[_Decoded]: + raise NotImplementedError + + def decode(self, packet: Packet) -> list[_Decoded]: + """Send one packet and return whatever is now ready (possibly empty, + e.g. while the codec buffers B-frames).""" + if self._drained: + raise RuntimeError( + "This decoder has been drained, and a codec that has been told " + "the stream ended ignores any further packet. Create a new " + "decoder to decode another stream." + ) + status = _blocks_packet_decoder_send_packet(self._handle, packet._handle) + if status < 0: + raise RuntimeError(f"Failed to send packet to decoder (status {status})") + return self._receive_ready_frames() + + def drain(self) -> list[_Decoded]: + """Tell the codec the stream ended, and return the frames it was still + holding on to.""" + _blocks_packet_decoder_send_eof(self._handle) + frames = self._receive_ready_frames() + self._drained = True + return frames + + def reset(self) -> None: + """Drop the codec's buffered state and start over. Needed after the + demuxer seeked, and after ``drain()``.""" + _blocks_packet_decoder_reset(self._handle) + self._drained = False + + +class VideoPacketDecoder(_BasePacketDecoder[RawFrame]): + """Decode building block: turns compressed :class:`Packet`\\ s into decoded + (YUV) :class:`RawFrame`\\ s. + + Built from a :class:`VideoDemuxer` (for its codec parameters) and stateful + (it holds the codec's reference-frame buffer). Passive and *not* + thread-safe: use one ``VideoPacketDecoder`` per thread. FFmpeg's internal + codec thread count is kept at 1 for now (not exposed); parallelism comes + from composing blocks on your own threads. + + ``device`` accepts a string or a ``torch.device``. It defaults to ``None``, + which means the current default device (see ``torch.set_default_device``). + """ + + def __init__(self, demuxer: VideoDemuxer, device: str | torch.device | None = None): + super().__init__(demuxer, convert_device_to_str(device)) + + def _receive_ready_frames(self) -> list[RawFrame]: frames = [] while True: handle, status, pts_seconds, duration_seconds, device, storage = ( @@ -84,7 +112,25 @@ def _receive_ready_video_frames(self) -> list[RawFrame]: ) return frames - def _receive_ready_audio_samples(self) -> list[RawAudioSamples]: + +class AudioPacketDecoder(_BasePacketDecoder[RawAudioSamples]): + """Decode building block: turns compressed :class:`Packet`\\ s into + :class:`RawAudioSamples`. + + Built from an :class:`AudioDemuxer` (for its codec parameters) and stateful: + a lossy codec's overlap-add state means the frames decoded right after a + seek are subtly wrong until it re-primes, so ``reset()`` is necessary but + not by itself sufficient - see :meth:`AudioDemuxer.seek`. Passive and *not* + thread-safe: use one ``AudioPacketDecoder`` per thread. + + There is no ``device`` parameter: audio is always decoded on the CPU, and + that doesn't change with ``torch.set_default_device``. + """ + + def __init__(self, demuxer: AudioDemuxer): + super().__init__(demuxer, "cpu") + + def _receive_ready_frames(self) -> list[RawAudioSamples]: samples = [] while True: data, status, pts_seconds, duration_seconds, sample_rate, sample_format = ( @@ -102,32 +148,3 @@ def _receive_ready_audio_samples(self) -> list[RawAudioSamples]: ) ) return samples - - def _receive_ready_frames(self) -> list[RawFrame] | list[RawAudioSamples]: - if self._is_audio: - return self._receive_ready_audio_samples() - return self._receive_ready_video_frames() - - def decode(self, packet: Packet) -> list[RawFrame] | list[RawAudioSamples]: - """Send one packet and return whatever is now ready (possibly empty, - e.g. while the codec buffers B-frames).""" - if self._drained: - raise RuntimeError( - "This PacketDecoder has been drained, and a codec that has been " - "told the stream ended ignores any further packet. Create a new " - "PacketDecoder to decode another stream." - ) - status = _blocks_packet_decoder_send_packet(self._handle, packet._handle) - if status < 0: - raise RuntimeError(f"Failed to send packet to decoder (status {status})") - return self._receive_ready_frames() - - def drain(self) -> list[RawFrame] | list[RawAudioSamples]: - _blocks_packet_decoder_send_eof(self._handle) - frames = self._receive_ready_frames() - self._drained = True - return frames - - def reset(self) -> None: - _blocks_packet_decoder_reset(self._handle) - self._drained = False diff --git a/test/test_decoders.py b/test/test_decoders.py index abfc314ff..d1606a106 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -43,12 +43,13 @@ ) from torchcodec.decoders._blocks import ( AudioDemuxer, + AudioPacketDecoder, ColorConverter, Packet, - PacketDecoder, RawAudioSamples, RawFrame, VideoDemuxer, + VideoPacketDecoder, ) from torchcodec.decoders._decoder_utils import _get_cuda_backend from torchcodec.decoders._image_decoders import _source_to_tensor @@ -3636,7 +3637,7 @@ class TestBlocks: @pytest.mark.parametrize("device", _block_devices()) def test_block_output_types(self, device): - # VideoDemuxer yields Packets, PacketDecoder yields RawFrames, and + # VideoDemuxer yields Packets, VideoPacketDecoder yields RawFrames, and # ColorConverter yields Frames with the expected shape/dtype. demuxer, decoder, converter = self._make_blocks(NASA_VIDEO.path, device) @@ -3711,7 +3712,7 @@ def drain(): @staticmethod def _make_blocks(path, device): demuxer = VideoDemuxer(path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) return demuxer, decoder, converter @@ -3762,7 +3763,7 @@ def _to_frame_batch(self, frames): def _assert_matches_video_decoder(got, ref, video): # We typically want exact equality but cannot achieve it on CUDA for HDR # videos that are downscaled to uint8: the VideoDecoder will ask NVDEC - # to output an 8bit surface while the PacketDecoder will decode on the + # to output an 8bit surface while the VideoPacketDecoder will decode on the # 16bit surface (by contract), so there are minor differences. if got.is_cuda and got.dtype == torch.uint8 and video in _HDR_VIDEOS: assert_tensor_close_on_at_least(got, ref, percentage=99, atol=2) @@ -3968,14 +3969,14 @@ def _first_frame(self, path, device): @pytest.mark.parametrize("device_str", _block_devices()) def test_device_none_default_device(self, device_str): - # PacketDecoder and ColorConverter default to device=None, which should + # VideoPacketDecoder and ColorConverter default to device=None, which should # respect both the torch.device() context manager and # torch.set_default_device(). def assert_first_frame_is_on_default_device(): # Note the absence of any device parameter. demuxer = VideoDemuxer(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) converter = ColorConverter() decoded = next(self._decode(decoder, self._demux(demuxer))) assert decoded.planes[0].device.type == device_str @@ -4079,7 +4080,7 @@ def test_no_rotation(self, device): @pytest.mark.needs_cuda @pytest.mark.parametrize("record_stream", (True, False)) def test_storage_record_stream(self, record_stream): - # Using PacketDecoder on one stream and consuming the frames on a + # Using VideoPacketDecoder on one stream and consuming the frames on a # different stream requires the user to call record_stream() on the # frame storage. # Without the record_stream() call the decoder's next frame may be @@ -4282,7 +4283,7 @@ def test_neutral_chroma_is_grayscale(self, video, device): ), ) def test_cpu_fallback_is_on_cuda(self, video, expected_pix_fmt): - # A CUDA PacketDecoder hands out CUDA frames even for the streams it has + # A CUDA VideoPacketDecoder hands out CUDA frames even for the streams it has # to decode on the CPU, and they're in an NVDEC surface format like any # other CUDA frame. assert VideoDecoder(video.path, device="cuda").cpu_fallback @@ -4479,7 +4480,7 @@ def test_seek_without_reset_yields_stale_frames(self, device): seconds = video_decoder.get_frame_at(keyframe_index).pts_seconds demuxer = VideoDemuxer(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) num_decoded = 0 for packet in demuxer: # decode a bit, so frames pile up in the codec num_decoded += len(decoder.decode(packet)) @@ -4744,7 +4745,7 @@ def test_seek_on_non_seekable_source_raises(self, tmp_path): ) try: demuxer = VideoDemuxer(fifo_path) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) num_decoded = 0 for packet in demuxer: # make sure the stream is really flowing num_decoded += len(decoder.decode(packet)) @@ -4765,7 +4766,7 @@ def test_decode_after_drain_raises(self, device): # ignores anything sent afterwards. Rather than silently decoding # nothing, say so. demuxer = VideoDemuxer(H265_VIDEO.path) - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) packet = demuxer.next_packet() decoder.decode(packet) decoder.drain() @@ -4927,7 +4928,7 @@ def test_scan_leaves_demuxer_at_the_start(self, video, device): demuxer = VideoDemuxer(video.path) index = demuxer.scan() - decoder = PacketDecoder(demuxer, device=device) + decoder = VideoPacketDecoder(demuxer, device=device) converter = ColorConverter(device=device) frames = list( self._convert(converter, self._decode(decoder, self._demux(demuxer))) @@ -4981,7 +4982,7 @@ def test_stream_index(self, stream_index): # nasa_13013.mp4 has two video streams, 0 and 3, of different sizes, # and 3 is the best one, i.e. the one used when nothing is requested. demuxer = VideoDemuxer(NASA_VIDEO.path, stream_index=stream_index) - decoder = PacketDecoder(demuxer) + decoder = VideoPacketDecoder(demuxer) converter = ColorConverter() got = [ converter.convert(raw_frame) @@ -5063,7 +5064,7 @@ def test_audio_demuxer_seek(self): @staticmethod def _decode_audio(asset, stream_index=None): demuxer = AudioDemuxer(asset.path, stream_index=stream_index) - decoder = PacketDecoder(demuxer) + decoder = AudioPacketDecoder(demuxer) chunks = [] for packet in demuxer: chunks += decoder.decode(packet) @@ -5138,28 +5139,48 @@ def test_audio_raw_samples_pts(self): assert pts == sorted(pts) assert pts[0] == pytest.approx(0, abs=1e-6) - def test_audio_decoder_output_type_follows_the_demuxer(self): - # Same PacketDecoder class, different output type: that's the whole - # reason it isn't split in two. - for demuxer_class, expected_type in ( - (AudioDemuxer, RawAudioSamples), - (VideoDemuxer, RawFrame), + def test_decoder_output_type_follows_the_demuxer(self): + # The two decoders are one class in C++; the split is a Python-level + # one, so that each has an exact output type and its own arguments. + for demuxer_class, decoder_class, expected_type in ( + (AudioDemuxer, AudioPacketDecoder, RawAudioSamples), + (VideoDemuxer, VideoPacketDecoder, RawFrame), ): demuxer = demuxer_class(NASA_VIDEO.path) - decoder = PacketDecoder(demuxer) + decoder = decoder_class(demuxer) # A codec needs more than one packet before it outputs anything. decoded = [] while not decoded: decoded = decoder.decode(demuxer.next_packet()) assert isinstance(decoded[0], expected_type) - @pytest.mark.parametrize("device", ("cuda", "cuda:0", torch.device("cuda"))) - def test_audio_decoder_non_cpu_device_raises(self, device): - # Not gated on CUDA being available: we reject the request before ever - # touching a device. - demuxer = AudioDemuxer(NASA_AUDIO_MP3.path) - with pytest.raises(ValueError, match="audio can only be decoded on the CPU"): - PacketDecoder(demuxer, device=device) + def test_audio_decoder_takes_no_device(self): + # Audio is CPU-only, so rather than accepting a device and rejecting + # everything but CPU, AudioPacketDecoder simply has no such argument. + with pytest.raises(TypeError, match="device"): + AudioPacketDecoder(AudioDemuxer(NASA_AUDIO_MP3.path), device="cuda") + + def test_audio_decoder_mpeg_ps_resync_after_seek(self): + # Seeking an MPEG program stream lands on a container-level byte + # offset, so the parser resumes mid-frame and the packets it rebuilds + # are AVERROR_INVALIDDATA until it resyncs. That's a property of the + # container, not of the codec: it applies to this file's audio stream + # exactly as it does to a video one. Without the resync handling this + # raises "Failed to send packet to decoder" on the very first packet. + asset = SINE_STEREO_MP2_MPEG_PS + demuxer = AudioDemuxer(asset.path) + decoder = AudioPacketDecoder(demuxer) + demuxer.seek(asset.duration_seconds / 2) + decoder.reset() + + chunks = [] + for packet in demuxer: + chunks += decoder.decode(packet) + chunks += decoder.drain() + + assert len(chunks) > 0 + num_samples = sum(chunk.num_samples for chunk in chunks) + assert 0 < num_samples < asset.duration_seconds * asset.sample_rate @needs_cuda def test_audio_decoder_ignores_non_cpu_default_device(self): From 2ef021892df0acd2889a1096f80ece73bac7d814 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 17:26:06 +0100 Subject: [PATCH 3/6] nits --- src/torchcodec/_core/PacketDecoder.cpp | 43 ++++++++++++-------------- src/torchcodec/_core/PacketDecoder.h | 4 +++ 2 files changed, 24 insertions(+), 23 deletions(-) diff --git a/src/torchcodec/_core/PacketDecoder.cpp b/src/torchcodec/_core/PacketDecoder.cpp index 652c15937..2cf8bc800 100644 --- a/src/torchcodec/_core/PacketDecoder.cpp +++ b/src/torchcodec/_core/PacketDecoder.cpp @@ -75,8 +75,6 @@ PacketDecoder::PacketDecoder( AVStream* stream = demuxer.active_stream(); time_base_ = stream->time_base; - // This is a property of the container, not of the codec: an audio stream in - // an MPEG program stream resyncs after a seek exactly like a video one does. is_mpeg_ps_ = std::string_view(demuxer.format_context()->iformat->name) == "mpeg"; @@ -314,27 +312,26 @@ torch::stable::Tensor audio_samples(const AVFrame& av_frame) { av_frame.extended_data[channel], bytes_per_channel); } - return samples; - } - - const uint8_t* src = av_frame.extended_data[0]; - int num_samples_int = static_cast(num_samples); - switch (bytes_per_sample) { - case 1: - deinterleave(src, dst, num_channels, num_samples_int); - break; - case 2: - deinterleave(src, dst, num_channels, num_samples_int); - break; - case 4: - deinterleave(src, dst, num_channels, num_samples_int); - break; - case 8: - deinterleave(src, dst, num_channels, num_samples_int); - break; - default: - STD_TORCH_CHECK( - false, "Unexpected sample width: ", bytes_per_sample, " bytes."); + } else { + const uint8_t* src = av_frame.extended_data[0]; + int num_samples_int = static_cast(num_samples); + switch (bytes_per_sample) { + case 1: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 2: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 4: + deinterleave(src, dst, num_channels, num_samples_int); + break; + case 8: + deinterleave(src, dst, num_channels, num_samples_int); + break; + default: + STD_TORCH_CHECK( + false, "Unexpected sample width: ", bytes_per_sample, " bytes."); + } } return samples; } diff --git a/src/torchcodec/_core/PacketDecoder.h b/src/torchcodec/_core/PacketDecoder.h index b8e3ea95a..71d50e9ea 100644 --- a/src/torchcodec/_core/PacketDecoder.h +++ b/src/torchcodec/_core/PacketDecoder.h @@ -103,6 +103,8 @@ struct FrameMetadata { double rotation_degrees = 0; }; +// TODO_API_BREAKDOWN CC P1 these should bet get_* + // Describes `av_frame` without touching its samples. Unlike frame_planes(), // this works for every pixel format, so callers can ask what a frame is before // asking for views they may not be able to get. @@ -124,6 +126,8 @@ FORCE_PUBLIC_VISIBILITY std::vector frame_planes( // interleave them, so neither is a [C, N] tensor as it stands. An audio frame // is a few kB, so normalizing here buys a uniform layout for the price of a // memcpy - and it means a converter can treat the result as planar-of-dtype. +// TODO_API_BREAKDOWN DESIGN P1: do we want to copy? Should we just keep the +// original layout? FORCE_PUBLIC_VISIBILITY torch::stable::Tensor audio_samples( const AVFrame& av_frame); From 083f841278cd4ea6905395ae2fb95b5d999b8929 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 17:30:29 +0100 Subject: [PATCH 4/6] nits --- test/test_decoders.py | 23 ----------------------- 1 file changed, 23 deletions(-) diff --git a/test/test_decoders.py b/test/test_decoders.py index d1606a106..9327cb09d 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -5155,8 +5155,6 @@ def test_decoder_output_type_follows_the_demuxer(self): assert isinstance(decoded[0], expected_type) def test_audio_decoder_takes_no_device(self): - # Audio is CPU-only, so rather than accepting a device and rejecting - # everything but CPU, AudioPacketDecoder simply has no such argument. with pytest.raises(TypeError, match="device"): AudioPacketDecoder(AudioDemuxer(NASA_AUDIO_MP3.path), device="cuda") @@ -5182,27 +5180,6 @@ def test_audio_decoder_mpeg_ps_resync_after_seek(self): num_samples = sum(chunk.num_samples for chunk in chunks) assert 0 < num_samples < asset.duration_seconds * asset.sample_rate - @needs_cuda - def test_audio_decoder_ignores_non_cpu_default_device(self): - # An unspecified device means CPU for audio, not whatever - # torch.get_default_device() happens to be. Compare against a plain - # run, so this also catches us silently decoding something else. - expected = self._decode_audio(NASA_AUDIO_MP3)[0].data - - with torch.device("cuda"): - got = self._decode_audio(NASA_AUDIO_MP3)[0].data - assert got.device.type == "cpu" - torch.testing.assert_close(got, expected, atol=0, rtol=0) - - original_device = torch.get_default_device() - try: - torch.set_default_device("cuda") - got = self._decode_audio(NASA_AUDIO_MP3)[0].data - finally: - torch.set_default_device(original_device) - assert got.device.type == "cpu" - torch.testing.assert_close(got, expected, atol=0, rtol=0) - @pytest.mark.parametrize("stream_index", (-1, 6, 1000)) def test_invalid_stream_index_raises(self, stream_index): with pytest.raises(RuntimeError, match="is not a valid stream"): From 6919cf4f7b23a6e6bc87c125bee4151327cc1e16 Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 17:33:20 +0100 Subject: [PATCH 5/6] nits --- test/test_decoders.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/test/test_decoders.py b/test/test_decoders.py index 9327cb09d..be4dc825c 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -5180,19 +5180,25 @@ def test_audio_decoder_mpeg_ps_resync_after_seek(self): num_samples = sum(chunk.num_samples for chunk in chunks) assert 0 < num_samples < asset.duration_seconds * asset.sample_rate + @pytest.mark.parametrize( + "demuxer_class", (VideoDemuxer, AudioDemuxer), ids=("video", "audio") + ) @pytest.mark.parametrize("stream_index", (-1, 6, 1000)) - def test_invalid_stream_index_raises(self, stream_index): + def test_invalid_stream_index_raises(self, demuxer_class, stream_index): with pytest.raises(RuntimeError, match="is not a valid stream"): - VideoDemuxer(NASA_VIDEO.path, stream_index=stream_index) + demuxer_class(NASA_VIDEO.path, stream_index=stream_index) - def test_bad_source_type_raises(self): + @pytest.mark.parametrize( + "demuxer_class", (VideoDemuxer, AudioDemuxer), ids=("video", "audio") + ) + def test_bad_source_type_raises(self, demuxer_class): with pytest.raises(TypeError, match="Unknown source type"): - VideoDemuxer(123) + demuxer_class(123) # user mistakenly forgets to specify binary reading when creating a # file-like object from open() with pytest.raises(TypeError, match="binary reading?"): - VideoDemuxer(open(NASA_VIDEO.path)) + demuxer_class(open(NASA_VIDEO.path)) # Small helpers to avoid having to always specify the same skip marks and decode_fn From 15311a05e242de4761e7ab7a2156024561917aad Mon Sep 17 00:00:00 2001 From: Nicolas Hug Date: Fri, 21 Aug 2026 17:43:01 +0100 Subject: [PATCH 6/6] More tests after seek --- test/test_decoders.py | 47 +++++++++++++++++++++++++++++++++++++++---- 1 file changed, 43 insertions(+), 4 deletions(-) diff --git a/test/test_decoders.py b/test/test_decoders.py index be4dc825c..914e3c1ce 100644 --- a/test/test_decoders.py +++ b/test/test_decoders.py @@ -5062,9 +5062,12 @@ def test_audio_demuxer_seek(self): # ===== Audio decoding: RawAudioSamples ===== @staticmethod - def _decode_audio(asset, stream_index=None): + def _decode_audio(asset, stream_index=None, seek_seconds=None): demuxer = AudioDemuxer(asset.path, stream_index=stream_index) decoder = AudioPacketDecoder(demuxer) + if seek_seconds is not None: + demuxer.seek(seek_seconds) + decoder.reset() chunks = [] for packet in demuxer: chunks += decoder.decode(packet) @@ -5117,11 +5120,39 @@ def test_audio_raw_samples_dtype_and_shape(self, asset, sample_format, dtype): NASA_AUDIO, ), ) - def test_audio_raw_samples_match_audio_decoder(self, asset): + @pytest.mark.parametrize("seek_fraction", (None, 1 / 3, 2 / 3)) + def test_audio_raw_samples_match_audio_decoder(self, asset, seek_fraction): # We hand out the true source samples: normalizing them the way FFmpeg # does reproduces AudioDecoder's output bit for bit. This is also what # pins the de-interleaving, most visibly on the 16-channel asset. - raw = torch.cat([chunk.data for chunk in self._decode_audio(asset)], dim=1) + # + # It holds after a seek too, but only once the caller has done the + # pre-roll these blocks don't do: a lossy codec decodes its first frames + # after a seek from a flushed state, so they come out subtly wrong - + # plausible, but not what whole-file decoding gives - until it + # re-primes. Dropping those frames is exactly what pre-rolling means, + # and everything from there on is bit exact again. Without the drop, + # mp3 and aac diverge over their first ~1000-1600 samples and match + # perfectly after that. + if seek_fraction is not None and asset is SINE_STEREO_MP2_MPEG_PS: + pytest.skip( + "MPEG-PS resync after a seek is unreliable for both the blocks " + "and AudioDecoder (which raises seeking this file to 2.8s), so " + "it can't tell us anything here. See " + "test_audio_decoder_mpeg_ps_resync_after_seek." + ) + + seek_seconds = ( + None if seek_fraction is None else asset.duration_seconds * seek_fraction + ) + chunks = self._decode_audio(asset, seek_seconds=seek_seconds) + if seek_seconds is not None: + # Same number of frames SingleStreamDecoder pre-rolls by, see + # Note [Audio pre-roll and post-roll]. + chunks = chunks[4:] + assert len(chunks) > 0 + + raw = torch.cat([chunk.data for chunk in chunks], dim=1) if raw.dtype == torch.uint8: got = (raw.to(torch.float32) - 128) / 128 @@ -5130,7 +5161,15 @@ def test_audio_raw_samples_match_audio_decoder(self, asset): else: got = raw.to(torch.float32) - expected = AudioDecoder(asset.path).get_all_samples().data + decoder = AudioDecoder(asset.path) + if seek_seconds is None: + expected = decoder.get_all_samples().data + else: + # Re-anchor on the first frame we kept, so both sides start on the + # same sample. + expected = decoder.get_samples_played_in_range( + start_seconds=chunks[0].pts_seconds + ).data torch.testing.assert_close(got, expected, atol=0, rtol=0) def test_audio_raw_samples_pts(self):