mirror of
https://github.com/sd-webui/stable-diffusion-webui.git
synced 2024-12-14 14:52:31 +03:00
9d95449624
* Metadata cleanup - Maintain metadata within UI This commit, when combined with Gradio 3.2.1b1+, maintains image metadata as an image is passed throughout the UI. For example, if you generate an image, send it to Image Lab, upscale it, fix faces, and then drag the resulting image back in to Image Lab, it will still remember the image generation parameters. When the image is saved, the metadata will be stripped from it if save-metadata is not enabled. If the image is saved by *dragging* out of the UI on to the filesystem it may maintain its metadata. Note: I have ran into UI responsiveness issues with upgrading Gradio. Seems there may be some Gradio queue management issues. *Without* the gradio update this commit will maintain current functionality, but will not keep meetadata when dragging an image between UI components. * Move ImageMetadata into its own file Cleans up webui, enables webui_streamlit et al to use it as well. * Fix typo
58 lines
2.0 KiB
Python
58 lines
2.0 KiB
Python
''' 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
|