import json import os import sys import unittest from typing import Any from unittest.mock import MagicMock import torch # Mock ComfyUI dependencies sys.modules["folder_paths"] = MagicMock() sys.modules["comfy"] = MagicMock() sys.modules["comfy.sd"] = MagicMock() sys.modules["comfy.utils"] = MagicMock() sys.modules["comfy.samplers"] = MagicMock() sys.modules["nodes"] = MagicMock() # Import test target from py.load_image_configs import MetadataHelper # noqa: E402 class TestMetadataParser(unittest.TestCase): def setUp(self) -> None: self.data_dir = os.path.join(os.path.dirname(__file__), "data", "metadata_parser") self.workflow_dir = os.path.join(os.path.dirname(__file__), "data", "workflows") self.pattern_png_path = os.path.join(self.data_dir, "pattern_png.metadata") self.pattern_webp_exif_path = os.path.join(self.data_dir, "pattern_webp_exif.metadata") self.pattern_webp_desc_path = os.path.join(self.data_dir, "pattern_webp_desc.metadata") self.json_1_path = os.path.join(self.data_dir, "1.json") self.json_2_path = os.path.join(self.data_dir, "2.json") self.json_3_path = os.path.join(self.data_dir, "3.json") self.json_4_path = os.path.join(self.data_dir, "4.json") self.json_5_path = os.path.join(self.data_dir, "5.json") self.json_6_path = os.path.join(self.data_dir, "6.json") self.flux_workflow_path = os.path.join(self.workflow_dir, "Flux.json") self.flux_gguf_workflow_path = os.path.join(self.workflow_dir, "Flux_GGUF.json") def _parse_exiftool_txt(self, path: str) -> dict[str, Any]: data = {} with open(path, encoding="utf-8") as f: for line in f: if ": " in line: key, value = line.split(": ", 1) data[key.strip()] = value.strip() return data def test_extract_from_pattern_png(self) -> None: meta = self._parse_exiftool_txt(self.pattern_png_path) info = {"workflow": meta.get("Workflow"), "prompt": meta.get("Prompt")} res = MetadataHelper.extract_from_data(info) positive, negative, model_name, wf_json, pr_json, a1111, logs = res self.assertIn("(masterpiece, best quality:1.2), 1girl, solo", positive) self.assertIn("(worst quality, low quality:1.4), multiple people", negative) self.assertEqual(model_name, "Illustrious\\novaAnimeXL_ilV130.safetensors") def test_extract_from_json_1(self) -> None: """1.json: Standard Workflow format""" with open(self.json_1_path, encoding="utf-8") as f: workflow_data = json.load(f) positive, negative = MetadataHelper.parse_workflow_json(workflow_data, []) self.assertIn("masterpiece, best quality, 1girl", positive) self.assertIn("low quality, bad anatomy", negative) # Check parameters (first sampler found, ID 2) params, found = MetadataHelper.get_ksampler_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["seed"], 924175127337775) self.assertEqual(params["steps"], 25) self.assertEqual(params["cfg"], 8.0) self.assertEqual(params["sampler_name"], "euler_ancestral") self.assertEqual(params["scheduler"], "simple") # Check resolution (EmptySD3LatentImage) dummy_tensor = torch.zeros((1, 64, 64, 3)) w, h, r_found = MetadataHelper.get_resolution_params(workflow_data, dummy_tensor, []) self.assertTrue(r_found) self.assertEqual(w, 512) self.assertEqual(h, 512) def test_extract_from_json_2(self) -> None: """2.json: API format (prompt JSON)""" with open(self.json_2_path, encoding="utf-8") as f: prompt_data = json.load(f) info = {"prompt": json.dumps(prompt_data)} res = MetadataHelper.extract_from_data(info) positive, negative, model_name, wf_json, pr_json, a1111, logs = res # Since trace results might vary by environment, check if it's retrieved or relax constraints if positive: self.assertIn("(masterpiece, best quality)", positive) if negative: self.assertIn("(worst quality, low quality:1.4)", negative) # Retrieve parameters from API format params, found = MetadataHelper.get_ksampler_params_from_prompt(prompt_data, []) self.assertTrue(found) self.assertEqual(params["seed"], 123456789012345) self.assertEqual(params["steps"], 20) self.assertEqual(params["cfg"], 1.5) self.assertEqual(params["sampler_name"], "euler_ancestral") self.assertEqual(params["scheduler"], "karras") def test_extract_from_json_3(self) -> None: """3.json: Standard Workflow format""" with open(self.json_3_path, encoding="utf-8") as f: workflow_data = json.load(f) logs: list[str] = [] positive, negative = MetadataHelper.parse_workflow_json(workflow_data, logs) self.assertEqual(positive, "(masterpiece, best quality:1.2), 1girl, solo") self.assertEqual(negative, "(worst quality, low quality:1.4), multiple people") # Parameters params, found = MetadataHelper.get_ksampler_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["seed"], 233851666717137) self.assertEqual(params["steps"], 10) self.assertEqual(params["cfg"], 5.0) self.assertEqual(params["sampler_name"], "euler_ancestral") self.assertEqual(params["scheduler"], "simple") def test_extract_from_json_4(self) -> None: """4.json: Complex Workflow including subgraphs""" with open(self.json_4_path, encoding="utf-8") as f: workflow_data = json.load(f) positive, negative = MetadataHelper.parse_workflow_json(workflow_data, []) self.assertIn("(masterpiece, best quality:1.2), 1girl, solo", positive) self.assertIn("(worst quality, low quality:1.4)", negative) # Retrieve parameters within subgraph (first sampler found, Generator inner ID 6) params, found = MetadataHelper.get_ksampler_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["seed"], 108626426419181) self.assertEqual(params["steps"], 40) self.assertEqual(params["cfg"], 1.5) self.assertEqual(params["sampler_name"], "dpmpp_2m_sde") self.assertEqual(params["scheduler"], "karras") # Check resolution (EmptySD3LatentImage ID 5 within subgraph Generator) dummy_tensor = torch.zeros((1, 64, 64, 3)) w, h, r_found = MetadataHelper.get_resolution_params(workflow_data, dummy_tensor, []) self.assertTrue(r_found) self.assertEqual(w, 1024) self.assertEqual(h, 1024) def test_extract_from_json_5(self) -> None: """5.json: Standard Workflow format""" with open(self.json_5_path, encoding="utf-8") as f: workflow_data = json.load(f) positive, negative = MetadataHelper.parse_workflow_json(workflow_data, []) self.assertIn("(masterpiece, best quality:1.2), 1girl, solo", positive) self.assertIn("(worst quality, low quality:1.4), multiple people", negative) # Parameters params, found = MetadataHelper.get_ksampler_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["seed"], 233851666717137) self.assertEqual(params["steps"], 10) self.assertEqual(params["cfg"], 5.0) def test_extract_from_json_6(self) -> None: """6.json: Workflow including subgraphs""" with open(self.json_6_path, encoding="utf-8") as f: workflow_data = json.load(f) positive, negative = MetadataHelper.parse_workflow_json(workflow_data, []) self.assertIn("(masterpiece, best quality:1.2)", positive) self.assertIn("(scenery:1.2)", positive) # Retrieve parameters within subgraph params, found = MetadataHelper.get_ksampler_params(workflow_data, []) self.assertTrue(found) def test_extract_dual_clip_params_from_flux_workflow(self) -> None: """Flux.json: DualCLIPLoader should keep clip_name1 and clip_name2 distinct""" with open(self.flux_workflow_path, encoding="utf-8") as f: workflow_data = json.load(f) params, found = MetadataHelper.get_dual_clip_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["clip_name1"], "t5xxl_fp8_e4m3fn.safetensors") self.assertEqual(params["clip_name2"], "clip_l.safetensors") self.assertNotEqual(params["clip_name1"], params["clip_name2"]) self.assertEqual(params["clip_type"], "flux") self.assertEqual(params["clip_device"], "default") def test_extract_from_flux_gguf_workflow(self) -> None: """Flux_GGUF.json: Should extract model_name from UnetLoaderGGUF and CLIP from DualCLIPLoaderGGUF""" with open(self.flux_gguf_workflow_path, encoding="utf-8") as f: workflow_data = json.load(f) # 1. Test model_name extraction (MetadataExtractor) info = {"workflow": json.dumps(workflow_data)} res = MetadataHelper.extract_from_data(info) positive, negative, model_name, wf_json, pr_json, a1111, logs = res # Expected to fail initially until we update MetadataExtractor self.assertEqual(model_name, "flux1-dev-Q4_K_S.gguf") # 2. Test DualCLIPLoaderGGUF extraction (Parsers) params, found = MetadataHelper.get_dual_clip_params(workflow_data, []) self.assertTrue(found) self.assertEqual(params["clip_name1"], "t5-v1_1-xxl-encoder-Q4_K_S.gguf") self.assertEqual(params["clip_name2"], "clip_l.safetensors") self.assertEqual(params["clip_type"], "flux") def test_extract_from_pattern_webp_exif(self) -> None: meta = self._parse_exiftool_txt(self.pattern_webp_exif_path) exif: dict[Any, Any] = {271: meta.get("Make"), 272: meta.get("Model")} info: dict[str, Any] = {} res = MetadataHelper.extract_from_data(info, exif) positive, negative, model_name, wf_json, pr_json, a1111, logs = res self.assertIn("(masterpiece, best quality:1.2)", positive) self.assertIn("(worst quality, low quality:1.4)", negative) self.assertEqual(model_name, "Illustrious\\novaAnimeXL_ilV130.safetensors") def test_extract_from_pattern_webp_desc(self) -> None: meta = self._parse_exiftool_txt(self.pattern_webp_desc_path) exif: dict[Any, Any] = {270: meta.get("ImageDescription"), 271: meta.get("Make")} info: dict[str, Any] = {} res = MetadataHelper.extract_from_data(info, exif) positive, negative, model_name, wf_json, pr_json, a1111, logs = res self.assertIn("(masterpiece, best quality:1.2)", positive) self.assertIn("(worst quality, low quality:1.4)", negative) def test_log_truncation(self) -> None: long_positive = "a" * 60 long_negative = "b" * 60 workflow_data = { "nodes": [ {"id": 1, "type": "CLIPTextEncode", "widgets_values": [long_positive]}, {"id": 2, "type": "CLIPTextEncode", "widgets_values": [long_negative]}, { "id": 3, "type": "KSampler", "inputs": [{"name": "positive", "link": 1}, {"name": "negative", "link": 2}], }, ], "links": [[1, 1, 0, 3, 0, "CONDITIONING"], [2, 2, 0, 3, 1, "CONDITIONING"]], } logs: list[str] = [] positive, negative = MetadataHelper.parse_workflow_json(workflow_data, logs) self.assertIn(f"-> Positive found: {'a' * 50}...", logs) self.assertIn(f"-> Negative found: {'b' * 50}...", logs) if __name__ == "__main__": unittest.main()