diff --git a/docs/source/api_doc/sd/metadata.rst b/docs/source/api_doc/sd/metadata.rst index ec1ba361dc..975ddb71cb 100644 --- a/docs/source/api_doc/sd/metadata.rst +++ b/docs/source/api_doc/sd/metadata.rst @@ -10,7 +10,7 @@ SDMetaData ------------------------- .. autoclass:: SDMetaData - :members: __str__, pnginfo + :members: __str__, text, pnginfo diff --git a/imgutils/sd/metadata.py b/imgutils/sd/metadata.py index f77e74c0b7..ab58ce44cf 100644 --- a/imgutils/sd/metadata.py +++ b/imgutils/sd/metadata.py @@ -84,6 +84,15 @@ def __str__(self): return self._sdmeta_text() def _sdmeta_text(self): + """ + Generate a formatted string representation of the metadata. + + This internal method is used by __str__ and other methods to create a consistent + string representation of the metadata. + + :return: A formatted string containing the metadata. + :rtype: str + """ with io.StringIO() as sio: print(self.prompt, file=sio) if self.neg_prompt: @@ -111,6 +120,26 @@ def _sdmeta_text(self): @property def text(self) -> str: + """ + Get the metadata as a formatted string. + + This property provides a convenient way to access the string representation + of the metadata without calling _sdmeta_text() directly. + + :return: A formatted string containing the metadata. + :rtype: str + + Example: + >>> metadata = SDMetaData( + ... prompt="A starry night", + ... neg_prompt="Daylight", + ... parameters={"Steps": 40, "Sampler": "Euler", "CFG scale": 8} + ... ) + >>> print(metadata.text) + A starry night + Negative prompt: Daylight + Steps: 40, Sampler: Euler, CFG scale: 8 + """ return self._sdmeta_text() @property @@ -242,27 +271,30 @@ def parse_sdmeta_from_text(x: str) -> SDMetaData: def get_sdmeta_from_image(image: ImageTyping) -> Optional[SDMetaData]: """ - Get metadata from a PNG image. + Extract and parse Stable Diffusion metadata from an image. + + This function attempts to read SD metadata from various sources within the image, + including PNG info, EXIF data, and GIF metadata. If found, it parses the metadata + into an SDMetaData object. - :param image: The input image. + :param image: The input image, which can be a file path, URL, or PIL Image object. :type image: ImageTyping - :return: A SDMetaData object containing the metadata if available, else None. + + :return: An SDMetaData object containing the parsed metadata if available, else None. :rtype: Optional[SDMetaData] - Examples:: + :raises: Various exceptions may be raised by the underlying image loading and + metadata reading functions. + + Example usage: >>> from imgutils.sd import get_sdmeta_from_image - >>> - >>> sd1 = get_sdmeta_from_image('sd_metadata_simple.png') - >>> sd1 - SDMetaData(prompt='(extremely delicate and beautiful), best quality, official art, global illumination, soft shadow, super detailed, Japanese light novel cover, 4K, metal_texture, (striped_background), super detailed background, more detailed, rich detailed, extremely detailed CG unity 8k wallpaper, ((unreal)), sci-fi,(fantasy),(masterpiece),(super delicate), (illustration), (extremely delicate and beautiful), anime coloring,\\n(silver_skin), ((high-cut silver_impossible_bodysuit), ((gem_on_chest)),(high-cut_silver_mechanical_leotard)),headgear,\\n(focus-on:1.1),(1_girl),((solo)),slim_waist,white hair, long hair, luminous yellow eyes,(medium_breast:1.2), (Indistinct_cameltoe:0.9), (flat_crotch:1.1),(coquettish), (squint:1.4),(evil_smile :1.35),(dark_persona), [open mouth: 0.7], standing,[wet:0.7],\\nslim_face, tall_girl,(mature),mature_face, (slim_figure), (slim_legs:1.1), (groin:1.1), ((bare_thighs)),', neg_prompt='EasyNegative, sketch, duplicate, ugly, huge eyes, text, logo, monochrome, worst face, (bad and mutated hands:1.3), (worst quality:2.0), (low quality:2.0), (blurry:2.0), horror, geometry, bad_prompt, (bad hands), (missing fingers), multiple limbs, bad anatomy, (interlocked fingers:1.2), Ugly Fingers, (extra digit and hands and fingers and legs and arms:1.4), ((2girl)), (deformed fingers:1.2), (long fingers:1.2),(bad-artist-anime), bad-artist, bad hand, blush, (lipstick),skindentation, tie, ((big_breast)), (nipple), thighhighs, pubic_hair, pussy, black and white,(3d), ((realistic)),blurry,nipple slip, (nipple), blush, head_out_of_frame,curvy,', parameters={'Steps': 20, 'Sampler': 'DDIM', 'CFG scale': 7, 'Seed': 3827064803, 'Size': (512, 848), 'Model hash': 'eb49192009', 'Model': 'AniDosMix', 'Clip skip': 2}) - >>> type(sd1) - - >>> - >>> sd2 = get_sdmeta_from_image('sd_metadata_complex.png') - >>> sd2 - SDMetaData(prompt='1girl, solo, blue eyes, black footwear, white hair, looking at viewer, shoes, full body, standing, bangs, indoors, wide sleeves, ahoge, dress, closed mouth, blush, long sleeves, potted plant, bag, plant, hair bun, window,,BlueArchive,', neg_prompt='Neg1,Negative,', parameters={'Steps': 20, 'Sampler': 'DPM++ 2M SDE Karras', 'CFG scale': 7, 'Seed': 2647703743, 'Size': (768, 768), 'Model hash': '72bd94132e', 'Model': 'CuteMix', 'Denoising strength': 0.7, 'ControlNet 0': 'preprocessor: openpose, model: control_v11p_sd15_openpose [cab727d4], weight: 1, starting/ending: (0, 1), resize mode: Crop and Resize, pixel perfect: False, control mode: Balanced, preprocessor params: (512, 64, 64)', 'Hires upscale': 2, 'Hires upscaler': 'Latent', 'TI hashes': 'Neg1: 339cc9210f70, Negative: 66a7279a88dd', 'Version': 'v1.5.1'}) - >>> type(sd2) - + >>> sd_meta = get_sdmeta_from_image('path/to/image.png') + >>> if sd_meta: + ... print(f"Prompt: {sd_meta.prompt}") + ... print(f"Negative prompt: {sd_meta.neg_prompt}") + ... print(f"Parameters: {sd_meta.parameters}") + ... else: + ... print("No SD metadata found in the image.") """ image = load_image(image, mode=None, force_background=None) pnginfo_text = (read_geninfo_parameters(image) or @@ -275,14 +307,38 @@ def get_sdmeta_from_image(image: ImageTyping) -> Optional[SDMetaData]: def _save_png_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs): + """ + Internal function to save a PNG image with SD metadata. + + :param image: The PIL Image object to save. + :param dst_file: The destination file path. + :param metadata: The SDMetaData object containing the metadata to save. + :param kwargs: Additional keyword arguments to pass to the PIL save function. + """ image.save(dst_file, pnginfo=metadata.pnginfo, **kwargs) def _save_exif_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs): + """ + Internal function to save an image with SD metadata in EXIF format. + + :param image: The PIL Image object to save. + :param dst_file: The destination file path. + :param metadata: The SDMetaData object containing the metadata to save. + :param kwargs: Additional keyword arguments to pass to the write_geninfo_exif function. + """ write_geninfo_exif(image, dst_file, metadata.text, **kwargs) def _save_gif_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs): + """ + Internal function to save a GIF image with SD metadata. + + :param image: The PIL Image object to save. + :param dst_file: The destination file path. + :param metadata: The SDMetaData object containing the metadata to save. + :param kwargs: Additional keyword arguments to pass to the write_geninfo_gif function. + """ write_geninfo_gif(image, dst_file, metadata.text, **kwargs) @@ -295,6 +351,36 @@ def _save_gif_with_sdmeta(image: Image.Image, dst_file: Union[str, os.PathLike], def save_image_with_sdmeta(image: ImageTyping, dst_file: Union[str, os.PathLike], metadata: SDMetaData, **kwargs): + """ + Save an image with Stable Diffusion metadata. + + This function saves the given image to the specified destination file, including + the provided SD metadata. The metadata is saved in a format appropriate for the + output image type (PNG, JPEG, WebP, or GIF). + + :param image: The input image, which can be a file path, URL, or PIL Image object. + :type image: ImageTyping + :param dst_file: The destination file path where the image will be saved. + :type dst_file: Union[str, os.PathLike] + :param metadata: The SD metadata to include with the image. + :type metadata: SDMetaData + :param kwargs: Additional keyword arguments to pass to the underlying save function. + + :raises SystemError: If the output file type is not supported for saving with metadata. + :raises: Various exceptions may be raised by the underlying image loading and + saving functions. + + Example usage: + >>> from imgutils.sd import get_sdmeta_from_image, save_image_with_sdmeta + >>> input_image = 'path/to/input.png' + >>> output_image = 'path/to/output.png' + >>> sd_meta = get_sdmeta_from_image(input_image) + >>> if sd_meta: + ... save_image_with_sdmeta(input_image, output_image, sd_meta) + ... print(f"Image saved with SD metadata to {output_image}") + ... else: + ... print("No SD metadata found in the input image.") + """ mimetype, _ = mimetypes.guess_type(str(dst_file)) if mimetype not in _FN_IMG_SAVE: raise SystemError(f'Not supported to save as a {mimetype!r} type, '