1221 lines
42 KiB
Python
1221 lines
42 KiB
Python
"""
|
|
Core Resolver Module
|
|
|
|
Integrates all components to provide high-level API for model linking.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import threading
|
|
from concurrent.futures import ThreadPoolExecutor
|
|
from dataclasses import replace
|
|
from typing import Any, Callable, Dict, List, Optional, Set, Tuple
|
|
from urllib.parse import unquote
|
|
|
|
from .contracts import (
|
|
MissingModel,
|
|
ModelMatch,
|
|
ModelReference,
|
|
Resolution,
|
|
ResolvedModel,
|
|
WorkflowAnalysisResult,
|
|
)
|
|
from .custom_nodes import (
|
|
custom_node_has_potential_model_reference,
|
|
should_skip_existing_custom_node_reference,
|
|
)
|
|
from .log_system import create_module_logger
|
|
|
|
log = create_module_logger(__name__)
|
|
|
|
from .local_hash_matches import collect_local_hash_matches_for_result
|
|
from .matcher import find_matches, strip_known_model_extension
|
|
from .metadata_model_utils import normalize_models
|
|
from .scanner import get_model_files, invalidate_model_files_cache
|
|
from .type_utils import MODEL_EXTENSIONS as _MODEL_EXTENSIONS
|
|
from .type_utils import (
|
|
as_dict,
|
|
as_list,
|
|
extract_sha256_from_metadata,
|
|
normalize_sha256,
|
|
prepare_remote_size_probe_url,
|
|
unique_ordered_strings,
|
|
)
|
|
from .workflow.analysis import identify_missing_models
|
|
from .workflow.inventory import get_workflow_model_inventory
|
|
from .workflow.references import (
|
|
get_model_widget_category_hint,
|
|
should_scan_as_model_reference,
|
|
)
|
|
from .workflow.traversal import (
|
|
iter_active_workflow_nodes_with_scope,
|
|
iter_workflow_nodes_with_scope,
|
|
)
|
|
from .workflow.widgets import NESTED_MODEL_KEYS, NODE_TYPE_TO_CATEGORY_HINTS
|
|
from .workflow_updater import update_workflow_nodes
|
|
|
|
# Regex patterns for URL extraction (matches HuggingFace and CivitAI URLs)
|
|
URL_PATTERN = re.compile(r'(https?://(?:huggingface\.co|civitai\.com)[^\s"\'<>\)\\]+)')
|
|
|
|
# Model file extensions to look for
|
|
MODEL_EXTENSIONS = tuple(_MODEL_EXTENSIONS)
|
|
|
|
|
|
from .path_utils import (
|
|
find_metadata_sidecar_path,
|
|
get_directory_entry_set,
|
|
get_filename_from_path,
|
|
get_path_identity,
|
|
get_path_key,
|
|
read_merged_model_metadata,
|
|
)
|
|
|
|
# Imported from .matcher
|
|
|
|
_LOCAL_HASH_MATCH_CACHE_LOCK = threading.Lock()
|
|
_LOCAL_HASH_MATCH_CACHE: Optional[Dict[str, List[ModelMatch]]] = None
|
|
_ACTIVE_DOWNLOAD_STATUSES = {"starting", "downloading", "paused", "cancelling"}
|
|
|
|
|
|
def invalidate_local_hash_match_cache() -> None:
|
|
"""Clear the in-memory SHA256 -> local model match index."""
|
|
global _LOCAL_HASH_MATCH_CACHE
|
|
with _LOCAL_HASH_MATCH_CACHE_LOCK:
|
|
_LOCAL_HASH_MATCH_CACHE = None
|
|
|
|
|
|
def invalidate_model_caches() -> None:
|
|
"""Clear all in-memory caches that depend on the local model inventory."""
|
|
invalidate_model_files_cache()
|
|
invalidate_local_hash_match_cache()
|
|
|
|
|
|
def _get_active_downloads_by_path() -> Dict[str, Dict[str, Any]]:
|
|
try:
|
|
from .download.api import get_all_progress
|
|
|
|
progress_items = get_all_progress()
|
|
except Exception:
|
|
return {}
|
|
|
|
active: Dict[str, Dict[str, Any]] = {}
|
|
for download_id, progress in progress_items.items():
|
|
if not isinstance(progress, dict):
|
|
continue
|
|
|
|
status = str(progress.get("status") or "").strip().lower()
|
|
if status not in _ACTIVE_DOWNLOAD_STATUSES:
|
|
continue
|
|
|
|
path = progress.get("path") or ""
|
|
if not path and progress.get("directory") and progress.get("filename"):
|
|
path = os.path.join(str(progress["directory"]), str(progress["filename"]))
|
|
|
|
path_key = get_path_key(path)
|
|
if not path_key:
|
|
continue
|
|
|
|
active[path_key] = {
|
|
"download_id": download_id,
|
|
"download_status": status,
|
|
"download_progress": progress.get("progress", 0),
|
|
"downloaded": progress.get("downloaded", 0),
|
|
"total_size": progress.get("total_size", 0),
|
|
}
|
|
|
|
return active
|
|
|
|
|
|
def annotate_local_matches_with_download_state(
|
|
matches: List[ModelMatch],
|
|
active_downloads_by_path: Optional[Dict[str, Dict[str, Any]]] = None,
|
|
) -> List[ModelMatch]:
|
|
active_downloads = (
|
|
active_downloads_by_path
|
|
if active_downloads_by_path is not None
|
|
else _get_active_downloads_by_path()
|
|
)
|
|
if not active_downloads:
|
|
return matches
|
|
|
|
enriched_matches: List[ModelMatch] = []
|
|
for match in matches:
|
|
model = match.model
|
|
candidate_paths = [
|
|
model.path,
|
|
model.extra_value("resolved_path"),
|
|
match.extra_value("path"),
|
|
match.extra_value("resolved_path"),
|
|
]
|
|
download_info = None
|
|
for candidate_path in candidate_paths:
|
|
path_key = get_path_key(candidate_path)
|
|
if path_key and path_key in active_downloads:
|
|
download_info = active_downloads[path_key]
|
|
break
|
|
|
|
if not download_info:
|
|
enriched_matches.append(match)
|
|
continue
|
|
|
|
download_fields = {
|
|
**download_info,
|
|
"is_downloading": True,
|
|
"downloading": True,
|
|
}
|
|
enriched_model = model.with_extra(**download_fields)
|
|
enriched_matches.append(
|
|
replace(match, model=enriched_model).with_extra(**download_fields)
|
|
)
|
|
|
|
return enriched_matches
|
|
|
|
|
|
def _is_local_hash_match_candidate(model: ResolvedModel) -> bool:
|
|
model_path = str(model.path or "").strip()
|
|
if not model_path:
|
|
return False
|
|
|
|
if os.path.isdir(model_path):
|
|
return True
|
|
|
|
filename = get_filename_from_path(model_path).lower()
|
|
if filename.endswith((".metadata.json", ".civitai.info")):
|
|
return False
|
|
|
|
file_ext = os.path.splitext(filename)[1].lower()
|
|
return file_ext in _MODEL_EXTENSIONS
|
|
|
|
|
|
def _build_local_hash_match_cache(
|
|
available_models: List[ResolvedModel],
|
|
) -> Dict[str, List[ModelMatch]]:
|
|
index: Dict[str, List[ModelMatch]] = {}
|
|
seen_entries = set()
|
|
|
|
normalized_models = normalize_models(available_models)
|
|
candidates = [
|
|
model
|
|
for model in normalized_models
|
|
if model.path and _is_local_hash_match_candidate(model)
|
|
]
|
|
if not candidates:
|
|
return index
|
|
|
|
dir_cache: Dict[str, Set[str]] = {}
|
|
for model in candidates:
|
|
model_path = model.path
|
|
if model_path:
|
|
directory = os.path.dirname(model_path)
|
|
if directory:
|
|
get_directory_entry_set(directory, cache=dir_cache)
|
|
|
|
def _inspect_model_metadata(
|
|
model: ResolvedModel,
|
|
) -> List[Tuple[str, str, ResolvedModel, str, str]]:
|
|
model_path = model.path
|
|
if not model_path:
|
|
return []
|
|
|
|
directory = os.path.dirname(model_path)
|
|
dir_entries = dir_cache.get(os.path.normcase(os.path.abspath(directory))) if directory else None
|
|
metadata_path = find_metadata_sidecar_path(model_path, dir_entries=dir_entries)
|
|
if not metadata_path:
|
|
return []
|
|
|
|
metadata = read_merged_model_metadata(model_path, None)
|
|
if not isinstance(metadata, dict) or not metadata:
|
|
log.debug(f"Could not read metadata sidecar for hash match: {metadata_path}")
|
|
return []
|
|
|
|
metadata_hashes = _extract_model_sha256_from_metadata(metadata, model)
|
|
if not metadata_hashes:
|
|
return []
|
|
|
|
model_filename = model.filename or get_filename_from_path(model_path)
|
|
model_identity = get_path_identity(model_path) or model_path
|
|
results: List[Tuple[str, str, ResolvedModel, str, str]] = []
|
|
for metadata_hash in metadata_hashes:
|
|
normalized_hash = normalize_sha256(metadata_hash)
|
|
if not normalized_hash:
|
|
continue
|
|
results.append(
|
|
(
|
|
normalized_hash,
|
|
model_identity,
|
|
model,
|
|
model_filename,
|
|
metadata_path,
|
|
)
|
|
)
|
|
return results
|
|
|
|
if len(candidates) <= 1:
|
|
inspected_results = [_inspect_model_metadata(candidates[0])]
|
|
else:
|
|
cpu_cores = os.cpu_count() or 2
|
|
max_workers = min(8, max(2, cpu_cores))
|
|
with ThreadPoolExecutor(max_workers=max_workers) as executor:
|
|
inspected_results = list(executor.map(_inspect_model_metadata, candidates))
|
|
|
|
for item_list in inspected_results:
|
|
for normalized_hash, model_identity, model, model_filename, metadata_path in item_list:
|
|
entry_key = (normalized_hash, model_identity)
|
|
if entry_key in seen_entries:
|
|
continue
|
|
seen_entries.add(entry_key)
|
|
|
|
model_with_metadata = model.with_extra(
|
|
sha256=normalized_hash,
|
|
metadata_path=metadata_path,
|
|
)
|
|
index.setdefault(normalized_hash, []).append(
|
|
ModelMatch.from_mapping(
|
|
{
|
|
"model": model_with_metadata,
|
|
"filename": model_filename,
|
|
"similarity": 1.0,
|
|
"confidence": 100.0,
|
|
"match_type": "hash",
|
|
"hash_match": True,
|
|
"hash_source": "metadata",
|
|
"sha256": normalized_hash,
|
|
"metadata_path": metadata_path,
|
|
}
|
|
)
|
|
)
|
|
|
|
return index
|
|
|
|
|
|
def _get_local_hash_match_cache(
|
|
force_rescan: bool = False,
|
|
) -> Dict[str, List[ModelMatch]]:
|
|
global _LOCAL_HASH_MATCH_CACHE
|
|
if force_rescan:
|
|
invalidate_local_hash_match_cache()
|
|
|
|
with _LOCAL_HASH_MATCH_CACHE_LOCK:
|
|
if _LOCAL_HASH_MATCH_CACHE is not None:
|
|
return _LOCAL_HASH_MATCH_CACHE
|
|
|
|
available_models = get_model_files(force_rescan=force_rescan)
|
|
index = _build_local_hash_match_cache(available_models)
|
|
|
|
with _LOCAL_HASH_MATCH_CACHE_LOCK:
|
|
if _LOCAL_HASH_MATCH_CACHE is None:
|
|
_LOCAL_HASH_MATCH_CACHE = index
|
|
return _LOCAL_HASH_MATCH_CACHE
|
|
|
|
|
|
def get_workflow_url_info_for_filename(
|
|
workflow_urls: Dict[str, Dict[str, Any]], filename: str
|
|
) -> Optional[Dict[str, Any]]:
|
|
if filename in workflow_urls:
|
|
return workflow_urls[filename]
|
|
|
|
filename_stem = strip_known_model_extension(get_filename_from_path(filename)).lower()
|
|
if not filename_stem:
|
|
return None
|
|
|
|
for workflow_filename, url_info in workflow_urls.items():
|
|
workflow_stem = strip_known_model_extension(
|
|
get_filename_from_path(workflow_filename)
|
|
).lower()
|
|
if workflow_stem == filename_stem:
|
|
return url_info
|
|
|
|
return None
|
|
|
|
|
|
def _workflow_hash_entry_candidates(entry: Any) -> List[Dict[str, Any]]:
|
|
if not isinstance(entry, dict):
|
|
return []
|
|
sha256 = normalize_sha256(extract_sha256_from_metadata(entry))
|
|
if not sha256:
|
|
return []
|
|
return [{**entry, "sha256": sha256}]
|
|
|
|
|
|
def _collect_workflow_hash_entries(value: Any) -> List[Dict[str, Any]]:
|
|
if not isinstance(value, dict):
|
|
return []
|
|
|
|
entries: List[Dict[str, Any]] = []
|
|
for key in (
|
|
"model_resolver_hashes",
|
|
"anomalous_hashes",
|
|
"model_hashes",
|
|
"hashes",
|
|
):
|
|
raw = value.get(key)
|
|
if isinstance(raw, dict):
|
|
for nested_key in ("models", "by_path", "by_node"):
|
|
nested = raw.get(nested_key)
|
|
if isinstance(nested, list):
|
|
for item in nested:
|
|
entries.extend(_workflow_hash_entry_candidates(item))
|
|
elif isinstance(nested, dict):
|
|
for item in nested.values():
|
|
entries.extend(_workflow_hash_entry_candidates(item))
|
|
entries.extend(_workflow_hash_entry_candidates(raw))
|
|
elif isinstance(raw, list):
|
|
for item in raw:
|
|
entries.extend(_workflow_hash_entry_candidates(item))
|
|
return entries
|
|
|
|
|
|
def extract_workflow_hash_metadata(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
|
|
"""Return embedded workflow SHA256 metadata indexed by path, filename, and node/widget."""
|
|
if not isinstance(workflow_json, dict):
|
|
return {}
|
|
|
|
entries: List[Dict[str, Any]] = []
|
|
entries.extend(_collect_workflow_hash_entries(workflow_json))
|
|
extra = workflow_json.get("extra")
|
|
if isinstance(extra, dict):
|
|
entries.extend(_collect_workflow_hash_entries(extra))
|
|
|
|
index: Dict[str, Dict[str, Any]] = {}
|
|
for entry in entries:
|
|
sha256 = normalize_sha256(entry.get("sha256"))
|
|
if not sha256:
|
|
continue
|
|
normalized = {**entry, "sha256": sha256}
|
|
for candidate in (
|
|
entry.get("path"),
|
|
entry.get("original_path"),
|
|
entry.get("filename"),
|
|
entry.get("name"),
|
|
entry.get("file_name"),
|
|
):
|
|
text = str(candidate or "").strip()
|
|
if not text:
|
|
continue
|
|
index[text] = normalized
|
|
index[get_filename_from_path(text)] = normalized
|
|
node_id = entry.get("node_id")
|
|
widget_index = entry.get("widget_index")
|
|
if node_id is not None and widget_index is not None:
|
|
index[f"{node_id}:{widget_index}"] = normalized
|
|
return index
|
|
|
|
|
|
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]
|
|
return None
|
|
|
|
|
|
def workflow_has_nodes(workflow_json: Dict[str, Any]) -> bool:
|
|
"""Return True when the active top-level workflow contains nodes."""
|
|
if not isinstance(workflow_json, dict):
|
|
return False
|
|
|
|
nodes = workflow_json.get("nodes")
|
|
return isinstance(nodes, list) and len(nodes) > 0
|
|
|
|
|
|
def node_has_potential_model_reference(node: Dict[str, Any]) -> bool:
|
|
"""Detect model-looking widget values without resolving paths or scanning disks."""
|
|
if not isinstance(node, dict):
|
|
return False
|
|
|
|
widgets_values = node.get("widgets_values")
|
|
if not isinstance(widgets_values, list) or not widgets_values:
|
|
return False
|
|
|
|
if custom_node_has_potential_model_reference(node):
|
|
return True
|
|
|
|
for idx, value in enumerate(widgets_values):
|
|
model_widget_category_hint = get_model_widget_category_hint(node, idx)
|
|
if should_scan_as_model_reference(
|
|
value, declared_model_widget=bool(model_widget_category_hint)
|
|
):
|
|
return True
|
|
|
|
if not isinstance(value, dict):
|
|
continue
|
|
|
|
for nested_key in NESTED_MODEL_KEYS:
|
|
nested_value = value.get(nested_key)
|
|
if (
|
|
isinstance(nested_value, str)
|
|
and should_scan_as_model_reference(
|
|
nested_value, declared_model_widget=True
|
|
)
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
|
|
def workflow_has_potential_model_references(workflow_json: Dict[str, Any]) -> bool:
|
|
"""Return True when active workflow nodes contain any model-looking values."""
|
|
return any(
|
|
node_has_potential_model_reference(context.node)
|
|
for context in iter_active_workflow_nodes_with_scope(workflow_json)
|
|
)
|
|
|
|
def normalize_workflow_download_url(url: str) -> str:
|
|
"""Convert workflow file-page URLs into direct download URLs when possible."""
|
|
normalized = prepare_remote_size_probe_url(url)
|
|
return normalized if normalized is not None else url
|
|
|
|
|
|
def workflow_url_points_to_file(url: str, filename: str) -> bool:
|
|
"""Return true when a URL appears to reference the specific model file."""
|
|
if not url or not filename:
|
|
return False
|
|
|
|
try:
|
|
decoded_url = unquote(url)
|
|
except Exception:
|
|
decoded_url = url
|
|
|
|
return filename in decoded_url or unquote(filename) in decoded_url
|
|
|
|
|
|
def _deduplicate_local_matches(matches: List[ModelMatch]) -> List[ModelMatch]:
|
|
"""Keep one highest-confidence match for each local model path."""
|
|
seen_absolute_paths = {}
|
|
deduplicated_matches = []
|
|
for match in matches:
|
|
model = match.model
|
|
absolute_path = model.path
|
|
path_identity = get_path_identity(absolute_path) if absolute_path else ""
|
|
dedupe_key = path_identity or os.path.normcase(
|
|
model.relative_path or match.filename
|
|
)
|
|
|
|
if dedupe_key not in seen_absolute_paths:
|
|
seen_absolute_paths[dedupe_key] = match
|
|
deduplicated_matches.append(match)
|
|
else:
|
|
existing_match = seen_absolute_paths[dedupe_key]
|
|
if match.confidence > existing_match.confidence:
|
|
idx = deduplicated_matches.index(existing_match)
|
|
deduplicated_matches[idx] = match
|
|
seen_absolute_paths[dedupe_key] = match
|
|
|
|
return deduplicated_matches
|
|
|
|
|
|
def search_local_matches(
|
|
target_for_matching: str,
|
|
category: Optional[str] = None,
|
|
similarity_threshold: float = 0.0,
|
|
max_matches_per_model: int = 10,
|
|
force_rescan: bool = False,
|
|
) -> List[ModelMatch]:
|
|
"""
|
|
Search local model files using the same matcher as workflow analysis.
|
|
|
|
Args:
|
|
target_for_matching: Filename/path to match against local files
|
|
category: Optional category hint to prioritize/filter candidates
|
|
similarity_threshold: Minimum similarity score (0.0 to 1.0)
|
|
max_matches_per_model: Maximum number of matches to return
|
|
|
|
Returns:
|
|
Deduplicated list of local matches sorted by similarity
|
|
"""
|
|
available_models = get_model_files(force_rescan=force_rescan)
|
|
available_models = normalize_models(available_models)
|
|
|
|
candidates = available_models
|
|
if category and category != "unknown":
|
|
candidates = [m for m in available_models if m.category == category]
|
|
candidates.extend(
|
|
[m for m in available_models if m.category != category]
|
|
)
|
|
|
|
matches = find_matches(
|
|
target_for_matching,
|
|
candidates,
|
|
threshold=similarity_threshold,
|
|
max_results=max_matches_per_model,
|
|
)
|
|
|
|
deduplicated_matches = _deduplicate_local_matches(matches)
|
|
return annotate_local_matches_with_download_state(deduplicated_matches)
|
|
|
|
|
|
def _collect_hashes_from_container(value: Any) -> List[str]:
|
|
h = extract_sha256_from_metadata(value)
|
|
return [h] if h else []
|
|
|
|
|
|
def _metadata_file_matches_model(
|
|
file_info: Dict[str, Any], model: ResolvedModel
|
|
) -> bool:
|
|
model_filename = str(
|
|
model.filename or get_filename_from_path(model.path)
|
|
).lower()
|
|
model_relative = str(model.relative_path or "").replace("\\", "/").lower()
|
|
model_stem = strip_known_model_extension(get_filename_from_path(model_filename)).lower()
|
|
|
|
candidates = [
|
|
file_info.get("name"),
|
|
file_info.get("filename"),
|
|
file_info.get("path"),
|
|
file_info.get("file_path"),
|
|
]
|
|
for candidate in candidates:
|
|
text = str(candidate or "").replace("\\", "/").lower()
|
|
basename = get_filename_from_path(text)
|
|
stem = strip_known_model_extension(basename).lower()
|
|
if text and model_relative and text == model_relative:
|
|
return True
|
|
if basename and basename == model_filename:
|
|
return True
|
|
if stem and model_stem and stem == model_stem:
|
|
return True
|
|
return False
|
|
|
|
|
|
def _extract_model_sha256_from_metadata(
|
|
metadata: Dict[str, Any], model: ResolvedModel
|
|
) -> List[str]:
|
|
if not isinstance(metadata, dict):
|
|
return []
|
|
|
|
hash_status = str(metadata.get("hash_status") or "").strip().lower()
|
|
if hash_status and hash_status != "completed":
|
|
return []
|
|
|
|
values: List[str] = []
|
|
values.extend(_collect_hashes_from_container(metadata))
|
|
values.extend(_collect_hashes_from_container(metadata.get("path_metadata")))
|
|
values.extend(_collect_hashes_from_container(metadata.get("file_info")))
|
|
values.extend(_collect_hashes_from_container(metadata.get("file")))
|
|
|
|
nested_file_lists = [
|
|
metadata.get("files"),
|
|
as_dict(metadata.get("selected_version")).get("files"),
|
|
as_dict(metadata.get("civitai")).get("files"),
|
|
]
|
|
for file_list in nested_file_lists:
|
|
for file_info in as_list(file_list):
|
|
if isinstance(file_info, dict) and _metadata_file_matches_model(file_info, model):
|
|
values.extend(_collect_hashes_from_container(file_info))
|
|
|
|
return [value for value in unique_ordered_strings(values) if normalize_sha256(value)]
|
|
|
|
|
|
|
|
def search_local_matches_by_hash(
|
|
sha256: str,
|
|
category: Optional[str] = None,
|
|
max_matches: int = 20,
|
|
force_rescan: bool = False,
|
|
) -> List[ModelMatch]:
|
|
"""
|
|
Find local models whose sidecar .metadata.json contains the given SHA256.
|
|
|
|
This intentionally does not hash model files. It only reads metadata sidecars
|
|
next to models already discovered by the scanner.
|
|
"""
|
|
normalized_hash = normalize_sha256(sha256)
|
|
if not normalized_hash:
|
|
return []
|
|
|
|
index = _get_local_hash_match_cache(force_rescan=force_rescan)
|
|
matches = list(index.get(normalized_hash, []))
|
|
|
|
if category and category != "unknown":
|
|
matches.sort(
|
|
key=lambda match: 0
|
|
if match.model.category == category
|
|
else 1
|
|
)
|
|
|
|
annotated_matches = annotate_local_matches_with_download_state(matches)
|
|
if max_matches > 0:
|
|
return annotated_matches[:max_matches]
|
|
return annotated_matches
|
|
|
|
|
|
def get_local_model_hash_metadata(
|
|
model_path: str,
|
|
model: Optional[ResolvedModel] = None,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Return SHA256 hashes already stored in sidecar metadata for a local model.
|
|
|
|
This intentionally does not hash model files. It only reads metadata sidecars
|
|
next to the selected local model, so callers can do quick hash comparisons.
|
|
"""
|
|
raw_path = str(model_path or "").strip()
|
|
if not raw_path:
|
|
return {"exists": False, "metadata_path": "", "hashes": [], "sha256": ""}
|
|
|
|
normalized_path = os.path.abspath(os.path.normpath(raw_path))
|
|
exists = os.path.exists(normalized_path)
|
|
file_size = 0
|
|
if exists and os.path.isfile(normalized_path):
|
|
try:
|
|
file_size = os.path.getsize(normalized_path)
|
|
except OSError:
|
|
file_size = 0
|
|
filename = get_filename_from_path(normalized_path)
|
|
model_info = (
|
|
model.with_updates(
|
|
path=normalized_path,
|
|
filename=model.filename or filename,
|
|
relative_path=model.relative_path or model.filename or filename,
|
|
)
|
|
if model is not None
|
|
else ResolvedModel(
|
|
path=normalized_path,
|
|
filename=filename,
|
|
relative_path=filename,
|
|
)
|
|
)
|
|
|
|
metadata_path = find_metadata_sidecar_path(normalized_path)
|
|
last_hash_status = ""
|
|
if metadata_path:
|
|
metadata = read_merged_model_metadata(normalized_path, None)
|
|
if isinstance(metadata, dict) and metadata:
|
|
last_hash_status = str(metadata.get("hash_status") or "").strip()
|
|
hashes = _extract_model_sha256_from_metadata(metadata, model_info)
|
|
if hashes:
|
|
return {
|
|
"exists": exists,
|
|
"metadata_path": metadata_path,
|
|
"hash_status": last_hash_status,
|
|
"hashes": hashes,
|
|
"sha256": hashes[0],
|
|
"size": file_size,
|
|
}
|
|
|
|
return {
|
|
"exists": exists,
|
|
"metadata_path": metadata_path,
|
|
"hash_status": last_hash_status,
|
|
"hashes": [],
|
|
"sha256": "",
|
|
"size": file_size,
|
|
}
|
|
|
|
|
|
def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str, Any]]:
|
|
"""
|
|
Extract model URLs from workflow JSON.
|
|
|
|
Sources:
|
|
1. node.properties.models array - contains {name, url, directory}
|
|
2. Regex extraction from workflow JSON string - finds HuggingFace/CivitAI URLs
|
|
|
|
Args:
|
|
workflow_json: Complete workflow JSON dictionary
|
|
|
|
Returns:
|
|
Dict mapping model filename -> {url, directory, source}
|
|
"""
|
|
url_map = {}
|
|
|
|
# Convert to string for regex search
|
|
workflow_str = json.dumps(workflow_json)
|
|
|
|
# Collect all serialized nodes, including nodes from subgraphs.
|
|
all_nodes = [
|
|
context.node for context in iter_workflow_nodes_with_scope(workflow_json)
|
|
]
|
|
|
|
# 1. Extract from node.properties.models (authoritative source)
|
|
for node in all_nodes:
|
|
node_type = node.get("type", "")
|
|
properties = node.get("properties", {})
|
|
models_list = properties.get("models", [])
|
|
|
|
for model_info in models_list:
|
|
if isinstance(model_info, dict):
|
|
name = model_info.get("name", "")
|
|
url = model_info.get("url", "")
|
|
directory = model_info.get("directory", "")
|
|
|
|
if name and name not in url_map:
|
|
url_map[name] = {
|
|
"url": normalize_workflow_download_url(url),
|
|
"model_url": url,
|
|
"directory": directory,
|
|
"node_type": node_type,
|
|
"source": "node_properties",
|
|
}
|
|
|
|
# 2. Extract URLs via regex from workflow JSON
|
|
urls_found = URL_PATTERN.findall(workflow_str)
|
|
|
|
# Clean URLs (remove trailing characters that may have been captured)
|
|
cleaned_urls = []
|
|
for url in urls_found:
|
|
url = url.split(")")[0].replace("\\n", "").replace("\n", "").strip()
|
|
if url:
|
|
cleaned_urls.append(url)
|
|
|
|
# 3. Extract model filenames via regex
|
|
model_pattern = re.compile(
|
|
r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx|gguf))", re.IGNORECASE
|
|
)
|
|
model_files_raw = model_pattern.findall(workflow_str)
|
|
|
|
# Clean and decode filenames
|
|
model_files = set()
|
|
model_name_map = {} # decoded -> original
|
|
|
|
for model in model_files_raw:
|
|
cleaned = model.strip()
|
|
if cleaned and cleaned[0].isalnum():
|
|
try:
|
|
decoded = unquote(cleaned)
|
|
except Exception:
|
|
decoded = cleaned
|
|
model_files.add(decoded)
|
|
model_name_map[decoded] = cleaned
|
|
|
|
# 4. Match URLs to model filenames
|
|
for model in model_files:
|
|
# Keep authoritative node.properties URLs only when they point to the file.
|
|
# Some workflows store a repo/model-page URL there and a concrete file URL
|
|
# elsewhere in the JSON; in that case the file URL is the usable download.
|
|
if (
|
|
model in url_map
|
|
and url_map[model].get("url")
|
|
and workflow_url_points_to_file(url_map[model].get("url"), model)
|
|
):
|
|
continue
|
|
|
|
original_name = model_name_map.get(model, model)
|
|
|
|
for url in cleaned_urls:
|
|
# Check decoded name in URL
|
|
if model in url:
|
|
if (
|
|
model not in url_map
|
|
or not url_map[model].get("url")
|
|
or not workflow_url_points_to_file(url_map[model].get("url"), model)
|
|
):
|
|
url_map[model] = {
|
|
"url": normalize_workflow_download_url(url),
|
|
"model_url": url,
|
|
"directory": url_map.get(model, {}).get("directory", ""),
|
|
"source": "regex",
|
|
}
|
|
else:
|
|
url_map[model]["url"] = normalize_workflow_download_url(url)
|
|
url_map[model]["model_url"] = url
|
|
url_map[model]["source"] = "regex"
|
|
break
|
|
# Check original (possibly URL-encoded) name in URL
|
|
if original_name in url:
|
|
if (
|
|
model not in url_map
|
|
or not url_map[model].get("url")
|
|
or not workflow_url_points_to_file(
|
|
url_map[model].get("url"), original_name
|
|
)
|
|
):
|
|
url_map[model] = {
|
|
"url": normalize_workflow_download_url(url),
|
|
"model_url": url,
|
|
"directory": url_map.get(model, {}).get("directory", ""),
|
|
"source": "regex",
|
|
}
|
|
else:
|
|
url_map[model]["url"] = normalize_workflow_download_url(url)
|
|
url_map[model]["model_url"] = url
|
|
url_map[model]["source"] = "regex"
|
|
break
|
|
# Check without extension
|
|
model_base = os.path.splitext(model)[0]
|
|
if model_base in url or unquote(model_base) in url:
|
|
if (
|
|
model not in url_map
|
|
or not url_map[model].get("url")
|
|
or not workflow_url_points_to_file(url_map[model].get("url"), model)
|
|
):
|
|
url_map[model] = {
|
|
"url": normalize_workflow_download_url(url),
|
|
"model_url": url,
|
|
"directory": url_map.get(model, {}).get("directory", ""),
|
|
"source": "regex",
|
|
}
|
|
else:
|
|
url_map[model]["url"] = normalize_workflow_download_url(url)
|
|
url_map[model]["model_url"] = url
|
|
url_map[model]["source"] = "regex"
|
|
break
|
|
|
|
return url_map
|
|
|
|
|
|
def analyze_and_find_matches(
|
|
workflow_json: Dict[str, Any],
|
|
similarity_threshold: float = 0.0,
|
|
max_matches_per_model: int = 10,
|
|
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
force_rescan: bool = False,
|
|
analysis_id: Optional[str] = None,
|
|
) -> WorkflowAnalysisResult:
|
|
"""
|
|
Main entry point: analyze workflow and find matches for missing models.
|
|
|
|
Args:
|
|
workflow_json: Complete workflow JSON dictionary
|
|
similarity_threshold: Minimum similarity score (0.0 to 1.0) for matches
|
|
max_matches_per_model: Maximum number of matches to return per missing model
|
|
force_rescan: If True, bypass the short-lived local model scan cache
|
|
analysis_id: Optional identifier used to correlate logs for one request
|
|
|
|
Returns:
|
|
Typed workflow analysis result. The HTTP route serializes it only after
|
|
applying route-level filtering and local download-source enrichment.
|
|
"""
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "extracting",
|
|
"message": "Extracting workflow model references...",
|
|
"current": 0,
|
|
"total": 0,
|
|
}
|
|
)
|
|
|
|
if (
|
|
not workflow_has_nodes(workflow_json)
|
|
or not workflow_has_potential_model_references(workflow_json)
|
|
):
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "completed",
|
|
"message": "Analysis complete",
|
|
"current": 0,
|
|
"total": 0,
|
|
}
|
|
)
|
|
return WorkflowAnalysisResult()
|
|
|
|
# Extract URLs from workflow (node.properties.models + regex)
|
|
analysis_context = f" (analysis_id={analysis_id})" if analysis_id else ""
|
|
workflow_urls = extract_workflow_urls(workflow_json)
|
|
log.debug(
|
|
f"Extracted {len(workflow_urls)} URLs from workflow{analysis_context}"
|
|
)
|
|
workflow_hashes = extract_workflow_hash_metadata(workflow_json)
|
|
log.debug(
|
|
f"Extracted {len(workflow_hashes)} workflow hash metadata keys"
|
|
f"{analysis_context}"
|
|
)
|
|
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "scanning",
|
|
"message": "Scanning local model index...",
|
|
"current": 0,
|
|
"total": 0,
|
|
}
|
|
)
|
|
|
|
inventory = get_workflow_model_inventory(
|
|
workflow_json,
|
|
force_rescan=force_rescan,
|
|
progress_callback=progress_callback,
|
|
analysis_id=analysis_id,
|
|
)
|
|
available_models = inventory.available_models
|
|
all_model_refs = inventory.model_refs
|
|
available_models = normalize_models(available_models)
|
|
all_model_refs = [
|
|
ref
|
|
if isinstance(ref, ModelReference)
|
|
else ModelReference.from_mapping(ref)
|
|
for ref in all_model_refs
|
|
]
|
|
available_models_by_category = {}
|
|
for model in available_models:
|
|
model_category = model.category
|
|
if model_category not in available_models_by_category:
|
|
available_models_by_category[model_category] = []
|
|
available_models_by_category[model_category].append(model)
|
|
|
|
ordered_candidates_cache: Dict[str, List[ResolvedModel]] = {}
|
|
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "identifying",
|
|
"message": "Identifying missing models...",
|
|
"current": 0,
|
|
"total": len(all_model_refs),
|
|
}
|
|
)
|
|
|
|
# Identify missing models
|
|
missing_models = identify_missing_models(all_model_refs, available_models)
|
|
resolved_model_refs = [
|
|
model_ref for model_ref in all_model_refs if model_ref.exists
|
|
]
|
|
|
|
# Enrich typed missing models with workflow URLs, hash metadata, and URN
|
|
# hints. Raw mappings are created only for helpers that still consume
|
|
# external workflow/custom-node metadata.
|
|
enriched_missing_models = []
|
|
for missing in missing_models:
|
|
updates: Dict[str, Any] = {}
|
|
original_path = missing.reference.original_path
|
|
filename = get_filename_from_path(original_path)
|
|
|
|
url_info = get_workflow_url_info_for_filename(workflow_urls, filename)
|
|
if url_info:
|
|
updates.update(
|
|
{
|
|
"workflow_url": url_info.get("url", ""),
|
|
"workflow_model_url": url_info.get("model_url", ""),
|
|
"workflow_directory": url_info.get("directory", ""),
|
|
"url_source": url_info.get("source", ""),
|
|
}
|
|
)
|
|
hash_info = get_workflow_hash_info_for_ref(
|
|
workflow_hashes,
|
|
missing.reference,
|
|
)
|
|
if hash_info:
|
|
updates.update(
|
|
{
|
|
"workflow_sha256": normalize_sha256(hash_info.get("sha256")),
|
|
"hash_lookup_source": hash_info.get("source")
|
|
or "workflow_metadata",
|
|
}
|
|
)
|
|
|
|
# Handle URNs: mark for async resolution by frontend. No sync CivitAI
|
|
# calls happen here; the frontend resolves them asynchronously.
|
|
if missing.extra_value("is_urn", False):
|
|
updates["needs_urn_resolve"] = True
|
|
urn = missing.extra_value("urn")
|
|
if urn:
|
|
updates.update(
|
|
{
|
|
"urn_model_id": urn.get("model_id"),
|
|
"urn_version_id": urn.get("version_id"),
|
|
"urn_type": urn.get("type", ""),
|
|
}
|
|
)
|
|
|
|
enriched_missing_models.append(missing.with_extra(**updates))
|
|
missing_models = enriched_missing_models
|
|
|
|
total_matching_models = len(missing_models)
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "matching",
|
|
"message": "Matching local models...",
|
|
"current": 0,
|
|
"total": total_matching_models,
|
|
}
|
|
)
|
|
|
|
local_match_cache: Dict[Tuple[str, str], List[ModelMatch]] = {}
|
|
active_downloads_by_path = _get_active_downloads_by_path()
|
|
|
|
def get_match_target(model_ref: MissingModel) -> Optional[str]:
|
|
reference = model_ref.reference
|
|
target_for_matching = (
|
|
reference.original_path
|
|
or model_ref.extra_value("expected_filename")
|
|
or model_ref.extra_value("name")
|
|
or model_ref.extra_value("filename")
|
|
or model_ref.extra_value("full_path")
|
|
or ""
|
|
)
|
|
|
|
# For URNs, prefer expected_filename for matching
|
|
expected_filename = model_ref.extra_value("expected_filename")
|
|
if model_ref.extra_value("is_urn") and expected_filename:
|
|
return expected_filename
|
|
|
|
if isinstance(target_for_matching, str) and target_for_matching.startswith(
|
|
"urn:air:"
|
|
):
|
|
# Don't fuzzy-match the full URN string against every local model.
|
|
# For worker-asset style URNs, try using the filename-like suffix after "@";
|
|
# otherwise skip local matching until the frontend resolves the URN async.
|
|
urn_suffix = (
|
|
target_for_matching.split("@", 1)[1]
|
|
if "@" in target_for_matching
|
|
else ""
|
|
)
|
|
urn_suffix_ext = os.path.splitext(urn_suffix)[1].lower()
|
|
if urn_suffix and urn_suffix_ext in MODEL_EXTENSIONS:
|
|
return urn_suffix
|
|
return None
|
|
|
|
return target_for_matching
|
|
|
|
def get_match_category(model_ref: MissingModel) -> str:
|
|
category = model_ref.reference.category
|
|
if not category or category == "unknown":
|
|
node_type = model_ref.reference.node_type
|
|
category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown")
|
|
return category or "unknown"
|
|
|
|
def get_candidates_for_category(category: str) -> List[ResolvedModel]:
|
|
if not category or category == "unknown":
|
|
return available_models
|
|
|
|
candidates = ordered_candidates_cache.get(category)
|
|
if candidates is None:
|
|
preferred = available_models_by_category.get(category, [])
|
|
others = [
|
|
m for m in available_models if m.category != category
|
|
]
|
|
candidates = preferred + others
|
|
ordered_candidates_cache[category] = candidates
|
|
|
|
return candidates
|
|
|
|
def find_local_matches_for_ref(model_ref: MissingModel) -> List[ModelMatch]:
|
|
target_for_matching = get_match_target(model_ref)
|
|
category = get_match_category(model_ref)
|
|
workflow_sha256 = normalize_sha256(
|
|
model_ref.extra_value("workflow_sha256")
|
|
)
|
|
hash_matches: List[ModelMatch] = []
|
|
if workflow_sha256:
|
|
hash_matches = collect_local_hash_matches_for_result(
|
|
workflow_sha256,
|
|
search_local_matches_by_hash=search_local_matches_by_hash,
|
|
category=category,
|
|
max_matches=max_matches_per_model,
|
|
force_rescan=False,
|
|
source=model_ref.extra_value("hash_lookup_source")
|
|
or "workflow_metadata",
|
|
filename=get_filename_from_path(
|
|
model_ref.reference.original_path or ""
|
|
),
|
|
)
|
|
|
|
if not target_for_matching:
|
|
return hash_matches
|
|
|
|
cache_key = (target_for_matching, category)
|
|
if cache_key in local_match_cache:
|
|
return _deduplicate_local_matches(hash_matches + local_match_cache[cache_key])
|
|
|
|
matches = find_matches(
|
|
target_for_matching,
|
|
get_candidates_for_category(category),
|
|
threshold=similarity_threshold,
|
|
max_results=max_matches_per_model,
|
|
)
|
|
deduplicated_matches = annotate_local_matches_with_download_state(
|
|
_deduplicate_local_matches(matches),
|
|
active_downloads_by_path,
|
|
)
|
|
local_match_cache[cache_key] = deduplicated_matches
|
|
return _deduplicate_local_matches(hash_matches + deduplicated_matches)
|
|
|
|
# Find matches for each missing model
|
|
missing_with_matches = []
|
|
total_missing = len(missing_models)
|
|
for index, missing in enumerate(missing_models, start=1):
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "matching",
|
|
"message": f"Analyzing model {index} of {total_missing}",
|
|
"current": index,
|
|
"total": total_matching_models,
|
|
"model_name": missing.extra_value("name")
|
|
or missing.reference.original_path,
|
|
}
|
|
)
|
|
|
|
name = missing.extra_value("name") or missing.reference.original_path
|
|
if should_skip_existing_custom_node_reference(missing.reference):
|
|
log.info(
|
|
f"Skipping existing custom-node model reference: {name}"
|
|
)
|
|
continue
|
|
|
|
missing_with_matches.append(
|
|
missing.with_matches(
|
|
find_local_matches_for_ref(missing)
|
|
)
|
|
)
|
|
|
|
# Existing models already have an exact local path. Fuzzy-matching every
|
|
# resolved reference against the full local model index is redundant and
|
|
# makes small workflow edits unnecessarily expensive.
|
|
resolved_references = tuple(resolved_model_refs)
|
|
result = WorkflowAnalysisResult(
|
|
missing_models=tuple(missing_with_matches),
|
|
resolved_models=resolved_references,
|
|
total_resolved=len(resolved_references),
|
|
total_missing=len(missing_with_matches),
|
|
total_models_analyzed=len(all_model_refs),
|
|
)
|
|
|
|
if progress_callback:
|
|
progress_callback(
|
|
{
|
|
"stage": "completed",
|
|
"message": "Analysis complete",
|
|
"current": total_matching_models,
|
|
"total": total_matching_models,
|
|
}
|
|
)
|
|
|
|
return result
|
|
|
|
|
|
def apply_resolution(
|
|
workflow_json: Dict[str, Any],
|
|
resolutions: List[Resolution],
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Apply model resolutions to workflow.
|
|
|
|
Args:
|
|
workflow_json: Workflow JSON dictionary (will be modified)
|
|
resolutions: Validated ``Resolution`` objects.
|
|
|
|
Returns:
|
|
Updated workflow JSON dictionary
|
|
"""
|
|
typed_resolutions = []
|
|
for resolution in resolutions:
|
|
try:
|
|
resolution.validate()
|
|
except (TypeError, ValueError) as exc:
|
|
log.warning(f"Skipping invalid workflow resolution: {exc}")
|
|
continue
|
|
|
|
typed_resolutions.append(resolution)
|
|
|
|
# Update workflow
|
|
updated_workflow = update_workflow_nodes(workflow_json, typed_resolutions)
|
|
|
|
return updated_workflow
|