From 584ecac04bfcb9db3da5077616718e8f760b623e Mon Sep 17 00:00:00 2001 From: StableLlama Date: Sun, 6 Sep 2026 18:19:35 +0200 Subject: [PATCH] Add (optionally) prompts to the images when saving them --- src/basic_data_handling/path_nodes.py | 181 +++++++++++++++++++++++--- tests/test_path_nodes.py | 129 ++++++++++++++++++ 2 files changed, 295 insertions(+), 15 deletions(-) diff --git a/src/basic_data_handling/path_nodes.py b/src/basic_data_handling/path_nodes.py index 078b33a..dc5a472 100644 --- a/src/basic_data_handling/path_nodes.py +++ b/src/basic_data_handling/path_nodes.py @@ -1005,12 +1005,135 @@ class PathSaveStringFile(ComfyNodeABC): return (False,) +def compose_prompt_text(prompt: str, negative_prompt: str) -> str: + """ + Build the generation-parameter text embedded into saved images. + + Follows the Stable Diffusion WebUI convention: the (optional) positive + prompt is written first, followed by an optional ``Negative prompt:`` + line:: + + + Negative prompt: + + Returns an empty string when neither value is provided, in which case no + metadata is embedded into the file. + """ + lines = [] + if prompt.strip(): + lines.append(prompt.strip()) + if negative_prompt.strip(): + lines.append(f"Negative prompt: {negative_prompt.strip()}") + return "\n".join(lines) + + +def build_png_info(metadata_text: str): + """ + Wrap ``metadata_text`` in a Pillow ``PngInfo`` container under the standard + ``parameters`` text-chunk key so it can be embedded in a PNG file. + + Returns ``None`` when there is no text to embed (or Pillow's PNG metadata + support is unavailable), in which case the image should be saved without + extra metadata. + """ + if not metadata_text: + return None + try: + from PIL import PngImagePlugin + except ModuleNotFoundError: + return None + pnginfo = PngImagePlugin.PngInfo() + pnginfo.add_text("parameters", metadata_text) + return pnginfo + + +def build_image_exif(metadata_text: str, include_description: bool = True): + """ + Build an EXIF block that stores ``metadata_text`` for formats without a + native text chunk (JPEG, WEBP, JXL). + + The payload is written into the EXIF ``UserComment`` field (tag 0x9286) of + the Exif IFD as ``UNICODE\0`` + UTF-16-BE, which matches what Stable + Diffusion WebUI / piexif based readers expect. When ``include_description`` + is true, the payload is also written as UTF-8 into the EXIF + ``ImageDescription`` field (tag 0x010E) of IFD0. + + Returns the EXIF bytes (starting with the ``Exif\0\0`` marker), or ``None`` + when there is no text to embed. + """ + if not metadata_text: + return None + try: + from PIL import ExifTags + except ModuleNotFoundError: + return None + Image, _ = _require_pillow() + exif = Image.Exif() + if include_description: + exif[0x010E] = metadata_text.encode("utf-8") + exif.get_ifd(ExifTags.IFD.Exif)[0x9286] = b"UNICODE\x00" + metadata_text.encode("utf-16-be") + return exif.tobytes() + + +def build_xmp_packet(metadata_text: str) -> bytes: + """ + Build an XMP packet storing ``metadata_text`` in the Dublin Core + ``dc:description`` tag, as expected for JPEG XL ``xml `` boxes. + """ + from xml.sax.saxutils import escape + body = escape(metadata_text) + packet = ( + '\n' + '\n' + '\n' + '\n' + '' + body + '\n' + '\n' + '\n' + '\n' + '' + ) + return packet.encode("utf-8") + + +def metadata_save_kwargs(metadata_text: str, fmt: str) -> dict: + """ + Return the extra keyword arguments that embed ``metadata_text`` when saving + an image in the (lower-case) format ``fmt``. + + Returns an empty dict when there is no text to embed or when the format + cannot carry text metadata. + """ + if not metadata_text: + return {} + if fmt == "png": + return {"pnginfo": build_png_info(metadata_text)} + if fmt in ("jpg", "jpeg"): + exif = build_image_exif(metadata_text, include_description=True) + return {"exif": exif} if exif is not None else {} + if fmt in ("webp", "jxl"): + exif = build_image_exif(metadata_text, include_description=False) + kwargs = {"exif": exif} if exif is not None else {} + if fmt == "jxl": + # EXIF and XMP boxes are only available in the JXL container format + kwargs["use_container"] = True + kwargs["xmp"] = build_xmp_packet(metadata_text) + return kwargs + return {} + + class PathSaveImageRGB(ComfyNodeABC): """ Saves an image to a file. This node takes an image tensor and saves it to the specified path. Supports various image formats like PNG, JPG, WEBP, JXL (if pillow-jxl is installed), etc. + + When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded + into the saved image as ``parameters`` metadata: in the PNG text chunk, in + the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in + the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata + ignore the prompts. """ @classmethod def INPUT_TYPES(cls): @@ -1023,6 +1146,8 @@ class PathSaveImageRGB(ComfyNodeABC): "format": (IO.STRING, {"default": "png"}), "quality": (IO.INT, {"default": 95, "min": 1, "max": 100}), "create_dirs": (IO.BOOLEAN, {"default": True}), + "prompt": (IO.STRING, {"default": ""}), + "negative_prompt": (IO.STRING, {"default": ""}), } } @@ -1033,7 +1158,8 @@ class PathSaveImageRGB(ComfyNodeABC): FUNCTION = "save_image" OUTPUT_NODE = True - def save_image(self, images, path: str, format: str = "png", quality: int = 95, create_dirs: bool = True): + def save_image(self, images, path: str, format: str = "png", quality: int = 95, + create_dirs: bool = True, prompt: str = "", negative_prompt: str = ""): if not path: print("Basic data handling: Save failed - no path specified") return (False,) @@ -1074,15 +1200,23 @@ class PathSaveImageRGB(ComfyNodeABC): # Create PIL image pil_img = Image.fromarray(img_np) - # Save the image - if format.lower() == "jpg" or format.lower() == "jpeg": - pil_img.save(path, format="JPEG", quality=quality) - elif format.lower() == "webp": - pil_img.save(path, format="WEBP", quality=quality) - elif format.lower() == "jxl" and has_jxl_support: + # Compose the prompt metadata to embed into the saved file + metadata_text = compose_prompt_text(prompt, negative_prompt) + fmt = format.lower() + + # Save the image, embedding prompt metadata where the format supports it + if fmt == "jpg" or fmt == "jpeg": + pil_img.save(path, format="JPEG", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "webp": + pil_img.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "jxl" and has_jxl_support: # JPEG XL specific options - pil_img.save(path, format="JXL", quality=quality) + pil_img.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "png": + pil_img.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt)) else: + if metadata_text: + print("Basic data handling: Prompt metadata is not supported for this format; skipping it.") pil_img.save(path, format=format.upper()) print(f"Basic data handling: Successfully saved image to {path}") @@ -1099,6 +1233,12 @@ class PathSaveImageRGBA(ComfyNodeABC): This node takes an image tensor and a mask tensor and saves them to the specified path as an image with transparency, where the mask defines the alpha channel. + + When ``prompt`` and/or ``negative_prompt`` are provided, they are embedded + into the saved image as ``parameters`` metadata: in the PNG text chunk, in + the EXIF ``UserComment`` (and ``ImageDescription`` for JPEG) fields, and in + the EXIF + XMP boxes for JPEG XL. Formats that cannot carry text metadata + ignore the prompts. """ @classmethod def INPUT_TYPES(cls): @@ -1113,6 +1253,8 @@ class PathSaveImageRGBA(ComfyNodeABC): "quality": (IO.INT, {"default": 95, "min": 1, "max": 100}), "invert_mask": (IO.BOOLEAN, {"default": False}), "create_dirs": (IO.BOOLEAN, {"default": True}), + "prompt": (IO.STRING, {"default": ""}), + "negative_prompt": (IO.STRING, {"default": ""}), } } @@ -1124,8 +1266,9 @@ class PathSaveImageRGBA(ComfyNodeABC): OUTPUT_NODE = True def save_image_with_mask(self, images, mask, path: str, format: str = "png", - quality: int = 95, invert_mask: bool = False, - create_dirs: bool = True): + quality: int = 95, invert_mask: bool = False, + create_dirs: bool = True, prompt: str = "", + negative_prompt: str = ""): if not path: print("Basic data handling: Save failed - no path specified") return (False,) @@ -1186,13 +1329,21 @@ class PathSaveImageRGBA(ComfyNodeABC): pil_img_rgba = pil_img.convert("RGBA") pil_img_rgba.putalpha(alpha_img) - # Save the image - if format.lower() == "webp": - pil_img_rgba.save(path, format="WEBP", quality=quality) - elif format.lower() == "jxl" and has_jxl_support: + # Compose the prompt metadata to embed into the saved file + metadata_text = compose_prompt_text(prompt, negative_prompt) + fmt = format.lower() + + # Save the image, embedding prompt metadata where the format supports it + if fmt == "webp": + pil_img_rgba.save(path, format="WEBP", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "jxl" and has_jxl_support: # JPEG XL supports alpha channel - pil_img_rgba.save(path, format="JXL", quality=quality) + pil_img_rgba.save(path, format="JXL", quality=quality, **metadata_save_kwargs(metadata_text, fmt)) + elif fmt == "png": + pil_img_rgba.save(path, format="PNG", **metadata_save_kwargs(metadata_text, fmt)) else: + if metadata_text: + print("Basic data handling: Prompt metadata is not supported for this format; skipping it.") pil_img_rgba.save(path, format=format.upper()) print(f"Basic data handling: Successfully saved image with mask to {path}") diff --git a/tests/test_path_nodes.py b/tests/test_path_nodes.py index 65b447d..3ed6fb6 100644 --- a/tests/test_path_nodes.py +++ b/tests/test_path_nodes.py @@ -15,6 +15,17 @@ from src.basic_data_handling.path_nodes import ( ) +def decode_user_comment(value): + """Decode an EXIF UserComment value like piexif/A1111 readers do.""" + if isinstance(value, bytes): + if value.startswith(b"ASCII\x00\x00\x00"): + return value[8:].decode("utf-8", errors="ignore") + if value.startswith(b"UNICODE\x00"): + return value[8:].decode("utf-16-be", errors="ignore") + return value.decode("utf-8", errors="ignore") + return value + + def test_path_join(): node = PathJoin() assert node.join_paths("folder", "file.txt") == (os.path.join("folder", "file.txt"),) @@ -459,6 +470,124 @@ def test_path_load_save_image_rgb(tmp_path, monkeypatch): load_node.load_image_rgb(str(tmp_path / "nonexistent.png")) +def test_path_save_image_rgb_prompt_metadata(tmp_path): + save_node = PathSaveImageRGB() + img_size = (16, 16) + red_img = torch.zeros(1, img_size[1], img_size[0], 3) + red_img[0, :, :, 0] = 1.0 # Red channel set to 1 + + # No metadata when neither prompt nor negative prompt is provided + plain_path = str(tmp_path / "plain") + assert save_node.save_image(red_img, plain_path) == (True,) + with Image.open(plain_path + ".png") as img: + img.load() + assert img.text.get("parameters") is None + + # Both prompt and negative prompt provided + full_path = str(tmp_path / "full") + assert save_node.save_image(red_img, full_path, prompt="a red square", + negative_prompt="blurry, low quality") == (True,) + with Image.open(full_path + ".png") as img: + img.load() + assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality" + + # Only the prompt is provided + pos_path = str(tmp_path / "pos_only") + assert save_node.save_image(red_img, pos_path, prompt="only positive") == (True,) + with Image.open(pos_path + ".png") as img: + img.load() + assert img.text.get("parameters") == "only positive" + + # Only the negative prompt is provided + neg_path = str(tmp_path / "neg_only") + assert save_node.save_image(red_img, neg_path, negative_prompt="only negative") == (True,) + with Image.open(neg_path + ".png") as img: + img.load() + assert img.text.get("parameters") == "Negative prompt: only negative" + + # JPEG embeds the prompt into the EXIF UserComment and ImageDescription fields + jpg_path = str(tmp_path / "jpeg_meta") + assert save_node.save_image(red_img, jpg_path, format="jpg", prompt="a red square", + negative_prompt="blurry, low quality") == (True,) + assert os.path.exists(jpg_path + ".jpg") + with Image.open(jpg_path + ".jpg") as img: + img.load() + exif = img.getexif() + expected = "a red square\nNegative prompt: blurry, low quality" + assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected + assert exif.get(0x010E) == expected + + # WEBP embeds the prompt into the EXIF UserComment field (no ImageDescription) + webp_path = str(tmp_path / "webp_meta") + assert save_node.save_image(red_img, webp_path, format="webp", prompt="a red square", + negative_prompt="blurry, low quality") == (True,) + with Image.open(webp_path + ".webp") as img: + img.load() + exif_bytes = img.info.get("exif", b"") + if exif_bytes.startswith(b"Exif\x00\x00"): + exif_bytes = exif_bytes[6:] + exif = Image.Exif() + exif.load(exif_bytes) + assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected + assert 0x010E not in exif + + +def test_path_save_image_rgba_prompt_metadata(tmp_path): + save_node = PathSaveImageRGBA() + img_size = (16, 16) + red_img = torch.zeros(1, img_size[1], img_size[0], 3) + red_img[0, :, :, 0] = 1.0 # Red channel set to 1 + mask = torch.zeros(1, img_size[1], img_size[0]) + + # No metadata when neither prompt nor negative prompt is provided + plain_path = str(tmp_path / "plain") + assert save_node.save_image_with_mask(red_img, mask, plain_path) == (True,) + with Image.open(plain_path + ".png") as img: + img.load() + assert img.text.get("parameters") is None + + # Both prompt and negative prompt provided + full_path = str(tmp_path / "full") + assert save_node.save_image_with_mask(red_img, mask, full_path, + prompt="a red square", + negative_prompt="blurry, low quality") == (True,) + with Image.open(full_path + ".png") as img: + img.load() + assert img.text.get("parameters") == "a red square\nNegative prompt: blurry, low quality" + + +def test_path_save_image_jxl_prompt_metadata(tmp_path): + pytest.importorskip("pillow_jxl") + save_node = PathSaveImageRGB() + img_size = (16, 16) + red_img = torch.zeros(1, img_size[1], img_size[0], 3) + red_img[0, :, :, 0] = 1.0 # Red channel set to 1 + + prompt = "a red square " + negative_prompt = "blurry, low quality" + jxl_path = str(tmp_path / "jxl_meta") + assert save_node.save_image(red_img, jxl_path, format="jxl", prompt=prompt, + negative_prompt=negative_prompt) == (True,) + assert os.path.exists(jxl_path + ".jxl") + + # JXL embeds the prompt into the EXIF UserComment field + with Image.open(jxl_path + ".jxl") as img: + img.load() + exif_bytes = img.info.get("exif", b"") + if exif_bytes.startswith(b"Exif\x00\x00"): + exif_bytes = exif_bytes[6:] + exif = Image.Exif() + exif.load(exif_bytes) + expected = "a red square \nNegative prompt: blurry, low quality" + assert decode_user_comment(exif.get_ifd(0x8769).get(0x9286)) == expected + + # JXL also embeds the prompt (XML-escaped) in an XMP xml box + raw = (tmp_path / "jxl_meta.jxl").read_bytes() + assert b"xml " in raw + assert b"" in raw + escaped = "a red square <lora:x:1.0>\nNegative prompt: blurry, low quality" + assert escaped.encode("utf-8") in raw + def test_path_normalize(): node = PathNormalize()