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:

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:

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)

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 and HDR tone mapping 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:

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 video_streams (each with index, width, height, fps) and audio_streams (each with index, sample_rate). Accepts the same inputs as MediaDecodeRequest. The GIL is released while probing.

MediaDecodeRequest(input: str | bytes)

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".

AudioStreamRequest

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.

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.

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.
• 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.

Contributing

Issues and pull requests are welcome. Project layout:

src/
  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
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