mirror of
https://github.com/sd-webui/stable-diffusion-webui.git
synced 2024-12-14 06:35:14 +03:00
78 lines
2.8 KiB
Python
78 lines
2.8 KiB
Python
# This file is part of sygil-webui (https://github.com/Sygil-Dev/sygil-webui/).
|
|
|
|
# Copyright 2022 Sygil-Dev team.
|
|
# This program is free software: you can redistribute it and/or modify
|
|
# it under the terms of the GNU Affero General Public License as published by
|
|
# the Free Software Foundation, either version 3 of the License, or
|
|
# (at your option) any later version.
|
|
|
|
# This program is distributed in the hope that it will be useful,
|
|
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
# GNU Affero General Public License for more details.
|
|
|
|
# You should have received a copy of the GNU Affero General Public License
|
|
# along with this program. If not, see <http://www.gnu.org/licenses/>.
|
|
""" Class to store image generation parameters to be stored as metadata in the image"""
|
|
from __future__ import annotations
|
|
from dataclasses import dataclass, asdict
|
|
from typing import Dict, Optional
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import copy
|
|
|
|
|
|
@dataclass
|
|
class ImageMetadata:
|
|
prompt: str = None
|
|
seed: str = None
|
|
width: str = None
|
|
height: str = None
|
|
steps: str = None
|
|
cfg_scale: str = None
|
|
normalize_prompt_weights: str = None
|
|
denoising_strength: str = None
|
|
GFPGAN: str = None
|
|
|
|
def as_png_info(self) -> PngInfo:
|
|
info = PngInfo()
|
|
for key, value in self.as_dict().items():
|
|
info.add_text(key, value)
|
|
return info
|
|
|
|
def as_dict(self) -> Dict[str, str]:
|
|
return {
|
|
f"SD:{key}": str(value)
|
|
for key, value in asdict(self).items()
|
|
if value is not None
|
|
}
|
|
|
|
@classmethod
|
|
def set_on_image(cls, image: Image, metadata: ImageMetadata) -> None:
|
|
"""Sets metadata on image, in both text form and as an ImageMetadata object"""
|
|
if metadata:
|
|
image.info = metadata.as_dict()
|
|
else:
|
|
metadata = ImageMetadata()
|
|
image.info["ImageMetadata"] = copy.copy(metadata)
|
|
|
|
@classmethod
|
|
def get_from_image(cls, image: Image) -> Optional[ImageMetadata]:
|
|
"""Gets metadata from an image, first looking for an ImageMetadata,
|
|
then if not found tries to construct one from the info"""
|
|
metadata = image.info.get("ImageMetadata", None)
|
|
if not metadata:
|
|
found_metadata = False
|
|
metadata = ImageMetadata()
|
|
for key, value in image.info.items():
|
|
if key.lower().startswith("sd:"):
|
|
key = key[3:]
|
|
if f"{key}" in metadata.__dict__:
|
|
metadata.__dict__[key] = value
|
|
found_metadata = True
|
|
if not found_metadata:
|
|
metadata = None
|
|
if not metadata:
|
|
print("Couldn't find metadata on image")
|
|
return metadata
|