From fec3c6bae88dd423e1c898548f9957fe53cfa983 Mon Sep 17 00:00:00 2001 From: CyberShadowX Date: Thu, 28 May 2026 13:56:21 +0800 Subject: [PATCH] Support USDU and EXIF UserComment metadata --- smart_metadata_reader/adapters/samplers.py | 29 ++- smart_metadata_reader/civitai_parser.py | 217 ++++++++++++++++++++ smart_metadata_reader/lora_extractor.py | 12 +- smart_metadata_reader/metadata_reader.py | 129 +++++++++++- smart_metadata_reader/models.py | 3 + smart_metadata_reader/sampler_selector.py | 85 ++++++-- smart_metadata_reader/settings_formatter.py | 3 + tests/test_metadata_reader.py | 202 ++++++++++++++++++ tests/test_sampler_selector.py | 58 ++++++ 9 files changed, 719 insertions(+), 19 deletions(-) create mode 100644 smart_metadata_reader/civitai_parser.py diff --git a/smart_metadata_reader/adapters/samplers.py b/smart_metadata_reader/adapters/samplers.py index 4f97b3b..921ca7a 100644 --- a/smart_metadata_reader/adapters/samplers.py +++ b/smart_metadata_reader/adapters/samplers.py @@ -12,9 +12,26 @@ SAMPLER_CLASS_TYPES = { "SamplerCustomAdvanced", } +SAMPLER_LIKE_INPUT_FIELDS = { + "positive", + "negative", + "model", + "vae", + "seed", + "noise_seed", + "steps", + "cfg", + "cfg_scale", + "sampler_name", + "sampler", + "scheduler", + "scheduler_name", + "denoise", +} + def is_sampler(node: NodeRecord) -> bool: - return node.class_type in SAMPLER_CLASS_TYPES + return node.class_type in SAMPLER_CLASS_TYPES or _is_usdu_sampler_like(node) def extract_sampler_settings(sampler: NodeRecord) -> dict[str, Any]: @@ -24,7 +41,7 @@ def extract_sampler_settings(sampler: NodeRecord) -> dict[str, Any]: "steps": ("steps",), "cfg": ("cfg", "cfg_scale"), "sampler_name": ("sampler_name", "sampler"), - "scheduler": ("scheduler", "schedule"), + "scheduler": ("scheduler", "scheduler_name", "schedule"), "denoise": ("denoise",), } @@ -40,3 +57,11 @@ def _first_existing_input(node: NodeRecord, field_names: tuple[str, ...]) -> Any if field_name in node.inputs: return node.inputs[field_name] return None + + +def _is_usdu_sampler_like(node: NodeRecord) -> bool: + lowered = node.class_type.lower().replace("_", " ") + looks_like_usdu = ("ultimate" in lowered and "upscale" in lowered) or "usdu" in lowered + if not looks_like_usdu: + return False + return any(field in node.inputs for field in SAMPLER_LIKE_INPUT_FIELDS) diff --git a/smart_metadata_reader/civitai_parser.py b/smart_metadata_reader/civitai_parser.py new file mode 100644 index 0000000..d7f1cc4 --- /dev/null +++ b/smart_metadata_reader/civitai_parser.py @@ -0,0 +1,217 @@ +from __future__ import annotations + +import json +import re +from typing import Any + +from .models import LoraInfo, ParseResult + + +SETTING_KEYS = ( + "Steps", + "Sampler", + "CFG scale", + "Seed", + "Size", + "Model type", + "Model", + "Created Date", + "Civitai resources", + "Civitai metadata", +) + + +def parse_exif_user_comment(comment: str, source_label: str = "EXIF UserComment") -> ParseResult: + civitai_metadata = _extract_labeled_json(comment, "Civitai metadata") + civitai_resources = _extract_labeled_json(comment, "Civitai resources") + source_format = _source_format(source_label, civitai_metadata, civitai_resources) + + positive, negative = _extract_prompts(comment) + if isinstance(civitai_metadata, dict): + positive = _string_value(civitai_metadata.get("prompt")) or positive + negative = _string_value(civitai_metadata.get("negativePrompt")) or negative + + model_name = _extract_checkpoint_name(civitai_resources) + if not model_name: + model_name = _extract_setting_value(comment, "Model type") or _extract_setting_value( + comment, "Model" + ) + + partial_result: dict[str, Any] = {"source_format": source_format} + if civitai_metadata is not None: + partial_result["civitai_metadata"] = civitai_metadata + if civitai_resources is not None: + partial_result["civitai_resources"] = civitai_resources + + return ParseResult( + positive=positive, + negative=negative, + seed=_int_or_default(_extract_setting_value(comment, "Seed"), -1), + steps=_int_or_default(_extract_setting_value(comment, "Steps"), 0), + cfg=_float_or_default(_extract_setting_value(comment, "CFG scale"), 0.0), + width=_width_from_size(_extract_setting_value(comment, "Size")), + height=_height_from_size(_extract_setting_value(comment, "Size")), + model_name=model_name, + sampler_name=_extract_setting_value(comment, "Sampler"), + loras=_extract_loras(civitai_resources), + partial_result=partial_result, + debug_trace=( + "EXIF_USER_COMMENT_FALLBACK:\n" + f"Parsed metadata from {source_label}." + ), + confidence=0.75, + status_message=f"{source_format} fallback", + ) + + +def _source_format( + source_label: str, + civitai_metadata: Any, + civitai_resources: Any, +) -> str: + if civitai_metadata is not None or civitai_resources is not None: + return f"{source_label} / Civitai metadata" + return source_label + + +def _extract_prompts(comment: str) -> tuple[str, str]: + settings_start = _settings_start(comment) + prompt_block = comment[:settings_start].strip(" \t\r\n,") if settings_start else comment.strip() + negative_match = re.search(r"(?im)^Negative prompt:\s*", prompt_block) + if negative_match is None: + return prompt_block, "" + positive = prompt_block[: negative_match.start()].strip() + negative = prompt_block[negative_match.end() :].strip() + return positive, negative + + +def _settings_start(comment: str) -> int: + match = re.search(r"(?im)(?:^|\n|,\s*)Steps:\s*", comment) + if match is None: + return 0 + return match.start() + + +def _extract_labeled_json(text: str, label: str) -> Any | None: + marker = f"{label}:" + index = text.find(marker) + if index < 0: + return None + + start = index + len(marker) + while start < len(text) and text[start].isspace(): + start += 1 + + try: + value, _end = json.JSONDecoder().raw_decode(text[start:]) + except json.JSONDecodeError: + return None + return value + + +def _extract_setting_value(text: str, key: str) -> str: + known_keys_pattern = "|".join(re.escape(known_key) for known_key in SETTING_KEYS) + pattern = ( + rf"(?:^|[\n,]\s*){re.escape(key)}:\s*" + rf"(.*?)" + rf"(?=,\s*(?:{known_keys_pattern}):|$)" + ) + match = re.search(pattern, text, flags=re.DOTALL) + if match is None: + return "" + return match.group(1).strip() + + +def _extract_checkpoint_name(resources: Any) -> str: + if not isinstance(resources, list): + return "" + for resource in resources: + if not isinstance(resource, dict): + continue + resource_type = str(resource.get("type", "")).lower() + if resource_type != "checkpoint": + continue + return _string_value(resource.get("modelName")) or _string_value( + resource.get("modelVersionName") + ) + return "" + + +def _extract_loras(resources: Any) -> list[LoraInfo]: + if not isinstance(resources, list): + return [] + + loras: list[LoraInfo] = [] + for index, resource in enumerate(resources): + if not isinstance(resource, dict): + continue + resource_type = str(resource.get("type", "")).lower() + if "lora" not in resource_type: + continue + name = _string_value(resource.get("modelName")) or _string_value( + resource.get("modelVersionName") + ) + if not name: + continue + weight = _number_or_none(resource.get("weight")) + loras.append( + LoraInfo( + node_id=f"civitai:{index}", + class_type="CivitaiResource", + lora_name=name, + strength_model=weight, + strength_clip=weight, + path=["EXIF UserComment", f"Civitai resources[{index}]"], + ) + ) + return loras + + +def _width_from_size(size: str) -> int: + width, _height = _parse_size(size) + return width + + +def _height_from_size(size: str) -> int: + _width, height = _parse_size(size) + return height + + +def _parse_size(size: str) -> tuple[int, int]: + match = re.match(r"\s*(\d+)\s*x\s*(\d+)\s*$", size, flags=re.IGNORECASE) + if match is None: + return 0, 0 + return int(match.group(1)), int(match.group(2)) + + +def _string_value(value: Any) -> str: + return value if isinstance(value, str) else "" + + +def _int_or_default(value: str | None, default: int) -> int: + if value is None: + return default + try: + return int(value) + except ValueError: + return default + + +def _float_or_default(value: str | None, default: float) -> float: + if value is None: + return default + try: + return float(value) + except ValueError: + return default + + +def _number_or_none(value: Any) -> float | None: + if value is None or isinstance(value, bool): + return None + if isinstance(value, (int, float)): + return float(value) + try: + return float(str(value)) + except ValueError: + return None diff --git a/smart_metadata_reader/lora_extractor.py b/smart_metadata_reader/lora_extractor.py index 06c8547..ff109d1 100644 --- a/smart_metadata_reader/lora_extractor.py +++ b/smart_metadata_reader/lora_extractor.py @@ -12,6 +12,7 @@ from .adapters.loaders import ( strength_model, vae_name, ) +from .adapters.samplers import is_sampler from .graph import GraphIndex from .models import LoraInfo, NodeRecord @@ -205,6 +206,11 @@ def _walk_output_chain_for_vae(graph: GraphIndex, output_node: NodeRecord) -> st continue visited.add(node.node_id) + if is_sampler(node): + name = _vae_name_from_node_input(graph, node) + if name: + return name + if _is_vae_decode(node): name = _vae_name_from_decode_node(graph, node) if name: @@ -221,7 +227,11 @@ def _walk_output_chain_for_vae(graph: GraphIndex, output_node: NodeRecord) -> st def _vae_name_from_decode_node(graph: GraphIndex, decode_node: NodeRecord) -> str: - target = graph.link_target(decode_node.inputs.get("vae")) + return _vae_name_from_node_input(graph, decode_node) + + +def _vae_name_from_node_input(graph: GraphIndex, node: NodeRecord) -> str: + target = graph.link_target(node.inputs.get("vae")) if target is None: return "" vae_node = graph.get_node(target[0]) diff --git a/smart_metadata_reader/metadata_reader.py b/smart_metadata_reader/metadata_reader.py index 284e5e6..09540b2 100644 --- a/smart_metadata_reader/metadata_reader.py +++ b/smart_metadata_reader/metadata_reader.py @@ -8,6 +8,7 @@ from PIL import Image from .a1111_parser import parse_a1111_parameters from .adapters.samplers import extract_sampler_settings +from .civitai_parser import parse_exif_user_comment from .graph import GraphIndex from .lora_extractor import extract_loras, extract_model_name, extract_vae_name from .models import MetadataBundle, ParseResult, PromptSegment @@ -26,11 +27,21 @@ def _metadata_value(info: dict[str, Any], key: str) -> str | None: return str(value) -def _source_format(prompt_raw: str | None, workflow_raw: str | None, parameters_raw: str | None) -> str: +def _source_format( + prompt_raw: str | None, + workflow_raw: str | None, + parameters_raw: str | None, + user_comment_raw: str | None, + image_description_raw: str | None, +) -> str: if prompt_raw or workflow_raw: return "ComfyUI prompt/workflow" if parameters_raw: return "A1111 parameters" + if user_comment_raw: + return "EXIF UserComment" + if image_description_raw: + return "EXIF ImageDescription" return "none" @@ -40,6 +51,9 @@ def read_metadata(image_path: str | Path) -> MetadataBundle: prompt_raw = _metadata_value(image.info, "prompt") workflow_raw = _metadata_value(image.info, "workflow") parameters_raw = _metadata_value(image.info, "parameters") + user_comment_raw = _exif_user_comment(image) + image_description_raw = _exif_text_tag(image, 270) + software_raw = _exif_text_tag(image, 305) width, height = image.size return MetadataBundle( @@ -49,7 +63,16 @@ def read_metadata(image_path: str | Path) -> MetadataBundle: prompt_raw=prompt_raw, workflow_raw=workflow_raw, parameters_raw=parameters_raw, - source_format=_source_format(prompt_raw, workflow_raw, parameters_raw), + source_format=_source_format( + prompt_raw, + workflow_raw, + parameters_raw, + user_comment_raw, + image_description_raw, + ), + user_comment_raw=user_comment_raw, + image_description_raw=image_description_raw, + software_raw=software_raw, ) @@ -84,6 +107,10 @@ def parse_metadata_bundle( fallback_reason="missing ComfyUI prompt JSON", ) + exif_text, exif_source = _exif_fallback_text(bundle) + if exif_text: + return _parse_exif_bundle(bundle, exif_text, exif_source) + return _failed_result( bundle=bundle, status_message="NO_METADATA", @@ -304,6 +331,22 @@ def _parse_a1111_bundle( return result +def _parse_exif_bundle( + bundle: MetadataBundle, + exif_text: str, + exif_source: str, +) -> ParseResult: + result = parse_exif_user_comment(exif_text, exif_source) + result.filename = bundle.filename + result.width = result.width or bundle.width + result.height = result.height or bundle.height + result.partial_result["fallback_reason"] = ( + "missing ComfyUI prompt/workflow and A1111 parameters metadata" + ) + result.setting = format_setting(result) + return result + + def _failed_result( bundle: MetadataBundle, status_message: str, @@ -325,6 +368,88 @@ def _failed_result( return result +def _exif_fallback_text(bundle: MetadataBundle) -> tuple[str | None, str]: + if bundle.user_comment_raw: + return bundle.user_comment_raw, "EXIF UserComment" + if bundle.image_description_raw: + return bundle.image_description_raw, "EXIF ImageDescription" + return None, "EXIF" + + +def _exif_user_comment(image: Image.Image) -> str | None: + exif = image.getexif() + value = exif.get(37510) + if value is None: + try: + exif_ifd = exif.get_ifd(34665) + except (AttributeError, KeyError, TypeError, ValueError): + exif_ifd = {} + value = exif_ifd.get(37510) if isinstance(exif_ifd, dict) else None + return _decode_exif_text(value) + + +def _exif_text_tag(image: Image.Image, tag: int) -> str | None: + return _decode_exif_text(image.getexif().get(tag)) + + +def _decode_exif_text(value: Any) -> str | None: + if value is None: + return None + if isinstance(value, str): + return _clean_decoded_exif_text(value) + if not isinstance(value, bytes): + return _clean_decoded_exif_text(str(value)) + + data = value + prefix = data[:8].upper() + if prefix.startswith(b"ASCII"): + payload = data[8:] + encodings = ("utf-8", "latin1") + elif prefix.startswith(b"UNICODE"): + payload = data[8:] + encodings = ("utf-16le", "utf-16be", "utf-8", "latin1") + elif prefix.startswith(b"JIS"): + payload = data[8:] + encodings = ("shift_jis", "utf-8", "latin1") + else: + payload = data + encodings = ("utf-8", "utf-16le", "utf-16be", "latin1") + + candidates: list[str] = [] + for encoding in encodings: + candidate_payload = payload + if encoding.startswith("utf-16") and len(candidate_payload) % 2: + candidate_payload += b"\x00" + try: + candidates.append(candidate_payload.decode(encoding)) + except UnicodeDecodeError: + continue + if not candidates: + candidates.append(payload.decode("latin1", errors="replace")) + + best = max(candidates, key=_readability_score) + return _clean_decoded_exif_text(best) + + +def _clean_decoded_exif_text(text: str) -> str | None: + cleaned = text.replace("\x00", "").strip() + return cleaned or None + + +def _readability_score(text: str) -> float: + if not text: + return 0.0 + cleaned = text.strip() + if not cleaned: + return 0.0 + printable = sum(1 for char in cleaned if char.isprintable() or char in "\r\n\t") + replacement = cleaned.count("\ufffd") + nuls = cleaned.count("\x00") + alnum = sum(1 for char in cleaned if char.isalnum()) + length = len(cleaned) + return (printable / length) + (alnum / max(length, 1) * 0.2) - replacement - nuls + + def _int_or_default(value: Any, default: int) -> int: try: return int(value) diff --git a/smart_metadata_reader/models.py b/smart_metadata_reader/models.py index 79301ea..2d0b21a 100644 --- a/smart_metadata_reader/models.py +++ b/smart_metadata_reader/models.py @@ -13,6 +13,9 @@ class MetadataBundle: workflow_raw: str | None parameters_raw: str | None source_format: str + user_comment_raw: str | None = None + image_description_raw: str | None = None + software_raw: str | None = None @dataclass diff --git a/smart_metadata_reader/sampler_selector.py b/smart_metadata_reader/sampler_selector.py index 0677f3d..4d788e5 100644 --- a/smart_metadata_reader/sampler_selector.py +++ b/smart_metadata_reader/sampler_selector.py @@ -29,14 +29,18 @@ class SamplerSelection: def select_final_sampler(graph: GraphIndex) -> SamplerSelection: - for output_node in _output_nodes(graph): - selection = _select_from_output_node(graph, output_node) + output_nodes = _output_nodes(graph) + output_chain_diagnostics: list[str] = [] + for output_node in output_nodes: + selection, diagnostic = _select_from_output_node(graph, output_node) + if diagnostic: + output_chain_diagnostics.append(diagnostic) if selection is not None: return selection fallback = _fallback_sampler(graph) if fallback is None: - raise ValueError("No sampler candidate found") + raise ValueError(_no_sampler_diagnostic(output_nodes, output_chain_diagnostics)) return SamplerSelection( sampler=fallback, @@ -69,7 +73,7 @@ def _is_output_image_node(node: NodeRecord) -> bool: def _select_from_output_node( graph: GraphIndex, output_node: NodeRecord, -) -> SamplerSelection | None: +) -> tuple[SamplerSelection | None, str]: queue: deque[tuple[NodeRecord, list[str]]] = deque() for field in OUTPUT_INPUT_FIELDS: target = graph.link_target(output_node.inputs.get(field)) @@ -86,24 +90,29 @@ def _select_from_output_node( ) visited: set[str] = set() + visited_chains: list[tuple[NodeRecord, list[str]]] = [] while queue: node, chain = queue.popleft() if node.node_id in visited: continue visited.add(node.node_id) + visited_chains.append((node, chain)) if is_sampler(node): trace_chain = " <- ".join(chain + [f"{node.class_type} {node.node_id}"]) - return SamplerSelection( - sampler=node, - confidence=1.0, - debug_trace=( - "SAMPLER_SELECTION:\n" - f"{trace_chain}\n" - f"Selected {node.class_type} {node.node_id} because it is " - f"upstream of final {output_node.class_type}." + return ( + SamplerSelection( + sampler=node, + confidence=1.0, + debug_trace=( + "SAMPLER_SELECTION:\n" + f"{trace_chain}\n" + f"Selected {node.class_type} {node.node_id} because it is " + f"upstream of final {output_node.class_type}." + ), + selected_by="output_chain", ), - selected_by="output_chain", + "", ) for field in UPSTREAM_FIELDS: @@ -119,7 +128,7 @@ def _select_from_output_node( chain + [f"{node.class_type} {node.node_id}.{field}"], ) ) - return None + return None, _format_output_chain_diagnostic(output_node, visited_chains) def _fallback_sampler(graph: GraphIndex) -> NodeRecord | None: @@ -139,3 +148,51 @@ def _fallback_sampler(graph: GraphIndex) -> NodeRecord | None: def _has_conditioning_inputs(node: NodeRecord) -> bool: return "positive" in node.inputs or "negative" in node.inputs + + +def _no_sampler_diagnostic( + output_nodes: list[NodeRecord], + output_chain_diagnostics: list[str], +) -> str: + lines = ["no supported sampler candidate found"] + if output_nodes: + output_summary = ", ".join( + f"{node.class_type} {node.node_id}" for node in output_nodes + ) + lines.append(f"final output nodes found: {output_summary}") + else: + lines.append("final output nodes found: none") + if output_chain_diagnostics: + lines.append("final output chain diagnostics:") + lines.extend(output_chain_diagnostics) + return "\n".join(lines) + + +def _format_output_chain_diagnostic( + output_node: NodeRecord, + visited_chains: list[tuple[NodeRecord, list[str]]], +) -> str: + lines = [f"- {output_node.class_type} {output_node.node_id}:"] + if not visited_chains: + lines.append(" no image/images upstream link found") + return "\n".join(lines) + + for node, chain in visited_chains: + trace_chain = " <- ".join(chain + [f"{node.class_type} {node.node_id}"]) + input_keys = ", ".join(sorted(node.inputs)) or "none" + lines.append(f" {trace_chain}; input keys: {input_keys}") + if _looks_like_unsupported_generator(node): + lines.append( + f" unsupported terminal generator node: " + f"{node.class_type} {node.node_id}; input keys: {input_keys}" + ) + return "\n".join(lines) + + +def _looks_like_unsupported_generator(node: NodeRecord) -> bool: + lowered = node.class_type.lower().replace("_", " ") + return ( + ("ultimate" in lowered and "upscale" in lowered) + or "usdu" in lowered + or "generator" in lowered + ) diff --git a/smart_metadata_reader/settings_formatter.py b/smart_metadata_reader/settings_formatter.py index e227074..c9652de 100644 --- a/smart_metadata_reader/settings_formatter.py +++ b/smart_metadata_reader/settings_formatter.py @@ -29,6 +29,9 @@ def format_setting(result: ParseResult) -> str: if unresolved_lines: lines.append("Unresolved:") lines.extend(unresolved_lines) + reason = result.partial_result.get("reason") + if isinstance(reason, str) and reason: + lines.append(f"Failure Reason: {reason}") return "\n".join(lines) diff --git a/tests/test_metadata_reader.py b/tests/test_metadata_reader.py index 72bfec7..9fb3b9f 100644 --- a/tests/test_metadata_reader.py +++ b/tests/test_metadata_reader.py @@ -4,6 +4,16 @@ from PIL import Image from PIL.PngImagePlugin import PngInfo +def _save_jpeg_with_user_comment(path, comment: str, prefix: bytes = b"ASCII\x00\x00\x00"): + exif = Image.Exif() + if prefix.startswith(b"UNICODE"): + payload = comment.encode("utf-16le") + else: + payload = comment.encode("utf-8") + exif[37510] = prefix + payload + Image.new("RGB", (32, 24), (20, 40, 60)).save(path, exif=exif) + + def test_read_metadata_extracts_png_text_chunks_and_dimensions(tmp_path): from smart_metadata_reader.metadata_reader import read_metadata @@ -171,3 +181,195 @@ def test_parse_a1111_scheduler_and_model_hash_metadata(): assert result.scheduler == "Karras" assert result.model_name == "forge_model" assert result.partial_result["model_hash"] == "abc123" + + +def test_parse_comfyui_ultimate_sd_upscale_sampler_like_node(): + from smart_metadata_reader.metadata_reader import parse_metadata_bundle + from smart_metadata_reader.models import MetadataBundle + + prompt = { + "1": { + "class_type": "CheckpointLoaderSimple", + "inputs": {"ckpt_name": "base_model.safetensors"}, + }, + "2": { + "class_type": "LoraLoader", + "inputs": { + "model": ["1", 0], + "clip": ["1", 1], + "lora_name": "upscale_detail.safetensors", + "strength_model": 0.5, + "strength_clip": 0.75, + }, + }, + "3": {"class_type": "VAELoader", "inputs": {"vae_name": "main_vae.safetensors"}}, + "4": {"class_type": "CLIPTextEncode", "inputs": {"text": "role_lora_trigger, 1girl, nsfw"}}, + "5": {"class_type": "CLIPTextEncode", "inputs": {"text": "low quality"}}, + "6": { + "class_type": "UltimateSDUpscaleNoUpscale", + "inputs": { + "model": ["2", 0], + "clip": ["2", 1], + "vae": ["3", 0], + "positive": ["4", 0], + "negative": ["5", 0], + "seed": 221, + "steps": 9, + "cfg_scale": 1.5, + "sampler": "Euler", + "scheduler_name": "normal", + "tile_width": 768, + "seam_fix_mode": "half tile", + "upscale_model": ["9", 0], + }, + }, + "7": {"class_type": "SaveImage", "inputs": {"images": ["6", 0]}}, + "9": {"class_type": "UpscaleModelLoader", "inputs": {"model_name": "4x.pth"}}, + } + bundle = MetadataBundle( + filename="usdu.png", + width=4800, + height=6400, + prompt_raw=json.dumps(prompt), + workflow_raw=None, + parameters_raw=None, + source_format="ComfyUI prompt/workflow", + ) + + result = parse_metadata_bundle(bundle) + + assert result.positive == "role_lora_trigger, 1girl, nsfw" + assert result.negative == "low quality" + assert result.seed == 221 + assert result.steps == 9 + assert result.cfg == 1.5 + assert result.sampler_name == "Euler" + assert result.scheduler == "normal" + assert result.model_name == "base_model.safetensors" + assert result.vae_name == "main_vae.safetensors" + assert [lora.lora_name for lora in result.loras] == ["upscale_detail.safetensors"] + assert "UltimateSDUpscaleNoUpscale 6" in result.debug_trace + assert "tile_width" not in result.positive + assert result.status_message == "OK" + + +def test_failed_comfyui_parse_includes_output_chain_diagnostics_in_setting(): + from smart_metadata_reader.metadata_reader import parse_metadata_bundle + from smart_metadata_reader.models import MetadataBundle + + prompt = { + "1": {"class_type": "UnsupportedGenerator", "inputs": {"image": ["3", 0]}}, + "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, + "3": {"class_type": "LoadImage", "inputs": {"image": "source.png"}}, + } + bundle = MetadataBundle( + filename="unsupported.png", + width=512, + height=768, + prompt_raw=json.dumps(prompt), + workflow_raw=json.dumps({"nodes": []}), + parameters_raw=None, + source_format="ComfyUI prompt/workflow", + ) + + result = parse_metadata_bundle(bundle) + + assert result.status_message == "FAILED" + assert "no supported sampler candidate found" in result.debug_trace + assert "SaveImage 2.images <- UnsupportedGenerator 1" in result.debug_trace + assert "Failure Reason:" in result.setting + assert "final output nodes found: SaveImage 2" in result.setting + + +def test_read_metadata_extracts_ascii_and_unicode_exif_user_comment(tmp_path): + from smart_metadata_reader.metadata_reader import read_metadata + + ascii_path = tmp_path / "ascii.jpeg" + unicode_path = tmp_path / "unicode.jpeg" + _save_jpeg_with_user_comment(ascii_path, "ASCII prompt\nSteps: 9, Seed: 1") + _save_jpeg_with_user_comment( + unicode_path, + "UNICODE prompt\nSteps: 10, Seed: 2", + prefix=b"UNICODE\x00", + ) + + ascii_bundle = read_metadata(ascii_path) + unicode_bundle = read_metadata(unicode_path) + + assert ascii_bundle.user_comment_raw == "ASCII prompt\nSteps: 9, Seed: 1" + assert unicode_bundle.user_comment_raw == "UNICODE prompt\nSteps: 10, Seed: 2" + assert ascii_bundle.source_format == "EXIF UserComment" + assert unicode_bundle.source_format == "EXIF UserComment" + + +def test_exif_user_comment_civitai_metadata_fallback(tmp_path): + from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata + + positive = ( + "A green tree frog wearing a tiny straw hat sitting on a lily pad with " + "a fishing rod made from a twig and spider silk." + ) + resources = [ + { + "type": "checkpoint", + "modelVersionId": 2442439, + "modelName": "Z Image Turbo", + "modelVersionName": "Turbo", + }, + { + "type": "lora", + "weight": 0.5, + "modelVersionId": 2460437, + "modelName": "Midjourney Luneva Cinematic Lora", + "modelVersionName": "LORA: R128-5000step", + }, + { + "type": "lora", + "weight": 0.8, + "modelVersionId": 2515203, + "modelName": "[ZIT] Detail Slider", + "modelVersionName": "v1.0", + }, + ] + civitai_metadata = { + "workflow": "txt2img", + "cfgScale": 1, + "steps": 9, + "seed": 2214950711, + "prompt": positive, + "resources": [ + {"modelVersionId": 2442439, "strength": 1, "type": "Checkpoint"}, + {"modelVersionId": 2460437, "strength": 0.5, "type": "LORA"}, + {"modelVersionId": 2515203, "strength": 0.8, "type": "LORA"}, + ], + } + user_comment = ( + f"{positive}\n" + "Steps: 9, Sampler: Euler, CFG scale: 1, Seed: 2214950711, " + "Size: 832x1216, Model type: Z-Image Turbo, " + f"Civitai resources: {json.dumps(resources)}, " + f"Civitai metadata: {json.dumps(civitai_metadata)}" + ) + image_path = tmp_path / "civitai.jpeg" + _save_jpeg_with_user_comment(image_path, user_comment, prefix=b"UNICODE\x00") + + bundle = read_metadata(image_path) + result = parse_metadata_bundle(bundle) + + assert bundle.user_comment_raw.startswith("A green tree frog") + assert result.positive == positive + assert result.negative == "" + assert result.seed == 2214950711 + assert result.steps == 9 + assert result.cfg == 1.0 + assert result.sampler_name == "Euler" + assert result.width == 832 + assert result.height == 1216 + assert result.model_name == "Z Image Turbo" + assert [(lora.lora_name, lora.strength_model) for lora in result.loras] == [ + ("Midjourney Luneva Cinematic Lora", 0.5), + ("[ZIT] Detail Slider", 0.8), + ] + assert result.partial_result["source_format"] == "EXIF UserComment / Civitai metadata" + assert result.status_message == "EXIF UserComment / Civitai metadata fallback" + assert "NO_METADATA" not in result.setting diff --git a/tests/test_sampler_selector.py b/tests/test_sampler_selector.py index b468dea..aa34bec 100644 --- a/tests/test_sampler_selector.py +++ b/tests/test_sampler_selector.py @@ -152,3 +152,61 @@ def test_extracts_ksampler_advanced_and_custom_field_variants(): "sampler_name": "custom_sampler", "scheduler": "sgm_uniform", } + + +def test_selects_ultimate_sd_upscale_sampler_like_node_upstream_of_save_image(): + prompt = { + "1": {"class_type": "CLIPTextEncode", "inputs": {"text": "positive"}}, + "2": {"class_type": "CLIPTextEncode", "inputs": {"text": "negative"}}, + "7": { + "class_type": "UltimateSDUpscale", + "inputs": { + "positive": ["1", 0], + "negative": ["2", 0], + "model": ["9", 0], + "vae": ["10", 0], + "seed": 123, + "steps": 12, + "cfg": 2.5, + "sampler_name": "euler", + "scheduler_name": "normal", + "tile_width": 768, + "seam_fix_mode": "band pass", + }, + }, + "8": {"class_type": "SaveImage", "inputs": {"images": ["7", 0]}}, + "9": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "model.safetensors"}}, + "10": {"class_type": "VAELoader", "inputs": {"vae_name": "vae.safetensors"}}, + } + + selection = select_sampler(prompt) + + assert selection.sampler.node_id == "7" + assert selection.sampler.class_type == "UltimateSDUpscale" + assert selection.selected_by == "output_chain" + assert "SaveImage 8.images <- UltimateSDUpscale 7" in selection.debug_trace + + +def test_unsupported_output_chain_error_includes_diagnostics(): + from smart_metadata_reader.sampler_selector import select_final_sampler + + prompt = { + "1": { + "class_type": "UnsupportedGenerator", + "inputs": {"image": ["3", 0], "tile_width": 768}, + }, + "2": {"class_type": "SaveImage", "inputs": {"images": ["1", 0]}}, + "3": {"class_type": "LoadImage", "inputs": {"image": "source.png"}}, + } + + try: + select_final_sampler(GraphIndex(prompt=prompt, workflow=None)) + except ValueError as exc: + message = str(exc) + else: + raise AssertionError("Expected unsupported output chain to fail") + + assert "no supported sampler candidate found" in message + assert "final output nodes found: SaveImage 2" in message + assert "SaveImage 2.images <- UnsupportedGenerator 1" in message + assert "input keys: image, tile_width" in message