Add (optionally) prompts to the images when saving them
This commit is contained in:
StableLlama
@@ -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::
|
||||
|
||||
<positive prompt>
|
||||
Negative prompt: <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 = (
|
||||
'<?xpacket begin="\ufeff" id="W5M0MpCehiHzreSzNTczkc9d"?>\n'
|
||||
'<x:xmpmeta xmlns:x="adobe:ns:meta/">\n'
|
||||
'<rdf:RDF xmlns:rdf="http://www.w3.org/1999/02/22-rdf-syntax-ns#">\n'
|
||||
'<rdf:Description xmlns:dc="http://purl.org/dc/elements/1.1/">\n'
|
||||
'<dc:description><rdf:Alt><rdf:li xml:lang="x-default">' + body + '</rdf:li></rdf:Alt></dc:description>\n'
|
||||
'</rdf:Description>\n'
|
||||
'</rdf:RDF>\n'
|
||||
'</x:xmpmeta>\n'
|
||||
'<?xpacket end="w"?>'
|
||||
)
|
||||
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}")
|
||||
|
||||
@@ -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 <lora:x:1.0>"
|
||||
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 <lora:x:1.0>\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"<dc:description>" 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()
|
||||
|
||||
Reference in New Issue
Block a user