342 lines
13 KiB
Python
342 lines
13 KiB
Python
import os
|
|
import tempfile
|
|
import unittest
|
|
|
|
from core.scanner import scan_directory
|
|
from core.workflow_analyzer import analyze_workflow_models, identify_missing_models
|
|
|
|
|
|
def _workflow_with_model(model_path):
|
|
return {
|
|
"nodes": [
|
|
{
|
|
"id": 1,
|
|
"type": "CheckpointLoaderSimple",
|
|
"widgets_values": [model_path],
|
|
"inputs": [],
|
|
"outputs": [],
|
|
}
|
|
]
|
|
}
|
|
|
|
|
|
class WorkflowAnalyzerCaseSensitivityTests(unittest.TestCase):
|
|
def test_folder_case_mismatch_is_reported_missing(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_dir = os.path.join(tmpdir, "QWEN")
|
|
os.makedirs(model_dir)
|
|
model_path = os.path.join(model_dir, "qwen_3_4b.safetensors")
|
|
with open(model_path, "wb"):
|
|
pass
|
|
|
|
available_models = [
|
|
{
|
|
"filename": "qwen_3_4b.safetensors",
|
|
"path": model_path,
|
|
"relative_path": r"QWEN\qwen_3_4b.safetensors",
|
|
"category": "checkpoints",
|
|
"base_directory": tmpdir,
|
|
}
|
|
]
|
|
|
|
refs = analyze_workflow_models(
|
|
_workflow_with_model(r"Qwen\qwen_3_4b.safetensors"),
|
|
available_models=available_models,
|
|
)
|
|
missing = identify_missing_models(refs, available_models)
|
|
|
|
self.assertFalse(refs[0]["exists"])
|
|
self.assertEqual(1, len(missing))
|
|
self.assertEqual(
|
|
r"Qwen\qwen_3_4b.safetensors", missing[0]["original_path"]
|
|
)
|
|
|
|
def test_exact_folder_case_is_not_reported_missing(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_dir = os.path.join(tmpdir, "QWEN")
|
|
os.makedirs(model_dir)
|
|
model_path = os.path.join(model_dir, "qwen_3_4b.safetensors")
|
|
with open(model_path, "wb"):
|
|
pass
|
|
|
|
available_models = [
|
|
{
|
|
"filename": "qwen_3_4b.safetensors",
|
|
"path": model_path,
|
|
"relative_path": r"QWEN\qwen_3_4b.safetensors",
|
|
"category": "checkpoints",
|
|
"base_directory": tmpdir,
|
|
}
|
|
]
|
|
|
|
refs = analyze_workflow_models(
|
|
_workflow_with_model(r"QWEN\qwen_3_4b.safetensors"),
|
|
available_models=available_models,
|
|
)
|
|
missing = identify_missing_models(refs, available_models)
|
|
|
|
self.assertTrue(refs[0]["exists"])
|
|
self.assertEqual([], missing)
|
|
|
|
|
|
class WorkflowAnalyzerCategoryHintTests(unittest.TestCase):
|
|
def test_upscale_model_loader_widget_named_model_name_stays_upscale(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 200,
|
|
"type": "UpscaleModelLoader",
|
|
"inputs": [],
|
|
"outputs": [
|
|
{
|
|
"name": "UPSCALE_MODEL",
|
|
"type": "UPSCALE_MODEL",
|
|
"links": [770],
|
|
}
|
|
],
|
|
"widgets": [{"name": "model_name"}],
|
|
"widgets_values": ["4x_NMKD-Siax_200k.pth"],
|
|
}
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
|
|
self.assertEqual(1, len(refs))
|
|
self.assertEqual("upscale_models", refs[0]["category"])
|
|
|
|
def test_upscale_output_hint_does_not_mark_every_text_widget_as_model(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 201,
|
|
"type": "CustomUpscaleLoader",
|
|
"outputs": [{"type": "UPSCALE_MODEL", "links": [771]}],
|
|
"widgets": [{"name": "model_name"}, {"name": "mode"}],
|
|
"widgets_values": ["4x_NMKD-Siax_200k.pth", "nearest"],
|
|
}
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
|
|
self.assertEqual(1, len(refs))
|
|
self.assertEqual("4x_NMKD-Siax_200k.pth", refs[0]["original_path"])
|
|
self.assertEqual("upscale_models", refs[0]["category"])
|
|
|
|
def test_impact_sam_loader_uses_sams_category(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 168,
|
|
"type": "SAMLoader",
|
|
"title": "SAMLoader",
|
|
"outputs": [{"type": "SAM_MODEL", "links": [590]}],
|
|
"widgets": [{"name": "model_name"}],
|
|
"widgets_values": ["sam_vit_b_01ec64.pth"],
|
|
}
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
|
|
self.assertEqual(1, len(refs))
|
|
self.assertEqual("sam_vit_b_01ec64.pth", refs[0]["original_path"])
|
|
self.assertEqual("sams", refs[0]["category"])
|
|
|
|
def test_impact_sam_loader_esam_option_is_not_missing_model(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 168,
|
|
"type": "SAMLoader",
|
|
"title": "SAMLoader",
|
|
"outputs": [{"type": "SAM_MODEL", "links": [590]}],
|
|
"widgets": [{"name": "model_name"}],
|
|
"widgets_values": ["ESAM"],
|
|
}
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
|
|
self.assertEqual([], refs)
|
|
|
|
def test_impact_ultralytics_detector_provider_uses_ultralytics_category(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 167,
|
|
"type": "UltralyticsDetectorProvider",
|
|
"title": "SEGM Detector",
|
|
"outputs": [
|
|
{"name": "BBOX_DETECTOR", "type": "BBOX_DETECTOR", "links": []},
|
|
{"name": "SEGM_DETECTOR", "type": "SEGM_DETECTOR", "links": [590]},
|
|
],
|
|
"widgets": [{"name": "model_name"}],
|
|
"widgets_values": ["segm/person_yolov8m-seg.pt"],
|
|
}
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
|
|
self.assertEqual(1, len(refs))
|
|
self.assertEqual("segm/person_yolov8m-seg.pt", refs[0]["original_path"])
|
|
self.assertEqual("ultralytics", refs[0]["category"])
|
|
|
|
def test_core_extra_loader_categories_match_comfyui_folder_paths(self):
|
|
cases = [
|
|
("DiffusersLoader", ["wan_diffusers"], [{"name": "model_path"}], "diffusers"),
|
|
("GLIGENLoader", ["gligen.safetensors"], [{"name": "gligen_name"}], "gligen"),
|
|
("AudioEncoderLoader", ["audio_encoder.safetensors"], [{"name": "audio_encoder_name"}], "audio_encoders"),
|
|
("LoadBackgroundRemovalModel", ["birefnet.pth"], [{"name": "bg_removal_name"}], "background_removal"),
|
|
("LoadDA3Model", ["depth_anything_3.pth", "default"], [{"name": "model_name"}, {"name": "weight_dtype"}], "geometry_estimation"),
|
|
("FrameInterpolationModelLoader", ["rife.pth"], [{"name": "model_name"}], "frame_interpolation"),
|
|
("LoadMediaPipeFaceLandmarker", ["face_landmarker.pth"], [{"name": "model_name"}], "detection"),
|
|
("ModelPatchLoader", ["qwen_patch.safetensors"], [{"name": "name"}], "model_patches"),
|
|
("LoadMoGeModel", ["moge.safetensors"], [{"name": "model_name"}], "geometry_estimation"),
|
|
("PhotoMakerLoader", ["photomaker.bin"], [{"name": "photomaker_model_name"}], "photomaker"),
|
|
("OpticalFlowLoader", ["raft_large.pth"], [{"name": "model_name"}], "optical_flow"),
|
|
]
|
|
|
|
for node_type, widget_values, widgets, expected_category in cases:
|
|
with self.subTest(node_type=node_type):
|
|
refs = analyze_workflow_models(
|
|
{
|
|
"nodes": [
|
|
{
|
|
"id": 300,
|
|
"type": node_type,
|
|
"widgets": widgets,
|
|
"widgets_values": widget_values,
|
|
"outputs": [{"links": [1]}],
|
|
}
|
|
]
|
|
},
|
|
available_models=[],
|
|
)
|
|
|
|
self.assertEqual(1, len(refs))
|
|
self.assertEqual(expected_category, refs[0]["category"])
|
|
|
|
def test_multi_clip_loader_indexes_are_text_encoders(self):
|
|
refs = analyze_workflow_models(
|
|
{
|
|
"nodes": [
|
|
{
|
|
"id": 301,
|
|
"type": "QuadrupleCLIPLoader",
|
|
"widgets_values": [
|
|
"clip_l.safetensors",
|
|
"clip_g.safetensors",
|
|
"t5xxl.safetensors",
|
|
"llama.safetensors",
|
|
],
|
|
"outputs": [{"links": [1]}],
|
|
}
|
|
]
|
|
},
|
|
available_models=[],
|
|
)
|
|
|
|
self.assertEqual(4, len(refs))
|
|
self.assertEqual({"text_encoders"}, {ref["category"] for ref in refs})
|
|
|
|
|
|
class ScannerFolderModelTests(unittest.TestCase):
|
|
def test_diffusers_folder_models_are_scanned_as_folder_entries(self):
|
|
with tempfile.TemporaryDirectory() as tmpdir:
|
|
model_dir = os.path.join(tmpdir, "wan_diffusers")
|
|
os.makedirs(model_dir)
|
|
with open(os.path.join(model_dir, "model_index.json"), "w", encoding="utf-8") as handle:
|
|
handle.write("{}")
|
|
ignored_dir = os.path.join(tmpdir, "not_a_diffusers_model")
|
|
os.makedirs(ignored_dir)
|
|
|
|
models = scan_directory(tmpdir, {"folder"}, "diffusers")
|
|
|
|
self.assertEqual(1, len(models))
|
|
self.assertEqual("wan_diffusers", models[0]["relative_path"])
|
|
self.assertEqual(model_dir, models[0]["path"])
|
|
|
|
|
|
class WorkflowCategoryHintsTests(unittest.TestCase):
|
|
def test_node_type_to_category_hints_is_populated(self):
|
|
from core.workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS
|
|
# Verify standard loader mappings are correctly generated
|
|
self.assertEqual("checkpoints", NODE_TYPE_TO_CATEGORY_HINTS.get("CheckpointLoaderSimple"))
|
|
self.assertEqual("checkpoints", NODE_TYPE_TO_CATEGORY_HINTS.get("CheckpointLoader"))
|
|
self.assertEqual("diffusion_models", NODE_TYPE_TO_CATEGORY_HINTS.get("UNETLoader"))
|
|
self.assertEqual("loras", NODE_TYPE_TO_CATEGORY_HINTS.get("LoraLoader"))
|
|
self.assertEqual("text_encoders", NODE_TYPE_TO_CATEGORY_HINTS.get("CLIPLoader"))
|
|
# Verify custom fallbacks are also correctly populated
|
|
self.assertEqual("loras", NODE_TYPE_TO_CATEGORY_HINTS.get("LoraLoaderV2"))
|
|
|
|
|
|
class WorkflowMissingReferenceGroupingTests(unittest.TestCase):
|
|
def test_power_lora_duplicate_missing_tracks_all_node_refs(self):
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 401,
|
|
"type": "Power Lora Loader (rgthree)",
|
|
"widgets_values": [
|
|
{"on": True, "lora": "missing_style.safetensors", "strength": 1.0}
|
|
],
|
|
},
|
|
{
|
|
"id": 402,
|
|
"type": "Power Lora Loader (rgthree)",
|
|
"widgets_values": [
|
|
{"on": True, "lora": "missing_style.safetensors", "strength": 0.5}
|
|
],
|
|
},
|
|
]
|
|
}
|
|
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
missing = identify_missing_models(refs, [])
|
|
|
|
self.assertEqual(2, len(refs))
|
|
self.assertEqual(1, len(missing))
|
|
self.assertEqual(2, missing[0]["reference_count"])
|
|
self.assertEqual(
|
|
[401, 402],
|
|
[ref["node_id"] for ref in missing[0]["all_node_refs"]],
|
|
)
|
|
self.assertEqual(
|
|
["lora", "lora"],
|
|
[ref["nested_key"] for ref in missing[0]["all_node_refs"]],
|
|
)
|
|
|
|
def test_same_filename_different_categories_are_not_merged(self):
|
|
refs = [
|
|
{
|
|
"node_id": 501,
|
|
"widget_index": 0,
|
|
"original_path": "shared_name.safetensors",
|
|
"category": "checkpoints",
|
|
"exists": False,
|
|
},
|
|
{
|
|
"node_id": 502,
|
|
"widget_index": 0,
|
|
"original_path": "shared_name.safetensors",
|
|
"category": "loras",
|
|
"exists": False,
|
|
},
|
|
]
|
|
|
|
missing = identify_missing_models(refs, [])
|
|
|
|
self.assertEqual(2, len(missing))
|
|
self.assertEqual(
|
|
{"checkpoints", "loras"},
|
|
{item["category"] for item in missing},
|
|
)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|