Load TensorFlow or PyTorch Models in FFmpeg SR

FFmpeg's sr (Super Resolution) filter applies deep learning models to upscale and enhance video frames using its internal Deep Neural Network (DNN) subsystem. This article explains how to prepare TensorFlow and PyTorch models for FFmpeg, ensure the proper backend dependencies are compiled, and construct the precise filter commands required to run inference on your media files.

1. Compile FFmpeg with DNN Backends

The sr filter requires FFmpeg to be compiled with the specific C++ runtime libraries corresponding to your model's framework. Standard static builds often omit these dependencies.

  • For TensorFlow: FFmpeg must be built with --enable-libtensorflow. You will need the C API version of libtensorflow.
  • For PyTorch: Modern FFmpeg releases support the LibTorch engine, requiring --enable-libtorch during the ./configure step. You must have the C++ LibTorch distribution installed on your system.

To check if your FFmpeg binary supports the necessary backend, run:

ffmpeg -filters | grep sr

Then check configuration flags using:

ffmpeg -buildconf

Verify that either --enable-libtensorflow or --enable-libtorch is present.

2. Prepare Your Model Files

FFmpeg does not accept raw Python script checkpoints; models must be serialized into deployment-ready formats.

TensorFlow (.pb)

TensorFlow models must be exported as frozen protocol buffer (.pb) files. You must also know the exact names of your input and output tensor nodes.

A typical export in TensorFlow produces a graph where nodes are identified by name:

  • Input node name example: x or input_tensor
  • Output node name example: y or output_tensor

PyTorch (.pt)

PyTorch models must be compiled using TorchScript via tracing or scripting before they can be loaded by LibTorch inside FFmpeg.

Example Python script to trace and save a PyTorch model:

import torch

# Load your custom or pre-trained super-resolution model
model = MySRModel().eval()

# Create dummy input matching the expected dimensions [batch, channels, height, width]
dummy_input = torch.rand(1, 3, 360, 640)

# Trace the model
traced_model = torch.jit.trace(model, dummy_input)

# Save as a TorchScript model
traced_model.save("model_sr.pt")

3. Load Models into the sr Filter

The sr filter accepts several parameters:

  • dnn_backend: The inference engine (tensorflow or torch).
  • model: Path to the exported model file.
  • scale: The upscale factor matching the model architecture (e.g., 2 for 2x, 4 for 4x).
  • input: The name of the input node/layer (required for TensorFlow).
  • output: The name of the output node/layer (required for TensorFlow).

TensorFlow Command Example

To process a video using a frozen TensorFlow graph:

ffmpeg -i input.mp4 -vf "sr=dnn_backend=tensorflow:model=model_sr.pb:input=input_node_name:output=output_node_name:scale=2" -c:a copy output.mp4

PyTorch Command Example

To run inference using a TorchScript .pt file via the LibTorch backend:

ffmpeg -i input.mp4 -vf "sr=dnn_backend=torch:model=model_sr.pt:scale=2" -c:a copy output.mp4

Notes on Memory and Performance

  • Input frame dimensions must be supported by the trained network's receptive field. If the network expects dynamic sizes, ensure the trace step allowed dynamic shapes.
  • If your model was trained on normalized color values (such as [0, 1] or RGB float conversions), insert a format filter before the sr filter (e.g., -vf "format=rgb24,sr=...") to prevent color distortion during model inference.