add back png info in image api

This commit is contained in:
evshiron 2022-11-02 22:37:45 +08:00
parent 172c4bc09f
commit e21fcd72fc

View File

@ -7,8 +7,9 @@ from fastapi import APIRouter, Depends, HTTPException
import modules.shared as shared import modules.shared as shared
from modules.api.models import * from modules.api.models import *
from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images
from modules.sd_samplers import all_samplers, sample_to_image, samples_to_image_grid from modules.sd_samplers import all_samplers
from modules.extras import run_extras, run_pnginfo from modules.extras import run_extras, run_pnginfo
from PIL import PngImagePlugin
def upscaler_to_index(name: str): def upscaler_to_index(name: str):
@ -31,9 +32,21 @@ def setUpscalers(req: dict):
def encode_pil_to_base64(image): def encode_pil_to_base64(image):
buffer = io.BytesIO() with io.BytesIO() as output_bytes:
image.save(buffer, format="png")
return base64.b64encode(buffer.getvalue()) # Copy any text-only metadata
use_metadata = False
metadata = PngImagePlugin.PngInfo()
for key, value in image.info.items():
if isinstance(key, str) and isinstance(value, str):
metadata.add_text(key, value)
use_metadata = True
image.save(
output_bytes, "PNG", pnginfo=(metadata if use_metadata else None)
)
bytes_data = output_bytes.getvalue()
return base64.b64encode(bytes_data)
class Api: class Api: