import json 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 _save_png_with_prompt(path, prompt, workflow=None): pnginfo = PngInfo() pnginfo.add_text("prompt", json.dumps(prompt)) if workflow is not None: pnginfo.add_text("workflow", json.dumps(workflow)) Image.new("RGB", (32, 24), (20, 40, 60)).save(path, pnginfo=pnginfo) def _simple_prompt_graph(positive="original positive prompt", negative="original negative prompt"): return { "1": {"class_type": "SaveImage", "inputs": {"images": ["2", 0]}}, "2": {"class_type": "VAEDecode", "inputs": {"samples": ["3", 0]}}, "3": { "class_type": "KSampler", "inputs": { "positive": ["4", 0], "negative": ["5", 0], "seed": 111, "steps": 12, "cfg": 5.5, "sampler_name": "euler", "scheduler": "normal", }, }, "4": {"class_type": "CLIPTextEncode", "inputs": {"text": positive}}, "5": {"class_type": "CLIPTextEncode", "inputs": {"text": negative}}, } def _upscale_prompt_graph(reader_image="003.png"): return { "22": { "class_type": "FaceDetailer", "inputs": { "image": ["24", 0], "positive": ["26", 0], "negative": ["27", 0], }, }, "24": { "class_type": "UltimateSDUpscale", "inputs": { "image": ["40", 0], "positive": ["26", 0], "negative": ["27", 0], "seed": 222, "steps": 9, "cfg": 1.5, "sampler_name": "Euler", "scheduler": "normal", "tile_width": 768, "seam_fix_mode": "half tile", }, }, "26": {"class_type": "CLIPTextEncode", "inputs": {"text": ["37", 0]}}, "27": {"class_type": "CLIPTextEncode", "inputs": {"text": ["38", 0]}}, "36": {"class_type": "SmartMetadataReader", "inputs": {"image": reader_image}}, "37": {"class_type": "ShowText|pysssss", "inputs": {"text": ["36", 2]}}, "38": {"class_type": "ShowText|pysssss", "inputs": {"text": ["36", 3]}}, "40": {"class_type": "LoadImage", "inputs": {"image": "source.png"}}, "90": { "class_type": "GeminiChatNode", "inputs": {"prompt": "unused Gemini prompt template"}, }, "91": {"class_type": "ShowText|pysssss", "inputs": {"text": ["90", 0]}}, "99": {"class_type": "PreviewImage", "inputs": {"images": ["22", 0]}}, } def _upscale_workflow(reader_image="003.png"): return { "nodes": [ { "id": 36, "type": "SmartMetadataReader", "inputs": {"image": reader_image}, "outputs": [ {"name": "image"}, {"name": "mask"}, {"name": "positive"}, {"name": "negative"}, {"name": "seed"}, {"name": "steps"}, {"name": "cfg"}, {"name": "width"}, {"name": "height"}, {"name": "model_name"}, {"name": "filename"}, {"name": "setting"}, ], }, { "id": 91, "type": "ShowText|pysssss", "inputs": {"text": ["90", 0]}, "widgets_values": ["wrong unused prompt"], }, ] } def test_read_metadata_extracts_png_text_chunks_and_dimensions(tmp_path): from smart_metadata_reader.metadata_reader import read_metadata image_path = tmp_path / "sample.png" prompt_graph = { "1": { "class_type": "CLIPTextEncode", "inputs": {"text": "positive trigger"}, } } workflow_graph = { "nodes": [ { "id": 1, "type": "CLIPTextEncode", "inputs": {"text": "positive trigger"}, } ] } pnginfo = PngInfo() pnginfo.add_text("prompt", json.dumps(prompt_graph)) pnginfo.add_text("workflow", json.dumps(workflow_graph)) pnginfo.add_text( "parameters", ( "positive\n" "Negative prompt: negative\n" "Steps: 20, Sampler: Euler, CFG scale: 7, Seed: 123, Size: 512x768" ), ) Image.new("RGBA", (64, 32), (255, 0, 0, 128)).save( image_path, pnginfo=pnginfo ) bundle = read_metadata(image_path) assert bundle.filename == "sample.png" assert bundle.width == 64 assert bundle.height == 32 assert json.loads(bundle.prompt_raw) == prompt_graph assert json.loads(bundle.workflow_raw) == workflow_graph assert "Negative prompt: negative" in bundle.parameters_raw assert bundle.source_format == "ComfyUI prompt/workflow" def test_load_image_and_mask_returns_batched_image_and_mask_shapes(tmp_path): from smart_metadata_reader.image_io import load_image_and_mask image_path = tmp_path / "transparent.png" Image.new("RGBA", (10, 12), (10, 20, 30, 128)).save(image_path) image, mask, width, height = load_image_and_mask(image_path) assert width == 10 assert height == 12 assert tuple(image.shape) == (1, 12, 10, 3) assert tuple(mask.shape) == (1, 12, 10) def test_parse_standard_a1111_parameters(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters parameters = ( "masterpiece, character trigger\n" "Negative prompt: blurry, low quality\n" "Steps: 25, Sampler: Euler a, CFG scale: 7.5, Seed: 12345, " "Size: 768x1024, Model: demo_model" ) result = parse_a1111_parameters(parameters) assert result.positive == "masterpiece, character trigger" assert result.negative == "blurry, low quality" assert result.steps == 25 assert result.sampler_name == "Euler a" assert result.cfg == 7.5 assert result.seed == 12345 assert result.width == 768 assert result.height == 1024 assert result.model_name == "demo_model" assert result.status_message == "A1111 parameters fallback" assert result.partial_result["source_format"] == "A1111 parameters" def test_parse_a1111_preserves_multiline_positive_and_negative_prompts(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters parameters = ( "masterpiece, best quality\n" "character trigger, detailed eyes\n" "Negative prompt: low quality\n" "bad hands, extra fingers\n" "Steps: 30, Sampler: DPM++ 2M, CFG scale: 6, Seed: 999, Size: 512x768" ) result = parse_a1111_parameters(parameters) assert result.positive == "masterpiece, best quality\ncharacter trigger, detailed eyes" assert result.negative == "low quality\nbad hands, extra fingers" assert result.steps == 30 assert result.seed == 999 def test_parse_a1111_preserves_lora_triggers_nsfw_and_user_tags(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters parameters = ( ", role_lora_trigger, nsfw, nude, user handwritten tag\n" "Negative prompt: low quality\n" "Steps: 20, Sampler: Euler, CFG scale: 7, Seed: 101, Size: 640x960" ) result = parse_a1111_parameters(parameters) assert result.positive == ( ", role_lora_trigger, nsfw, nude, user handwritten tag" ) assert result.negative == "low quality" def test_parse_a1111_without_negative_prompt_keeps_positive_and_empty_negative(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters parameters = ( "solo, dramatic lighting\n" "Steps: 12, Sampler: Euler, CFG scale: 4.5, Seed: 42, Size: 320x448" ) result = parse_a1111_parameters(parameters) assert result.positive == "solo, dramatic lighting" assert result.negative == "" assert result.steps == 12 assert result.seed == 42 def test_parse_a1111_missing_partial_parameters_uses_defaults(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters result = parse_a1111_parameters("simple prompt only") assert result.positive == "simple prompt only" assert result.negative == "" assert result.seed == -1 assert result.steps == 0 assert result.cfg == 0.0 assert result.width == 0 assert result.height == 0 assert result.model_name == "" def test_parse_a1111_scheduler_and_model_hash_metadata(): from smart_metadata_reader.a1111_parser import parse_a1111_parameters parameters = ( "positive text\n" "Negative prompt: negative text\n" "Steps: 25, Sampler: DPM++ 2M, Scheduler: Karras, CFG scale: 7.25, " "Seed: 123, Size: 768x1024, Model hash: abc123, Model: forge_model" ) result = parse_a1111_parameters(parameters) 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 def test_nested_smart_metadata_reader_outputs_resolve_referenced_image_prompts(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata original_path = tmp_path / "003.png" _save_png_with_prompt(original_path, _simple_prompt_graph()) upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt( upscale_path, _upscale_prompt_graph(), workflow=_upscale_workflow(), ) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "original positive prompt" assert result.negative == "original negative prompt" assert result.status_message != "FAILED" assert not (result.status_message == "OK" and result.positive == "" and result.negative == "") assert "Resolved SmartMetadataReader positive output from referenced image 003.png" in result.debug_trace assert "Resolved SmartMetadataReader negative output from referenced image 003.png" in result.debug_trace def test_unconnected_showtext_workflow_cache_does_not_pollute_nested_reader_prompt(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata original_path = tmp_path / "003.png" _save_png_with_prompt(original_path, _simple_prompt_graph()) upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt( upscale_path, _upscale_prompt_graph(), workflow=_upscale_workflow(), ) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "original positive prompt" assert "wrong unused prompt" not in result.positive assert "wrong unused prompt" not in result.negative def test_nested_smart_metadata_reader_falls_back_to_current_output_contract_without_workflow_outputs(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata original_path = tmp_path / "003.png" _save_png_with_prompt(original_path, _simple_prompt_graph()) upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt(upscale_path, _upscale_prompt_graph()) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "original positive prompt" assert result.negative == "original negative prompt" def test_nested_smart_metadata_reader_prefers_workflow_output_names_over_indices(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata original_path = tmp_path / "003.png" _save_png_with_prompt(original_path, _simple_prompt_graph()) prompt = _upscale_prompt_graph() prompt["37"]["inputs"]["text"] = ["36", 3] prompt["38"]["inputs"]["text"] = ["36", 2] workflow = _upscale_workflow() workflow["nodes"][0]["outputs"][2] = {"name": "negative"} workflow["nodes"][0]["outputs"][3] = {"name": "positive"} upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt(upscale_path, prompt, workflow=workflow) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "original positive prompt" assert result.negative == "original negative prompt" def test_nested_smart_metadata_reader_missing_image_returns_partial_unresolved(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt( upscale_path, _upscale_prompt_graph(reader_image="missing.png"), workflow=_upscale_workflow(reader_image="missing.png"), ) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "" assert result.negative == "" assert result.status_message == "PARTIAL" assert result.confidence <= 0.6 assert result.partial_result["unresolved"][0]["class_type"] == "SmartMetadataReader" assert result.partial_result["unresolved"][0]["output_index"] == 2 assert result.partial_result["unresolved"][0]["resolved_output_role"] == "positive" assert result.partial_result["unresolved"][0]["referenced_image_filename"] == "missing.png" assert "referenced image not found" in result.partial_result["unresolved"][0]["reason"] assert "Unresolved:" in result.setting def test_nested_smart_metadata_reader_rejects_unsafe_image_paths(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata for unsafe_path in ("../secret.png", "C:\\Users\\xxx\\secret.png"): upscale_path = tmp_path / f"upscaled-{unsafe_path[0].encode().hex()}.png" _save_png_with_prompt( upscale_path, _upscale_prompt_graph(reader_image=unsafe_path), workflow=_upscale_workflow(reader_image=unsafe_path), ) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "" assert result.status_message == "PARTIAL" assert result.partial_result["unresolved"][0]["class_type"] == "SmartMetadataReader" assert "unsafe referenced image path" in result.partial_result["unresolved"][0]["reason"] def test_nested_smart_metadata_reader_respects_actual_connected_output_role(tmp_path): from smart_metadata_reader.metadata_reader import parse_metadata_bundle, read_metadata original_path = tmp_path / "003.png" _save_png_with_prompt(original_path, _simple_prompt_graph()) prompt = _upscale_prompt_graph() prompt["37"]["inputs"]["text"] = ["36", 3] upscale_path = tmp_path / "upscaled.png" _save_png_with_prompt(upscale_path, prompt, workflow=_upscale_workflow()) result = parse_metadata_bundle(read_metadata(upscale_path)) assert result.positive == "original negative prompt" assert "SmartMetadataReader output negative used in positive chain" in result.debug_trace