fix: preserve alpha when upscaling RGBA images with ESRGAN (#2029)

This commit is contained in:
leejet 2026-09-23 02:23:10 +08:00 committed by GitHub
parent 28b454bda1
commit c92d73c408
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 24 additions and 2 deletions

View File

@ -2,6 +2,8 @@
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. 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.
- Specify the model path using the `--upscale-model PATH` parameter. example: - Specify the model path using the `--upscale-model PATH` parameter. example:
```bash ```bash

View File

@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path,
sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) { sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_tensor) {
sd::ParallelScope tensor_scope(&tensor_executor); 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");
return {};
}
const bool has_alpha = input_tensor.shape()[2] == 4;
sd::Tensor<float> rgb;
if (has_alpha) {
rgb = sd::ops::slice(input_tensor, 2, 0, 3);
}
const sd::Tensor<float>& model_input = has_alpha ? rgb : input_tensor;
sd::Tensor<float> upscaled; sd::Tensor<float> upscaled;
const int scale = esrgan_upscaler->config.scale; const int scale = esrgan_upscaler->config.scale;
if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) { if (tile_size <= 0 || (input_tensor.shape()[0] <= tile_size && input_tensor.shape()[1] <= tile_size)) {
upscaled = esrgan_upscaler->compute(n_threads, input_tensor); upscaled = esrgan_upscaler->compute(n_threads, model_input);
} else { } else {
auto on_processing = [&](const sd::Tensor<float>& input_tile) -> sd::Tensor<float> { 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 = esrgan_upscaler->compute(n_threads, input_tile);
@ -125,7 +137,7 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
return output_tile; return output_tile;
}; };
upscaled = process_tiles_2d(input_tensor, upscaled = process_tiles_2d(model_input,
static_cast<int>(input_tensor.shape()[0] * scale), static_cast<int>(input_tensor.shape()[0] * scale),
static_cast<int>(input_tensor.shape()[1] * scale), static_cast<int>(input_tensor.shape()[1] * scale),
scale, scale,
@ -141,6 +153,14 @@ sd::Tensor<float> UpscalerGGML::upscale_tensor(const sd::Tensor<float>& input_te
LOG_ERROR("esrgan compute failed"); LOG_ERROR("esrgan compute failed");
return {}; return {};
} }
if (has_alpha) {
auto alpha = sd::ops::slice(input_tensor, 2, 3, 4);
auto alpha_shape = alpha.shape();
alpha_shape[0] = upscaled.shape()[0];
alpha_shape[1] = upscaled.shape()[1];
alpha = sd::ops::interpolate(alpha, alpha_shape, sd::ops::InterpolateMode::Bilinear);
upscaled = sd::ops::concat(upscaled, alpha, 2);
}
return upscaled; return upscaled;
} }