diff --git a/docs/esrgan.md b/docs/esrgan.md index 708f49bad7..4068ef11c7 100644 --- a/docs/esrgan.md +++ b/docs/esrgan.md @@ -2,7 +2,9 @@ You can use ESRGAN—such as the model [RealESRGAN_x4plus_anime_6B.pth](https://github.com/xinntao/Real-ESRGAN/releases/download/v0.2.2.4/RealESRGAN_x4plus_anime_6B.pth)—to upscale the generated images and improve their overall resolution and clarity. -RGBA images, including Qwen Image 2.1 output, keep their alpha channel during model upscaling and hires fix. ESRGAN processes the RGB channels; the alpha channel is resized with bilinear interpolation and recombined with the upscaled image. +MoSR models with the PixelShuffle upsampler, such as [4xNomos2_hq_mosr.pth](https://github.com/Phhofm/models/releases/download/4xNomos2_hq_mosr/4xNomos2_hq_mosr.pth), are also supported. The architecture is detected from the weights, so they are used the same way as ESRGAN models. + +RGBA images, including Qwen Image 2.1 output, keep their alpha channel during model upscaling and hires fix. The upscaler processes the RGB channels; the alpha channel is resized with bilinear interpolation and recombined with the upscaled image. - Specify the model path using the `--upscale-model PATH` parameter. example: diff --git a/src/model/upscaler/mosr.hpp b/src/model/upscaler/mosr.hpp new file mode 100644 index 0000000000..0aa158c81c --- /dev/null +++ b/src/model/upscaler/mosr.hpp @@ -0,0 +1,337 @@ +#ifndef __SD_MODEL_UPSCALER_MOSR_HPP__ +#define __SD_MODEL_UPSCALER_MOSR_HPP__ + +#include +#include +#include +#include +#include +#include + +#include "core/ggml_extend.h" +#include "core/ggml_runner.h" +#include "core/util.h" +#include "model/common/ggml_block.hpp" + +/* + =================================== MoSR =================================== + References: + https://github.com/umzi2/MoSR + https://github.com/chaiNNer-org/spandrel/blob/main/libs/spandrel/spandrel/architectures/MoSR/__arch/mosr_arch.py + + Only the PixelShuffle ("ps") upsampler is supported. +*/ + +struct MoSRConfig { + int scale = 4; + int n_block = 24; + int in_ch = 3; + int dim = 64; + int hidden = 96; + int conv_ch = 64; + int kernel_size = 7; + + bool is_valid() const { + // The upscaler feeds RGB tiles, so only 3-channel models are usable. + return in_ch == 3 && scale > 0 && n_block > 0 && conv_ch <= hidden; + } + + static MoSRConfig detect_from_weights(const String2TensorStorage& tensor_storage_map, + const std::string& prefix = "") { + MoSRConfig config; + auto find_weight = [&](const std::string& suffix) -> const TensorStorage* { + std::string name = prefix.empty() ? suffix : prefix + "." + suffix; + auto iter = tensor_storage_map.find(name); + if (iter == tensor_storage_map.end() || iter->second.n_dims != 4) { + return nullptr; + } + return &iter->second; + }; + + const TensorStorage* first = find_weight("gblocks.0.weight"); + const TensorStorage* fc1 = find_weight("gblocks.1.fc1.weight"); + const TensorStorage* conv = find_weight("gblocks.1.conv.weight"); + const TensorStorage* up = find_weight("upsampler.0.weight"); + if (first == nullptr || fc1 == nullptr || conv == nullptr || up == nullptr) { + // Missing tensors, or the unsupported "gps"/"dys" upsamplers. + config.n_block = 0; + return config; + } + + int seq_len = 0; + const std::string gblocks_prefix = prefix.empty() ? "gblocks." : prefix + ".gblocks."; + for (const auto& [name, _] : tensor_storage_map) { + if (!starts_with(name, gblocks_prefix)) { + continue; + } + size_t pos = name.find('.', gblocks_prefix.size()); + if (pos == std::string::npos) { + continue; + } + try { + int idx = std::stoi(name.substr(gblocks_prefix.size(), pos - gblocks_prefix.size())); + seq_len = std::max(seq_len, idx + 1); + } catch (...) { + } + } + + // gblocks = [conv_in, n_block x GatedCNNBlock, conv, Mish, conv, Mish, conv_1x1] + config.n_block = seq_len - 6; + config.in_ch = (int)first->ne[2]; + config.dim = (int)first->ne[3]; + config.hidden = (int)fc1->ne[3] / 2; + config.conv_ch = (int)conv->ne[3]; + config.kernel_size = (int)conv->ne[0]; + config.scale = (int)std::lround(std::sqrt((double)up->ne[3] / config.in_ch)); + if (config.scale * config.scale * config.in_ch != up->ne[3]) { + config.scale = 0; + } + + LOG_VERBOSE("mosr: scale = %d, n_block = %d, in_ch = %d, dim = %d, hidden = %d, conv_ch = %d, kernel_size = %d", + config.scale, + config.n_block, + config.in_ch, + config.dim, + config.hidden, + config.conv_ch, + config.kernel_size); + return config; + } +}; + +static inline ggml_tensor* mosr_mish(ggml_context* ctx, ggml_tensor* x) { + return ggml_mul(ctx, x, ggml_tanh(ctx, ggml_softplus(ctx, x))); +} + +class MoSRLayerNorm2d : public UnaryBlock { +protected: + int64_t channels; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + params["weight"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + } + +public: + explicit MoSRLayerNorm2d(int64_t channels) + : channels(channels) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + // [N, C, H, W] -> [N, H, W, C] so the norm reduces over channels. + x = ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 1, 2, 0, 3)); + x = ggml_ext_layer_norm(ctx->ggml_ctx, x, params["weight"], params["bias"], 1e-6f); + return ggml_ext_cont(ctx->ggml_ctx, ggml_permute(ctx->ggml_ctx, x, 2, 0, 1, 3)); + } +}; + +class MoSRDepthwiseConv2d : public UnaryBlock { +protected: + int64_t channels; + int kernel_size; + + void init_params(ggml_context* ctx, const String2TensorStorage& tensor_storage_map = {}, const std::string prefix = "") override { + // The direct depthwise kernels of the CUDA and Vulkan backends need an F32 kernel. + params["weight"] = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, kernel_size, kernel_size, 1, channels); + params["bias"] = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, channels); + } + +public: + MoSRDepthwiseConv2d(int64_t channels, int kernel_size) + : channels(channels), kernel_size(kernel_size) {} + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) override { + const int pad = kernel_size / 2; + x = ggml_conv_2d_dw_direct(ctx->ggml_ctx, params["weight"], ggml_ext_cont(ctx->ggml_ctx, x), 1, 1, pad, pad, 1, 1); + auto b = ggml_reshape_4d(ctx->ggml_ctx, params["bias"], 1, 1, channels, 1); + return ggml_add_inplace(ctx->ggml_ctx, x, b); + } +}; + +class MoSRGatedCNNBlock : public GGMLBlock { +protected: + int hidden; + int conv_ch; + +public: + MoSRGatedCNNBlock(int dim, int hidden, int conv_ch, int kernel_size) + : hidden(hidden), conv_ch(conv_ch) { + blocks["norm"] = std::shared_ptr(new MoSRLayerNorm2d(dim)); + blocks["fc1"] = std::shared_ptr(new Conv2d(dim, hidden * 2, {3, 3}, {1, 1}, {1, 1})); + blocks["conv"] = std::shared_ptr(new MoSRDepthwiseConv2d(conv_ch, kernel_size)); + blocks["fc2"] = std::shared_ptr(new Conv2d(hidden, dim, {3, 3}, {1, 1}, {1, 1})); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [n, dim, h, w] + // return: [n, dim, h, w] + auto norm = std::dynamic_pointer_cast(blocks["norm"]); + auto fc1 = std::dynamic_pointer_cast(blocks["fc1"]); + auto conv = std::dynamic_pointer_cast(blocks["conv"]); + auto fc2 = std::dynamic_pointer_cast(blocks["fc2"]); + + auto h = fc1->forward(ctx, norm->forward(ctx, x)); + + auto channel_view = [&](int64_t offset, int64_t count) { + return ggml_view_4d(ctx->ggml_ctx, h, h->ne[0], h->ne[1], count, h->ne[3], + h->nb[1], h->nb[2], h->nb[3], offset * h->nb[2]); + }; + auto g = channel_view(0, hidden); + auto i = channel_view(hidden, hidden - conv_ch); + auto c = conv->forward(ctx, channel_view(2 * hidden - conv_ch, conv_ch)); + if (hidden > conv_ch) { + c = ggml_concat(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, i), c, 2); + } + + h = ggml_mul(ctx->ggml_ctx, c, mosr_mish(ctx->ggml_ctx, ggml_ext_cont(ctx->ggml_ctx, g))); + h = mosr_mish(ctx->ggml_ctx, fc2->forward(ctx, h)); + return ggml_add(ctx->ggml_ctx, h, ggml_scale_bias(ctx->ggml_ctx, x, 1.f, -0.5f)); + } +}; + +class MoSRConvBlock : public GGMLBlock { +public: + MoSRConvBlock(int in_ch, int out_ch) { + blocks["block.0"] = std::shared_ptr(new Conv2d(in_ch, out_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["block.2"] = std::shared_ptr(new Conv2d(out_ch, out_ch, {3, 3}, {1, 1}, {1, 1})); + blocks["conv11"] = std::shared_ptr(new Conv2d(in_ch, out_ch, {1, 1})); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + auto conv0 = std::dynamic_pointer_cast(blocks["block.0"]); + auto conv2 = std::dynamic_pointer_cast(blocks["block.2"]); + auto conv11 = std::dynamic_pointer_cast(blocks["conv11"]); + + auto out = mosr_mish(ctx->ggml_ctx, conv0->forward(ctx, x)); + out = mosr_mish(ctx->ggml_ctx, conv2->forward(ctx, out)); + return ggml_add(ctx->ggml_ctx, out, conv11->forward(ctx, x)); + } +}; + +class MoSRNet : public GGMLBlock { +protected: + MoSRConfig config; + + static std::string gblock_name(int idx) { + return "gblocks." + std::to_string(idx); + } + + // torch.nn.PixelShuffle: [n, c*r*r, h, w] -> [n, c, h*r, w*r] + static ggml_tensor* pixel_shuffle(ggml_context* ctx, ggml_tensor* x, int r) { + const int64_t w = x->ne[0]; + const int64_t h = x->ne[1]; + const int64_t c = x->ne[2] / (r * r); + const int64_t n = x->ne[3]; + x = ggml_reshape_4d(ctx, x, w, h, r, r * c * n); // [n*c*r, r, h, w] + x = ggml_cont(ctx, ggml_permute(ctx, x, 1, 2, 0, 3)); // [n*c*r, h, w, r] + x = ggml_reshape_4d(ctx, x, r * w, h, r, c * n); // [n*c, r, h, w*r] + x = ggml_cont(ctx, ggml_permute(ctx, x, 0, 2, 1, 3)); // [n*c, h, r, w*r] + return ggml_reshape_4d(ctx, x, r * w, r * h, c, n); + } + +public: + explicit MoSRNet(MoSRConfig config) + : config(std::move(config)) { + const int dim = this->config.dim; + const int nb = this->config.n_block; + + blocks[gblock_name(0)] = std::shared_ptr(new Conv2d(this->config.in_ch, dim, {3, 3}, {1, 1}, {1, 1})); + for (int i = 1; i <= nb; i++) { + blocks[gblock_name(i)] = std::shared_ptr(new MoSRGatedCNNBlock(dim, + this->config.hidden, + this->config.conv_ch, + this->config.kernel_size)); + } + blocks[gblock_name(nb + 1)] = std::shared_ptr(new Conv2d(dim, dim * 2, {3, 3}, {1, 1}, {1, 1})); + blocks[gblock_name(nb + 3)] = std::shared_ptr(new Conv2d(dim * 2, dim, {3, 3}, {1, 1}, {1, 1})); + blocks[gblock_name(nb + 5)] = std::shared_ptr(new Conv2d(dim, dim, {1, 1})); + blocks["shortcut"] = std::shared_ptr(new MoSRConvBlock(this->config.in_ch, dim)); + blocks["upsampler.0"] = std::shared_ptr(new Conv2d(dim, + this->config.in_ch * this->config.scale * this->config.scale, + {3, 3}, + {1, 1}, + {1, 1})); + } + + ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) { + // x: [n, in_ch, h, w] + // return: [n, in_ch, h*scale, w*scale] + const int nb = config.n_block; + auto conv_in = std::dynamic_pointer_cast(blocks[gblock_name(0)]); + auto conv_mid1 = std::dynamic_pointer_cast(blocks[gblock_name(nb + 1)]); + auto conv_mid2 = std::dynamic_pointer_cast(blocks[gblock_name(nb + 3)]); + auto conv_out = std::dynamic_pointer_cast(blocks[gblock_name(nb + 5)]); + auto shortcut = std::dynamic_pointer_cast(blocks["shortcut"]); + auto upsampler = std::dynamic_pointer_cast(blocks["upsampler.0"]); + + auto skip = shortcut->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(skip, "mosr.shortcut", "skip"); + + auto feat = conv_in->forward(ctx, x); + sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.prelude", "feat"); + for (int i = 1; i <= nb; i++) { + auto block = std::dynamic_pointer_cast(blocks[gblock_name(i)]); + feat = block->forward(ctx, feat); + sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.gblocks." + std::to_string(i), "feat"); + } + feat = mosr_mish(ctx->ggml_ctx, conv_mid1->forward(ctx, feat)); + feat = mosr_mish(ctx->ggml_ctx, conv_mid2->forward(ctx, feat)); + feat = conv_out->forward(ctx, feat); + feat = ggml_add(ctx->ggml_ctx, feat, ggml_scale_bias(ctx->ggml_ctx, skip, 1.f, -0.5f)); + sd::ggml_graph_cut::mark_graph_cut(feat, "mosr.body.out", "feat"); + + auto out = pixel_shuffle(ctx->ggml_ctx, upsampler->forward(ctx, feat), config.scale); + sd::ggml_graph_cut::mark_graph_cut(out, "mosr.final", "out"); + return out; + } +}; + +struct MoSR : public GGMLRunner { + MoSRConfig config; + std::unique_ptr mosr_net; + + MoSR(ggml_backend_t backend, + const String2TensorStorage& tensor_storage_map = {}, + std::shared_ptr weight_manager = nullptr) + : GGMLRunner(backend, weight_manager), + config(MoSRConfig::detect_from_weights(tensor_storage_map)) { + if (!config.is_valid()) { + return; + } + mosr_net = std::make_unique(config); + mosr_net->init(params_ctx, tensor_storage_map, ""); + } + + std::string get_desc() override { + return "mosr"; + } + + void get_param_tensors(std::map& tensors) { + if (!mosr_net) { + return; + } + + mosr_net->get_param_tensors(tensors); + } + + ggml_cgraph* build_graph(const sd::Tensor& x_tensor) { + if (!mosr_net) + return nullptr; + constexpr int kGraphNodes = 1 << 16; // 65k + ggml_cgraph* gf = new_graph_custom(kGraphNodes); + ggml_tensor* x = make_input(x_tensor); + + auto runner_ctx = get_context(); + ggml_tensor* out = mosr_net->forward(&runner_ctx, x); + ggml_build_forward_expand(gf, out); + return gf; + } + + sd::Tensor compute(const int n_threads, + const sd::Tensor& x) { + auto get_graph = [&]() -> ggml_cgraph* { return build_graph(x); }; + auto result = restore_trailing_singleton_dims(GGMLRunner::compute(get_graph, n_threads, false), x.dim()); + return result; + } +}; + +#endif // __SD_MODEL_UPSCALER_MOSR_HPP__ diff --git a/src/upscaler.cpp b/src/upscaler.cpp index dc3498e646..dfe5e4b22f 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -25,13 +25,38 @@ UpscalerGGML::~UpscalerGGML() { // ModelManager holds raw ggml tensor pointers owned by the runner context. model_manager.reset(); esrgan_upscaler.reset(); + mosr_upscaler.reset(); } void UpscalerGGML::set_max_graph_vram_bytes(size_t max_vram_bytes) { max_graph_vram_bytes = max_vram_bytes; + if (runner() != nullptr) { + runner()->set_max_graph_vram_bytes(max_vram_bytes); + } +} + +GGMLRunner* UpscalerGGML::runner() { + if (mosr_upscaler) { + return mosr_upscaler.get(); + } + return esrgan_upscaler.get(); +} + +int UpscalerGGML::get_scale() { + if (mosr_upscaler) { + return mosr_upscaler->config.scale; + } if (esrgan_upscaler) { - esrgan_upscaler->set_max_graph_vram_bytes(max_vram_bytes); + return esrgan_upscaler->config.scale; + } + return 1; +} + +sd::Tensor UpscalerGGML::compute(const sd::Tensor& input_tensor) { + if (mosr_upscaler) { + return mosr_upscaler->compute(n_threads, input_tensor); } + return esrgan_upscaler->compute(n_threads, input_tensor); } bool UpscalerGGML::load_from_file(const std::string& esrgan_path, @@ -82,20 +107,33 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, } model_loader.set_wtype_override(model_data_type); LOG_INFO("Upscaler weight type: %s", ggml_type_name(model_data_type)); - esrgan_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), - model_loader.get_tensor_storage_map(), + const auto& tensor_storage_map = model_loader.get_tensor_storage_map(); + std::map tensors; + if (tensor_storage_map.count("gblocks.0.weight") > 0) { + mosr_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), + tensor_storage_map, model_manager); - if (esrgan_upscaler == nullptr || esrgan_upscaler->rrdb_net == nullptr) { - LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); - return false; + if (mosr_upscaler->mosr_net == nullptr) { + LOG_ERROR("init mosr model from metadata failed: '%s'", esrgan_path.c_str()); + return false; + } + mosr_upscaler->get_param_tensors(tensors); + } else { + esrgan_upscaler = std::make_shared(backend_for(SDBackendModule::UPSCALER), + tensor_storage_map, + model_manager); + if (esrgan_upscaler->rrdb_net == nullptr) { + LOG_ERROR("init esrgan model from metadata failed: '%s'", esrgan_path.c_str()); + return false; + } + esrgan_upscaler->get_param_tensors(tensors); } - esrgan_upscaler->set_max_graph_vram_bytes(max_graph_vram_bytes); + LOG_INFO("Upscaler architecture: %s", runner()->get_desc().c_str()); + runner()->set_max_graph_vram_bytes(max_graph_vram_bytes); if (direct) { - esrgan_upscaler->set_conv2d_direct_enabled(true); + runner()->set_conv2d_direct_enabled(true); } - std::map tensors; - esrgan_upscaler->get_param_tensors(tensors); if (!model_manager->set_loader(model_loader) || !model_manager->register_param_tensors(ModelComponent::Upscaler, std::move(tensors), @@ -103,7 +141,7 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, backend_for(SDBackendModule::UPSCALER), params_backend_for(SDBackendModule::UPSCALER)) || !model_manager->validate_registered_tensors()) { - LOG_ERROR("register esrgan tensors with model manager failed"); + LOG_ERROR("register upscaler tensors with model manager failed"); return false; } return true; @@ -113,7 +151,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te sd::ParallelScope tensor_scope(&tensor_executor); if (input_tensor.empty() || input_tensor.dim() != 4 || (input_tensor.shape()[2] != 3 && input_tensor.shape()[2] != 4)) { - LOG_ERROR("esrgan expects a 4D RGB or RGBA image tensor"); + LOG_ERROR("upscaler expects a 4D RGB or RGBA image tensor"); return {}; } @@ -124,14 +162,14 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te } const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; sd::Tensor upscaled; - const int scale = esrgan_upscaler->config.scale; + const int scale = get_scale(); if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { - upscaled = esrgan_upscaler->compute(n_threads, model_input); + upscaled = compute(model_input); } else { auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { - auto output_tile = esrgan_upscaler->compute(n_threads, input_tile); + auto output_tile = compute(input_tile); if (output_tile.empty()) { - LOG_ERROR("esrgan compute failed while processing a tile"); + LOG_ERROR("upscaler compute failed while processing a tile"); return {}; } return output_tile; @@ -148,9 +186,9 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te false, on_processing); } - esrgan_upscaler->runner_end(); + runner()->runner_end(); if (upscaled.empty()) { - LOG_ERROR("esrgan compute failed"); + LOG_ERROR("upscaler compute failed"); return {}; } if (has_alpha) { @@ -168,7 +206,7 @@ sd_image_t UpscalerGGML::upscale(sd_image_t input_image, uint32_t upscale_factor sd::ParallelScope tensor_scope(&tensor_executor); // upscale_factor, unused for RealESRGAN_x4plus_anime_6B.pth sd_image_t upscaled_image = {0, 0, 0, nullptr}; - const int scale = esrgan_upscaler->config.scale; + const int scale = get_scale(); int output_width = (int)input_image.width * scale; int output_height = (int)input_image.height * scale; LOG_INFO("upscaling from (%i x %i) to (%i x %i)", @@ -256,10 +294,10 @@ bool upscale(upscaler_ctx_t* upscaler_ctx, } int get_upscale_factor(upscaler_ctx_t* upscaler_ctx) { - if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr || upscaler_ctx->upscaler->esrgan_upscaler == nullptr) { + if (upscaler_ctx == nullptr || upscaler_ctx->upscaler == nullptr) { return 1; } - return upscaler_ctx->upscaler->esrgan_upscaler->config.scale; + return upscaler_ctx->upscaler->get_scale(); } int get_upscaler_model_scale(const char* model_path) { @@ -272,8 +310,12 @@ int get_upscaler_model_scale(const char* model_path) { return 0; } const auto& tensors = loader.get_tensor_storage_map(); - auto first = tensors.find("conv_first.weight"); - auto last = tensors.find("conv_last.weight"); + if (tensors.count("gblocks.0.weight") > 0) { + MoSRConfig config = MoSRConfig::detect_from_weights(tensors); + return config.is_valid() ? config.scale : 0; + } + auto first = tensors.find("conv_first.weight"); + auto last = tensors.find("conv_last.weight"); if (first == tensors.end() || last == tensors.end() || tensors.count("body.0.rdb1.conv1.weight") == 0 || first->second.n_dims != 4 || last->second.n_dims != 4 || diff --git a/src/upscaler.h b/src/upscaler.h index c23e9d0f25..b93065fcc9 100644 --- a/src/upscaler.h +++ b/src/upscaler.h @@ -4,6 +4,7 @@ #include "core/ggml_extend_backend.h" #include "core/tensor.hpp" #include "model/upscaler/esrgan.hpp" +#include "model/upscaler/mosr.hpp" #include "model_manager.h" #include "stable-diffusion.h" @@ -15,6 +16,7 @@ struct UpscalerGGML { std::shared_ptr model_manager; ggml_type model_data_type = GGML_TYPE_F16; std::shared_ptr esrgan_upscaler; + std::shared_ptr mosr_upscaler; std::string esrgan_path; int n_threads; sd::ParallelExecutor tensor_executor; @@ -34,6 +36,9 @@ struct UpscalerGGML { bool load_from_file(const std::string& esrgan_path, int n_threads); void set_max_graph_vram_bytes(size_t max_vram_bytes); + GGMLRunner* runner(); + int get_scale(); + sd::Tensor compute(const sd::Tensor& input_tensor); sd::Tensor upscale_tensor(const sd::Tensor& input_tensor); sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor); };