mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-23 20:14:08 -03:00
feat: add image metadata loader with native LoRA Manager integration
Add Load Image Metadata (LoraManager) to extract reusable prompts, model references, LoRA stacks, and sampling settings from images. Prefer saved A1111-style parameters by default, with optional workflow and subgraph sampler selection. Resolve local model and LoRA names, report missing resources, and recover extraction failures with explicit defaults and readable diagnostics. Include parser, resource-resolution, and node regression tests, plus usage documentation.
This commit is contained in:
Martial Michel
@@ -0,0 +1,430 @@
|
||||
import json
|
||||
import sys
|
||||
import types
|
||||
from pathlib import Path
|
||||
|
||||
import piexif
|
||||
import piexif.helper
|
||||
import pytest
|
||||
from PIL import Image, PngImagePlugin
|
||||
|
||||
from py.nodes.load_image_metadata import LoadImageMetadataLM, MetadataError, resolve_resource
|
||||
from py.utils.exif_utils import ExifUtils
|
||||
|
||||
|
||||
PARAMETERS = 'cat <lora:style:0.7:0.2>\nNegative prompt: blur\nSteps: 25, Sampler: Euler, Schedule type: Normal, CFG scale: 6.5, Seed: 18446744073709551615, Size: 768x1024, Model: base'
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def runtime(tmp_path, monkeypatch):
|
||||
import comfy
|
||||
import folder_paths
|
||||
import nodes
|
||||
|
||||
image_path = tmp_path / "input.png"
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("parameters", PARAMETERS)
|
||||
Image.new("RGB", (16, 24)).save(image_path, pnginfo=info)
|
||||
model = tmp_path / "base.safetensors"
|
||||
lora = tmp_path / "style.safetensors"
|
||||
model.touch()
|
||||
lora.touch()
|
||||
library = ([{"file_path": str(model), "sub_type": "checkpoint"}], [str(tmp_path)], [{"file_path": str(lora)}], [str(tmp_path)])
|
||||
monkeypatch.setattr(LoadImageMetadataLM, "_library", staticmethod(lambda: library))
|
||||
monkeypatch.setattr(folder_paths, "get_annotated_filepath", lambda name: str(image_path), raising=False)
|
||||
monkeypatch.setattr(folder_paths, "exists_annotated_filepath", lambda name: image_path.exists(), raising=False)
|
||||
pixels = types.SimpleNamespace(shape=(1, 24, 16, 3))
|
||||
mask = object()
|
||||
class LoadImage:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {"required": {"image": (["input.png"], {"image_upload": True})}}
|
||||
|
||||
def load_image(self, name):
|
||||
return pixels, mask
|
||||
monkeypatch.setattr(nodes, "LoadImage", LoadImage, raising=False)
|
||||
samplers = types.ModuleType("comfy.samplers")
|
||||
samplers.KSampler = types.SimpleNamespace(SAMPLERS=["euler", "dpmpp_2m"], SCHEDULERS=["normal", "karras"])
|
||||
monkeypatch.setitem(sys.modules, "comfy.samplers", samplers)
|
||||
monkeypatch.setattr(comfy, "samplers", samplers, raising=False)
|
||||
return image_path, library, pixels, mask
|
||||
|
||||
|
||||
def test_full_node_contract_with_real_png_metadata(runtime):
|
||||
_, library, pixels, mask = runtime
|
||||
result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert len(result) == len(LoadImageMetadataLM.RETURN_TYPES)
|
||||
assert result[:4] == (pixels, mask, "cat", "blur")
|
||||
assert result[5] == [(library[2][0]["file_path"], .7, .2)]
|
||||
assert result[7:15] == (2**64 - 1, 25, 6.5, "euler", "normal", 768, 1024, 1.0)
|
||||
assert "Resolved 1 LoRA" in result[15]
|
||||
assert LoadImageMetadataLM.INPUT_TYPES()["required"]["image"][1]["image_upload"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("extension", ["webp", "jpg"])
|
||||
def test_exif_parameters_from_real_image(runtime, extension):
|
||||
image_path, *_ = runtime
|
||||
exif = piexif.dump({"Exif": {piexif.ExifIFD.UserComment: piexif.helper.UserComment.dump(PARAMETERS, encoding="unicode")}})
|
||||
alternate = image_path.with_suffix("." + extension)
|
||||
Image.new("RGB", (16, 24)).save(alternate, exif=exif)
|
||||
fields = ExifUtils._load_structured_metadata(str(alternate))
|
||||
assert "Steps: 25" in fields["parameters"]
|
||||
|
||||
|
||||
def test_missing_lora_strict_or_explicit_skip(runtime):
|
||||
runtime[1][2].clear()
|
||||
strict_result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert strict_result[5] == []
|
||||
assert "LoRA: style | model weight: 0.7 | CLIP weight: 0.2" in strict_result[17]
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="use_defaults")
|
||||
assert result[5] == []
|
||||
assert "Skipped LoRA" in result[15]
|
||||
|
||||
|
||||
def test_overrides_replace_loras_and_preserve_large_seed(runtime):
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json=json.dumps({"seed": 2**64 - 2, "loras": [], "positive": "changed"}))
|
||||
assert result[2] == "changed"
|
||||
assert result[5] == []
|
||||
assert result[7] == 2**64 - 2
|
||||
|
||||
|
||||
def test_no_metadata_can_be_inspected_with_defaults(runtime):
|
||||
Image.new("RGB", (16, 24)).save(runtime[0])
|
||||
assert LoadImageMetadataLM().load_metadata("input.png")[12:14] == (1024, 1024)
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="use_defaults")
|
||||
assert result[12:14] == (1024, 1024)
|
||||
assert "No model resolved" in result[15]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("override", [{"seed": -1}, {"steps": 2.5}, {"cfg": float("nan")}, {"sampler_name": "made_up"}, {"positive": ["1", 0]}, {"unknown": 1}])
|
||||
def test_invalid_override_rejected(runtime, override):
|
||||
with pytest.raises((MetadataError, ValueError)):
|
||||
LoadImageMetadataLM().load_metadata("input.png", overrides_json=json.dumps(override))
|
||||
|
||||
|
||||
def test_duplicate_basenames_require_path(tmp_path):
|
||||
items = []
|
||||
for folder in ("a", "b"):
|
||||
directory = tmp_path / folder
|
||||
directory.mkdir()
|
||||
path = directory / "same.safetensors"
|
||||
path.touch()
|
||||
items.append({"file_path": str(path)})
|
||||
with pytest.raises(MetadataError, match="Ambiguous"):
|
||||
resolve_resource("same", items, [str(tmp_path)])
|
||||
assert resolve_resource("b/same.safetensors", items, [str(tmp_path)]) == items[1]
|
||||
assert resolve_resource("b/same", items, [str(tmp_path)]) == items[1]
|
||||
|
||||
|
||||
def test_file_hash_detects_replacement_and_accepts_all_inputs(runtime):
|
||||
before = LoadImageMetadataLM.IS_CHANGED("input.png", sampler_node_id="", missing_settings="strict", overrides_json="{}")
|
||||
Image.new("RGB", (32, 32)).save(runtime[0])
|
||||
assert before != LoadImageMetadataLM.IS_CHANGED("input.png")
|
||||
|
||||
|
||||
def test_comfy_webp_exif_prompt_fields(runtime):
|
||||
image_path, *_ = runtime
|
||||
graph = {"1": {"class_type": "KSampler", "inputs": {"seed": 42}}}
|
||||
exif = piexif.dump({"0th": {
|
||||
piexif.ImageIFD.Make: "prompt:" + json.dumps(graph),
|
||||
piexif.ImageIFD.Model: 'workflow:{"nodes": []}',
|
||||
}})
|
||||
alternate = image_path.with_suffix(".webp")
|
||||
Image.new("RGB", (16, 24)).save(alternate, exif=exif)
|
||||
fields = ExifUtils._load_structured_metadata(str(alternate))
|
||||
assert json.loads(fields["prompt"]) == graph
|
||||
assert json.loads(fields["workflow"]) == {"nodes": []}
|
||||
|
||||
|
||||
|
||||
def test_report_preserves_extracted_names_without_catalog(runtime):
|
||||
runtime[1][0].clear()
|
||||
runtime[1][2].clear()
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="use_defaults")
|
||||
payload = json.loads(result[15].split("\n\n", 1)[1])
|
||||
assert result[4:7] == ("", [], "")
|
||||
assert payload["source_resources"]["checkpoint_name"] == "base"
|
||||
assert payload["source_resources"]["loras"] == [["style", .7, .2]]
|
||||
|
||||
|
||||
# These user-provided images are optional local integration fixtures, not assets
|
||||
# required by the public test suite.
|
||||
_SAMPLE_PNGS = sorted((Path(__file__).resolve().parents[2] / "_tmp").glob("*.png"))
|
||||
_SAMPLE_PNGS = [path for path in _SAMPLE_PNGS if path.stem.endswith("_")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("sample", _SAMPLE_PNGS or [pytest.param(None, marks=pytest.mark.skip(reason="No local PNG samples"))], ids=lambda path: path.name if path else "no-samples")
|
||||
def test_local_png_node_without_catalog(runtime, monkeypatch, sample):
|
||||
import comfy.samplers
|
||||
import folder_paths
|
||||
|
||||
runtime[1][0].clear()
|
||||
runtime[1][2].clear()
|
||||
monkeypatch.setattr(folder_paths, "get_annotated_filepath", lambda name: str(sample))
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SAMPLERS", ["euler", "euler_ancestral", "er_sde"])
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SCHEDULERS", ["normal", "simple", "sgm_uniform"])
|
||||
result = LoadImageMetadataLM().load_metadata(sample.name, missing_settings="use_defaults")
|
||||
payload = json.loads(result[15].split("\n\n", 1)[1])
|
||||
assert result[2] and result[3]
|
||||
assert "<lora:" not in result[2]
|
||||
assert result[7] == int(sample.stem.split("_")[-3])
|
||||
assert result[4:7] == ("", [], "")
|
||||
assert "Default " not in result[15]
|
||||
assert "Replaced unsupported" not in result[15]
|
||||
assert payload["source_resources"]["checkpoint_name"] in sample.name
|
||||
expected_count = 0 if any(name in sample.name for name in ("hyphoria", "pieModelsAnima")) else 1
|
||||
assert len(payload["source_resources"]["loras"]) == expected_count
|
||||
|
||||
|
||||
@pytest.mark.parametrize("chunk_type", [b"tEXt", b"zTXt", b"iTXt"])
|
||||
def test_png_metadata_after_pixel_data_is_read(runtime, chunk_type):
|
||||
import struct
|
||||
import zlib
|
||||
|
||||
image_path = runtime[0]
|
||||
Image.new("RGB", (16, 24)).save(image_path)
|
||||
original = image_path.read_bytes()
|
||||
encoded = PARAMETERS.encode("utf-8")
|
||||
if chunk_type == b"zTXt":
|
||||
payload = b"parameters\0\0" + zlib.compress(encoded)
|
||||
elif chunk_type == b"iTXt":
|
||||
payload = b"parameters\0\0\0\0\0" + encoded
|
||||
else:
|
||||
payload = b"parameters\0" + encoded
|
||||
chunk = (struct.pack(">I", len(payload)) + chunk_type + payload
|
||||
+ struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF))
|
||||
# Place metadata immediately before IEND, after all pixel data.
|
||||
image_path.write_bytes(original[:-12] + chunk + original[-12:])
|
||||
result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert result[2:4] == ("cat", "blur")
|
||||
assert result[4] == "base.safetensors"
|
||||
assert result[7] == 2**64 - 1
|
||||
|
||||
|
||||
def test_missing_metadata_report_identifies_actual_file(runtime):
|
||||
Image.new("RGB", (16, 24)).save(runtime[0])
|
||||
message = LoadImageMetadataLM().load_metadata("input.png")[15]
|
||||
assert str(runtime[0]) in message
|
||||
assert "Format: PNG" in message
|
||||
assert "metadata keys: (none)" in message
|
||||
assert "settings were not extracted" in message
|
||||
|
||||
|
||||
|
||||
def test_readable_report_contains_settings_prompts_and_missing_resources(runtime):
|
||||
runtime[1][0].clear()
|
||||
runtime[1][2].clear()
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="use_defaults")
|
||||
readable = result[16]
|
||||
assert LoadImageMetadataLM.RETURN_NAMES[16] == "readable_report"
|
||||
assert "Checkpoint recorded in image: base" in readable
|
||||
assert "No local model resolved." in readable
|
||||
assert "Seed: 18446744073709551615" in readable
|
||||
assert "Sampler: euler" in readable
|
||||
assert "Size: 768 × 1024" in readable
|
||||
assert "style (model: 0.7, CLIP: 0.2)" in readable
|
||||
assert "Resolved locally: 0 of 1 requested entries." in readable
|
||||
assert "POSITIVE PROMPT\ncat" in readable
|
||||
assert "NEGATIVE PROMPT\nblur" in readable
|
||||
assert "WARNING" in readable
|
||||
assert json.loads(result[15].split("\n\n", 1)[1])["seed"] == 2**64 - 1
|
||||
|
||||
|
||||
|
||||
def test_empty_metadata_starter_respects_overrides_and_indexed_model(runtime):
|
||||
Image.new("RGB", (16, 24)).save(runtime[0])
|
||||
base = runtime[0].parent / "sd_xl_base_1.0.safetensors"
|
||||
base.touch()
|
||||
runtime[1][0].append({"file_path": str(base), "sub_type": "checkpoint"})
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json='{"seed": 123, "positive": "custom prompt", "width": 768}')
|
||||
assert result[2] == "custom prompt"
|
||||
assert result[4] == base.name
|
||||
assert result[7] == 123
|
||||
assert result[12:14] == (768, 1024)
|
||||
assert result[5] == []
|
||||
|
||||
|
||||
def test_user_example_png_runs_with_saved_strict_setting(runtime, monkeypatch):
|
||||
import folder_paths
|
||||
|
||||
path = Path(__file__).resolve().parents[2] / "_tmp" / "example.png"
|
||||
if not path.exists():
|
||||
pytest.skip("No local example.png fixture")
|
||||
monkeypatch.setattr(folder_paths, "get_annotated_filepath", lambda name: str(path))
|
||||
assert not any(ExifUtils._load_structured_metadata(str(path)).values())
|
||||
runtime[1][0].clear()
|
||||
runtime[1][2].clear()
|
||||
result = LoadImageMetadataLM().load_metadata("example.png", missing_settings="strict")
|
||||
assert "glass bottle" in result[2]
|
||||
assert result[3] == "text, watermark"
|
||||
assert result[4:7] == ("", [], "")
|
||||
assert result[7:15] == (0, 20, 7.0, "euler", "normal", 1024, 1024, 1.0)
|
||||
assert "starter preset" in result[16]
|
||||
|
||||
|
||||
|
||||
def test_missing_files_includes_model_and_lora_in_strict_mode(runtime):
|
||||
runtime[1][0].clear()
|
||||
runtime[1][2].clear()
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="strict")
|
||||
assert result[4:7] == ("", [], "")
|
||||
assert "Model: base" in result[17]
|
||||
assert "LoRA: style | model weight: 0.7 | CLIP weight: 0.2" in result[17]
|
||||
assert LoadImageMetadataLM.RETURN_NAMES[17] == "missing_files"
|
||||
|
||||
|
||||
def test_missing_files_keeps_valid_stack_entries(runtime):
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json=json.dumps({"loras": [["style", .7, .2], ["missing", -.5, 0]]}))
|
||||
assert result[5] == [(runtime[1][2][0]["file_path"], .7, .2)]
|
||||
assert "LoRA: missing | model weight: -0.5 | CLIP weight: 0" in result[17]
|
||||
assert "LoRA: style" not in result[17]
|
||||
assert LoadImageMetadataLM().load_metadata("input.png")[17] == ""
|
||||
|
||||
|
||||
@pytest.mark.parametrize("subtype", ["checkpoint", "diffusion_model"])
|
||||
def test_generic_model_name_resolves_both_model_categories(runtime, subtype):
|
||||
runtime[1][0][0]["sub_type"] = subtype
|
||||
result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert result[4] == "base.safetensors"
|
||||
assert result[17] == ""
|
||||
assert subtype in result[16]
|
||||
assert LoadImageMetadataLM.RETURN_NAMES[4:7] == ("model_name", "lora_stack", "lora_stack_text")
|
||||
assert result[6] == f"{runtime[1][2][0]['file_path']} | model weight: 0.7 | CLIP weight: 0.2"
|
||||
|
||||
|
||||
def test_duplicate_model_names_across_categories_require_path(runtime):
|
||||
directory = runtime[0].parent / "unet"
|
||||
directory.mkdir()
|
||||
model = directory / "base.safetensors"
|
||||
model.touch()
|
||||
runtime[1][0].append({"file_path": str(model), "sub_type": "diffusion_model"})
|
||||
# The exact root-relative name wins when present.
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json='{"model_name":"unet/base.safetensors"}')
|
||||
assert result[4] == "unet/base.safetensors"
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json='{"model_name":"old/base.safetensors"}')
|
||||
assert result[4] == ""
|
||||
assert "Ambiguous" in result[17]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("key", ["model_name", "checkpoint_name", "unet_name"])
|
||||
def test_model_override_aliases(runtime, key):
|
||||
runtime[1][0][0]["sub_type"] = "diffusion_model"
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", overrides_json=json.dumps({key: "base.safetensors"}))
|
||||
assert result[4] == "base.safetensors"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("policy", ["strict", "use_defaults"])
|
||||
def test_unsupported_sampler_returns_defaults_and_error(runtime, policy):
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("prompt", json.dumps({"1": {"class_type": "CustomSampler", "inputs": {}}}))
|
||||
Image.new("RGB", (16, 24)).save(runtime[0], pnginfo=info)
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings=policy)
|
||||
assert result[7:15] == (0, 20, 7.0, "euler", "normal", 1024, 1024, 1.0)
|
||||
assert "glass bottle" in result[2]
|
||||
assert result[5] == []
|
||||
assert "❌ ERROR" in result[16]
|
||||
assert "supported sampler IDs: none" in result[16]
|
||||
assert "⚙️ SAMPLING" in result[16]
|
||||
|
||||
|
||||
def test_unsupported_graph_uses_valid_parameters_before_defaults(runtime):
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("prompt", json.dumps({"1": {"class_type": "CustomSampler", "inputs": {}}}))
|
||||
info.add_text("parameters", PARAMETERS)
|
||||
Image.new("RGB", (16, 24)).save(runtime[0], pnginfo=info)
|
||||
result = LoadImageMetadataLM().load_metadata("input.png", missing_settings="strict", prefer_saved_image_metadata=False)
|
||||
assert result[2] == "cat"
|
||||
assert result[7] == 2**64 - 1
|
||||
assert result[8] == 25
|
||||
assert "recovered saved generation parameters" in result[16]
|
||||
assert "❌ ERROR" in result[16]
|
||||
|
||||
|
||||
def test_invalid_extracted_number_preserves_other_settings(runtime):
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("parameters", PARAMETERS.replace("CFG scale: 6.5", "CFG scale: nan"))
|
||||
Image.new("RGB", (16, 24)).save(runtime[0], pnginfo=info)
|
||||
result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert result[9] == 7.0
|
||||
assert result[8] == 25
|
||||
assert "ERROR: Invalid cfg" in result[16]
|
||||
|
||||
|
||||
|
||||
def test_actual_custom_sampler_png_uses_saved_parameters(runtime, monkeypatch):
|
||||
import comfy.samplers
|
||||
import folder_paths
|
||||
|
||||
path = Path(__file__).resolve().parents[2] / "_tmp" / "20260613-122517_S4_unnamedaANIMA_v10_617459040116303.png"
|
||||
if not path.exists():
|
||||
pytest.skip("No local custom sampler PNG")
|
||||
monkeypatch.setattr(folder_paths, "get_annotated_filepath", lambda name: str(path))
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SAMPLERS", ["euler", "er_sde"])
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SCHEDULERS", ["normal", "simple"])
|
||||
result = LoadImageMetadataLM().load_metadata(path.name, missing_settings="strict", prefer_saved_image_metadata=False)
|
||||
assert result[7:15] == (617459040116303, 30, 4.0, "er_sde", "simple", 1664, 1088, 1.0)
|
||||
assert result[2]
|
||||
assert "❌ ERROR" in result[16]
|
||||
assert "recovered saved generation parameters" in result[16]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ["1481:1783", "1481/1783", "1481", "1783"])
|
||||
def test_actual_png_subgraph_sampler_selection(runtime, monkeypatch, selector):
|
||||
import comfy.samplers
|
||||
import folder_paths
|
||||
|
||||
path = Path(__file__).resolve().parents[2] / "_tmp" / "20260613-122517_S4_unnamedaANIMA_v10_617459040116303.png"
|
||||
if not path.exists():
|
||||
pytest.skip("No local custom sampler PNG")
|
||||
monkeypatch.setattr(folder_paths, "get_annotated_filepath", lambda name: str(path))
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SAMPLERS", ["euler", "er_sde"])
|
||||
monkeypatch.setattr(comfy.samplers.KSampler, "SCHEDULERS", ["normal", "simple"])
|
||||
result = LoadImageMetadataLM().load_metadata(path.name, sampler_node_id=selector, prefer_saved_image_metadata=False)
|
||||
assert result[7:12] == (617459040116303, 30, 4.0, "er_sde", "simple")
|
||||
assert "sampler 1481:1783" in result[16]
|
||||
assert "Detail Daemon" in result[16]
|
||||
assert "recovered saved generation parameters" not in result[16]
|
||||
|
||||
|
||||
def test_source_preference_flag_defaults_true(runtime):
|
||||
assert LoadImageMetadataLM.INPUT_TYPES()["required"]["prefer_saved_image_metadata"][1]["default"] is True
|
||||
info = PngImagePlugin.PngInfo()
|
||||
info.add_text("parameters", PARAMETERS)
|
||||
info.add_text("prompt", json.dumps({"1": {"class_type": "CustomSampler", "inputs": {}}}))
|
||||
Image.new("RGB", (16, 24)).save(runtime[0], pnginfo=info)
|
||||
result = LoadImageMetadataLM().load_metadata("input.png")
|
||||
assert result[7] == 2**64 - 1
|
||||
assert "saved image generation parameters (preferred)" in result[16]
|
||||
assert "❌ ERROR" not in result[16]
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["Kroma.v2.1", "Kroma.v2.1.safetensors", " Kroma.v2.1 "])
|
||||
def test_model_resolution_preserves_dotted_extensionless_names(tmp_path, name):
|
||||
directory = tmp_path / "Krea 2"
|
||||
directory.mkdir()
|
||||
path = directory / "Kroma.v2.1.safetensors"
|
||||
path.touch()
|
||||
item = {"file_path": str(path)}
|
||||
assert resolve_resource(name, [item], [str(tmp_path)]) == item
|
||||
|
||||
|
||||
def test_model_resolution_accepts_unique_catalog_model_name(tmp_path):
|
||||
path = tmp_path / "local-renamed.safetensors"
|
||||
path.touch()
|
||||
item = {"file_path": str(path), "model_name": "Kroma catalog name"}
|
||||
assert resolve_resource("Kroma catalog name", [item], [str(tmp_path)]) == item
|
||||
|
||||
|
||||
def test_catalog_alias_ambiguity_and_stale_entries(tmp_path):
|
||||
items = []
|
||||
for name in ("a", "b"):
|
||||
path = tmp_path / (name + ".safetensors")
|
||||
path.touch()
|
||||
items.append({"file_path": str(path), "model_name": "Kroma"})
|
||||
with pytest.raises(MetadataError, match="Ambiguous"):
|
||||
resolve_resource("Kroma", items, [str(tmp_path)])
|
||||
items.append({"file_path": str(tmp_path / "absent.safetensors"), "model_name": "missing"})
|
||||
with pytest.raises(MetadataError, match="could not be matched"):
|
||||
resolve_resource("missing", items, [str(tmp_path)])
|
||||
assert resolve_resource("a.safetensors", items, [str(tmp_path)]) == items[0]
|
||||
@@ -200,3 +200,19 @@ def test_lora_loader_qwen_model_raises_clear_error_when_helper_import_fails(monk
|
||||
[],
|
||||
lora_stack=[("stack_qwen.safetensors", 0.6, 0.1)],
|
||||
)
|
||||
|
||||
|
||||
def test_stack_entry_keeps_resolved_absolute_path(monkeypatch):
|
||||
from py.nodes.lora_loader import _collect_stack_entries
|
||||
|
||||
seen = []
|
||||
|
||||
def resolve(name):
|
||||
seen.append(name)
|
||||
return name, ["trigger"]
|
||||
|
||||
monkeypatch.setattr("py.nodes.lora_loader.get_lora_info_absolute", resolve)
|
||||
result = _collect_stack_entries([("/models/b/same.safetensors", .7, .3)])
|
||||
assert seen == ["/models/b/same.safetensors"]
|
||||
assert result[0]["absolute_path"] == "/models/b/same.safetensors"
|
||||
assert result[0]["clip_strength"] == .3
|
||||
|
||||
@@ -0,0 +1,250 @@
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
from py.utils.generation_metadata import (
|
||||
GraphReader, MetadataError, extract_generation_metadata, parse_parameters, split_lora_tags,
|
||||
)
|
||||
|
||||
|
||||
def graph():
|
||||
return {
|
||||
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "base.safetensors"}},
|
||||
"2": {"class_type": "CLIPTextEncode", "inputs": {"text": "ugly monster, (detail:1.2)", "clip": ["1", 1]}},
|
||||
"3": {"class_type": "CLIPTextEncode", "inputs": {"text": "sunshine", "clip": ["1", 1]}},
|
||||
"4": {"class_type": "EmptyLatentImage", "inputs": {"width": 768, "height": 1024}},
|
||||
"5": {"class_type": "KSampler", "inputs": {"model": ["1", 0], "positive": ["2", 0], "negative": ["3", 0], "latent_image": ["4", 0], "seed": 18446744073709551615, "steps": 25, "cfg": 6.5, "sampler_name": "euler", "scheduler": "normal", "denoise": 1}},
|
||||
}
|
||||
|
||||
|
||||
def test_traces_polarity_without_content_heuristics():
|
||||
result = GraphReader(graph()).read("")
|
||||
assert result.values["positive"] == "ugly monster, (detail:1.2)"
|
||||
assert result.values["negative"] == "sunshine"
|
||||
assert result.values["seed"] == 2**64 - 1
|
||||
assert result.values["width"] == 768
|
||||
assert not result.issues
|
||||
|
||||
|
||||
def test_multiple_samplers_require_selection_and_do_not_mix():
|
||||
data = graph()
|
||||
data["6"] = {"class_type": "KSampler", "inputs": {**data["5"]["inputs"], "seed": 42}}
|
||||
with pytest.raises(MetadataError, match="5, 6"):
|
||||
GraphReader(data).read("")
|
||||
assert GraphReader(data).read("6").values["seed"] == 42
|
||||
|
||||
|
||||
def test_model_lora_order_repeated_entries_and_clip_strength():
|
||||
data = graph()
|
||||
data["6"] = {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "lora_name": "same.safetensors", "strength_model": .7, "strength_clip": .3}}
|
||||
data["7"] = {"class_type": "Lora Loader (LoraManager)", "inputs": {"model": ["6", 0], "loras": {"__value__": [{"name": "same", "active": True, "strength": .4, "clipStrength": 0}, {"name": "disabled", "active": False}]}}}
|
||||
data["5"]["inputs"]["model"] = ["7", 0]
|
||||
result = GraphReader(data).read("")
|
||||
assert result.loras == [("same.safetensors", .7, .3), ("same", .4, 0)]
|
||||
|
||||
|
||||
def test_linked_primitive_and_cycle_detection():
|
||||
data = graph()
|
||||
data["6"] = {"class_type": "PrimitiveInt", "inputs": {"value": 123}}
|
||||
data["5"]["inputs"]["seed"] = ["6", 0]
|
||||
assert GraphReader(data).read("").values["seed"] == 123
|
||||
data["6"]["inputs"]["value"] = ["6", 0]
|
||||
assert "Cyclic" in GraphReader(data).read("").issues["seed"]
|
||||
|
||||
|
||||
def test_unsupported_conditioning_is_not_silently_flattened():
|
||||
data = graph()
|
||||
data["2"]["class_type"] = "ConditioningCombine"
|
||||
assert "Unsupported conditioning" in GraphReader(data).read("").issues["positive"]
|
||||
|
||||
|
||||
def test_parameters_sampler_mapping_and_clean_prompts():
|
||||
result = parse_parameters('portrait (detail:1.2) <lora:style:0.7:0.2>\nsecond line\nNegative prompt: blur\nmore blur\nSteps: 25, Sampler: DPM++ 2M Karras, CFG scale: 7, Seed: 123, Size: 512x768, Model: base')
|
||||
assert result.values["sampler_name"] == "dpmpp_2m"
|
||||
assert result.values["scheduler"] == "karras"
|
||||
assert result.values["negative"] == "blur\nmore blur"
|
||||
clean, loras = split_lora_tags(result.values["positive"])
|
||||
assert clean == "portrait (detail:1.2) \nsecond line"
|
||||
assert loras == []
|
||||
assert result.loras == [("style", .7, .2)]
|
||||
|
||||
|
||||
def test_unspecified_a1111_scheduler_requires_decision():
|
||||
result = parse_parameters("cat\nSteps: 20, Sampler: Euler a, Seed: 1, CFG scale: 7")
|
||||
assert result.values["sampler_name"] == "euler_ancestral"
|
||||
assert "scheduler" in result.issues
|
||||
|
||||
|
||||
@pytest.mark.parametrize("value", ["<lora:foo:nan>", "<lora:foo:1e999>", "<lora:foo:bad>"])
|
||||
def test_bad_lora_strength(value):
|
||||
with pytest.raises(ValueError):
|
||||
split_lora_tags(value)
|
||||
|
||||
|
||||
def test_malformed_and_missing_metadata():
|
||||
with pytest.raises(MetadataError, match="Malformed"):
|
||||
extract_generation_metadata({"prompt": "{"})
|
||||
with pytest.raises(MetadataError, match="no supported"):
|
||||
extract_generation_metadata({})
|
||||
assert extract_generation_metadata({"comment": json.dumps(graph())}).values["steps"] == 25
|
||||
|
||||
|
||||
def test_core_ui_workflow_fallback():
|
||||
workflow = {"nodes": [
|
||||
{"id": 1, "type": "CheckpointLoaderSimple", "widgets_values": ["base.safetensors"]},
|
||||
{"id": 2, "type": "CLIPTextEncode", "widgets_values": ["positive"]},
|
||||
{"id": 3, "type": "CLIPTextEncode", "widgets_values": ["negative"]},
|
||||
{"id": 4, "type": "KSampler", "widgets_values": [42, "fixed", 20, 7, "euler", "normal", 1], "inputs": [
|
||||
{"name": "model", "link": 1}, {"name": "positive", "link": 2}, {"name": "negative", "link": 3}]},
|
||||
], "links": [[1, 1, 0, 4, 0, "MODEL"], [2, 2, 0, 4, 1, "CONDITIONING"], [3, 3, 0, 4, 2, "CONDITIONING"]]}
|
||||
result = extract_generation_metadata({"workflow": json.dumps(workflow)})
|
||||
assert result.values["positive"] == "positive"
|
||||
assert result.values["seed"] == 42
|
||||
assert "UI workflow fallback" in result.notes[0]
|
||||
|
||||
|
||||
def test_stack_combiner_uses_numeric_order():
|
||||
data = {str(i): {"class_type": "Lora Stacker (LoraManager)", "inputs": {"loras": [{"name": str(i), "strength": 1, "active": True}]}} for i in (1, 2, 10)}
|
||||
data["20"] = {"class_type": "Lora Stack Combiner (LoraManager)", "inputs": {"lora_stack10": ["10", 0], "lora_stack2": ["2", 0], "lora_stack1": ["1", 0]}}
|
||||
assert [entry[0] for entry in GraphReader(data).stack(["20", 0])] == ["1", "2", "10"]
|
||||
|
||||
|
||||
def test_model_and_clip_lora_mismatch_requires_override():
|
||||
data = graph()
|
||||
data["6"] = {"class_type": "LoraLoader", "inputs": {"model": ["1", 0], "clip": ["1", 1], "lora_name": "style", "strength_model": .7, "strength_clip": .3}}
|
||||
data["5"]["inputs"]["model"] = ["6", 0]
|
||||
assert "different LoRAs" in GraphReader(data).read("").issues["loras"]
|
||||
data["2"]["inputs"]["clip"] = ["6", 1]
|
||||
data["3"]["inputs"]["clip"] = ["6", 1]
|
||||
assert not GraphReader(data).read("").issues
|
||||
|
||||
|
||||
def test_malformed_sampler_inputs():
|
||||
data = graph()
|
||||
data["5"]["inputs"] = None
|
||||
with pytest.raises(MetadataError, match="Malformed sampler"):
|
||||
GraphReader(data).read("")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("label,sampler,scheduler", [
|
||||
("Euler a SGM Uniform", "euler_ancestral", "sgm_uniform"),
|
||||
("Euler simple", "euler", "simple"),
|
||||
("Euler Normal", "euler", "normal"),
|
||||
("er_sde simple", "er_sde", "simple"),
|
||||
])
|
||||
def test_combined_sampler_scheduler_labels(label, sampler, scheduler):
|
||||
result = parse_parameters(f"cat\nSteps: 30, Sampler: {label}, Seed: 42, CFG scale: 5")
|
||||
assert result.values["sampler_name"] == sampler
|
||||
assert result.values["scheduler"] == scheduler
|
||||
assert not result.issues
|
||||
|
||||
|
||||
def test_multiline_settings_and_single_resource_weight():
|
||||
result = parse_parameters('cat\nNegative prompt: blur\nSteps: 30, Sampler: Euler Normal, Seed: 42, CFG scale: 5, Clip skip: 0, extra text,\nmore text\n, Model: example, Hashes: {"model":"123", "LORA:style, special":"456"}, Civitai resources: [{"air":"urn:model"}, {"air":"urn:lora", "weight":0.74}]')
|
||||
assert result.values["checkpoint_name"] == "example"
|
||||
assert result.values["negative"] == "blur"
|
||||
assert result.loras == [("style, special", .74, .74)]
|
||||
assert result.resource_hints[0]["hash"] == "456"
|
||||
|
||||
|
||||
def test_multiple_resource_weights_are_not_paired_by_order():
|
||||
result = parse_parameters('cat\nSteps: 20, Sampler: Euler Normal, Hashes: {"LORA:first":"aaa","LORA:second":"bbb"}, Civitai resources: [{"weight":0.5},{"weight":0.8}]')
|
||||
assert result.loras == []
|
||||
assert "loras" in result.issues
|
||||
assert [item["name"] for item in result.resource_hints] == ["first", "second"]
|
||||
|
||||
|
||||
|
||||
def test_duplicate_tags_with_single_authoritative_resource():
|
||||
result = parse_parameters('cat <lora:style:0.45> <lora:style:0.45>\nSteps: 10, Sampler: Euler simple, Hashes: {"LORA:style":"abc"}, Civitai resources: [{"weight":0.45}]')
|
||||
assert result.loras == [("style", .45, .45)]
|
||||
assert "<lora:" not in result.values["positive"]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("selector", ["outer:inner:5", "outer/inner/5", "outer:inner", "5", ""])
|
||||
def test_qualified_subgraph_sampler_selection(selector):
|
||||
original = graph()
|
||||
expanded = {}
|
||||
for key, node in original.items():
|
||||
inputs = {name: ["outer:inner:" + value[0], value[1]] if isinstance(value, list) else value for name, value in node["inputs"].items()}
|
||||
expanded["outer:inner:" + key] = {**node, "inputs": inputs}
|
||||
result = GraphReader(expanded).read(selector)
|
||||
assert result.values["seed"] == 2**64 - 1
|
||||
assert "outer:inner:5" in result.notes[0]
|
||||
assert not result.issues
|
||||
|
||||
|
||||
def test_subgraph_leaf_selection_rejects_ambiguity():
|
||||
reader = GraphReader({
|
||||
"10:5": {"class_type": "KSampler", "inputs": {}},
|
||||
"20:5": {"class_type": "KSampler", "inputs": {}},
|
||||
})
|
||||
with pytest.raises(MetadataError, match="10:5, 20:5"):
|
||||
reader.read("5")
|
||||
assert reader.select_sampler("20") == "20:5"
|
||||
|
||||
|
||||
def test_standard_custom_sampler_pipeline():
|
||||
data = graph()
|
||||
old = data["5"]["inputs"]
|
||||
data["noise"] = {"class_type": "RandomNoise", "inputs": {"noise_seed": 123}}
|
||||
data["guider"] = {"class_type": "CFGGuider", "inputs": {key: old[key] for key in ("model", "positive", "negative", "cfg")}}
|
||||
data["schedule"] = {"class_type": "BasicScheduler", "inputs": {"steps": 28, "scheduler": "karras", "denoise": .6}}
|
||||
data["sampler"] = {"class_type": "KSamplerSelect", "inputs": {"sampler_name": "euler"}}
|
||||
data["5"] = {"class_type": "SamplerCustomAdvanced", "inputs": {"noise": ["noise", 0], "guider": ["guider", 0], "sigmas": ["schedule", 0], "sampler": ["sampler", 0], "latent_image": old["latent_image"]}}
|
||||
result = GraphReader(data).read("5")
|
||||
assert not result.issues
|
||||
assert result.values["seed"] == 123
|
||||
assert result.values["steps"] == 28
|
||||
assert result.values["denoise"] == .6
|
||||
assert result.values["positive"] == "ugly monster, (detail:1.2)"
|
||||
|
||||
|
||||
def test_saved_metadata_is_preferred_and_workflow_can_be_selected():
|
||||
fields = {
|
||||
"prompt": json.dumps(graph()),
|
||||
"parameters": "saved prompt\nSteps: 12, Sampler: Euler Normal, CFG scale: 4, Seed: 42, Model: saved",
|
||||
}
|
||||
result = extract_generation_metadata(fields, "not-a-node")
|
||||
assert result.values["seed"] == "42"
|
||||
assert result.values["positive"] == "saved prompt"
|
||||
assert any("ignored" in note for note in result.notes)
|
||||
result = extract_generation_metadata(fields, "5", prefer_saved_image_metadata=False)
|
||||
assert result.values["seed"] == 2**64 - 1
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [2, 4])
|
||||
def test_muted_or_bypassed_api_sampler_is_not_selected(mode):
|
||||
data = graph()
|
||||
data["6"] = {"class_type": "KSampler", "mode": mode, "inputs": {**data["5"]["inputs"], "seed": 123}}
|
||||
reader = GraphReader(data)
|
||||
assert reader.read("").values["seed"] == 2**64 - 1
|
||||
with pytest.raises(MetadataError, match="muted, bypassed"):
|
||||
reader.read("6")
|
||||
|
||||
|
||||
@pytest.mark.parametrize("mode", [2, 4])
|
||||
@pytest.mark.parametrize("inactive_parent", [False, True])
|
||||
def test_workflow_modes_exclude_nested_api_sampler(mode, inactive_parent):
|
||||
data = graph()
|
||||
sampler = data.pop("5")
|
||||
data["10:20:5"] = sampler
|
||||
data["30:5"] = {**sampler, "inputs": {**sampler["inputs"], "seed": 123}}
|
||||
workflow = {
|
||||
"nodes": [{"id": 10, "type": "outer", "mode": mode if inactive_parent else 0}, {"id": 30, "type": "active"}],
|
||||
"definitions": {"subgraphs": [
|
||||
{"id": "outer", "nodes": [{"id": 20, "type": "inner"}]},
|
||||
{"id": "inner", "nodes": [{"id": 5, "type": "KSampler", "mode": 0 if inactive_parent else mode}]},
|
||||
{"id": "active", "nodes": [{"id": 5, "type": "KSampler"}]},
|
||||
]},
|
||||
}
|
||||
fields = {"prompt": json.dumps(data), "workflow": json.dumps(workflow)}
|
||||
assert extract_generation_metadata(fields, prefer_saved_image_metadata=False).values["seed"] == 123
|
||||
with pytest.raises(MetadataError, match="muted, bypassed"):
|
||||
extract_generation_metadata(fields, "10:20:5", prefer_saved_image_metadata=False)
|
||||
|
||||
|
||||
def test_invalid_preferred_parameters_recover_workflow():
|
||||
result = extract_generation_metadata({"parameters": "invalid", "prompt": json.dumps(graph())})
|
||||
assert result.values["seed"] == 2**64 - 1
|
||||
assert any("ERROR: Saved image metadata" in note for note in result.notes)
|
||||
@@ -427,3 +427,22 @@ def test_get_lora_info_not_found_returns_original(mock_lora_scanner):
|
||||
|
||||
assert path == "nonexistent"
|
||||
assert triggers == []
|
||||
|
||||
|
||||
def test_get_lora_info_absolute_preserves_exact_stack_path(mock_lora_scanner):
|
||||
mock_lora_scanner([
|
||||
{"file_name": "same", "folder": "a", "file_path": "/models/a/same.safetensors", "civitai": {"trainedWords": ["wrong"]}},
|
||||
{"file_name": "same", "folder": "b", "file_path": "/models/b/same.safetensors", "civitai": {"trainedWords": ["right"]}},
|
||||
])
|
||||
assert get_lora_info_absolute("/models/b/same.safetensors") == (
|
||||
"/models/b/same.safetensors", ["right"]
|
||||
)
|
||||
|
||||
|
||||
def test_get_lora_info_absolute_does_not_substitute_missing_absolute_path(mock_lora_scanner):
|
||||
mock_lora_scanner([
|
||||
{"file_name": "same", "folder": "a", "file_path": "/models/a/same.safetensors"},
|
||||
])
|
||||
assert get_lora_info_absolute("/models/missing/same.safetensors") == (
|
||||
"/models/missing/same.safetensors", []
|
||||
)
|
||||
|
||||
Reference in New Issue
Block a user