feat(server): add generation metadata to png images (#1217)

This commit is contained in:
Wagner Bruna
2026-03-31 12:06:27 -03:00
committed by GitHub
parent 4d5232083f
commit 4fe7a35939
5 changed files with 110 additions and 70 deletions

View File

@@ -220,13 +220,24 @@ std::string extract_and_remove_sd_cpp_extra_args(std::string& text) {
enum class ImageFormat { JPEG,
PNG };
static int stbi_ext_write_png_to_func(stbi_write_func* func, void* context, int x, int y, int comp, const void* data, int stride_bytes, const char* parameters) {
int len;
unsigned char* png = stbi_write_png_to_mem((const unsigned char*)data, stride_bytes, x, y, comp, &len, parameters);
if (png == NULL)
return 0;
func(context, png, len);
STBIW_FREE(png);
return 1;
}
std::vector<uint8_t> write_image_to_vector(
ImageFormat format,
const uint8_t* image,
int width,
int height,
int channels,
int quality = 90) {
std::string params = "",
int quality = 90) {
std::vector<uint8_t> buffer;
auto write_func = [&buffer](void* context, void* data, int size) {
@@ -249,7 +260,7 @@ std::vector<uint8_t> write_image_to_vector(
result = stbi_write_jpg_to_func(c_func, &ctx, width, height, channels, image, quality);
break;
case ImageFormat::PNG:
result = stbi_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels);
result = stbi_ext_write_png_to_func(c_func, &ctx, width, height, channels, image, width * channels, params.size() > 0 ? params.c_str() : nullptr);
break;
default:
throw std::runtime_error("invalid image format");
@@ -497,11 +508,15 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
if (results[i].data == nullptr) {
continue;
}
auto image_bytes = write_image_to_vector(output_format == "jpeg" ? ImageFormat::JPEG : ImageFormat::PNG,
std::string params = gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, gen_params, gen_params.seed + i)
: "";
auto image_bytes = write_image_to_vector(output_format == "jpeg" ? ImageFormat::JPEG : ImageFormat::PNG,
results[i].data,
results[i].width,
results[i].height,
results[i].channel,
params,
output_compression);
if (image_bytes.empty()) {
LOG_ERROR("write image to mem failed");
@@ -747,11 +762,15 @@ void register_openai_api_endpoints(httplib::Server& svr, ServerRuntime& rt) {
for (int i = 0; i < num_results; i++) {
if (results[i].data == nullptr)
continue;
auto image_bytes = write_image_to_vector(output_format == "jpeg" ? ImageFormat::JPEG : ImageFormat::PNG,
std::string params = gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, gen_params, gen_params.seed + i)
: "";
auto image_bytes = write_image_to_vector(output_format == "jpeg" ? ImageFormat::JPEG : ImageFormat::PNG,
results[i].data,
results[i].width,
results[i].height,
results[i].channel,
params,
output_compression);
std::string b64 = base64_encode(image_bytes);
json item;
@@ -1062,11 +1081,15 @@ void register_sdapi_endpoints(httplib::Server& svr, ServerRuntime& rt) {
continue;
}
auto image_bytes = write_image_to_vector(ImageFormat::PNG,
results[i].data,
results[i].width,
results[i].height,
results[i].channel);
std::string params = gen_params.embed_image_metadata
? get_image_params(*runtime->ctx_params, gen_params, gen_params.seed + i)
: "";
auto image_bytes = write_image_to_vector(ImageFormat::PNG,
results[i].data,
results[i].width,
results[i].height,
results[i].channel,
params);
if (image_bytes.empty()) {
LOG_ERROR("write image to mem failed");