diff options
Diffstat (limited to 'modules/api')
| -rw-r--r-- | modules/api/api.py | 35 | ||||
| -rw-r--r-- | modules/api/models.py | 32 |
2 files changed, 40 insertions, 27 deletions
diff --git a/modules/api/api.py b/modules/api/api.py index 8a7ab2f5..112000b8 100644 --- a/modules/api/api.py +++ b/modules/api/api.py @@ -10,6 +10,7 @@ from modules.api.models import * from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img, process_images from modules.sd_samplers import all_samplers from modules.extras import run_extras, run_pnginfo +from PIL import PngImagePlugin from modules.sd_models import checkpoints_list from modules.realesrgan_model import get_realesrgan_models from typing import List @@ -34,9 +35,21 @@ def setUpscalers(req: dict): def encode_pil_to_base64(image): - buffer = io.BytesIO() - image.save(buffer, format="png") - return base64.b64encode(buffer.getvalue()) + with io.BytesIO() as output_bytes: + + # 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: @@ -205,7 +218,7 @@ class Api: shared.state.interrupt() return {} - + def get_config(self): options = {} for key in shared.opts.data.keys(): @@ -214,9 +227,9 @@ class Api: options.update({key: shared.opts.data.get(key, shared.opts.data_labels.get(key).default)}) else: options.update({key: shared.opts.data.get(key, None)}) - + return options - + def set_config(self, req: OptionsModel): # currently req has all options fields even if you send a dict like { "send_seed": false }, which means it will # overwrite all options with default values. @@ -237,13 +250,13 @@ class Api: def get_upscalers(self): upscalers = [] - + for upscaler in shared.sd_upscalers: u = upscaler.scaler upscalers.append({"name":u.name, "model_name":u.model_name, "model_path":u.model_path, "model_url":u.model_url}) - + return upscalers - + def get_sd_models(self): return [{"title":x.title, "model_name":x.model_name, "hash":x.hash, "filename": x.filename, "config": x.config} for x in checkpoints_list.values()] @@ -255,11 +268,11 @@ class Api: def get_realesrgan_models(self): return [{"name":x.name,"path":x.data_path, "scale":x.scale} for x in get_realesrgan_models(None)] - + def get_promp_styles(self): styleList = [] for k in shared.prompt_styles.styles: - style = shared.prompt_styles.styles[k] + style = shared.prompt_styles.styles[k] styleList.append({"name":style[0], "prompt": style[1], "negative_prompr": style[2]}) return styleList diff --git a/modules/api/models.py b/modules/api/models.py index 2ae75f43..f89da1ff 100644 --- a/modules/api/models.py +++ b/modules/api/models.py @@ -1,11 +1,11 @@ import inspect from pydantic import BaseModel, Field, create_model -from typing import Any, Optional, Union +from typing import Any, Optional from typing_extensions import Literal from inflection import underscore from modules.processing import StableDiffusionProcessingTxt2Img, StableDiffusionProcessingImg2Img from modules.shared import sd_upscalers, opts, parser -from typing import List +from typing import Dict, List API_NOT_ALLOWED = [ "self", @@ -185,22 +185,22 @@ _options = vars(parser)['_option_string_actions'] for key in _options: if(_options[key].dest != 'help'): flag = _options[key] - _type = str - if(_options[key].default != None): _type = type(_options[key].default) + _type = str + if _options[key].default is not None: _type = type(_options[key].default) flags.update({flag.dest: (_type,Field(default=flag.default, description=flag.help))}) FlagsModel = create_model("Flags", **flags) class SamplerItem(BaseModel): name: str = Field(title="Name") - aliases: list[str] = Field(title="Aliases") - options: dict[str, str] = Field(title="Options") + aliases: List[str] = Field(title="Aliases") + options: Dict[str, str] = Field(title="Options") class UpscalerItem(BaseModel): name: str = Field(title="Name") - model_name: str | None = Field(title="Model Name") - model_path: str | None = Field(title="Path") - model_url: str | None = Field(title="URL") + model_name: Optional[str] = Field(title="Model Name") + model_path: Optional[str] = Field(title="Path") + model_url: Optional[str] = Field(title="URL") class SDModelItem(BaseModel): title: str = Field(title="Title") @@ -211,23 +211,23 @@ class SDModelItem(BaseModel): class HypernetworkItem(BaseModel): name: str = Field(title="Name") - path: str | None = Field(title="Path") + path: Optional[str] = Field(title="Path") class FaceRestorerItem(BaseModel): name: str = Field(title="Name") - cmd_dir: str | None = Field(title="Path") + cmd_dir: Optional[str] = Field(title="Path") class RealesrganItem(BaseModel): name: str = Field(title="Name") - path: str | None = Field(title="Path") - scale: int | None = Field(title="Scale") + path: Optional[str] = Field(title="Path") + scale: Optional[int] = Field(title="Scale") class PromptStyleItem(BaseModel): name: str = Field(title="Name") - prompt: str | None = Field(title="Prompt") - negative_prompt: str | None = Field(title="Negative Prompt") + prompt: Optional[str] = Field(title="Prompt") + negative_prompt: Optional[str] = Field(title="Negative Prompt") class ArtistItem(BaseModel): name: str = Field(title="Name") score: float = Field(title="Score") - category: str = Field(title="Category")
\ No newline at end of file + category: str = Field(title="Category") |