avtensor

High-performance media decoding straight into PyTorch tensors.

avtensor is a Rust library with Python bindings that decodes video, audio and images into torch.Tensors using FFmpeg. It is built for ML data-loading pipelines: data moves directly from FFmpeg frames into tensors. Inputs can be local files, gs:// / s3:// objects or any HTTP(S) URL.

• Decode-time transforms inside FFmpeg: resize, frame-rate resampling, audio resampling, EBU R128 loudness normalization
• Color-correct output: YUV→RGB follows the stream's color metadata, and HDR sources are tone-mapped to sRGB
• Windowed decode of a start/end time range
• NVDEC hardware decode, optionally with GPU-resident output tensors

Comparison with torchcodec

torchcodec is PyTorch's official media decoder. Both decode through FFmpeg and produce bit-identical RGB output on CPU decode. Some differences as of torchcodec 0.14:

	avtensor	torchcodec
cloud inputs (gs://, s3://)	native	—
video + audio	one decode pass	separate decoders
HDR	tone-mapped to SDR	float32 output, no tone mapping (beta)
encoding	—	audio + video

Benchmarks

Median wall time on a 30 s 1080p H.264 clip (Xeon 8481C + H100, identical FFmpeg builds). Whole-clip decode times are within ±1% of each other; decode-time resize is 1.7–6× faster in avtensor. Full results, methodology, and reproduction steps in benchmarks/:

Installation

Install from PyPI with pip install avtensor. The package is distributed as source and compiles at install time against your environment's PyTorch (libtorch) and FFmpeg >= 7.1 (shared build):

export LIBTORCH_USE_PYTORCH=1
export LIBTORCH_STATIC=0
export FFMPEG_PKG_CONFIG_PATH=/opt/ffmpeg/lib/pkgconfig
export LD_LIBRARY_PATH=/opt/ffmpeg/lib:$LD_LIBRARY_PATH
export LIBRARY_PATH=$LD_LIBRARY_PATH:$LIBRARY_PATH

pip install avtensor

To build a reusable wheel for deployment, use maturin: maturin build --release -o wheelhouse/.

Requirements:

dependency	version	notes
Python	>= 3.8 (CPython)	wheel is built per Python minor version
PyTorch	2.13 by default (tch 0.26 in Cargo.lock); 2.11 via tch 0.24	linked dynamically at build time
FFmpeg	>= 7.1, shared libraries	needs zscale/tonemap filters (libzimg) for HDR input and for color_conversion="accurate" (the float32 default; without libzimg, float32 falls back to "fast" with a warning)
Rust	stable toolchain	build-time only

Quickstart

import avtensor
from avtensor import (
    AudioStreamRequest,
    MediaDecodeRequest,
    StreamType,
    VideoStreamRequest,
)

# Decode a video with its audio track.
request = MediaDecodeRequest("clip.mp4")
request.video_stream = VideoStreamRequest()
request.audio_streams = [AudioStreamRequest()]

streams = avtensor.decode_asset(request)

for stream in streams:
    if stream["stream_type"] == StreamType.Video:
        frames = stream["data"]      # uint8 Tensor, shape [T, C, H, W] (RGB)
        print(f"video: {frames.shape} @ {stream['fps']} fps")
    elif stream["stream_type"] == StreamType.Audio:
        samples = stream["data"]     # float32 Tensor, shape [C, T]
        print(f"audio: {samples.shape} @ {stream['sample_rate']} Hz")

Resize and resample at decode time

request = MediaDecodeRequest(
    "clip.mp4",
    video_stream=VideoStreamRequest(
        width=512,
        height=288,
        fps=24.0,          # frame-rate resampling
    ),
    audio_streams=[AudioStreamRequest(sample_rate=16000)],  # audio resampling
)

streams = avtensor.decode_asset(request)

By default fps keeps the source frame nearest each output timestamp, which spaces frames unevenly for non-integer ratios (30 -> 24 fps keeps 4 of every 5 frames). fps_conversion="blend" cross-fades neighbouring frames instead, and fps_conversion="motion" interpolates along estimated motion (about 1.25x the cost of the default):

VideoStreamRequest(fps=24.0, fps_conversion="motion")

Both interpolating modes detect hard scene cuts and use the nearer source frame for an output frame that falls between two shots, so no frame mixes two unrelated shots. The detector looks for a one-frame spike in frame difference against recent motion and confirms it with local block correlation, so fast pans and lighting changes don't count as cuts. scene_cuts="off" interpolates across cuts instead. Frame count and timestamps are the same as "nearest" in every mode.

Decode a time window (seek)

request = MediaDecodeRequest(
    "movie.mp4",
    start_time=42.0,   # seconds
    end_time=48.0,
    video_stream=VideoStreamRequest(),
    audio_streams=[AudioStreamRequest()],
)

streams = avtensor.decode_asset(request)

Decode from cloud storage

# Google Cloud Storage (Application Default Credentials):
request = MediaDecodeRequest("gs://my-bucket/path/clip.mp4")

# S3 (standard AWS credential chain):
request = MediaDecodeRequest("s3://my-bucket/path/clip.mp4")

# Any HTTP(S) URL FFmpeg can read — e.g. a presigned URL:
request = MediaDecodeRequest("https://my-bucket.s3.amazonaws.com/clip.mp4?X-Amz-...")

See Reading from cloud storage for credentials, custom endpoints and S3-compatible providers.

In-memory assets and probing

bytes input decodes without touching the filesystem, and probe_asset returns the stream layout without decoding:

data: bytes = fetch_asset()  # e.g. from a queue or cache

meta = avtensor.probe_asset(data)
print(meta["video_streams"][0])  # {'index': 0, 'width': 1920, 'height': 1080, 'fps': 30.0}

request = MediaDecodeRequest(data, video_stream=VideoStreamRequest())
streams = avtensor.decode_asset(request)

High bit-depth sources (float32 output)

By default frames are quantized to 8-bit RGB. For 10/12-bit sources (HDR masters, ProRes, 10-bit H.264/HEVC), request float32 output: FFmpeg converts to planar float and avtensor returns float32 tensors in [0, 1] with the source's full precision. Unlike uint8 output, float32 tensors are contiguous in NCHW (the NHWC order is the view).

request = MediaDecodeRequest(
    "10bit_master.mov",
    video_stream=VideoStreamRequest(dtype="float32"),
)
(video,) = avtensor.decode_asset(request)
video["data"].dtype  # torch.float32, values in [0, 1]

float32 output is 4× the memory of uint8 (a 30 s 1080p clip is ~22 GB); decode a start_time/end_time window for long assets.

Images

Images decode as a single-frame video stream:

request = MediaDecodeRequest("photo.jpg")
request.video_stream = VideoStreamRequest()

(image,) = avtensor.decode_asset(request)
image["data"].shape  # [1, 3, H, W]

GPU-accelerated decoding (NVDEC)

request.video_stream = VideoStreamRequest(hardware_acceleration=True)

Bitstream decoding runs on the GPU's NVDEC engine (FFmpeg's *_cuvid decoders); frames return to system memory, so filters and tensor conversion are unchanged. A requested downscale runs on the GPU before the transfer, shrinking the GPU→CPU copy. On 1080p production shots decoded to 256×144 this measured ~1.7× faster with ~3.4× less CPU than software decode.

Caveats: a GPU has a limited number of decode engines, so many concurrent decodes can saturate them (software decode scales with cores instead). NVDEC H.264 supports only 4:2:0 chroma. Requires an NVIDIA GPU and an FFmpeg build with cuvid support.

GPU-resident output

With device, frames never leave the GPU: the returned tensor is CUDA-resident, and NV12 → RGB conversion runs on the device (NPP), so the GPU→CPU transfer and CPU color conversion disappear. device implies hardware_acceleration.

request.video_stream = VideoStreamRequest(device="cuda")  # tensor lands on cuda:0

Constraints on this path:

• width/height must both be set to an even, strictly smaller size (NVDEC's own scaler) or both left unset.
• fps resampling, HDR tone mapping and autorotate of a rotated stream are unavailable; frames never reach the CPU filter graph.
• Values may differ from the CPU path by ±1 (NPP vs swscale rounding).

The CUDA runtime and NPP libraries must be present at run time. They are dlopened, so CPU-only deployments carry no CUDA dependency.

Loudness normalization

Audio can be loudness-normalized (FFmpeg loudnorm, EBU R128) during decode:

from avtensor import LoudnessNormalization

norm = LoudnessNormalization(
    integrated_loudness_target=-18.0,  # LUFS
    true_peak_level_target=-1.0,       # dBTP
    loudness_range_target=7.0,         # LU
)

request = MediaDecodeRequest(
    "clip.mp4",
    audio_streams=[AudioStreamRequest(loudness_normalization=norm)],
)

API reference

decode_asset(request: MediaDecodeRequest) -> list[DecodeResult]

Decodes the requested streams. Returns one result dict per decoded stream:

key	video streams	audio streams
data	uint8 Tensor [T, C, H, W], RGB (float32 in [0, 1], NCHW-contiguous, with dtype="float32"; [T, H, W, C] with dimension_order="NHWC")	float32 Tensor [C, T]
stream_type	StreamType.Video	StreamType.Audio
stream_index	index of the stream in the container	index of the stream in the container
fps	output frame rate	—
pts	float64 Tensor [T], presentation timestamp of each frame in seconds	—
sample_rate	—	output sample rate in Hz

The GIL is released for the duration of the decode.

probe_asset(input: str | bytes) -> MediaMetadata

Returns the asset's stream layout without decoding it: a dict with format_name, duration_s, video_streams (each with index, coded width/height, fps, codec_name, duration_s, duration_source, nb_frames, rotation, sample_aspect_ratio) and audio_streams (each with index, sample_rate, channels, codec_name, duration_s, duration_source). Stream durations fall back to the container duration (duration_source="container", e.g. Matroska/WebM). sample_aspect_ratio is reported but not applied by the decoder. rotation is the stream's display matrix (counterclockwise degrees); the decoder applies it only with VideoStreamRequest(autorotate=True). It is the stream-level matrix only: orientation carried per frame (JPEG EXIF Orientation, H.264/HEVC display-orientation SEI) is reported as rotation=None and not applied, even though the ffmpeg CLI applies it. Accepts the same inputs as MediaDecodeRequest. The GIL is released while probing.

plan_asset(request: MediaDecodeRequest) -> list[PlannedStream]

Returns, per stream, the buffer decode_asset(request) would allocate, without decoding: stream_type, stream_index, shape (in the returned layout), dtype, fps or sample_rate, and the duration_s / duration_source being decoded. Stream selection, the seek window, fps, width/height, autorotate and dtype resolve through the same code as decode_asset. The decoded tensor is trimmed to what was actually decoded, so shape is an upper bound on its length, except for video with fps_mode="cli": its output grows past the plan when the CLI emits more frames (a file with no stored duration plans 1 frame), up to the output size limit. Decoders are not opened, so codec and hardware errors only surface at decode time.

MediaDecodeRequest(input: str | bytes)

attribute	type	meaning
input	str | bytes	local path, gs:// / s3:// URI, HTTP(S) URL or the raw bytes of an in-memory asset
start_time	float | None	decode window start, in seconds
end_time	float | None	decode window end, in seconds
video_stream	VideoStreamRequest | None	request the video stream (None = skip video)
audio_streams	list[AudioStreamRequest] | None	audio streams to decode (None = skip audio)
max_output_bytes	int | None	largest output buffer per stream, in bytes (see Output size limit)

All request classes accept their attributes as keyword arguments (VideoStreamRequest(width=512, height=288)). Nested request objects are held by reference, so mutating request.video_stream (or an element of request.audio_streams) after assignment is reflected in decode_asset.

VideoStreamRequest

All attributes default to None, meaning "keep the source value".

attribute	type	meaning
index	int | None	select a specific video stream
width, height	int | None	rescale output frames (a lone width or height keeps the source's coded aspect ratio; sample aspect ratio is not applied)
fps	float | None	resample to this frame rate
number_of_threads	int | None	FFmpeg decoder threads (default: 1; 0 = auto). Also sizes the RGB conversion on the direct conversion path; other decodes convert in the filter graph with its own pool of up to 16 threads
hardware_acceleration	bool | None	decode on the GPU's NVDEC engine (see GPU-accelerated decoding)
dimension_order	str | None	"NCHW" (default, [T, C, H, W], a non-contiguous view) or "NHWC" ([T, H, W, C], contiguous)
device	str | None	"cuda" / "cuda:N": keep frames on the GPU, implies hardware_acceleration (see GPU-accelerated decoding)
dtype	str | None	"uint8" (default) or "float32" ([0, 1], NCHW-contiguous, preserves 10/12-bit source depth)
color_conversion	str | None	"fast" (libswscale default; default for uint8) or "accurate" (zimg: exact float32, round-to-nearest uint8, resizes after the RGB conversion; default for float32 when FFmpeg has libzimg). Not supported with device="cuda"
autorotate	bool | None	apply the stream's display-matrix rotation as the ffmpeg CLI's -autorotate does (90°/270° swap the output width and height); width/height refer to the rotated frame. Default False. Stream-level matrix only: per-frame orientation (JPEG EXIF, H.264/HEVC SEI) is not applied
fps_mode	str | None	"filter" (default) or "cli": the frames ffmpeg -i <input> -f rawvideo - outputs, at the stream's guessed rate, with the CLI's duplication/dropping and MPEG-TS/PS/HLS timestamp-discontinuity correction. Whole-file decodes only, without fps, device or hardware_acceleration=True. Unlike the CLI it autorotates only with autorotate=True.

AudioStreamRequest

attribute	type	meaning
index	int | None	select a specific audio stream
sample_rate	int | None	resample to this rate
loudness_normalization	LoudnessNormalization | None	apply EBU R128 loudness normalization

LoudnessNormalization

Mirrors FFmpeg's loudnorm filter: integrated_loudness_target (LUFS), true_peak_level_target (dBTP), loudness_range_target (LU), the measured_* variants for two-pass normalization, offset_gain, linear and dual_mono.

Output size limit

Each stream's output buffer is preallocated at a size computed from the file's metadata (duration, frame rate, dimensions, sample rate), which an untrusted file controls. To bound that allocation, a request whose buffer for any one stream would exceed 16 GiB fails before any buffer is allocated, for example 30 s of uncropped 4K uint8 video (about 22 GB). plan_asset applies the same check.

• The limit is per stream: video plus N audio streams can total (1 + N)× it.
• Set it per request with MediaDecodeRequest(max_output_bytes=...), or for the whole process with AVTENSOR_MAX_OUTPUT_TENSOR_BYTES (bytes, or with a K/M/G/T binary suffix such as 32G); the request's value wins. Prefer raising it for the requests that need it over a huge process-wide value, which lets a file with spoofed metadata allocate that much.
• 0 refuses every non-empty stream; it does not mean unlimited.
• An invalid AVTENSOR_MAX_OUTPUT_TENSOR_BYTES keeps the 16 GiB default, and the error names the ignored value.

Reading from cloud storage

gs:// and s3:// objects are fetched with streaming reads and concurrent range requests on seek, buffered in memory (AVTENSOR_MAX_CLOUD_OBJECT_BYTES caps the buffer; default 16 GiB). HTTP(S) inputs go through FFmpeg's own protocol layer.

Google Cloud Storage (gs://)

Authentication uses Application Default Credentials. Set GCS_ENDPOINT to point at a different GCS-API-compatible endpoint (defaults to https://storage.googleapis.com).

Amazon S3 and S3-compatible stores (s3://)

Credentials and region come from the standard AWS provider chain (AWS_ACCESS_KEY_ID/AWS_SECRET_ACCESS_KEY, AWS_PROFILE, IMDS, ...). S3-compatible providers work by overriding the endpoint:

export AWS_ENDPOINT_URL_S3=https://objects.example.com

# MinIO and other path-style-only stores additionally need:
export AVTENSOR_S3_FORCE_PATH_STYLE=1

S3 Configuration can also be passed explicitly per call with an S3Config, instead of through the environment.

import avtensor

s3_cfg = avtensor.S3Config(
    endpoint_url="https://objects.example.com",
    region="us-east-1",
    access_key_id=...,
    secret_access_key=...,
    force_path_style=True,
)

request = avtensor.MediaDecodeRequest("s3://bucket/key.mp4", ...)
avtensor.decode_asset(request, s3_config=s3_cfg)
avtensor.probe_asset("s3://bucket/key.mp4", s3_config=s3_cfg)

Fields set on an S3Config are authoritative — the environment is not consulted for them; unset fields fall back to the standard AWS environment.

Anything else: presigned HTTP(S) URLs

Any provider that can issue a presigned or public HTTP(S) URL works without provider-specific support — the URL is handed to FFmpeg's http protocol.

Color handling

• YUV→RGB conversion uses the colorspace and range tagged on the stream, rather than letting FFmpeg guess from the resolution.
• HDR sources — HLG (arib-std-b67) or PQ (smpte2084) transfer or BT.2020 primaries — are automatically tone-mapped (zscale → tonemap=hable → zscale) so the returned RGB tensors are sRGB/BT.709, consistent with SDR sources.
• HDR streams (PQ/HLG transfer or BT.2020 primaries) missing a matrix or primaries tag are treated as BT.2020 (non-constant luminance), in both hdr_modes. Untagged SDR keeps swscale's BT.601 default.
• float32 output is clamped to [0, 1]. zimg does not clip out-of-gamut YUV, and tone-mapped HDR can exceed 1.0 (in either color_conversion mode).
• color_conversion="accurate" converts with zimg and resizes after the conversion to RGB (for tone-mapped HDR, in the first, linearizing step). zimg upsamples subsampled chroma with bicubic rather than swscale's bilinear. That is more accurate on natural content, but can ring slightly at saturated hard edges. Odd-size 4:2:0/4:2:2 frames are upsampled to 4:4:4 first, because zimg rejects them otherwise.

Direct conversion path

When no filtering is needed, avtensor converts decoded frames straight into the output tensor with swscale instead of going through an FFmpeg filter graph. The output is identical; it saves a copy per frame, and its conversion threads follow number_of_threads. A video stream takes this path only when all of these hold:

• software decoding (no hardware_acceleration or device);
• uint8 (rgb24) output, with no width/height resize and no fps;
• an output width that is a multiple of 32 (swscale's vectorized RGB24 writer needs it on the tensor's tightly packed rows), so e.g. 1920 or 1280 wide qualify and 720 or 854 do not;
• an SDR source, or an HDR source with hdr_mode="raw" whose YUV matrix is BT.709, BT.601, FCC, SMPTE 240M, BT.2020 NCL or CL, RGB, or untagged. These are the matrices the scale filter accepts, plus BT.2020 CL, which only this path can convert: a BT.2020 CL source that needs the filter graph (a resize, fps, a width that isn't a multiple of 32, or AVTENSOR_DISABLE_DIRECT_PATH) fails.

Everything else converts in the filter graph, whose thread pool is up to 16 threads whatever number_of_threads is. AVTENSOR_DISABLE_DIRECT_PATH=1 forces the filter graph for every stream.

Troubleshooting

• Build fails with "Cannot find a libtorch install" — set LIBTORCH_USE_PYTORCH=1 and make sure python -c "import torch" works in the active environment.
• Build fails with FFmpeg/pkg-config errors — set FFMPEG_PKG_CONFIG_PATH to the directory containing libavcodec.pc etc., and make sure the FFmpeg is a shared (not static) build, version >= 7.1.
• ImportError: libavcodec.so.61: cannot open shared object file — the FFmpeg libraries must be on LD_LIBRARY_PATH at runtime, matching the version the wheel was built against.
• Hardware decode fails with CUDA_ERROR_NOT_SUPPORTED — the stream's format exceeds NVDEC's capabilities; for H.264 only 4:2:0 chroma is supported (4:4:4 will not decode).
• HDR decode fails with a filter error — your FFmpeg lacks the zscale filter (libzimg); use a build that includes it.
• color_conversion="accurate" needs FFmpeg's zscale filter — same cause: use an FFmpeg built with libzimg, or pass color_conversion="fast". Without libzimg, float32 decodes fall back to "fast" (swscale, scaled by 255/256) and log a warning.
• s3:// works with AWS but not with your S3-compatible store — check whether the store requires path-style addressing (AVTENSOR_S3_FORCE_PATH_STYLE=1) and that AWS_ENDPOINT_URL_S3 is set.
• Resident memory grows when decoded tensors are freed on a different thread than the one that decoded them — torch's CPU allocator is not always plain malloc (the aarch64 Linux wheels embed mimalloc, which keeps blocks freed cross-thread mapped and resident until the allocating thread exits). avtensor therefore allocates its output buffers with the system allocator and hands them to torch via from_blob with a deleter, so a free is an immediate munmap on any thread. AVTENSOR_OUTPUT_ALLOCATOR=torch switches back to torch.empty-style allocation.
• Filling large output buffers is slow — on Linux, set AVTENSOR_OUTPUT_HUGE_PAGES=1 to advise output buffers of 64 MiB or more to use transparent huge pages (THP), which roughly halves the cost of filling them on hosts with plenty of free memory. It is off by default: under the common THP settings enabled=madvise and defrag=madvise, a fault in an advised buffer may stall on memory reclaim and compaction when free memory is low (for example, a data loader whose page cache fills RAM), which can make decodes several times slower instead. Under enabled=always, unadvised memory already gets huge pages when they are cheap. If you enable it, watch compact_stall and thp_fault_fallback in /proc/vmstat. The advice is skipped on kernels whose huge pages are larger than 2 MiB (such as 512 MiB on 64 KiB-page arm64).

Contributing

Issues and pull requests are welcome. Project layout:

src/
  alloc.rs       output-buffer allocation (system allocator + torch from_blob)
  decoder/       demuxing, decoding, filter graphs (mod.rs), cloud AVIO reader (io.rs)
  ffi/           PyO3 bindings: request/response types, decode_asset
  util/          gcs/s3 URI handling, memory, test media generation
csrc/            C++ shim for the one libtorch call tch-rs does not wrap (built by build.rs)
avtensor.pyi     Python type stubs, shipped with the wheel

Set the environment variables from Installation first — builds and tests fail without them. CI enforces all of the following:

make test              # cargo test --no-default-features (required PyO3 workaround)
cargo fmt --check
cargo clippy --no-default-features --all-targets -- -D warnings

# Python tooling (benchmarks/):
ruff check benchmarks/ && ruff format --check benchmarks/
ty check

# Line coverage (65% floor):
cargo llvm-cov --no-default-features --summary-only

The test suite is self-contained: all test media is generated locally with FFmpeg at test time. A performance-regression gate also runs on every PR. While debugging, RUST_LOG=debug RUST_BACKTRACE=1 enables native logs and backtraces; errors are forwarded to Python logging by default.

Cargo.lock is pinned deliberately — wheels link against FFmpeg and libtorch, and downstream deployments validate specific versions. Bump dependencies intentionally rather than accepting automated lockfile updates.

License

Apache 2.0