| @@ -0,0 +1,597 @@ |
| 1 | +/* |
| 2 | + * Copyright (c) 2024 |
| 3 | + * |
| 4 | + * This file is part of FFmpeg. |
| 5 | + * |
| 6 | + * FFmpeg is free software; you can redistribute it and/or |
| 7 | + * modify it under the terms of the GNU Lesser General Public |
| 8 | + * License as published by the Free Software Foundation; either |
| 9 | + * version 2.1 of the License, or (at your option) any later version. |
| 10 | + * |
| 11 | + * FFmpeg is distributed in the hope that it will be useful, |
| 12 | + * but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 13 | + * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU |
| 14 | + * Lesser General Public License for more details. |
| 15 | + * |
| 16 | + * You should have received a copy of the GNU Lesser General Public |
| 17 | + * License along with FFmpeg; if not, write to the Free Software |
| 18 | + * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA |
| 19 | + */ |
| 20 | + |
| 21 | +/** |
| 22 | + * @file |
| 23 | + * DNN Torch backend implementation. |
| 24 | + */ |
| 25 | + |
| 26 | +#include <torch/torch.h> |
| 27 | +#include <torch/script.h> |
| 28 | + |
| 29 | +extern "C" { |
| 30 | +#include "../internal.h" |
| 31 | +#include "dnn_io_proc.h" |
| 32 | +#include "dnn_backend_common.h" |
| 33 | +#include "libavutil/opt.h" |
| 34 | +#include "queue.h" |
| 35 | +#include "safe_queue.h" |
| 36 | +} |
| 37 | + |
| 38 | +typedef struct THOptions{ |
| 39 | + char *device_name; |
| 40 | + int optimize; |
| 41 | +} THOptions; |
| 42 | + |
| 43 | +typedef struct THContext { |
| 44 | + const AVClass *c_class; |
| 45 | + THOptions options; |
| 46 | +} THContext; |
| 47 | + |
| 48 | +typedef struct THModel { |
| 49 | + THContext ctx; |
| 50 | + DNNModel *model; |
| 51 | + torch::jit::Module *jit_model; |
| 52 | + SafeQueue *request_queue; |
| 53 | + Queue *task_queue; |
| 54 | + Queue *lltask_queue; |
| 55 | +} THModel; |
| 56 | + |
| 57 | +typedef struct THInferRequest { |
| 58 | + torch::Tensor *output; |
| 59 | + torch::Tensor *input_tensor; |
| 60 | +} THInferRequest; |
| 61 | + |
| 62 | +typedef struct THRequestItem { |
| 63 | + THInferRequest *infer_request; |
| 64 | + LastLevelTaskItem *lltask; |
| 65 | + DNNAsyncExecModule exec_module; |
| 66 | +} THRequestItem; |
| 67 | + |
| 68 | + |
| 69 | +#define OFFSET(x) offsetof(THContext, x) |
| 70 | +#define FLAGS AV_OPT_FLAG_FILTERING_PARAM |
| 71 | +static const AVOption dnn_th_options[] = { |
| 72 | + { "device", "device to run model", OFFSET(options.device_name), AV_OPT_TYPE_STRING, { .str = "cpu" }, 0, 0, FLAGS }, |
| 73 | + { "optimize", "turn on graph executor optimization", OFFSET(options.optimize), AV_OPT_TYPE_INT, { .i64 = 0 }, 0, 1, FLAGS}, |
| 74 | + { NULL } |
| 75 | +}; |
| 76 | + |
| 77 | +AVFILTER_DEFINE_CLASS(dnn_th); |
| 78 | + |
| 79 | +static int extract_lltask_from_task(TaskItem *task, Queue *lltask_queue) |
| 80 | +{ |
| 81 | + THModel *th_model = (THModel *)task->model; |
| 82 | + THContext *ctx = &th_model->ctx; |
| 83 | + LastLevelTaskItem *lltask = (LastLevelTaskItem *)av_malloc(sizeof(*lltask)); |
| 84 | + if (!lltask) { |
| 85 | + av_log(ctx, AV_LOG_ERROR, "Failed to allocate memory for LastLevelTaskItem\n"); |
| 86 | + return AVERROR(ENOMEM); |
| 87 | + } |
| 88 | + task->inference_todo = 1; |
| 89 | + task->inference_done = 0; |
| 90 | + lltask->task = task; |
| 91 | + if (ff_queue_push_back(lltask_queue, lltask) < 0) { |
| 92 | + av_log(ctx, AV_LOG_ERROR, "Failed to push back lltask_queue.\n"); |
| 93 | + av_freep(&lltask); |
| 94 | + return AVERROR(ENOMEM); |
| 95 | + } |
| 96 | + return 0; |
| 97 | +} |
| 98 | + |
| 99 | +static void th_free_request(THInferRequest *request) |
| 100 | +{ |
| 101 | + if (!request) |
| 102 | + return; |
| 103 | + if (request->output) { |
| 104 | + delete(request->output); |
| 105 | + request->output = NULL; |
| 106 | + } |
| 107 | + if (request->input_tensor) { |
| 108 | + delete(request->input_tensor); |
| 109 | + request->input_tensor = NULL; |
| 110 | + } |
| 111 | + return; |
| 112 | +} |
| 113 | + |
| 114 | +static inline void destroy_request_item(THRequestItem **arg) |
| 115 | +{ |
| 116 | + THRequestItem *item; |
| 117 | + if (!arg || !*arg) { |
| 118 | + return; |
| 119 | + } |
| 120 | + item = *arg; |
| 121 | + th_free_request(item->infer_request); |
| 122 | + av_freep(&item->infer_request); |
| 123 | + av_freep(&item->lltask); |
| 124 | + ff_dnn_async_module_cleanup(&item->exec_module); |
| 125 | + av_freep(arg); |
| 126 | +} |
| 127 | + |
| 128 | +static void dnn_free_model_th(DNNModel **model) |
| 129 | +{ |
| 130 | + THModel *th_model; |
| 131 | + if (!model || !*model) |
| 132 | + return; |
| 133 | + |
| 134 | + th_model = (THModel *) (*model)->model; |
| 135 | + while (ff_safe_queue_size(th_model->request_queue) != 0) { |
| 136 | + THRequestItem *item = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue); |
| 137 | + destroy_request_item(&item); |
| 138 | + } |
| 139 | + ff_safe_queue_destroy(th_model->request_queue); |
| 140 | + |
| 141 | + while (ff_queue_size(th_model->lltask_queue) != 0) { |
| 142 | + LastLevelTaskItem *item = (LastLevelTaskItem *)ff_queue_pop_front(th_model->lltask_queue); |
| 143 | + av_freep(&item); |
| 144 | + } |
| 145 | + ff_queue_destroy(th_model->lltask_queue); |
| 146 | + |
| 147 | + while (ff_queue_size(th_model->task_queue) != 0) { |
| 148 | + TaskItem *item = (TaskItem *)ff_queue_pop_front(th_model->task_queue); |
| 149 | + av_frame_free(&item->in_frame); |
| 150 | + av_frame_free(&item->out_frame); |
| 151 | + av_freep(&item); |
| 152 | + } |
| 153 | + ff_queue_destroy(th_model->task_queue); |
| 154 | + delete th_model->jit_model; |
| 155 | + av_opt_free(&th_model->ctx); |
| 156 | + av_freep(&th_model); |
| 157 | + av_freep(model); |
| 158 | +} |
| 159 | + |
| 160 | +static int get_input_th(void *model, DNNData *input, const char *input_name) |
| 161 | +{ |
| 162 | + input->dt = DNN_FLOAT; |
| 163 | + input->order = DCO_RGB; |
| 164 | + input->layout = DL_NCHW; |
| 165 | + input->dims[0] = 1; |
| 166 | + input->dims[1] = 3; |
| 167 | + input->dims[2] = -1; |
| 168 | + input->dims[3] = -1; |
| 169 | + return 0; |
| 170 | +} |
| 171 | + |
| 172 | +static void deleter(void *arg) |
| 173 | +{ |
| 174 | + av_freep(&arg); |
| 175 | +} |
| 176 | + |
| 177 | +static int fill_model_input_th(THModel *th_model, THRequestItem *request) |
| 178 | +{ |
| 179 | + LastLevelTaskItem *lltask = NULL; |
| 180 | + TaskItem *task = NULL; |
| 181 | + THInferRequest *infer_request = NULL; |
| 182 | + DNNData input = { 0 }; |
| 183 | + THContext *ctx = &th_model->ctx; |
| 184 | + int ret, width_idx, height_idx, channel_idx; |
| 185 | + |
| 186 | + lltask = (LastLevelTaskItem *)ff_queue_pop_front(th_model->lltask_queue); |
| 187 | + if (!lltask) { |
| 188 | + ret = AVERROR(EINVAL); |
| 189 | + goto err; |
| 190 | + } |
| 191 | + request->lltask = lltask; |
| 192 | + task = lltask->task; |
| 193 | + infer_request = request->infer_request; |
| 194 | + |
| 195 | + ret = get_input_th(th_model, &input, NULL); |
| 196 | + if ( ret != 0) { |
| 197 | + goto err; |
| 198 | + } |
| 199 | + width_idx = dnn_get_width_idx_by_layout(input.layout); |
| 200 | + height_idx = dnn_get_height_idx_by_layout(input.layout); |
| 201 | + channel_idx = dnn_get_channel_idx_by_layout(input.layout); |
| 202 | + input.dims[height_idx] = task->in_frame->height; |
| 203 | + input.dims[width_idx] = task->in_frame->width; |
| 204 | + input.data = av_malloc(input.dims[height_idx] * input.dims[width_idx] * |
| 205 | + input.dims[channel_idx] * sizeof(float)); |
| 206 | + if (!input.data) |
| 207 | + return AVERROR(ENOMEM); |
| 208 | + infer_request->input_tensor = new torch::Tensor(); |
| 209 | + infer_request->output = new torch::Tensor(); |
| 210 | + |
| 211 | + switch (th_model->model->func_type) { |
| 212 | + case DFT_PROCESS_FRAME: |
| 213 | + input.scale = 255; |
| 214 | + if (task->do_ioproc) { |
| 215 | + if (th_model->model->frame_pre_proc != NULL) { |
| 216 | + th_model->model->frame_pre_proc(task->in_frame, &input, th_model->model->filter_ctx); |
| 217 | + } else { |
| 218 | + ff_proc_from_frame_to_dnn(task->in_frame, &input, ctx); |
| 219 | + } |
| 220 | + } |
| 221 | + break; |
| 222 | + default: |
| 223 | + avpriv_report_missing_feature(NULL, "model function type %d", th_model->model->func_type); |
| 224 | + break; |
| 225 | + } |
| 226 | + *infer_request->input_tensor = torch::from_blob(input.data, |
| 227 | + {1, input.dims[channel_idx], input.dims[height_idx], input.dims[width_idx]}, |
| 228 | + deleter, torch::kFloat32); |
| 229 | + return 0; |
| 230 | + |
| 231 | +err: |
| 232 | + th_free_request(infer_request); |
| 233 | + return ret; |
| 234 | +} |
| 235 | + |
| 236 | +static int th_start_inference(void *args) |
| 237 | +{ |
| 238 | + THRequestItem *request = (THRequestItem *)args; |
| 239 | + THInferRequest *infer_request = NULL; |
| 240 | + LastLevelTaskItem *lltask = NULL; |
| 241 | + TaskItem *task = NULL; |
| 242 | + THModel *th_model = NULL; |
| 243 | + THContext *ctx = NULL; |
| 244 | + std::vector<torch::jit::IValue> inputs; |
| 245 | + torch::NoGradGuard no_grad; |
| 246 | + |
| 247 | + if (!request) { |
| 248 | + av_log(NULL, AV_LOG_ERROR, "THRequestItem is NULL\n"); |
| 249 | + return AVERROR(EINVAL); |
| 250 | + } |
| 251 | + infer_request = request->infer_request; |
| 252 | + lltask = request->lltask; |
| 253 | + task = lltask->task; |
| 254 | + th_model = (THModel *)task->model; |
| 255 | + ctx = &th_model->ctx; |
| 256 | + |
| 257 | + if (ctx->options.optimize) |
| 258 | + torch::jit::setGraphExecutorOptimize(true); |
| 259 | + else |
| 260 | + torch::jit::setGraphExecutorOptimize(false); |
| 261 | + |
| 262 | + if (!infer_request->input_tensor || !infer_request->output) { |
| 263 | + av_log(ctx, AV_LOG_ERROR, "input or output tensor is NULL\n"); |
| 264 | + return DNN_GENERIC_ERROR; |
| 265 | + } |
| 266 | + inputs.push_back(*infer_request->input_tensor); |
| 267 | + |
| 268 | + *infer_request->output = th_model->jit_model->forward(inputs).toTensor(); |
| 269 | + |
| 270 | + return 0; |
| 271 | +} |
| 272 | + |
| 273 | +static void infer_completion_callback(void *args) { |
| 274 | + THRequestItem *request = (THRequestItem*)args; |
| 275 | + LastLevelTaskItem *lltask = request->lltask; |
| 276 | + TaskItem *task = lltask->task; |
| 277 | + DNNData outputs = { 0 }; |
| 278 | + THInferRequest *infer_request = request->infer_request; |
| 279 | + THModel *th_model = (THModel *)task->model; |
| 280 | + torch::Tensor *output = infer_request->output; |
| 281 | + |
| 282 | + c10::IntArrayRef sizes = output->sizes(); |
| 283 | + outputs.order = DCO_RGB; |
| 284 | + outputs.layout = DL_NCHW; |
| 285 | + outputs.dt = DNN_FLOAT; |
| 286 | + if (sizes.size() == 4) { |
| 287 | + // 4 dimensions: [batch_size, channel, height, width] |
| 288 | + // this format of data is normally used for video frame SR |
| 289 | + outputs.dims[0] = sizes.at(0); // N |
| 290 | + outputs.dims[1] = sizes.at(1); // C |
| 291 | + outputs.dims[2] = sizes.at(2); // H |
| 292 | + outputs.dims[3] = sizes.at(3); // W |
| 293 | + } else { |
| 294 | + avpriv_report_missing_feature(&th_model->ctx, "Support of this kind of model"); |
| 295 | + goto err; |
| 296 | + } |
| 297 | + |
| 298 | + switch (th_model->model->func_type) { |
| 299 | + case DFT_PROCESS_FRAME: |
| 300 | + if (task->do_ioproc) { |
| 301 | + outputs.scale = 255; |
| 302 | + outputs.data = output->data_ptr(); |
| 303 | + if (th_model->model->frame_post_proc != NULL) { |
| 304 | + th_model->model->frame_post_proc(task->out_frame, &outputs, th_model->model->filter_ctx); |
| 305 | + } else { |
| 306 | + ff_proc_from_dnn_to_frame(task->out_frame, &outputs, &th_model->ctx); |
| 307 | + } |
| 308 | + } else { |
| 309 | + task->out_frame->width = outputs.dims[dnn_get_width_idx_by_layout(outputs.layout)]; |
| 310 | + task->out_frame->height = outputs.dims[dnn_get_height_idx_by_layout(outputs.layout)]; |
| 311 | + } |
| 312 | + break; |
| 313 | + default: |
| 314 | + avpriv_report_missing_feature(&th_model->ctx, "model function type %d", th_model->model->func_type); |
| 315 | + goto err; |
| 316 | + } |
| 317 | + task->inference_done++; |
| 318 | + av_freep(&request->lltask); |
| 319 | +err: |
| 320 | + th_free_request(infer_request); |
| 321 | + |
| 322 | + if (ff_safe_queue_push_back(th_model->request_queue, request) < 0) { |
| 323 | + destroy_request_item(&request); |
| 324 | + av_log(&th_model->ctx, AV_LOG_ERROR, "Unable to push back request_queue when failed to start inference.\n"); |
| 325 | + } |
| 326 | +} |
| 327 | + |
| 328 | +static int execute_model_th(THRequestItem *request, Queue *lltask_queue) |
| 329 | +{ |
| 330 | + THModel *th_model = NULL; |
| 331 | + LastLevelTaskItem *lltask; |
| 332 | + TaskItem *task = NULL; |
| 333 | + int ret = 0; |
| 334 | + |
| 335 | + if (ff_queue_size(lltask_queue) == 0) { |
| 336 | + destroy_request_item(&request); |
| 337 | + return 0; |
| 338 | + } |
| 339 | + |
| 340 | + lltask = (LastLevelTaskItem *)ff_queue_peek_front(lltask_queue); |
| 341 | + if (lltask == NULL) { |
| 342 | + av_log(NULL, AV_LOG_ERROR, "Failed to get LastLevelTaskItem\n"); |
| 343 | + ret = AVERROR(EINVAL); |
| 344 | + goto err; |
| 345 | + } |
| 346 | + task = lltask->task; |
| 347 | + th_model = (THModel *)task->model; |
| 348 | + |
| 349 | + ret = fill_model_input_th(th_model, request); |
| 350 | + if ( ret != 0) { |
| 351 | + goto err; |
| 352 | + } |
| 353 | + if (task->async) { |
| 354 | + avpriv_report_missing_feature(&th_model->ctx, "LibTorch async"); |
| 355 | + } else { |
| 356 | + ret = th_start_inference((void *)(request)); |
| 357 | + if (ret != 0) { |
| 358 | + goto err; |
| 359 | + } |
| 360 | + infer_completion_callback(request); |
| 361 | + return (task->inference_done == task->inference_todo) ? 0 : DNN_GENERIC_ERROR; |
| 362 | + } |
| 363 | + |
| 364 | +err: |
| 365 | + th_free_request(request->infer_request); |
| 366 | + if (ff_safe_queue_push_back(th_model->request_queue, request) < 0) { |
| 367 | + destroy_request_item(&request); |
| 368 | + } |
| 369 | + return ret; |
| 370 | +} |
| 371 | + |
| 372 | +static int get_output_th(void *model, const char *input_name, int input_width, int input_height, |
| 373 | + const char *output_name, int *output_width, int *output_height) |
| 374 | +{ |
| 375 | + int ret = 0; |
| 376 | + THModel *th_model = (THModel*) model; |
| 377 | + THContext *ctx = &th_model->ctx; |
| 378 | + TaskItem task = { 0 }; |
| 379 | + THRequestItem *request = NULL; |
| 380 | + DNNExecBaseParams exec_params = { |
| 381 | + .input_name = input_name, |
| 382 | + .output_names = &output_name, |
| 383 | + .nb_output = 1, |
| 384 | + .in_frame = NULL, |
| 385 | + .out_frame = NULL, |
| 386 | + }; |
| 387 | + ret = ff_dnn_fill_gettingoutput_task(&task, &exec_params, th_model, input_height, input_width, ctx); |
| 388 | + if ( ret != 0) { |
| 389 | + goto err; |
| 390 | + } |
| 391 | + |
| 392 | + ret = extract_lltask_from_task(&task, th_model->lltask_queue); |
| 393 | + if ( ret != 0) { |
| 394 | + av_log(ctx, AV_LOG_ERROR, "unable to extract last level task from task.\n"); |
| 395 | + goto err; |
| 396 | + } |
| 397 | + |
| 398 | + request = (THRequestItem*) ff_safe_queue_pop_front(th_model->request_queue); |
| 399 | + if (!request) { |
| 400 | + av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n"); |
| 401 | + ret = AVERROR(EINVAL); |
| 402 | + goto err; |
| 403 | + } |
| 404 | + |
| 405 | + ret = execute_model_th(request, th_model->lltask_queue); |
| 406 | + *output_width = task.out_frame->width; |
| 407 | + *output_height = task.out_frame->height; |
| 408 | + |
| 409 | +err: |
| 410 | + av_frame_free(&task.out_frame); |
| 411 | + av_frame_free(&task.in_frame); |
| 412 | + return ret; |
| 413 | +} |
| 414 | + |
| 415 | +static THInferRequest *th_create_inference_request(void) |
| 416 | +{ |
| 417 | + THInferRequest *request = (THInferRequest *)av_malloc(sizeof(THInferRequest)); |
| 418 | + if (!request) { |
| 419 | + return NULL; |
| 420 | + } |
| 421 | + request->input_tensor = NULL; |
| 422 | + request->output = NULL; |
| 423 | + return request; |
| 424 | +} |
| 425 | + |
| 426 | +static DNNModel *dnn_load_model_th(const char *model_filename, DNNFunctionType func_type, const char *options, AVFilterContext *filter_ctx) |
| 427 | +{ |
| 428 | + DNNModel *model = NULL; |
| 429 | + THModel *th_model = NULL; |
| 430 | + THRequestItem *item = NULL; |
| 431 | + THContext *ctx; |
| 432 | + |
| 433 | + model = (DNNModel *)av_mallocz(sizeof(DNNModel)); |
| 434 | + if (!model) { |
| 435 | + return NULL; |
| 436 | + } |
| 437 | + |
| 438 | + th_model = (THModel *)av_mallocz(sizeof(THModel)); |
| 439 | + if (!th_model) { |
| 440 | + av_freep(&model); |
| 441 | + return NULL; |
| 442 | + } |
| 443 | + th_model->model = model; |
| 444 | + model->model = th_model; |
| 445 | + th_model->ctx.c_class = &dnn_th_class; |
| 446 | + ctx = &th_model->ctx; |
| 447 | + //parse options |
| 448 | + av_opt_set_defaults(ctx); |
| 449 | + if (av_opt_set_from_string(ctx, options, NULL, "=", "&") < 0) { |
| 450 | + av_log(ctx, AV_LOG_ERROR, "Failed to parse options \"%s\"\n", options); |
| 451 | + return NULL; |
| 452 | + } |
| 453 | + |
| 454 | + c10::Device device = c10::Device(ctx->options.device_name); |
| 455 | + if (!device.is_cpu()) { |
| 456 | + av_log(ctx, AV_LOG_ERROR, "Not supported device:\"%s\"\n", ctx->options.device_name); |
| 457 | + goto fail; |
| 458 | + } |
| 459 | + |
| 460 | + try { |
| 461 | + th_model->jit_model = new torch::jit::Module; |
| 462 | + (*th_model->jit_model) = torch::jit::load(model_filename); |
| 463 | + } catch (const c10::Error& e) { |
| 464 | + av_log(ctx, AV_LOG_ERROR, "Failed to load torch model\n"); |
| 465 | + goto fail; |
| 466 | + } |
| 467 | + |
| 468 | + th_model->request_queue = ff_safe_queue_create(); |
| 469 | + if (!th_model->request_queue) { |
| 470 | + goto fail; |
| 471 | + } |
| 472 | + |
| 473 | + item = (THRequestItem *)av_mallocz(sizeof(THRequestItem)); |
| 474 | + if (!item) { |
| 475 | + goto fail; |
| 476 | + } |
| 477 | + item->lltask = NULL; |
| 478 | + item->infer_request = th_create_inference_request(); |
| 479 | + if (!item->infer_request) { |
| 480 | + av_log(NULL, AV_LOG_ERROR, "Failed to allocate memory for Torch inference request\n"); |
| 481 | + goto fail; |
| 482 | + } |
| 483 | + item->exec_module.start_inference = &th_start_inference; |
| 484 | + item->exec_module.callback = &infer_completion_callback; |
| 485 | + item->exec_module.args = item; |
| 486 | + |
| 487 | + if (ff_safe_queue_push_back(th_model->request_queue, item) < 0) { |
| 488 | + goto fail; |
| 489 | + } |
| 490 | + item = NULL; |
| 491 | + |
| 492 | + th_model->task_queue = ff_queue_create(); |
| 493 | + if (!th_model->task_queue) { |
| 494 | + goto fail; |
| 495 | + } |
| 496 | + |
| 497 | + th_model->lltask_queue = ff_queue_create(); |
| 498 | + if (!th_model->lltask_queue) { |
| 499 | + goto fail; |
| 500 | + } |
| 501 | + |
| 502 | + model->get_input = &get_input_th; |
| 503 | + model->get_output = &get_output_th; |
| 504 | + model->options = NULL; |
| 505 | + model->filter_ctx = filter_ctx; |
| 506 | + model->func_type = func_type; |
| 507 | + return model; |
| 508 | + |
| 509 | +fail: |
| 510 | + if (item) { |
| 511 | + destroy_request_item(&item); |
| 512 | + av_freep(&item); |
| 513 | + } |
| 514 | + dnn_free_model_th(&model); |
| 515 | + return NULL; |
| 516 | +} |
| 517 | + |
| 518 | +static int dnn_execute_model_th(const DNNModel *model, DNNExecBaseParams *exec_params) |
| 519 | +{ |
| 520 | + THModel *th_model = (THModel *)model->model; |
| 521 | + THContext *ctx = &th_model->ctx; |
| 522 | + TaskItem *task; |
| 523 | + THRequestItem *request; |
| 524 | + int ret = 0; |
| 525 | + |
| 526 | + ret = ff_check_exec_params(ctx, DNN_TH, model->func_type, exec_params); |
| 527 | + if (ret != 0) { |
| 528 | + av_log(ctx, AV_LOG_ERROR, "exec parameter checking fail.\n"); |
| 529 | + return ret; |
| 530 | + } |
| 531 | + |
| 532 | + task = (TaskItem *)av_malloc(sizeof(TaskItem)); |
| 533 | + if (!task) { |
| 534 | + av_log(ctx, AV_LOG_ERROR, "unable to alloc memory for task item.\n"); |
| 535 | + return AVERROR(ENOMEM); |
| 536 | + } |
| 537 | + |
| 538 | + ret = ff_dnn_fill_task(task, exec_params, th_model, 0, 1); |
| 539 | + if (ret != 0) { |
| 540 | + av_freep(&task); |
| 541 | + av_log(ctx, AV_LOG_ERROR, "unable to fill task.\n"); |
| 542 | + return ret; |
| 543 | + } |
| 544 | + |
| 545 | + ret = ff_queue_push_back(th_model->task_queue, task); |
| 546 | + if (ret < 0) { |
| 547 | + av_freep(&task); |
| 548 | + av_log(ctx, AV_LOG_ERROR, "unable to push back task_queue.\n"); |
| 549 | + return ret; |
| 550 | + } |
| 551 | + |
| 552 | + ret = extract_lltask_from_task(task, th_model->lltask_queue); |
| 553 | + if (ret != 0) { |
| 554 | + av_log(ctx, AV_LOG_ERROR, "unable to extract last level task from task.\n"); |
| 555 | + return ret; |
| 556 | + } |
| 557 | + |
| 558 | + request = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue); |
| 559 | + if (!request) { |
| 560 | + av_log(ctx, AV_LOG_ERROR, "unable to get infer request.\n"); |
| 561 | + return AVERROR(EINVAL); |
| 562 | + } |
| 563 | + |
| 564 | + return execute_model_th(request, th_model->lltask_queue); |
| 565 | +} |
| 566 | + |
| 567 | +static DNNAsyncStatusType dnn_get_result_th(const DNNModel *model, AVFrame **in, AVFrame **out) |
| 568 | +{ |
| 569 | + THModel *th_model = (THModel *)model->model; |
| 570 | + return ff_dnn_get_result_common(th_model->task_queue, in, out); |
| 571 | +} |
| 572 | + |
| 573 | +static int dnn_flush_th(const DNNModel *model) |
| 574 | +{ |
| 575 | + THModel *th_model = (THModel *)model->model; |
| 576 | + THRequestItem *request; |
| 577 | + |
| 578 | + if (ff_queue_size(th_model->lltask_queue) == 0) |
| 579 | + // no pending task need to flush |
| 580 | + return 0; |
| 581 | + |
| 582 | + request = (THRequestItem *)ff_safe_queue_pop_front(th_model->request_queue); |
| 583 | + if (!request) { |
| 584 | + av_log(&th_model->ctx, AV_LOG_ERROR, "unable to get infer request.\n"); |
| 585 | + return AVERROR(EINVAL); |
| 586 | + } |
| 587 | + |
| 588 | + return execute_model_th(request, th_model->lltask_queue); |
| 589 | +} |
| 590 | + |
| 591 | +extern const DNNModule ff_dnn_backend_torch = { |
| 592 | + .load_model = dnn_load_model_th, |
| 593 | + .execute_model = dnn_execute_model_th, |
| 594 | + .get_result = dnn_get_result_th, |
| 595 | + .flush = dnn_flush_th, |
| 596 | + .free_model = dnn_free_model_th, |
| 597 | +}; |