PR #25466 opened by Raja-89 URL: https://code.ffmpeg.org/FFmpeg/FFmpeg/pulls/25466 Patch URL: https://code.ffmpeg.org/FFmpeg/FFmpeg/pulls/25466.patch
Add support for loading PyTorch 2 Export (.pt2) models through the AOTInductor runtime (AOTIModelPackageLoader). TorchScript (.pt) models remain supported but are deprecated, and a warning is printed when one is loaded. Requires libtorch >= 2.8.0. The .pt2 package must be generated with the same PyTorch version that FFmpeg is built and run against. Signed-off-by: Raja-89 <[email protected]> >From 2e3aaa2cd57f5cc7e887ad79a3871cce040f4405 Mon Sep 17 00:00:00 2001 From: Raja-89 <[email protected]> Date: Sun, 25 Oct 2026 20:30:18 +0530 Subject: [PATCH] avfilter/dnn/dnn_backend_torch: add PT2E (AOTInductor) model loading Add support for loading PyTorch 2 Export (.pt2) models through the AOTInductor runtime (AOTIModelPackageLoader). TorchScript (.pt) models remain supported but are deprecated, and a warning is printed when one is loaded. Requires libtorch >= 2.8.0. The .pt2 package must be generated with the same PyTorch version that FFmpeg is built and run against. Signed-off-by: Raja-89 <[email protected]> --- doc/filters.texi | 13 ++- libavfilter/dnn/dnn_backend_torch.cpp | 127 +++++++++++++++++++++----- 2 files changed, 116 insertions(+), 24 deletions(-) diff --git a/doc/filters.texi b/doc/filters.texi index 37f9123df49..1c8efcd653d 100644 --- a/doc/filters.texi +++ b/doc/filters.texi @@ -12306,8 +12306,13 @@ need to build and install the OpenVINO for C library (see be needed if the header files and libraries are not installed into system path) @item torch -Libtorch backend. To enable this backend you need to build and install Libtorch -for C++ library. Please download cxx11 ABI version (see +Libtorch backend. To enable this backend you need to build and install the +LibTorch C++ library. The backend accepts TorchScript @file{.pt} models and +AOTInductor @file{.pt2} packages. TorchScript models remain supported but are +deprecated; use @file{.pt2} for new models. Support for @file{.pt2} packages +requires LibTorch 2.8.0 or newer, and the package should be generated with the +same LibTorch version that FFmpeg is built and run against. +Please download the cxx11 ABI version (see @url{https://pytorch.org/get-started/locally}) and configure FFmpeg with @code{--enable-libtorch --extra-cflags=-I/libtorch_root/libtorch/include @@ -12373,6 +12378,10 @@ must either be consumed by a D3D12 encoder or downloaded explicitly with Set path to model file specifying network architecture and its parameters. Note that different backends use different file formats. TensorFlow, OpenVINO, Libtorch, and ONNX Runtime backends can load files for only their respective formats. +For the Libtorch backend, a @file{.pt2} package is compiled for a specific +device at export time, so a package exported for CPU runs on CPU only. A +package exported with static shapes accepts only the input shape used at +export time. @item input Set the input name of the dnn network. Required for the TensorFlow backend; diff --git a/libavfilter/dnn/dnn_backend_torch.cpp b/libavfilter/dnn/dnn_backend_torch.cpp index 59ab0a9c729..eb7b41b9d13 100644 --- a/libavfilter/dnn/dnn_backend_torch.cpp +++ b/libavfilter/dnn/dnn_backend_torch.cpp @@ -25,11 +25,25 @@ #include <torch/torch.h> #include <torch/script.h> +#include <string.h> +#include <utility> + +/* AOTInductor (.pt2) support targets the multi-runner API introduced in + * LibTorch 2.8. */ +#if (TORCH_VERSION_MAJOR > 2 || \ + (TORCH_VERSION_MAJOR == 2 && TORCH_VERSION_MINOR >= 8)) && \ + __has_include(<torch/csrc/inductor/aoti_package/model_package_loader.h>) +#define PT2E_AVAILABLE 1 +#include <torch/csrc/inductor/aoti_package/model_package_loader.h> +#else +#define PT2E_AVAILABLE 0 +#endif extern "C" { #include "config.h" #include "dnn_io_proc.h" #include "dnn_backend_common.h" +#include "libavutil/avstring.h" #include "libavutil/opt.h" #include "libavutil/mem.h" #include "libavutil/cpu.h" @@ -43,10 +57,24 @@ extern "C" { #include "safe_queue.h" } +/* Return 1 if model_filename ends with ".pt2". */ +static int th_is_pt2_model(const char *model_filename) +{ + const char *ext; + if (!model_filename) + return 0; + ext = strrchr(model_filename, '.'); + return ext && !av_strcasecmp(ext, ".pt2"); +} + typedef struct THModel { DNNModel model; DnnContext *ctx; torch::jit::Module *jit_model; +#if PT2E_AVAILABLE + torch::inductor::AOTIModelPackageLoader *aoti_model; +#endif + int is_aoti; SafeQueue *request_queue; Queue *task_queue; Queue *lltask_queue; @@ -144,6 +172,10 @@ static void dnn_free_model_th(DNNModel **model) if (th_model->task_queue) ff_queue_destroy(th_model->task_queue); +#if PT2E_AVAILABLE + if (th_model->aoti_model) + delete th_model->aoti_model; +#endif if (th_model->jit_model) delete th_model->jit_model; @@ -376,7 +408,6 @@ static int th_start_inference(void *args) TaskItem *task = NULL; THModel *th_model = NULL; DnnContext *ctx = NULL; - std::vector<torch::jit::IValue> inputs; torch::NoGradGuard no_grad; if (!request) { @@ -389,23 +420,50 @@ static int th_start_inference(void *args) th_model = (THModel *)task->model; ctx = th_model->ctx; - if (ctx->torch_option.optimize) - torch::jit::setGraphExecutorOptimize(true); - else - torch::jit::setGraphExecutorOptimize(false); - if (!infer_request->input_tensor || !infer_request->output) { av_log(ctx, AV_LOG_ERROR, "input or output tensor is NULL\n"); return DNN_GENERIC_ERROR; } - // Transfer tensor to the same device as model - const char *device_name = ctx->device ? ctx->device : "cpu"; - c10::Device device(device_name); - if (infer_request->input_tensor->device() != device) - *infer_request->input_tensor = infer_request->input_tensor->to(device); - inputs.push_back(*infer_request->input_tensor); - *infer_request->output = th_model->jit_model->forward(inputs).toTensor(); + if (th_model->is_aoti) { +#if PT2E_AVAILABLE + /* AOTInductor path: device is baked at export time, + * no .to(device) needed for CPU. */ + std::vector<at::Tensor> aoti_inputs = { + *infer_request->input_tensor}; + try { + auto outputs = + th_model->aoti_model->boxed_run(std::move(aoti_inputs)); + if (outputs.empty()) { + av_log(ctx, AV_LOG_ERROR, + "AOTI model returned no outputs\n"); + return DNN_GENERIC_ERROR; + } + /* boxed_run() transfers input ownership to the selected runner; + * each runner owns its execution state, so no output clone is needed. */ + *infer_request->output = outputs[0]; + } catch (const std::exception &e) { + av_log(ctx, AV_LOG_ERROR, + "AOTI inference failed: %s\n", e.what()); + return DNN_GENERIC_ERROR; + } +#endif + } else { + std::vector<torch::jit::IValue> inputs; + if (ctx->torch_option.optimize) + torch::jit::setGraphExecutorOptimize(true); + else + torch::jit::setGraphExecutorOptimize(false); + /* Transfer tensor to the same device as model */ + const char *device_name = ctx->device ? ctx->device : "cpu"; + c10::Device device(device_name); + if (infer_request->input_tensor->device() != device) + *infer_request->input_tensor = + infer_request->input_tensor->to(device); + inputs.push_back(*infer_request->input_tensor); + *infer_request->output = + th_model->jit_model->forward(inputs).toTensor(); + } return 0; } @@ -637,19 +695,44 @@ static DNNModel *dnn_load_model_th(DnnContext *ctx, DNNFunctionType func_type, A goto fail; } - try { - th_model->jit_model = new torch::jit::Module; - (*th_model->jit_model) = torch::jit::load(ctx->model_filename); - th_model->jit_model->to(device); - } catch (const c10::Error& e) { - av_log(ctx, AV_LOG_ERROR, "Failed to load torch model\n"); - goto fail; - } - if (ctx->nireq <= 0) { ctx->nireq = av_cpu_count() / 2 + 1; } + if (th_is_pt2_model(ctx->model_filename)) { +#if PT2E_AVAILABLE + try { + th_model->aoti_model = + new torch::inductor::AOTIModelPackageLoader( + ctx->model_filename, "model", false, ctx->nireq); + th_model->is_aoti = 1; + } catch (const std::exception &e) { + av_log(ctx, AV_LOG_ERROR, + "Failed to load AOTI model package: %s\n", e.what()); + goto fail; + } + av_log(ctx, AV_LOG_INFO, + "Loaded AOTInductor model: %s\n", ctx->model_filename); +#else + av_log(ctx, AV_LOG_ERROR, + ".pt2 models require libtorch >= 2.8 with " + "AOTInductor support\n"); + goto fail; +#endif + } else { + av_log(ctx, AV_LOG_WARNING, + "TorchScript models are deprecated. " + "Consider exporting to .pt2 (AOTInductor).\n"); + try { + th_model->jit_model = new torch::jit::Module; + (*th_model->jit_model) = torch::jit::load(ctx->model_filename); + th_model->jit_model->to(device); + } catch (const c10::Error& e) { + av_log(ctx, AV_LOG_ERROR, "Failed to load torch model\n"); + goto fail; + } + } + th_model->request_queue = ff_safe_queue_create(); if (!th_model->request_queue) { goto fail; -- 2.54.0 _______________________________________________ ffmpeg-devel mailing list -- [email protected] To unsubscribe send an email to [email protected]
