trending/stable-diffusion.cpp
1
0
Fork 0
mirror of https://github.com/leejet/stable-diffusion.cpp.git synced 2026-10-27 07:14:04 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat: add MoSR upscaler support (#2128)

Browse Source
This commit is contained in:
Gabriele Lozupone
2026-10-12 22:38:38 +08:00
committed by GitHub
parent 2988060a43
commit 792bcd344c
4 changed files with 411 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/upscaler.cpp
+66 -24
View File
@@ -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 ||