Will Miao
749c71976a
Newly created workflow entries used to be hardcoded to "Unknown". The base model is now resolved at creation time only (Create LM-managed workflow and first-time scans); write-back never re-resolves it. Resolution order: 1. generation-time hint: on create with a prompt context, the primary checkpoint from the metadata collector (primary-sampler tracing plus lm_marker_role marks) outranks static analysis when it resolves in the local checkpoint library 2. static votes: majority base model of the referenced checkpoints in the local library, falling back to the referenced LoRAs 3. CivitAI URN fallback for one-off user actions (never scans) Workflow analysis now extracts structured loader model references from both API-format and UI-format graphs, including loaders whose widget values are serialized as arrays.
208 lines
7.3 KiB
Python
208 lines
7.3 KiB
Python
"""Tests for py.utils.workflow_analysis: sniffing and graph extraction."""
|
|
|
|
import json
|
|
|
|
from py.utils.workflow_analysis import (
|
|
WORKFLOW_FORMAT_API,
|
|
WORKFLOW_FORMAT_UI,
|
|
analyze_workflow,
|
|
analyze_workflow_file,
|
|
load_workflow_file,
|
|
sniff_workflow_text,
|
|
)
|
|
|
|
|
|
UI_WORKFLOW = {
|
|
"nodes": [
|
|
{"id": 1, "type": "CheckpointLoaderSimple", "widgets_values": ["sd_xl_base_1.0.safetensors"]},
|
|
{"id": 2, "type": "LoraLoader", "widgets_values": ["civitai:12345@67890", 1.0, 1.0]},
|
|
{"id": 3, "type": "KSampler", "widgets_values": [42]},
|
|
],
|
|
"links": [],
|
|
"groups": [],
|
|
"version": 0.4,
|
|
}
|
|
|
|
API_WORKFLOW = {
|
|
"1": {"class_type": "CheckpointLoaderSimple", "inputs": {"ckpt_name": "v1-5-pruned.safetensors"}},
|
|
"2": {"class_type": "LoraLoader", "inputs": {"lora_name": "civitai:111@222", "strength_model": 1.0}},
|
|
"3": {"class_type": "SaveImage", "inputs": {"filename_prefix": "ComfyUI"}},
|
|
}
|
|
|
|
|
|
def test_sniff_ui_format():
|
|
assert sniff_workflow_text(json.dumps(UI_WORKFLOW)) == WORKFLOW_FORMAT_UI
|
|
|
|
|
|
def test_sniff_api_format():
|
|
assert sniff_workflow_text(json.dumps(API_WORKFLOW)) == WORKFLOW_FORMAT_API
|
|
|
|
|
|
def test_sniff_rejects_non_workflow_json():
|
|
assert sniff_workflow_text(json.dumps({"setting": True, "other": 1})) is None
|
|
assert sniff_workflow_text(json.dumps(["a", "b"])) is None
|
|
assert sniff_workflow_text("not json at all") is None
|
|
# A sidecar-shaped document (flat metadata fields) is not a workflow.
|
|
assert sniff_workflow_text(json.dumps({"file_name": "x", "sha256": "abc"})) is None
|
|
|
|
|
|
def test_analyze_ui_format_extracts_classes_and_refs():
|
|
analysis = analyze_workflow(UI_WORKFLOW)
|
|
assert analysis is not None
|
|
assert analysis.workflow_format == WORKFLOW_FORMAT_UI
|
|
assert analysis.node_count == 3
|
|
assert analysis.node_classes == ["CheckpointLoaderSimple", "KSampler", "LoraLoader"]
|
|
assert analysis.model_refs == ["sd_xl_base_1.0.safetensors"]
|
|
assert analysis.civitai_refs == ["civitai:12345@67890"]
|
|
|
|
|
|
def test_analyze_api_format_extracts_classes_and_refs():
|
|
analysis = analyze_workflow(API_WORKFLOW)
|
|
assert analysis is not None
|
|
assert analysis.workflow_format == WORKFLOW_FORMAT_API
|
|
assert analysis.node_count == 3
|
|
assert "SaveImage" in analysis.node_classes
|
|
assert analysis.model_refs == ["v1-5-pruned.safetensors"]
|
|
assert analysis.civitai_refs == ["civitai:111@222"]
|
|
|
|
|
|
def test_analyze_rejects_non_workflow():
|
|
assert analyze_workflow({"foo": "bar"}) is None
|
|
|
|
|
|
def test_analyze_skips_prompt_text_with_extension_like_suffix(tmp_path):
|
|
# Prompt strings contain spaces and must not become model refs.
|
|
workflow = {
|
|
"1": {
|
|
"class_type": "CLIPTextEncode",
|
|
"inputs": {"text": "a beautiful landscape, best quality.safetensors style"},
|
|
}
|
|
}
|
|
analysis = analyze_workflow(workflow)
|
|
assert analysis is not None
|
|
assert analysis.model_refs == []
|
|
|
|
|
|
def test_analyze_workflow_file_roundtrip(tmp_path):
|
|
path = tmp_path / "flow.json"
|
|
path.write_text(json.dumps(API_WORKFLOW), encoding="utf-8")
|
|
analysis = analyze_workflow_file(str(path))
|
|
assert analysis is not None
|
|
assert analysis.workflow_format == WORKFLOW_FORMAT_API
|
|
|
|
|
|
def test_load_workflow_file_rejects_sidecar(tmp_path):
|
|
sidecar = tmp_path / "flow.metadata.json"
|
|
sidecar.write_text(json.dumps({"file_name": "flow", "sha256": "x"}), encoding="utf-8")
|
|
assert load_workflow_file(str(sidecar)) is None
|
|
|
|
|
|
def test_load_workflow_file_handles_missing(tmp_path):
|
|
assert load_workflow_file(str(tmp_path / "nope.json")) is None
|
|
|
|
|
|
def test_analyze_api_format_extracts_primary_model_refs():
|
|
workflow = {
|
|
"1": {
|
|
"class_type": "CheckpointLoaderSimple",
|
|
"inputs": {"ckpt_name": "My Model v2.safetensors"},
|
|
},
|
|
"2": {
|
|
"class_type": "UNETLoader",
|
|
"inputs": {"unet_name": "flux1-dev.safetensors", "weight_dtype": "default"},
|
|
},
|
|
# LoRA loaders and encoders are not base-model loaders.
|
|
"3": {"class_type": "LoraLoader", "inputs": {"lora_name": "detail.safetensors"}},
|
|
"4": {"class_type": "CLIPTextEncode", "inputs": {"text": "a cat.safetensors"}},
|
|
}
|
|
analysis = analyze_workflow(workflow)
|
|
assert analysis is not None
|
|
# Loader-attributed refs accept file names with spaces, which the generic
|
|
# model_refs scan rejects to keep prompt text out.
|
|
assert analysis.primary_model_refs == [
|
|
"My Model v2.safetensors",
|
|
"flux1-dev.safetensors",
|
|
]
|
|
# The space-carrying name stays out of the generic refs; the LoRA and the
|
|
# prompt-like string behave as before.
|
|
assert analysis.model_refs == ["detail.safetensors", "flux1-dev.safetensors"]
|
|
|
|
|
|
def test_analyze_ui_format_extracts_primary_model_refs():
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "CheckpointLoaderSimple",
|
|
"widgets_values": ["SDXL/My Model.safetensors"],
|
|
},
|
|
{
|
|
"id": 2,
|
|
"type": "UNETLoader",
|
|
"widgets_values": ["flux dev.safetensors", "default"],
|
|
},
|
|
# LoRA loaders are not base-model loaders.
|
|
{
|
|
"id": 3,
|
|
"type": "LoraLoader",
|
|
"widgets_values": ["style.safetensors", 1.0, 1.0],
|
|
},
|
|
# A model name fed through a non-loader node (e.g. a converted
|
|
# widget's PrimitiveNode source) is not attributed to a loader.
|
|
{
|
|
"id": 4,
|
|
"type": "PrimitiveNode",
|
|
"widgets_values": ["wired checkpoint.safetensors"],
|
|
},
|
|
],
|
|
"links": [],
|
|
}
|
|
analysis = analyze_workflow(workflow)
|
|
assert analysis is not None
|
|
assert analysis.primary_model_refs == [
|
|
"SDXL/My Model.safetensors",
|
|
"flux dev.safetensors",
|
|
]
|
|
|
|
|
|
def test_analyze_primary_refs_empty_without_loader_nodes():
|
|
analysis = analyze_workflow({"1": {"class_type": "KSampler", "inputs": {"seed": 1}}})
|
|
assert analysis is not None
|
|
assert analysis.primary_model_refs == []
|
|
|
|
|
|
def test_analyze_recognizes_lora_manager_display_name_loaders():
|
|
"""LoRA Manager nodes register their display NAME ("Checkpoint Loader
|
|
(LoraManager)") as the workflow JSON type — not the Python class name."""
|
|
ui_workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "Checkpoint Loader (LoraManager)",
|
|
"widgets_values": ["Illustrious/anime/waiIllustriousSDXL_v140.safetensors"],
|
|
},
|
|
{
|
|
"id": 2,
|
|
"type": "Unet Loader (LoraManager)",
|
|
"widgets_values": ["Qwen/qwenImageEdit_fp8.safetensors", "default"],
|
|
},
|
|
],
|
|
"links": [],
|
|
}
|
|
analysis = analyze_workflow(ui_workflow)
|
|
assert analysis is not None
|
|
assert analysis.primary_model_refs == [
|
|
"Illustrious/anime/waiIllustriousSDXL_v140.safetensors",
|
|
"Qwen/qwenImageEdit_fp8.safetensors",
|
|
]
|
|
|
|
api_workflow = {
|
|
"1": {
|
|
"class_type": "Checkpoint Loader (LoraManager)",
|
|
"inputs": {"ckpt_name": "sd_xl_base_1.0.safetensors", "base_model": "Any"},
|
|
},
|
|
}
|
|
analysis = analyze_workflow(api_workflow)
|
|
assert analysis is not None
|
|
assert analysis.primary_model_refs == ["sd_xl_base_1.0.safetensors"]
|