diff --git a/docs/esrgan.md b/docs/esrgan.md index 39a97605..708f49ba 100644 --- a/docs/esrgan.md +++ b/docs/esrgan.md @@ -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. +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: ```bash diff --git a/src/upscaler.cpp b/src/upscaler.cpp index 341c76ec..3c72dbf3 100644 --- a/src/upscaler.cpp +++ b/src/upscaler.cpp @@ -111,10 +111,22 @@ bool UpscalerGGML::load_from_file(const std::string& esrgan_path, sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_tensor) { 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 rgb; + if (has_alpha) { + rgb = sd::ops::slice(input_tensor, 2, 0, 3); + } + const sd::Tensor& model_input = has_alpha ? rgb : input_tensor; sd::Tensor upscaled; const int scale = esrgan_upscaler->config.scale; 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 { auto on_processing = [&](const sd::Tensor& input_tile) -> sd::Tensor { auto output_tile = esrgan_upscaler->compute(n_threads, input_tile); @@ -125,7 +137,7 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te return output_tile; }; - upscaled = process_tiles_2d(input_tensor, + upscaled = process_tiles_2d(model_input, static_cast(input_tensor.shape()[0] * scale), static_cast(input_tensor.shape()[1] * scale), scale, @@ -141,6 +153,14 @@ sd::Tensor UpscalerGGML::upscale_tensor(const sd::Tensor& input_te LOG_ERROR("esrgan compute failed"); 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; }