Skip to content

Commit e21fcd7

Browse files
committed
add back png info in image api
1 parent 172c4bc commit e21fcd7

File tree

1 file changed

+17
-4
lines changed

1 file changed

+17
-4
lines changed

modules/api/api.py

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,9 @@
77
import modules.shared as shared
88
from modules.api.models import *
99
from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
10-
from modules.sd_samplers import all_samplers, sample_to_image, samples_to_image_grid
10+
from modules.sd_samplers import all_samplers
1111
from modules.extras import run_extras, run_pnginfo
12+
from PIL import PngImagePlugin
1213

1314

1415
def upscaler_to_index(name: str):
@@ -31,9 +32,21 @@ def setUpscalers(req: dict):
3132

3233

3334
def encode_pil_to_base64(image):
34-
buffer = io.BytesIO()
35-
image.save(buffer, format="png")
36-
return base64.b64encode(buffer.getvalue())
35+
with io.BytesIO() as output_bytes:
36+
37+
# Copy any text-only metadata
38+
use_metadata = False
39+
metadata = PngImagePlugin.PngInfo()
40+
for key, value in image.info.items():
41+
if isinstance(key, str) and isinstance(value, str):
42+
metadata.add_text(key, value)
43+
use_metadata = True
44+
45+
image.save(
46+
output_bytes, "PNG", pnginfo=(metadata if use_metadata else None)
47+
)
48+
bytes_data = output_bytes.getvalue()
49+
return base64.b64encode(bytes_data)
3750

3851

3952
class Api:

0 commit comments

Comments
 (0)