Dariusz L
55197621eb
Add custom node adapter for ComfyUI_SmartModelLoader_Nova suite: - Support 'Smart Model Loader [SML2]', 'CLIP Loader [SML2]', and 'LoRA Stack [SML2]' - Parse and serialize canonical JSON configuration widget - Extract diffusion models, VAE, CLIP text encoders, and LoRAs with strengths - Enable model resolution by hash using per-resource SHA256 stored in the node - Update model paths and SHA256 checksums in canonical JSON format during relinking - Add unit and end-to-end integration tests
392 lines
14 KiB
Python
392 lines
14 KiB
Python
"""Tests for Smart Model Loader Nova [SML2] custom node adapter."""
|
|
|
|
import json
|
|
import unittest
|
|
|
|
from core.contracts import (
|
|
CustomNodeMetadata,
|
|
ModelReference,
|
|
Resolution,
|
|
ResolvedModel,
|
|
)
|
|
from core.custom_nodes.registry import (
|
|
get_custom_node_model_adapter,
|
|
get_custom_node_resolution_metadata,
|
|
)
|
|
from core.custom_nodes.smart_model_loader_nova import (
|
|
ADAPTER_ID,
|
|
CONFIG_WIDGET_INDEX,
|
|
_dump_configuration,
|
|
_load_configuration,
|
|
analyze_references,
|
|
has_potential_reference,
|
|
update_model_path,
|
|
)
|
|
from core.resolver import get_workflow_hash_info_for_ref
|
|
from core.workflow.analysis import analyze_workflow_models
|
|
from core.workflow_updater import update_model_path as update_workflow_model_path
|
|
|
|
|
|
class SmartModelLoaderNovaTests(unittest.TestCase):
|
|
def test_adapter_registration(self):
|
|
for node_type in (
|
|
"Smart Model Loader [SML2]",
|
|
"CLIP Loader [SML2]",
|
|
"LoRA Stack [SML2]",
|
|
):
|
|
with self.subTest(node_type=node_type):
|
|
adapter = get_custom_node_model_adapter(node_type)
|
|
self.assertIsNotNone(adapter)
|
|
self.assertEqual(ADAPTER_ID, adapter.adapter_id)
|
|
|
|
def test_load_and_dump_configuration(self):
|
|
raw_dict = {"model": {"kind": "unet"}, "version": 1}
|
|
config, stored_as_string = _load_configuration(raw_dict)
|
|
self.assertFalse(stored_as_string)
|
|
self.assertEqual(raw_dict, config)
|
|
|
|
raw_json = '{"version":1,"model":{"kind":"unet"}}'
|
|
config, stored_as_string = _load_configuration(raw_json)
|
|
self.assertTrue(stored_as_string)
|
|
self.assertEqual("unet", config["model"]["kind"])
|
|
|
|
dumped = _dump_configuration(config)
|
|
self.assertEqual('{"model":{"kind":"unet"},"version":1}', dumped)
|
|
|
|
def test_analyze_references_extracts_all_resource_types(self):
|
|
config = {
|
|
"model": {
|
|
"kind": "unet",
|
|
"resource": {
|
|
"name": "Qwen/diffusion_model.safetensors",
|
|
"sha256": "a" * 64,
|
|
},
|
|
},
|
|
"clip": {
|
|
"source": "external",
|
|
"encoders": [
|
|
{
|
|
"id": "clip-1",
|
|
"resource": {
|
|
"name": "text_encoders/qwen_clip.safetensors",
|
|
"sha256": "b" * 64,
|
|
},
|
|
}
|
|
],
|
|
},
|
|
"vae": {
|
|
"source": "external",
|
|
"resource": {
|
|
"name": "vae/qwen_vae.safetensors",
|
|
},
|
|
},
|
|
"loras": {
|
|
"entries": [
|
|
{
|
|
"id": "lora-1",
|
|
"resource": {
|
|
"name": "loras/detail_slider.safetensors",
|
|
},
|
|
"strength": 1.5,
|
|
"enabled": True,
|
|
}
|
|
]
|
|
},
|
|
"version": 1,
|
|
}
|
|
|
|
node = {
|
|
"id": 10,
|
|
"type": "Smart Model Loader [SML2]",
|
|
"title": "Smart Loader",
|
|
"widgets_values": [_dump_configuration(config)],
|
|
}
|
|
|
|
available = [
|
|
ResolvedModel(
|
|
path="/models/vae/qwen_vae.safetensors",
|
|
filename="qwen_vae.safetensors",
|
|
category="vae",
|
|
relative_path="vae/qwen_vae.safetensors",
|
|
)
|
|
]
|
|
|
|
refs = analyze_references(
|
|
node,
|
|
available_models=available,
|
|
is_active=True,
|
|
get_widget_name_hint=lambda n, idx: "configuration",
|
|
)
|
|
|
|
self.assertIsNotNone(refs)
|
|
self.assertEqual(4, len(refs))
|
|
|
|
by_cat = {ref.category: ref for ref in refs}
|
|
|
|
# 1. Diffusion model
|
|
unet_ref = by_cat["diffusion_models"]
|
|
self.assertEqual("Qwen/diffusion_model.safetensors", unet_ref.original_path)
|
|
self.assertEqual("a" * 64, unet_ref.extra_value("workflow_sha256"))
|
|
self.assertEqual("sml2_configuration", unet_ref.extra_value("hash_lookup_source"))
|
|
self.assertFalse(unet_ref.exists)
|
|
|
|
# 2. Text encoder (clip)
|
|
clip_ref = by_cat["text_encoders"]
|
|
self.assertEqual("text_encoders/qwen_clip.safetensors", clip_ref.original_path)
|
|
self.assertEqual("b" * 64, clip_ref.extra_value("workflow_sha256"))
|
|
|
|
# 3. VAE (matched from available_models)
|
|
vae_ref = by_cat["vae"]
|
|
self.assertEqual("vae/qwen_vae.safetensors", vae_ref.original_path)
|
|
self.assertTrue(vae_ref.exists)
|
|
self.assertEqual("/models/vae/qwen_vae.safetensors", vae_ref.extra_value("full_path"))
|
|
|
|
# 4. LoRA
|
|
lora_ref = by_cat["loras"]
|
|
self.assertEqual("loras/detail_slider.safetensors", lora_ref.original_path)
|
|
self.assertEqual(1.5, lora_ref.extra_value("strength"))
|
|
self.assertTrue(lora_ref.extra_value("active"))
|
|
|
|
def test_analyze_references_skips_baked_and_none(self):
|
|
config = {
|
|
"model": {
|
|
"kind": "unet",
|
|
"resource": {"name": "None"},
|
|
},
|
|
"clip": {
|
|
"source": "baked",
|
|
"encoders": [{"resource": {"name": "ignored.safetensors"}}],
|
|
},
|
|
"vae": {
|
|
"source": "baked",
|
|
"resource": {"name": "ignored_vae.safetensors"},
|
|
},
|
|
}
|
|
node = {
|
|
"id": 11,
|
|
"type": "Smart Model Loader [SML2]",
|
|
"widgets_values": [_dump_configuration(config)],
|
|
}
|
|
refs = analyze_references(
|
|
node,
|
|
is_active=True,
|
|
get_widget_name_hint=lambda n, idx: "configuration",
|
|
)
|
|
self.assertEqual([], refs)
|
|
|
|
def test_has_potential_reference(self):
|
|
valid_node = {
|
|
"widgets_values": [
|
|
json.dumps({"model": {"resource": {"name": "model.safetensors"}}})
|
|
]
|
|
}
|
|
empty_node = {
|
|
"widgets_values": [
|
|
json.dumps({"model": {"resource": {"name": "None"}}})
|
|
]
|
|
}
|
|
invalid_node = {"widgets_values": ["not-json"]}
|
|
|
|
self.assertTrue(has_potential_reference(valid_node))
|
|
self.assertFalse(has_potential_reference(empty_node))
|
|
self.assertFalse(has_potential_reference(invalid_node))
|
|
|
|
def test_update_model_path_rewrites_resource_and_hash(self):
|
|
config = {
|
|
"model": {
|
|
"kind": "unet",
|
|
"resource": {
|
|
"name": "AuthorFolder/qwen_unet.safetensors",
|
|
"sha256": "1" * 64,
|
|
},
|
|
},
|
|
"version": 1,
|
|
}
|
|
node = {
|
|
"id": 12,
|
|
"type": "Smart Model Loader [SML2]",
|
|
"widgets_values": [_dump_configuration(config)],
|
|
}
|
|
|
|
resolved = ResolvedModel(
|
|
path="C:/models/diffusion_models/Local/qwen_unet.safetensors",
|
|
filename="qwen_unet.safetensors",
|
|
category="diffusion_models",
|
|
relative_path="Local/qwen_unet.safetensors",
|
|
extra={"sha256": "2" * 64},
|
|
)
|
|
metadata = CustomNodeMetadata(
|
|
adapter_id=ADAPTER_ID,
|
|
original_identity="AuthorFolder/qwen_unet.safetensors",
|
|
)
|
|
|
|
success = update_model_path(node, CONFIG_WIDGET_INDEX, resolved, metadata)
|
|
self.assertTrue(success)
|
|
|
|
updated_config, is_str = _load_configuration(node["widgets_values"][0])
|
|
self.assertTrue(is_str)
|
|
resource = updated_config["model"]["resource"]
|
|
self.assertEqual("Local/qwen_unet.safetensors", resource["name"])
|
|
self.assertEqual("2" * 64, resource["sha256"])
|
|
|
|
def test_update_model_path_removes_stale_hash_when_resolved_lacks_hash(self):
|
|
config = {
|
|
"model": {
|
|
"kind": "unet",
|
|
"resource": {
|
|
"name": "AuthorFolder/qwen_unet.safetensors",
|
|
"sha256": "1" * 64,
|
|
},
|
|
},
|
|
}
|
|
node = {
|
|
"id": 13,
|
|
"type": "Smart Model Loader [SML2]",
|
|
"widgets_values": [_dump_configuration(config)],
|
|
}
|
|
resolved = ResolvedModel(
|
|
path="C:/models/diffusion_models/other.safetensors",
|
|
filename="other.safetensors",
|
|
category="diffusion_models",
|
|
relative_path="other.safetensors",
|
|
)
|
|
metadata = CustomNodeMetadata(
|
|
adapter_id=ADAPTER_ID,
|
|
original_identity="AuthorFolder/qwen_unet.safetensors",
|
|
)
|
|
|
|
success = update_model_path(node, CONFIG_WIDGET_INDEX, resolved, metadata)
|
|
self.assertTrue(success)
|
|
|
|
updated_config, _ = _load_configuration(node["widgets_values"][0])
|
|
resource = updated_config["model"]["resource"]
|
|
self.assertEqual("other.safetensors", resource["name"])
|
|
self.assertNotIn("sha256", resource)
|
|
|
|
def test_get_workflow_hash_info_for_ref_custom_node_fallback(self):
|
|
ref = ModelReference.from_mapping(
|
|
{
|
|
"node_id": 14,
|
|
"original_path": "model.safetensors",
|
|
"name": "model.safetensors",
|
|
"workflow_sha256": "c" * 64,
|
|
"hash_lookup_source": "sml2_configuration",
|
|
}
|
|
)
|
|
|
|
# 1. With empty workflow_hashes: falls back to node's own hash
|
|
hash_info = get_workflow_hash_info_for_ref({}, ref)
|
|
self.assertIsNotNone(hash_info)
|
|
self.assertEqual("c" * 64, hash_info["sha256"])
|
|
self.assertEqual("sml2_configuration", hash_info["source"])
|
|
|
|
# 2. When workflow_hashes has a match, workflow_hashes takes priority
|
|
workflow_hashes = {
|
|
"model.safetensors": {"sha256": "d" * 64, "source": "workflow_metadata"}
|
|
}
|
|
priority_info = get_workflow_hash_info_for_ref(workflow_hashes, ref)
|
|
self.assertEqual("d" * 64, priority_info["sha256"])
|
|
self.assertEqual("workflow_metadata", priority_info["source"])
|
|
|
|
def test_workflow_analyzer_and_updater_end_to_end(self):
|
|
"""Verify the full pipeline: analyze_workflow_models -> resolution metadata -> update_model_path."""
|
|
config = {
|
|
"model": {
|
|
"kind": "unet",
|
|
"resource": {
|
|
"name": "Qwen Image/qwen_image_2.1_int8_convrot.safetensors",
|
|
"sha256": "cb74" * 16,
|
|
},
|
|
},
|
|
"clip": {
|
|
"source": "external",
|
|
"type": "qwen_image",
|
|
"encoders": [
|
|
{
|
|
"id": "clip-1",
|
|
"resource": {
|
|
"name": "qwen3-vl-8b-heretic-1.3.0_int8_convrot.safetensors",
|
|
"sha256": "443f" * 16,
|
|
},
|
|
}
|
|
],
|
|
},
|
|
"vae": {
|
|
"source": "external",
|
|
"resource": {
|
|
"name": "qwen_image_2.1_vae_bf16.safetensors",
|
|
},
|
|
},
|
|
"loras": {
|
|
"entries": [
|
|
{
|
|
"id": "lora-1",
|
|
"resource": {
|
|
"name": "qwen/Pornmaster_QI2.1_Breasts_Slider_V1.safetensors",
|
|
},
|
|
"strength": 5.0,
|
|
}
|
|
]
|
|
},
|
|
"version": 1,
|
|
}
|
|
workflow = {
|
|
"nodes": [
|
|
{
|
|
"id": 42,
|
|
"type": "Smart Model Loader [SML2]",
|
|
"widgets_values": [_dump_configuration(config)],
|
|
}
|
|
]
|
|
}
|
|
|
|
# 1. Workflow analyzer extracts references via custom node adapter
|
|
refs = analyze_workflow_models(workflow, available_models=[])
|
|
self.assertEqual(4, len(refs))
|
|
by_cat = {ref.category: ref for ref in refs}
|
|
|
|
unet_ref = by_cat["diffusion_models"]
|
|
self.assertEqual(
|
|
"Qwen Image/qwen_image_2.1_int8_convrot.safetensors",
|
|
unet_ref.original_path,
|
|
)
|
|
self.assertEqual(ADAPTER_ID, unet_ref.extra_value("custom_node_adapter"))
|
|
self.assertEqual("cb74" * 16, unet_ref.extra_value("workflow_sha256"))
|
|
|
|
# 2. Verify resolution metadata extraction for workflow updater
|
|
metadata = get_custom_node_resolution_metadata(unet_ref)
|
|
self.assertEqual(ADAPTER_ID, metadata.adapter_id)
|
|
self.assertEqual(
|
|
"Qwen Image/qwen_image_2.1_int8_convrot.safetensors",
|
|
metadata.original_identity,
|
|
)
|
|
|
|
# 3. Test end-to-end update through core.workflow_updater.update_model_path
|
|
resolution = Resolution(
|
|
node_id=42,
|
|
widget_index=CONFIG_WIDGET_INDEX,
|
|
resolved_path="diffusion_models/qwen_fixed.safetensors",
|
|
category="diffusion_models",
|
|
resolved_model=ResolvedModel(
|
|
path="/models/diffusion_models/qwen_fixed.safetensors",
|
|
filename="qwen_fixed.safetensors",
|
|
category="diffusion_models",
|
|
relative_path="qwen_fixed.safetensors",
|
|
extra={"sha256": "1111" * 16},
|
|
),
|
|
custom_node_metadata=metadata,
|
|
)
|
|
updated = update_workflow_model_path(workflow, resolution)
|
|
self.assertTrue(updated)
|
|
|
|
updated_config, _ = _load_configuration(
|
|
workflow["nodes"][0]["widgets_values"][0]
|
|
)
|
|
res = updated_config["model"]["resource"]
|
|
self.assertEqual("qwen_fixed.safetensors", res["name"])
|
|
self.assertEqual("1111" * 16, res["sha256"])
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|