From dd6f9360035a45f363efa10b21763b247b9821ae Mon Sep 17 00:00:00 2001 From: StableLlama Date: Fri, 16 Oct 2026 18:28:00 +0200 Subject: [PATCH] Make saved image Comfy compliant --- CHANGELOG.md | 8 + src/basic_data_handling/path_nodes.py | 210 +++++++++++++++++++++++--- tests/test_path_nodes.py | 167 ++++++++++++++++++++ 3 files changed, 364 insertions(+), 21 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5ce29be..ab6c4e8 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -22,6 +22,14 @@ Fixed: and the reduction is reported; the image itself is always written. Should embedding metadata ever fail for another reason, the image is written again without metadata instead of not at all. +- "save IMAGE to file" and "save IMAGE+MASK to file": the `parameters` metadata + now ends with the sampler settings of the workflow (`Steps: 8, Sampler: euler, + Schedule type: simple, CFG scale: 1, Seed: 51, Size: 1024x1024`). Sites like + Civitai only read the prompt out of a `parameters` text that contains such a + line; without it the whole text is discarded and the explicitly set `prompt` + was not extracted at all. The settings are read from the sampler node of the + workflow and nothing is invented - when the graph does not provide them, the + text stays the plain prompt as before. ## [2.1.0] - 2026-09-19 diff --git a/src/basic_data_handling/path_nodes.py b/src/basic_data_handling/path_nodes.py index e554adc..602a825 100644 --- a/src/basic_data_handling/path_nodes.py +++ b/src/basic_data_handling/path_nodes.py @@ -1084,28 +1084,152 @@ class PathSaveStringFile(ComfyNodeABC): return (False,) -def compose_prompt_text(prompt: str, negative_prompt: str) -> str: +def compose_prompt_text(prompt: str, negative_prompt: str, generation_parameters: 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:: + prompt is written first, followed by an optional ``Negative prompt:`` line + and the optional ``generation_parameters`` line (see + ``build_generation_parameters``):: Negative prompt: + Steps: 8, Sampler: euler, ... - Returns an empty string when neither value is provided, in which case no - metadata is embedded into the file. + Returns an empty string when nothing 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()}") + if generation_parameters: + lines.append(generation_parameters) return "\n".join(lines) +def _is_node_link(value: Any) -> bool: + """True for a ComfyUI ``[node_id, output_index]`` input link.""" + return isinstance(value, (list, tuple)) and len(value) == 2 and isinstance(value[0], str) and isinstance(value[1], int) + + +def _static_input_value(value: Any, workflow_prompt: Any) -> Any: + """ + Return a sampler setting as the value the graph was run with. + + An input that is not driven by a link is serialised as its plain value: that + is the value. A driven one is a link instead, and what it carries is unknown + in general - but when the link points at a node consisting of nothing but a + ``value`` input (ComfyUI's primitives, this pack's ``IntCreate``/ + ``FloatCreate``, ...) then that ``value`` *is* the setting the user typed, so + it is taken. Anything computed stays unknown (``None``) instead of being + guessed, and the same holds for everything that is not a single value (a + list input feeding a sampler is a batch, not a setting). + """ + if not _is_node_link(value): + return value + if not isinstance(workflow_prompt, dict): + return None + node = workflow_prompt.get(value[0]) + inputs = node.get("inputs") if isinstance(node, dict) else None + if isinstance(inputs, dict) and len(inputs) == 1: + source = inputs.get("value") + if isinstance(source, (str, int, float)): + return source + return None + + +def _ancestor_node_ids(node_id: str, workflow_prompt: dict) -> set[str]: + """All node ids reachable from ``node_id`` by following its input links.""" + seen: set[str] = set() + stack = [node_id] + while stack: + node = workflow_prompt.get(stack.pop()) + inputs = node.get("inputs") if isinstance(node, dict) else None + if not isinstance(inputs, dict): + continue + for value in inputs.values(): + if _is_node_link(value) and value[0] not in seen: + seen.add(value[0]) + stack.append(value[0]) + return seen + + +def _format_setting(value: Any) -> str: + """Format one sampler setting the way A1111 writes it (``1.0`` -> ``1``).""" + if isinstance(value, float) and value.is_integer(): + return str(int(value)) + return str(value) + + +def build_generation_parameters(workflow_prompt: Any, width: int = 0, height: int = 0) -> str: + """ + Compose the A1111-style generation-parameters line of a save. + + Readers of the ``parameters`` text look for a line holding the sampler + settings (``Steps:``, ``Sampler:``, ``CFG scale:``, ...) and treat the whole + text as unreadable when there is none - the explicitly set prompt would then + not be extracted at all (which is what Civitai did with the saved files). + The settings are therefore read from the sampler node of ComfyUI's own + workflow prompt (``PROMPT``), so the line describes the graph that produced + the image, e.g.:: + + Steps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51, Size: 1024x1024 + + Only settings the graph really provides are written, and only when they are + known without guessing (see ``_static_input_value``). Returns an empty + string when the three settings the line is built around (``steps``, ``cfg``, + ``sampler_name``) are not all available; the ``parameters`` text then stays + the plain prompt text. + """ + if not isinstance(workflow_prompt, dict): + return "" + + samplers: dict[str, dict] = {} + for node_id, node in workflow_prompt.items(): + inputs = node.get("inputs") if isinstance(node, dict) else None + if isinstance(inputs, dict) and all(key in inputs for key in ("steps", "cfg", "sampler_name")): + samplers[str(node_id)] = inputs + if not samplers: + return "" + + # A graph may hold several samplers (hires fix and other multi-pass + # workflows). The pass that ran last describes the file, exactly like A1111 + # reports the settings of its second pass: that is a sampler whose output no + # other sampler consumed (it is no other sampler's ancestor), and among those + # the one with the deepest dependency chain. + ancestors = {node_id: _ancestor_node_ids(node_id, workflow_prompt) for node_id in samplers} + consumed_by_a_later_pass = set().union(*ancestors.values()) + last_pass = [node_id for node_id in samplers if node_id not in consumed_by_a_later_pass] + chosen = max(last_pass or list(samplers), key=lambda node_id: len(ancestors[node_id])) + inputs = samplers[chosen] + + steps = _static_input_value(inputs["steps"], workflow_prompt) + cfg = _static_input_value(inputs["cfg"], workflow_prompt) + sampler = _static_input_value(inputs["sampler_name"], workflow_prompt) + if steps is None or cfg is None or not sampler: + return "" + + settings = [f"Steps: {_format_setting(steps)}", f"Sampler: {sampler}"] + scheduler = _static_input_value(inputs.get("scheduler"), workflow_prompt) + if scheduler: + settings.append(f"Schedule type: {scheduler}") + settings.append(f"CFG scale: {_format_setting(cfg)}") + seed = _static_input_value(inputs.get("seed"), workflow_prompt) + if seed is None: + seed = _static_input_value(inputs.get("noise_seed"), workflow_prompt) + if seed is not None: + settings.append(f"Seed: {_format_setting(seed)}") + if width and height: + settings.append(f"Size: {width}x{height}") + denoise = _static_input_value(inputs.get("denoise"), workflow_prompt) + if isinstance(denoise, (int, float)) and not isinstance(denoise, bool) and float(denoise) != 1.0: + settings.append(f"Denoising strength: {_format_setting(denoise)}") + return ", ".join(settings) + + def _metadata_disabled() -> bool: """ True when ComfyUI was started with ``--disable-metadata``. @@ -1145,7 +1269,7 @@ def build_workflow_metadata(workflow_prompt: Any, workflow_extra_pnginfo: Any) - def compose_save_metadata( - prompt: str, negative_prompt: str, workflow_prompt: Any, workflow_extra_pnginfo: Any + prompt: str, negative_prompt: str, workflow_prompt: Any, workflow_extra_pnginfo: Any, width: int = 0, height: int = 0 ) -> tuple[str, dict[str, str]]: """ Build both metadata channels of an image save. @@ -1154,13 +1278,21 @@ def compose_save_metadata( ``parameters`` text of the explicit prompt(s) and the ``{key: json}`` mapping of ComfyUI's workflow prompt info. + The prompt text is completed by the sampler settings of the workflow (see + ``build_generation_parameters``): without that line readers like Civitai + discard the whole ``parameters`` text as unreadable, so the explicitly set + prompt would not be extracted. + ComfyUI's ``--disable-metadata`` flag suppresses *all* metadata, exactly like the built-in save node does, so a run started with it writes files without either channel. """ if _metadata_disabled(): return "", {} - return compose_prompt_text(prompt, negative_prompt), build_workflow_metadata(workflow_prompt, workflow_extra_pnginfo) + generation_parameters = build_generation_parameters(workflow_prompt, width, height) + return compose_prompt_text(prompt, negative_prompt, generation_parameters), build_workflow_metadata( + workflow_prompt, workflow_extra_pnginfo + ) def build_png_info(metadata_text: str, workflow_metadata: Optional[dict[str, str]] = None): @@ -1776,9 +1908,12 @@ class PathSaveImageRGB(ComfyNodeABC): Two independent metadata channels are embedded, so both stay usable: with ``prompt``/``negative_prompt`` the A1111-style ``parameters`` text is written (PNG text chunk, EXIF ``UserComment``, XMP for JPEG XL) which image viewers - display and which is easy to extract afterwards, and in addition ComfyUI's - hidden workflow prompt info (``prompt``/``workflow`` chunks) is stored, which - is what makes a saved file drag-and-drop back into ComfyUI as a workflow. + display and which is easy to extract afterwards - completed by the + ``Steps: ...`` line of the workflow's sampler settings, without which readers + like Civitai discard the whole text as unreadable - and in addition + ComfyUI's hidden workflow prompt info (``prompt``/``workflow`` chunks) is + stored, which is what makes a saved file drag-and-drop back into ComfyUI as + a workflow. When a batch is saved from a data list the workflow necessarily contains every prompt of the graph; the prompt of the individual image is the one in ``parameters``. Formats that cannot carry text metadata ignore both. @@ -1810,8 +1945,20 @@ class PathSaveImageRGB(ComfyNodeABC): ), "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (jpg/webp/jxl)."}), "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}), - "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}), - "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}), + "prompt": ( + IO.STRING, + { + "default": "", + "tooltip": 'Optional positive prompt embedded as "parameters" metadata, together with the sampler settings of the workflow (readers like Civitai need them to extract the prompt).', + }, + ), + "negative_prompt": ( + IO.STRING, + { + "default": "", + "tooltip": 'Optional negative prompt embedded as "parameters" metadata, together with the sampler settings of the workflow (readers like Civitai need them to extract the prompt).', + }, + ), # Appended last on purpose: the node stores widget values # positionally, so adding before the existing widgets would shift # old workflows and mis-assign their saved values. @@ -1876,8 +2023,11 @@ class PathSaveImageRGB(ComfyNodeABC): paths, create_dirs = _plan_save_paths(path, fmt, use_prefix_mode, batch, width, height) # Compose both metadata channels: the explicit prompt(s) of this save - # (the A1111-style "parameters" text) and ComfyUI's workflow prompt info. - metadata_text, workflow_metadata = compose_save_metadata(prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo) + # (the A1111-style "parameters" text, completed by the sampler settings + # of the workflow) and ComfyUI's workflow prompt info. + metadata_text, workflow_metadata = compose_save_metadata( + prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo, width, height + ) # Reporting the written files makes the frontend preview them and list # them with the generated images (same payload as the built-in node). @@ -1925,9 +2075,12 @@ class PathSaveImageRGBA(ComfyNodeABC): Two independent metadata channels are embedded, so both stay usable: with ``prompt``/``negative_prompt`` the A1111-style ``parameters`` text is written (PNG text chunk, EXIF ``UserComment``, XMP for JPEG XL) which image viewers - display and which is easy to extract afterwards, and in addition ComfyUI's - hidden workflow prompt info (``prompt``/``workflow`` chunks) is stored, which - is what makes a saved file drag-and-drop back into ComfyUI as a workflow. + display and which is easy to extract afterwards - completed by the + ``Steps: ...`` line of the workflow's sampler settings, without which readers + like Civitai discard the whole text as unreadable - and in addition + ComfyUI's hidden workflow prompt info (``prompt``/``workflow`` chunks) is + stored, which is what makes a saved file drag-and-drop back into ComfyUI as + a workflow. When a batch is saved from a data list the workflow necessarily contains every prompt of the graph; the prompt of the individual image is the one in ``parameters``. Formats that cannot carry text metadata ignore both. @@ -1961,8 +2114,20 @@ class PathSaveImageRGBA(ComfyNodeABC): "quality": (IO.INT, {"default": 95, "min": 1, "max": 100, "tooltip": "Quality for lossy formats (webp/jxl)."}), "invert_mask": (IO.BOOLEAN, {"default": False, "tooltip": "Invert the mask before using it as the alpha channel."}), "create_dirs": (IO.BOOLEAN, {"default": True, "tooltip": "Create missing parent directories (plain path mode only)."}), - "prompt": (IO.STRING, {"default": "", "tooltip": "Optional positive prompt embedded as parameters metadata."}), - "negative_prompt": (IO.STRING, {"default": "", "tooltip": "Optional negative prompt embedded as parameters metadata."}), + "prompt": ( + IO.STRING, + { + "default": "", + "tooltip": 'Optional positive prompt embedded as "parameters" metadata, together with the sampler settings of the workflow (readers like Civitai need them to extract the prompt).', + }, + ), + "negative_prompt": ( + IO.STRING, + { + "default": "", + "tooltip": 'Optional negative prompt embedded as "parameters" metadata, together with the sampler settings of the workflow (readers like Civitai need them to extract the prompt).', + }, + ), # Appended last on purpose: the node stores widget values # positionally, so adding before the existing widgets would shift # old workflows and mis-assign their saved values. @@ -2038,8 +2203,11 @@ class PathSaveImageRGBA(ComfyNodeABC): mask_frames = mask.cpu() # Compose both metadata channels: the explicit prompt(s) of this save - # (the A1111-style "parameters" text) and ComfyUI's workflow prompt info. - metadata_text, workflow_metadata = compose_save_metadata(prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo) + # (the A1111-style "parameters" text, completed by the sampler settings + # of the workflow) and ComfyUI's workflow prompt info. + metadata_text, workflow_metadata = compose_save_metadata( + prompt, negative_prompt, workflow_prompt, workflow_extra_pnginfo, width, height + ) # Reporting the written files makes the frontend preview them and list # them with the generated images (same payload as the built-in node). diff --git a/tests/test_path_nodes.py b/tests/test_path_nodes.py index 5503e3e..15c4300 100644 --- a/tests/test_path_nodes.py +++ b/tests/test_path_nodes.py @@ -953,6 +953,173 @@ def test_path_save_image_embeds_workflow_and_explicit_prompt(tmp_path, monkeypat assert json.loads(img.text["workflow"]) == extra_pnginfo["workflow"] +#: The sampler settings an A1111 infotext reader knows. Readers (Civitai and +#: similar) scan the ``parameters`` text from its end for a line holding at least +#: three of them - ``Sampler`` among them - and discard the whole text when there +#: is no such line, so the prompt would not be extracted either. +_A1111_SAMPLER_KEYS = ("Sampler", "CFG scale", "Seed", "Steps", "ENSD", "Schedule type") + + +def read_as_a1111_parameters(parameters_text): + """Return the sampler settings a reader finds in a ``parameters`` text.""" + for line in reversed(parameters_text.splitlines()): + settings = {} + for item in line.split(","): + key, separator, value = item.partition(":") + if separator and key.strip() in _A1111_SAMPLER_KEYS: + settings[key.strip()] = value.strip() + if len(settings) >= 3 and "Sampler" in settings: + return settings + return None + + +def _sampler_workflow(**overrides): + """An API workflow prompt holding one KSampler with literal settings.""" + settings = {"steps": 8, "cfg": 1.0, "sampler_name": "euler", "scheduler": "simple", "seed": 51, "denoise": 1.0} + settings.update(overrides) + return {"3": {"class_type": "KSampler", "inputs": {"model": ["10", 0], "positive": ["6", 0], **settings}}} + + +def test_path_save_image_parameters_text_carries_the_sampler_settings(tmp_path): + """The explicit prompt is only extracted by readers like Civitai when the + ``parameters`` text ends with a ``Steps: ...`` line, so the save node adds the + sampler settings of the workflow below the prompt.""" + save_node = PathSaveImageRGB() + out = str(tmp_path / "a1111_meta") + + result = save_node.save_image( + _batch_rgb(1), + out, + prompt="the prompt of this image", + negative_prompt="bad", + workflow_prompt=_sampler_workflow(), + workflow_extra_pnginfo={"workflow": {"nodes": [{"id": 3}]}}, + ) + assert save_success(result) is True + + with Image.open(out + ".png") as img: + img.load() + parameters = img.text["parameters"] + + assert parameters == ( + "the prompt of this image\nNegative prompt: bad\n" + "Steps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51, Size: 16x16" + ) + # Everything above the settings line is the prompt, exactly what a reader + # of A1111 parameters text expects. + assert read_as_a1111_parameters(parameters) == { + "Steps": "8", + "Sampler": "euler", + "Schedule type": "simple", + "CFG scale": "1", + "Seed": "51", + } + + +def test_path_save_image_with_mask_also_carries_the_sampler_settings(tmp_path): + """Both save nodes share the metadata composition.""" + save_node = PathSaveImageRGBA() + out = str(tmp_path / "a1111_rgba") + + result = save_node.save_image_with_mask( + _batch_rgb(1), torch.zeros(1, 16, 16), out, prompt="explicit", workflow_prompt=_sampler_workflow() + ) + + assert save_success(result) is True + with Image.open(out + ".png") as img: + img.load() + assert img.text["parameters"].endswith("\nSteps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51, Size: 16x16") + + +def test_path_save_image_sampler_settings_follow_value_links(tmp_path): + """A sampler widget is often driven by a plain value node (this pack's + ``IntCreate`` and the like): such a link carries the setting, while a device + that computes a value stays unknown and is left out rather than guessed.""" + workflow_prompt = { + "3": { + "class_type": "KSampler", + "inputs": {"steps": ["5", 0], "cfg": 3.5, "sampler_name": "dpmpp_2m", "seed": ["7", 0], "scheduler": "karras"}, + }, + "5": {"class_type": "Basic data handling: IntCreate", "inputs": {"value": "20"}}, + "7": {"class_type": "RandomInt", "inputs": {"min": 0, "max": 100}}, + } + out = str(tmp_path / "linked") + + result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt) + + assert save_success(result) is True + with Image.open(out + ".png") as img: + img.load() + # "Steps: 20" comes from the value node, the seed cannot be resolved and + # is therefore missing - which the line itself has to survive. + assert img.text["parameters"] == "explicit\nSteps: 20, Sampler: dpmpp_2m, Schedule type: karras, CFG scale: 3.5, Size: 16x16" + assert read_as_a1111_parameters(img.text["parameters"]) == { + "Steps": "20", + "Sampler": "dpmpp_2m", + "Schedule type": "karras", + "CFG scale": "3.5", + } + + +def test_path_save_image_uses_the_last_pass_of_a_multi_sampler_workflow(tmp_path): + """A hires fix runs two samplers; like A1111 (which reports its second pass) + the settings of the pass that produced the file are the ones written.""" + workflow_prompt = { + "1": _sampler_workflow()["3"], + "2": {"class_type": "LatentUpscale", "inputs": {"samples": ["1", 0]}}, + "3": { + "class_type": "KSampler", + "inputs": {"steps": 20, "cfg": 7.0, "sampler_name": "dpmpp_2m", "scheduler": "karras", "seed": 52, "latent_image": ["2", 0]}, + }, + } + out = str(tmp_path / "second_pass") + + result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt) + + assert save_success(result) is True + with Image.open(out + ".png") as img: + img.load() + assert ( + img.text["parameters"] == "explicit\nSteps: 20, Sampler: dpmpp_2m, Schedule type: karras, CFG scale: 7, Seed: 52, Size: 16x16" + ) + + +def test_path_save_image_without_sampler_settings_keeps_the_plain_prompt(tmp_path): + """A workflow whose settings are spread over several nodes gets no settings + line at all - the text stays the plain prompt of the previous versions.""" + workflow_prompt = { + "3": {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["4", 0], "guider": ["5", 0], "sampler": ["6", 0]}}, + "4": {"class_type": "BasicScheduler", "inputs": {"steps": 20, "scheduler": "simple"}}, + } + out = str(tmp_path / "no_settings") + + result = PathSaveImageRGB().save_image(_batch_rgb(1), out, prompt="explicit", workflow_prompt=workflow_prompt) + + assert save_success(result) is True + with Image.open(out + ".png") as img: + img.load() + assert img.text["parameters"] == "explicit" + assert read_as_a1111_parameters(img.text["parameters"]) is None + + +def test_build_generation_parameters_reads_only_what_the_graph_provides(): + """The settings line is derived from the workflow and never invented.""" + import src.basic_data_handling.path_nodes as pn + + assert pn.build_generation_parameters(None) == "" + assert pn.build_generation_parameters("not a workflow") == "" + assert pn.build_generation_parameters({}) == "" + # The image size is part of the line only when it is known ... + assert pn.build_generation_parameters(_sampler_workflow()) == "Steps: 8, Sampler: euler, Schedule type: simple, CFG scale: 1, Seed: 51" + # ... a sampling that only partly denoises is reported like A1111 does ... + assert pn.build_generation_parameters(_sampler_workflow(denoise=0.6)).endswith("Denoising strength: 0.6") + # ... and without the settings the line is built around there is no line. + for missing in ("steps", "cfg", "sampler_name"): + workflow_prompt = _sampler_workflow() + del workflow_prompt["3"]["inputs"][missing] + assert pn.build_generation_parameters(workflow_prompt) == "" + + def test_path_save_image_workflow_metadata_in_exif(tmp_path, monkeypatch): """Formats without a native text chunk keep the explicit prompt in UserComment and store the workflow in the fields ComfyUI uses for them."""