Support Smart Model Loader Nova [SML2] nodes (#6)
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
This commit is contained in:
Dariusz L
@@ -12,10 +12,14 @@ from .lora_manager import ADAPTER as LORA_MANAGER_ADAPTER
|
||||
from .rgthree_power_lora_loader import (
|
||||
ADAPTER as RGTHREE_POWER_LORA_LOADER_ADAPTER,
|
||||
)
|
||||
from .smart_model_loader_nova import (
|
||||
ADAPTER as SMART_MODEL_LOADER_NOVA_ADAPTER,
|
||||
)
|
||||
|
||||
CUSTOM_NODE_MODEL_ADAPTERS = (
|
||||
LORA_MANAGER_ADAPTER,
|
||||
RGTHREE_POWER_LORA_LOADER_ADAPTER,
|
||||
SMART_MODEL_LOADER_NOVA_ADAPTER,
|
||||
)
|
||||
_ADAPTERS_BY_ID = {
|
||||
adapter.adapter_id: adapter for adapter in CUSTOM_NODE_MODEL_ADAPTERS
|
||||
|
||||
@@ -0,0 +1,307 @@
|
||||
"""Backend adapter for ComfyUI Smart Model Loader Nova ([SML2]) loader nodes.
|
||||
|
||||
These nodes keep every model choice inside one JSON "configuration" widget
|
||||
(a canonical, key-sorted JSON string), for example::
|
||||
|
||||
{"model": {"kind": "unet", "resource": {"name": "...", "sha256": "..."}},
|
||||
"clip": {"source": "external", "encoders": [{"id": "...", "resource": {...}}]},
|
||||
"vae": {"source": "external", "resource": {...}},
|
||||
"loras": {"entries": [{"id": "...", "resource": {...}, "strength": 1}]}, ...}
|
||||
|
||||
The CLIP Loader and LoRA Stack variants use the same resource objects under
|
||||
``clip.encoders`` and ``entries``. Each resource becomes one reference on
|
||||
widget 0; updates rewrite the matching resource name and keep the string
|
||||
canonical so the node's own editor reads it back unchanged.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import os
|
||||
from typing import Any, Callable, Dict, Iterator, List, Optional, Tuple
|
||||
|
||||
from ..contracts import CustomNodeMetadata, ModelReference, ResolvedModel
|
||||
from ..log_system import create_module_logger
|
||||
from ..type_utils import normalize_sha256
|
||||
from .base import CustomNodeModelAdapter
|
||||
|
||||
log = create_module_logger(__name__)
|
||||
|
||||
ADAPTER_ID = "smart-model-loader-nova"
|
||||
NODE_TYPES = (
|
||||
"Smart Model Loader [SML2]",
|
||||
"CLIP Loader [SML2]",
|
||||
"LoRA Stack [SML2]",
|
||||
)
|
||||
CONFIG_WIDGET_INDEX = 0
|
||||
|
||||
# model.kind -> Smart Model Loader folder role (core/model_loader/configuration.py MODES)
|
||||
MODEL_KIND_ROLES = {
|
||||
"checkpoint": "checkpoints",
|
||||
"unet": "diffusion_models",
|
||||
"gguf": "diffusion_models_gguf",
|
||||
"nunchaku-flux": "diffusion_models",
|
||||
"nunchaku-qwen": "diffusion_models",
|
||||
"nunchaku-zimage": "diffusion_models",
|
||||
}
|
||||
# role -> (resolver category, ComfyUI folder keys searched for the file)
|
||||
ROLE_FOLDERS = {
|
||||
"checkpoints": ("checkpoints", ("checkpoints",)),
|
||||
"diffusion_models": ("diffusion_models", ("diffusion_models", "unet")),
|
||||
"diffusion_models_gguf": (
|
||||
"diffusion_models",
|
||||
("diffusion_models_gguf", "unet_gguf", "diffusion_models", "unet"),
|
||||
),
|
||||
"text_encoders": ("text_encoders", ("text_encoders", "clip")),
|
||||
"vae": ("vae", ("vae",)),
|
||||
"loras": ("loras", ("loras",)),
|
||||
}
|
||||
CATEGORY_ALIASES = {
|
||||
"unet": "diffusion_models",
|
||||
"clip": "text_encoders",
|
||||
"diffusion_models_gguf": "diffusion_models",
|
||||
"unet_gguf": "diffusion_models",
|
||||
}
|
||||
|
||||
|
||||
def _load_configuration(value: Any) -> Tuple[Optional[Dict[str, Any]], bool]:
|
||||
"""Return (configuration, stored_as_string)."""
|
||||
if isinstance(value, dict):
|
||||
return value, False
|
||||
if isinstance(value, str) and value.strip().startswith("{"):
|
||||
try:
|
||||
parsed = json.loads(value)
|
||||
except ValueError:
|
||||
return None, True
|
||||
return (parsed if isinstance(parsed, dict) else None), True
|
||||
return None, False
|
||||
|
||||
|
||||
def _dump_configuration(config: Dict[str, Any]) -> str:
|
||||
# Same canonical form as the node editor (sorted keys, no whitespace).
|
||||
return json.dumps(config, sort_keys=True, separators=(",", ":"), ensure_ascii=False)
|
||||
|
||||
|
||||
def _resource_name(resource: Any) -> str:
|
||||
if not isinstance(resource, dict):
|
||||
return ""
|
||||
name = str(resource.get("name") or "").strip()
|
||||
return "" if name == "None" else name
|
||||
|
||||
|
||||
def _iter_resources(config: Dict[str, Any]) -> Iterator[Tuple[Dict[str, Any], str, Dict[str, Any]]]:
|
||||
"""Yield (resource, role, entry) for every model file the configuration names."""
|
||||
model = config.get("model")
|
||||
if isinstance(model, dict) and isinstance(model.get("resource"), dict):
|
||||
role = MODEL_KIND_ROLES.get(str(model.get("kind") or ""), "diffusion_models")
|
||||
yield model["resource"], role, model
|
||||
|
||||
for key in ("vae", "audioVae"):
|
||||
section = config.get(key)
|
||||
if (
|
||||
isinstance(section, dict)
|
||||
and section.get("source") == "external"
|
||||
and isinstance(section.get("resource"), dict)
|
||||
):
|
||||
yield section["resource"], "vae", section
|
||||
|
||||
clip = config.get("clip")
|
||||
if isinstance(clip, dict) and clip.get("source") != "baked":
|
||||
if isinstance(clip.get("resource"), dict):
|
||||
yield clip["resource"], "text_encoders", clip
|
||||
for entry in clip.get("encoders") or []:
|
||||
if isinstance(entry, dict) and isinstance(entry.get("resource"), dict):
|
||||
yield entry["resource"], "text_encoders", entry
|
||||
|
||||
loras = config.get("loras")
|
||||
lora_entries = (
|
||||
loras.get("entries")
|
||||
if isinstance(loras, dict)
|
||||
else (loras if isinstance(loras, list) else config.get("entries"))
|
||||
)
|
||||
for entry in lora_entries or []:
|
||||
if isinstance(entry, dict) and isinstance(entry.get("resource"), dict):
|
||||
yield entry["resource"], "loras", entry
|
||||
|
||||
|
||||
def _find_full_path(
|
||||
role: str,
|
||||
name: str,
|
||||
available_models: Optional[List[ResolvedModel]] = None,
|
||||
) -> Optional[str]:
|
||||
try:
|
||||
import folder_paths
|
||||
except ImportError:
|
||||
folder_paths = None
|
||||
|
||||
if folder_paths is not None:
|
||||
for folder in ROLE_FOLDERS.get(role, ("", ()))[1]:
|
||||
try:
|
||||
full_path = folder_paths.get_full_path(folder, name)
|
||||
except Exception:
|
||||
full_path = None
|
||||
if full_path and os.path.exists(full_path):
|
||||
return full_path
|
||||
|
||||
if available_models:
|
||||
category = ROLE_FOLDERS.get(role, ("diffusion_models", ()))[0]
|
||||
norm_name = name.replace("\\", "/").lower()
|
||||
for model in available_models:
|
||||
if model.category == category:
|
||||
model_rel = (model.relative_path or "").replace("\\", "/").lower()
|
||||
model_fn = (model.filename or "").replace("\\", "/").lower()
|
||||
if norm_name in (model_rel, model_fn):
|
||||
return model.path
|
||||
|
||||
return None
|
||||
|
||||
|
||||
def analyze_references(
|
||||
node: Dict[str, Any],
|
||||
available_models: Optional[List[ResolvedModel]] = None,
|
||||
*,
|
||||
is_active: bool,
|
||||
get_widget_name_hint: Callable[[Dict[str, Any], int], str],
|
||||
) -> Optional[List[ModelReference]]:
|
||||
"""Extract every model file named inside the configuration widget."""
|
||||
widgets_values = node.get("widgets_values")
|
||||
if not isinstance(widgets_values, list) or not widgets_values:
|
||||
return None
|
||||
config, _ = _load_configuration(widgets_values[CONFIG_WIDGET_INDEX])
|
||||
if config is None:
|
||||
return None
|
||||
|
||||
node_title = str(node.get("title", "") or "").strip()
|
||||
references: List[ModelReference] = []
|
||||
for resource, role, entry in _iter_resources(config):
|
||||
name = _resource_name(resource)
|
||||
if not name:
|
||||
continue
|
||||
category, folder_keys = ROLE_FOLDERS.get(role, ("diffusion_models", ()))
|
||||
full_path = _find_full_path(role, name, available_models=available_models)
|
||||
reference = {
|
||||
"node_id": node.get("id"),
|
||||
"node_type": node.get("type", ""),
|
||||
"widget_index": CONFIG_WIDGET_INDEX,
|
||||
"widget_name": get_widget_name_hint(node, CONFIG_WIDGET_INDEX),
|
||||
"original_path": name,
|
||||
"name": name,
|
||||
"node_title": node_title,
|
||||
"category": category,
|
||||
"category_hints": [category],
|
||||
"folder_key_hints": list(folder_keys),
|
||||
"full_path": full_path,
|
||||
"exists": bool(full_path),
|
||||
"is_urn": False,
|
||||
"custom_node_adapter": ADAPTER_ID,
|
||||
"custom_node_original_identity": name,
|
||||
"sml2_role": role,
|
||||
"connected": bool(is_active),
|
||||
}
|
||||
sha256 = normalize_sha256(resource.get("sha256"))
|
||||
if sha256:
|
||||
# Lets the resolver find the same file by hash under any name or folder.
|
||||
reference["workflow_sha256"] = sha256
|
||||
reference["hash_lookup_source"] = "sml2_configuration"
|
||||
if role == "loras":
|
||||
try:
|
||||
reference["strength"] = float(entry.get("strength", 1.0))
|
||||
except (TypeError, ValueError):
|
||||
reference["strength"] = 1.0
|
||||
reference["active"] = entry.get("enabled", entry.get("active", True)) is not False
|
||||
references.append(ModelReference.from_mapping(reference))
|
||||
return references
|
||||
|
||||
|
||||
def has_potential_reference(node: Dict[str, Any]) -> bool:
|
||||
widgets_values = node.get("widgets_values")
|
||||
if not isinstance(widgets_values, list) or not widgets_values:
|
||||
return False
|
||||
config, _ = _load_configuration(widgets_values[CONFIG_WIDGET_INDEX])
|
||||
if config is None:
|
||||
return False
|
||||
return any(_resource_name(resource) for resource, _, _ in _iter_resources(config))
|
||||
|
||||
|
||||
def _category_matches(role: str, resolved_category: str) -> bool:
|
||||
if not resolved_category:
|
||||
return True
|
||||
wanted = CATEGORY_ALIASES.get(resolved_category, resolved_category)
|
||||
return ROLE_FOLDERS.get(role, ("",))[0] == wanted
|
||||
|
||||
|
||||
def update_model_path(
|
||||
node: Dict[str, Any],
|
||||
widget_index: int,
|
||||
resolved_model: Optional[ResolvedModel],
|
||||
metadata: CustomNodeMetadata,
|
||||
) -> Optional[bool]:
|
||||
"""Point the matching resource at the resolved file, keeping the JSON canonical."""
|
||||
if metadata.adapter_id != ADAPTER_ID or widget_index != CONFIG_WIDGET_INDEX:
|
||||
return None
|
||||
original_name = str(metadata.original_identity or "").strip()
|
||||
if not original_name or resolved_model is None:
|
||||
return False
|
||||
new_name = resolved_model.relative_path or resolved_model.filename
|
||||
if not new_name:
|
||||
return False
|
||||
|
||||
widgets_values = node.get("widgets_values")
|
||||
if not isinstance(widgets_values, list) or not widgets_values:
|
||||
return False
|
||||
config, stored_as_string = _load_configuration(widgets_values[CONFIG_WIDGET_INDEX])
|
||||
if config is None:
|
||||
log.warning("Smart Model Loader Nova configuration could not be parsed")
|
||||
return False
|
||||
|
||||
resolved_sha = normalize_sha256((resolved_model.extra or {}).get("sha256"))
|
||||
updated = 0
|
||||
for resource, role, _ in _iter_resources(config):
|
||||
current = _resource_name(resource)
|
||||
if current.replace("\\", "/").lower() != original_name.replace("\\", "/").lower():
|
||||
continue
|
||||
if not _category_matches(role, resolved_model.category):
|
||||
continue
|
||||
resource["name"] = new_name
|
||||
if resolved_sha:
|
||||
resource["sha256"] = resolved_sha
|
||||
else:
|
||||
# A different file may have been chosen; a stale hash would fail integrity checks.
|
||||
resource.pop("sha256", None)
|
||||
updated += 1
|
||||
|
||||
if not updated:
|
||||
log.warning(f"'{original_name}' was not found in the Smart Model Loader Nova configuration")
|
||||
return False
|
||||
|
||||
widgets_values[CONFIG_WIDGET_INDEX] = (
|
||||
_dump_configuration(config) if stored_as_string else config
|
||||
)
|
||||
log.info(f"Updated Smart Model Loader Nova model: {original_name} -> {new_name}")
|
||||
return True
|
||||
|
||||
|
||||
def should_skip_existing(reference: ModelReference) -> bool:
|
||||
return reference.exists is True
|
||||
|
||||
|
||||
def adapt_loaded_model(
|
||||
reference: ModelReference,
|
||||
model_name: str,
|
||||
strength: Any,
|
||||
) -> Tuple[str, Any]:
|
||||
return (
|
||||
reference.extra_value("name", model_name),
|
||||
reference.extra_value("strength", strength),
|
||||
)
|
||||
|
||||
|
||||
ADAPTER = CustomNodeModelAdapter(
|
||||
adapter_id=ADAPTER_ID,
|
||||
node_types=NODE_TYPES,
|
||||
analyze_references=analyze_references,
|
||||
has_potential_reference=has_potential_reference,
|
||||
update_model_path=update_model_path,
|
||||
should_skip_existing=should_skip_existing,
|
||||
adapt_loaded_model=adapt_loaded_model,
|
||||
)
|
||||
+21
-14
@@ -405,20 +405,27 @@ def extract_workflow_hash_metadata(workflow_json: Dict[str, Any]) -> Dict[str, D
|
||||
def get_workflow_hash_info_for_ref(
|
||||
workflow_hashes: Dict[str, Dict[str, Any]], model_ref: ModelReference
|
||||
) -> Optional[Dict[str, Any]]:
|
||||
if not workflow_hashes:
|
||||
return None
|
||||
original_path = model_ref.original_path
|
||||
# Node/widget positions are reused after edits, and multi-model widgets
|
||||
# share one position. Only model names/paths can identify a saved hash.
|
||||
candidates = [
|
||||
original_path,
|
||||
get_filename_from_path(original_path),
|
||||
model_ref.extra_value("filename", "") or "",
|
||||
model_ref.extra_value("name", "") or "",
|
||||
]
|
||||
for candidate in candidates:
|
||||
if candidate and candidate in workflow_hashes:
|
||||
return workflow_hashes[candidate]
|
||||
if workflow_hashes:
|
||||
original_path = model_ref.original_path
|
||||
# Node/widget positions are reused after edits, and multi-model widgets
|
||||
# share one position. Only model names/paths can identify a saved hash.
|
||||
candidates = [
|
||||
original_path,
|
||||
get_filename_from_path(original_path),
|
||||
model_ref.extra_value("filename", "") or "",
|
||||
model_ref.extra_value("name", "") or "",
|
||||
]
|
||||
for candidate in candidates:
|
||||
if candidate and candidate in workflow_hashes:
|
||||
return workflow_hashes[candidate]
|
||||
|
||||
# Custom-node adapters may carry the hash stored inside the node itself.
|
||||
node_sha256 = normalize_sha256(model_ref.extra_value("workflow_sha256"))
|
||||
if node_sha256:
|
||||
return {
|
||||
"sha256": node_sha256,
|
||||
"source": model_ref.extra_value("hash_lookup_source") or "custom_node",
|
||||
}
|
||||
return None
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,391 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user