diff --git a/include/libvideo2x/libvideo2x.h b/include/libvideo2x/libvideo2x.h index 5b315e6..6ae6376 100644 --- a/include/libvideo2x/libvideo2x.h +++ b/include/libvideo2x/libvideo2x.h @@ -58,7 +58,7 @@ struct RealESRGANConfig { #ifdef _WIN32 const wchar_t *model_path; #else - const char *model_path; + const char *model_name; #endif }; diff --git a/include/libvideo2x/realesrgan_filter.h b/include/libvideo2x/realesrgan_filter.h index 88e3695..f7089fa 100644 --- a/include/libvideo2x/realesrgan_filter.h +++ b/include/libvideo2x/realesrgan_filter.h @@ -1,8 +1,6 @@ #ifndef REALSRGAN_FILTER_H #define REALSRGAN_FILTER_H -#include - extern "C" { #include } @@ -17,9 +15,11 @@ class RealesrganFilter : public Filter { int gpuid; bool tta_mode; int scaling_factor; - const std::filesystem::path model_path; - const std::filesystem::path custom_model_param_path; - const std::filesystem::path custom_model_bin_path; +#ifdef _WIN32 + const std::wstring model_name; +#else + const std::string model_name; +#endif AVRational in_time_base; AVRational out_time_base; AVPixelFormat out_pix_fmt; @@ -30,9 +30,11 @@ class RealesrganFilter : public Filter { int gpuid = 0, bool tta_mode = false, int scaling_factor = 4, - const std::filesystem::path model = std::filesystem::path("realesr-animevideov3"), - const std::filesystem::path custom_model_param_path = std::filesystem::path(), - const std::filesystem::path custom_model_bin_path = std::filesystem::path() +#ifdef _WIN32 + const std::wstring model_name = L"realesr-animevideov3" +#else + const std::string model_name = "realesr-animevideov3" +#endif ); // Destructor diff --git a/src/libplacebo_filter.cpp b/src/libplacebo_filter.cpp index c3d9edc..e862a82 100644 --- a/src/libplacebo_filter.cpp +++ b/src/libplacebo_filter.cpp @@ -49,7 +49,7 @@ int LibplaceboFilter::init(AVCodecContext *dec_ctx, AVCodecContext *enc_ctx, AVB // Check if the shader file exists if (!std::filesystem::exists(shader_full_path)) { - spdlog::error("libplacebo shader file not found: {}", shader_full_path.string()); + spdlog::error("libplacebo shader file not found: '{}'", shader_path.string()); return -1; } diff --git a/src/libvideo2x.cpp b/src/libvideo2x.cpp index ee1438f..34d3a84 100644 --- a/src/libvideo2x.cpp +++ b/src/libvideo2x.cpp @@ -439,13 +439,13 @@ extern "C" int process_video( }; } else if (filter_config->filter_type == FILTER_REALESRGAN) { const auto &config = filter_config->config.realesrgan; - if (!config.model_path) { + if (!config.model_name) { spdlog::error("Model name must be provided for the RealESRGAN filter"); cleanup(); return -1; } filter = new RealesrganFilter{ - config.gpuid, config.tta_mode, config.scaling_factor, config.model_path + config.gpuid, config.tta_mode, config.scaling_factor, config.model_name }; } else { spdlog::error("Unknown filter type"); diff --git a/src/realesrgan_filter.cpp b/src/realesrgan_filter.cpp index 2699851..992be23 100644 --- a/src/realesrgan_filter.cpp +++ b/src/realesrgan_filter.cpp @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -13,17 +14,18 @@ RealesrganFilter::RealesrganFilter( int gpuid, bool tta_mode, int scaling_factor, - const std::filesystem::path model_path, - const std::filesystem::path custom_model_param_path, - const std::filesystem::path custom_model_bin_path +#ifdef _WIN32 + const std::wstring model_name +#else + const std::string model_name +#endif ) : realesrgan(nullptr), gpuid(gpuid), tta_mode(tta_mode), scaling_factor(scaling_factor), - model_path(std::move(model_path)), - custom_model_param_path(std::move(custom_model_param_path)), - custom_model_bin_path(std::move(custom_model_bin_path)) {} + model_name(std::move(model_name)) { +} RealesrganFilter::~RealesrganFilter() { if (realesrgan) { @@ -37,22 +39,17 @@ int RealesrganFilter::init(AVCodecContext *dec_ctx, AVCodecContext *enc_ctx, AVB std::filesystem::path model_param_path; std::filesystem::path model_bin_path; - if (!model_path.empty()) { - // Find the model paths by model name if provided - // TODO: ensure this works with wide strings on Windows - model_param_path = std::filesystem::path("models") / "realesrgan" / - (model_path.string() + "-x" + std::to_string(scaling_factor) + ".param"); - model_bin_path = std::filesystem::path("models") / "realesrgan" / - (model_path.string() + "-x" + std::to_string(scaling_factor) + ".bin"); - } else if (!custom_model_param_path.empty() && !custom_model_bin_path.empty()) { - // Use the custom model paths if provided - model_param_path = custom_model_param_path; - model_bin_path = custom_model_bin_path; - } else { - // Neither model name nor custom model paths provided - spdlog::error("Model or model paths must be provided for RealESRGAN filter"); - return -1; - } +#ifdef _WIN32 + std::wstring param_file_name = model_name + L"-x" + std::to_wstring(scaling_factor) + L".param"; + std::wstring bin_file_name = model_name + L"-x" + std::to_wstring(scaling_factor) + L".bin"; +#else + std::string param_file_name = model_name + "-x" + std::to_string(scaling_factor) + ".param"; + std::string bin_file_name = model_name + "-x" + std::to_string(scaling_factor) + ".bin"; +#endif + + // Find the model paths by model name if provided + model_param_path = std::filesystem::path("models") / "realesrgan" / param_file_name; + model_bin_path = std::filesystem::path("models") / "realesrgan" / bin_file_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_path); @@ -60,11 +57,11 @@ int RealesrganFilter::init(AVCodecContext *dec_ctx, AVCodecContext *enc_ctx, AVB // Check if the model files exist if (!std::filesystem::exists(model_param_full_path)) { - spdlog::error("RealESRGAN model param file not found: {}", model_param_full_path.string()); + spdlog::error("RealESRGAN model param file not found: {}", model_param_path.string()); return -1; } if (!std::filesystem::exists(model_bin_full_path)) { - spdlog::error("RealESRGAN model bin file not found: {}", model_bin_full_path.string()); + spdlog::error("RealESRGAN model bin file not found: {}", model_bin_path.string()); return -1; }