Commit 792bcd3 for stable-diffusion.cpp
commit 792bcd344c3fa4d14143b240d5f6cbb324583b18
Author: Gabriele Lozupone <54721912+GabrieleLozupone@users.noreply.github.com>
Date: Mon Oct 12 16:38:38 2026 +0200
feat: add MoSR upscaler support (#2128)
diff --git a/docs/esrgan.md b/docs/esrgan.md
index 708f49b..4068ef1 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 0000000..0aa158c
--- /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 <algorithm>
+#include <cmath>
+#include <map>
+#include <string>
+#include <utility>
+#include <vector>
+
+#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<GGMLBlock>(new MoSRLayerNorm2d(dim));
+ blocks["fc1"] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, hidden * 2, {3, 3}, {1, 1}, {1, 1}));
+ blocks["conv"] = std::shared_ptr<GGMLBlock>(new MoSRDepthwiseConv2d(conv_ch, kernel_size));
+ blocks["fc2"] = std::shared_ptr<GGMLBlock>(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<MoSRLayerNorm2d>(blocks["norm"]);
+ auto fc1 = std::dynamic_pointer_cast<Conv2d>(blocks["fc1"]);
+ auto conv = std::dynamic_pointer_cast<MoSRDepthwiseConv2d>(blocks["conv"]);
+ auto fc2 = std::dynamic_pointer_cast<Conv2d>(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<GGMLBlock>(new Conv2d(in_ch, out_ch, {3, 3}, {1, 1}, {1, 1}));
+ blocks["block.2"] = std::shared_ptr<GGMLBlock>(new Conv2d(out_ch, out_ch, {3, 3}, {1, 1}, {1, 1}));
+ blocks["conv11"] = std::shared_ptr<GGMLBlock>(new Conv2d(in_ch, out_ch, {1, 1}));
+ }
+
+ ggml_tensor* forward(GGMLRunnerContext* ctx, ggml_tensor* x) {
+ auto conv0 = std::dynamic_pointer_cast<Conv2d>(blocks["block.0"]);
+ auto conv2 = std::dynamic_pointer_cast<Conv2d>(blocks["block.2"]);
+ auto conv11 = std::dynamic_pointer_cast<Conv2d>(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<GGMLBlock>(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<GGMLBlock>(new MoSRGatedCNNBlock(dim,
+ this->config.hidden,
+ this->config.conv_ch,
+ this->config.kernel_size));
+ }
+ blocks[gblock_name(nb + 1)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim * 2, {3, 3}, {1, 1}, {1, 1}));
+ blocks[gblock_name(nb + 3)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim * 2, dim, {3, 3}, {1, 1}, {1, 1}));
+ blocks[gblock_name(nb + 5)] = std::shared_ptr<GGMLBlock>(new Conv2d(dim, dim, {1, 1}));
+ blocks["shortcut"] = std::shared_ptr<GGMLBlock>(new MoSRConvBlock(this->config.in_ch, dim));
+ blocks["upsampler.0"] = std::shared_ptr<GGMLBlock>(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<Conv2d>(blocks[gblock_name(0)]);
+ auto conv_mid1 = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 1)]);
+ auto conv_mid2 = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 3)]);
+ auto conv_out = std::dynamic_pointer_cast<Conv2d>(blocks[gblock_name(nb + 5)]);
+ auto shortcut = std::dynamic_pointer_cast<MoSRConvBlock>(blocks["shortcut"]);
+ auto upsampler = std::dynamic_pointer_cast<Conv2d>(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<MoSRGatedCNNBlock>(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<MoSRNet> mosr_net;
+
+ MoSR(ggml_backend_t backend,
+ const String2TensorStorage& tensor_storage_map = {},
+ std::shared_ptr<RunnerWeightManager> 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<MoSRNet>(config);
+ mosr_net->init(params_ctx, tensor_storage_map, "");
+ }
+
+ std::string get_desc() override {
+ return "mosr";
+ }
+
+ void get_param_tensors(std::map<std::string, ggml_tensor*>& tensors) {
+ if (!mosr_net) {
+ return;
+ }
+
+ mosr_net->get_param_tensors(tensors);
+ }
+
+ ggml_cgraph* build_graph(const sd::Tensor<float>& 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<float> compute(const int n_threads,
+ const sd::Tensor<float>& 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 dc3498e..dfe5e4b 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<float> UpscalerGGML::compute(const sd::Tensor<float>& 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<ESRGAN>(backend_for(SDBackendModule::UPSCALER),
- model_loader.get_tensor_storage_map(),
+ const auto& tensor_storage_map = model_loader.get_tensor_storage_map();
+ std::map<std::string, ggml_tensor*> tensors;
+ if (tensor_storage_map.count("gblocks.0.weight") > 0) {
+ mosr_upscaler = std::make_shared<MoSR>(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<ESRGAN>(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<std::string, ggml_tensor*> 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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& 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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
}
const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
sd::Tensor<float> 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<float>& input_tile) -> sd::Tensor<float> {
- 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<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& 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 c23e9d0..b93065f 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<ModelManager> model_manager;
ggml_type model_data_type = GGML_TYPE_F16;
std::shared_ptr<ESRGAN> esrgan_upscaler;
+ std::shared_ptr<MoSR> 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<float> compute(const sd::Tensor<float>& input_tensor);
sd::Tensor<float> upscale_tensor(const sd::Tensor<float>& input_tensor);
sd_image_t upscale(sd_image_t input_image, uint32_t upscale_factor);
};