mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-09-23 20:14:08 -03:00
Merge pull request #1121 from mmartial/loader
Add Load Image Metadata node for reusing generation settings
This commit is contained in:
pixelpaws
@@ -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)
|
||||
@@ -562,3 +562,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