feat: add MoSR upscaler support (#2128)

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__

src/upscaler.cpp

@@ -25,15 +25,40 @@ 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 (esrgan_upscaler) {
-        esrgan_upscaler->set_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) {
+        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,
                                   int n_threads) {
     sd::ParallelScope tensor_scope(&tensor_executor);
@@ -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 ||

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);
 };