feat(rife): add support for frame interpolation and RIFE (#1244)

* feat: add RIFE files and processor/interpolator abstractions * feat: add `rife` as processor option * feat: add frame interpolation math except first frame * feat: complete motion interpolation and add scene detection * feat: improve Vulkan device validation * fix: fix casting issues and variable names * refactor: improve error-checking; add abstractions and factories * refactor: improve readability of the frames processor * docs: update changelog Signed-off-by: k4yt3x <i@k4yt3x.com>
2026-02-14 09:14:53 +08:00 · 2024-12-01 09:55:56 +00:00
parent 2fc89e3883
commit 627f3d84a4
84 changed files with 4914 additions and 615 deletions
--- a/src/interpolator_rife.cpp
+++ b/src/interpolator_rife.cpp
@@ -0,0 +1,121 @@
+#include "interpolator_rife.h"
+
+#include <cstdio>
+#include <filesystem>
+
+#include <spdlog/spdlog.h>
+
+#include "conversions.h"
+#include "fsutils.h"
+
+InterpolatorRIFE::InterpolatorRIFE(
+    int gpuid,
+    bool tta_mode,
+    bool tta_temporal_mode,
+    bool uhd_mode,
+    int num_threads,
+    bool rife_v2,
+    bool rife_v4,
+    const StringType model_name
+)
+    : rife_(nullptr),
+      gpuid_(gpuid),
+      tta_mode_(tta_mode),
+      tta_temporal_mode_(tta_temporal_mode),
+      uhd_mode_(uhd_mode),
+      num_threads_(num_threads),
+      rife_v2_(rife_v2),
+      rife_v4_(rife_v4),
+      model_name_(std::move(model_name)) {}
+
+InterpolatorRIFE::~InterpolatorRIFE() {
+    if (rife_) {
+        delete rife_;
+        rife_ = nullptr;
+    }
+}
+
+int InterpolatorRIFE::init(AVCodecContext *dec_ctx, AVCodecContext *enc_ctx, AVBufferRef *) {
+    // Construct the model directory path using std::filesystem
+    std::filesystem::path model_param_dir;
+
+    // Find the model paths by model name if provided
+    model_param_dir = std::filesystem::path(STR("models")) / STR("rife") / model_name_;
+
+    // Get the full paths using a function that possibly modifies or validates the path
+    std::filesystem::path model_param_full_path = find_resource_file(model_param_dir);
+
+    // Check if the model files exist
+    if (!std::filesystem::exists(model_param_full_path)) {
+        spdlog::error("RIFE model param directory not found: {}", model_param_dir.u8string());
+        return -1;
+    }
+
+    // Create a new RIFE instance
+    rife_ = new RIFE(
+        gpuid_, tta_mode_, tta_temporal_mode_, uhd_mode_, num_threads_, rife_v2_, rife_v4_
+    );
+
+    // Store the time bases
+    in_time_base_ = dec_ctx->time_base;
+    out_time_base_ = enc_ctx->time_base;
+    out_pix_fmt_ = enc_ctx->pix_fmt;
+
+    // Load the model
+    if (rife_->load(model_param_full_path) != 0) {
+        spdlog::error("Failed to load RIFE model");
+        return -1;
+    }
+
+    return 0;
+}
+
+int InterpolatorRIFE::interpolate(
+    AVFrame *prev_frame,
+    AVFrame *in_frame,
+    AVFrame **out_frame,
+    float time_step
+) {
+    int ret;
+
+    ncnn::Mat in_mat1 = avframe_to_ncnn_mat(prev_frame);
+    if (in_mat1.empty()) {
+        spdlog::error("Failed to convert AVFrame to ncnn::Mat");
+        return -1;
+    }
+
+    ncnn::Mat in_mat2 = avframe_to_ncnn_mat(in_frame);
+    if (in_mat2.empty()) {
+        spdlog::error("Failed to convert AVFrame to ncnn::Mat");
+        return -1;
+    }
+
+    // Allocate space for output ncnn::Mat
+    ncnn::Mat out_mat = ncnn::Mat(in_mat2.w, in_mat2.h, static_cast<size_t>(3), 3);
+
+    ret = rife_->process(in_mat1, in_mat2, time_step, out_mat);
+    if (ret != 0) {
+        spdlog::error("RIFE processing failed");
+        return ret;
+    }
+
+    // Convert ncnn::Mat to AVFrame
+    *out_frame = ncnn_mat_to_avframe(out_mat, out_pix_fmt_);
+
+    // Rescale PTS to encoder's time base
+    (*out_frame)->pts = av_rescale_q(in_frame->pts, in_time_base_, out_time_base_);
+
+    // Return the processed frame to the caller
+    return ret;
+}
+
+void InterpolatorRIFE::get_output_dimensions(
+    const ProcessorConfig *,
+    int in_width,
+    int in_height,
+    int &out_width,
+    int &out_height
+) const {
+    out_width = in_width;
+    out_height = in_height;
+}