Dariusz L
0837964874
Add support for the gguf model format across the codebase and improve handling of nested/dict-type widget model references. Changes include: - Add ".gguf" to model extension lists in core/linker.py, core/scanner.py, core/sources/huggingface.py, and core/workflow_analyzer.py so gguf files are detected and suggested. - Update HuggingFace search/repo listing to consider gguf files and improve some request/response handling and logging. - Improve scanner implementation: formatting/PEP8 cleanup, robust extension checks, normalized MODEL_EXTENSIONS set, safer path handling, and clearer logging. - Add NESTED_MODEL_KEYS and logic in workflow_analyzer to detect model filenames stored inside dict-type widget values (e.g. rgthree Power Lora Loader) and record a nested_key for resolution. - Propagate nested_key through apply_resolution mapping and update update_model_path to write resolved paths back into nested dict keys when present (instead of only replacing the top-level widget value). - Update web UI (web/linker.js) regexes to strip/display gguf extensions in model name formatting. These changes allow gguf models to be discovered, resolved and updated in workflows, and enable support for nodes that store model references inside nested dict fields.
397 lines
14 KiB
Python
397 lines
14 KiB
Python
"""
|
|
Core Linker Module
|
|
|
|
Integrates all components to provide high-level API for model linking.
|
|
"""
|
|
|
|
import os
|
|
import re
|
|
import json
|
|
import logging
|
|
from typing import Dict, Any, List, Optional, Tuple
|
|
from urllib.parse import unquote
|
|
|
|
from .scanner import get_model_files
|
|
from .workflow_analyzer import analyze_workflow_models, identify_missing_models
|
|
from .matcher import find_matches
|
|
from .workflow_updater import update_workflow_nodes
|
|
from .sources.civitai import resolve_urn
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# 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 = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx", ".gguf")
|
|
|
|
|
|
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 nodes including from subgraphs
|
|
all_nodes = list(workflow_json.get("nodes", []))
|
|
definitions = workflow_json.get("definitions", {})
|
|
subgraphs = definitions.get("subgraphs", [])
|
|
for subgraph in subgraphs:
|
|
subgraph_nodes = subgraph.get("nodes", [])
|
|
all_nodes.extend(subgraph_nodes)
|
|
|
|
# 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": 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:
|
|
# Skip if already found in node.properties.models
|
|
if model in url_map and url_map[model].get("url"):
|
|
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:
|
|
url_map[model] = {"url": url, "directory": "", "source": "regex"}
|
|
elif not url_map[model].get("url"):
|
|
url_map[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:
|
|
url_map[model] = {"url": url, "directory": "", "source": "regex"}
|
|
elif not url_map[model].get("url"):
|
|
url_map[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:
|
|
url_map[model] = {"url": url, "directory": "", "source": "regex"}
|
|
elif not url_map[model].get("url"):
|
|
url_map[model]["url"] = url
|
|
url_map[model]["source"] = "regex"
|
|
break
|
|
|
|
return url_map
|
|
|
|
|
|
def parse_huggingface_url(url: str) -> Tuple[Optional[str], Optional[str]]:
|
|
"""
|
|
Extract HuggingFace repo and path from URL.
|
|
|
|
Args:
|
|
url: HuggingFace URL
|
|
|
|
Returns:
|
|
Tuple of (repo_id, file_path) or (None, None) if not valid
|
|
"""
|
|
if not url or "huggingface.co" not in url:
|
|
return None, None
|
|
|
|
# Pattern: https://huggingface.co/user/repo/resolve/main/path/to/file.safetensors
|
|
match = re.match(
|
|
r"https?://huggingface\.co/([^/]+/[^/]+)/(?:resolve|blob)/[^/]+/(.+)", url
|
|
)
|
|
if match:
|
|
return match.group(1), match.group(2)
|
|
|
|
return None, None
|
|
|
|
|
|
def analyze_and_find_matches(
|
|
workflow_json: Dict[str, Any],
|
|
similarity_threshold: float = 0.0,
|
|
max_matches_per_model: int = 10,
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
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
|
|
|
|
Returns:
|
|
Dictionary with analysis results:
|
|
{
|
|
'missing_models': [
|
|
{
|
|
'node_id': node ID,
|
|
'node_type': node type,
|
|
'widget_index': widget index,
|
|
'original_path': original path from workflow,
|
|
'category': model category,
|
|
'workflow_url': URL from workflow if found,
|
|
'workflow_directory': directory from workflow if found,
|
|
'matches': [
|
|
{
|
|
'model': model dict from scanner,
|
|
'filename': model filename,
|
|
'similarity': similarity score (0.0-1.0),
|
|
'confidence': confidence percentage (0-100)
|
|
},
|
|
...
|
|
]
|
|
},
|
|
...
|
|
],
|
|
'total_missing': count of missing models,
|
|
'total_models_analyzed': count of all models in workflow
|
|
}
|
|
"""
|
|
# Extract URLs from workflow (node.properties.models + regex)
|
|
workflow_urls = extract_workflow_urls(workflow_json)
|
|
logger.debug(f"Extracted {len(workflow_urls)} URLs from workflow")
|
|
|
|
# Analyze workflow to find all model references
|
|
all_model_refs = analyze_workflow_models(workflow_json)
|
|
|
|
# Get available models
|
|
available_models = get_model_files()
|
|
|
|
# Identify missing models
|
|
missing_models = identify_missing_models(all_model_refs, available_models)
|
|
|
|
# Enrich missing models with workflow URLs
|
|
for missing in missing_models:
|
|
original_path = missing.get("original_path", "")
|
|
filename = os.path.basename(original_path)
|
|
|
|
if filename in workflow_urls:
|
|
url_info = workflow_urls[filename]
|
|
missing["workflow_url"] = url_info.get("url", "")
|
|
missing["workflow_directory"] = url_info.get("directory", "")
|
|
missing["url_source"] = url_info.get("source", "")
|
|
|
|
# Handle URNs: fetch expected filename from CivitAI
|
|
for missing in missing_models:
|
|
if missing.get("is_urn"):
|
|
urn = missing["urn"]
|
|
model_info = resolve_urn(urn["model_id"], urn["version_id"])
|
|
if model_info:
|
|
missing["civitai_info"] = {
|
|
"model_name": model_info.get("model_name"),
|
|
"version_name": model_info.get("version_name"),
|
|
}
|
|
if model_info:
|
|
missing["civitai_info"] = model_info
|
|
missing["expected_filename"] = model_info["expected_filename"]
|
|
logger.debug(
|
|
f"URN {missing['original_path']} → {missing['expected_filename']}"
|
|
)
|
|
|
|
# Add URN type to missing model for search functionality
|
|
if "urn" in missing:
|
|
missing["urn_type"] = missing["urn"].get("type", "")
|
|
|
|
# Find matches for each missing model
|
|
missing_with_matches = []
|
|
for missing in missing_models:
|
|
# Skip LoraManager lorAs that already exist locally (exists=True means no linking needed)
|
|
is_lora_v2 = missing.get("is_lora_v2")
|
|
exists = missing.get("exists")
|
|
name = missing.get("name") or missing.get("original_path", "")
|
|
logger.debug(f"Checking {name}: is_lora_v2={is_lora_v2}, exists={exists}")
|
|
|
|
if is_lora_v2 and exists:
|
|
logger.info(f"Skipping LoraManager lora {name} - already exists locally")
|
|
continue
|
|
|
|
target_for_matching = missing.get("original_path", "")
|
|
|
|
# For URNs, prefer expected_filename for matching
|
|
if missing.get("is_urn") and missing.get("expected_filename"):
|
|
target_for_matching = missing["expected_filename"]
|
|
|
|
# Filter available models by category if known
|
|
category = missing.get("category")
|
|
if not category or category == "unknown":
|
|
from .workflow_analyzer import NODE_TYPE_TO_CATEGORY_HINTS
|
|
|
|
node_type = missing.get("node_type", "")
|
|
category = NODE_TYPE_TO_CATEGORY_HINTS.get(node_type, "unknown")
|
|
|
|
candidates = available_models
|
|
if category and category != "unknown":
|
|
candidates = [m for m in available_models if m.get("category") == category]
|
|
candidates.extend(
|
|
[m for m in available_models if m.get("category") != category]
|
|
)
|
|
|
|
# Find matches
|
|
matches = find_matches(
|
|
target_for_matching,
|
|
candidates,
|
|
threshold=similarity_threshold,
|
|
max_results=max_matches_per_model,
|
|
)
|
|
|
|
# Deduplicate matches by absolute path
|
|
seen_absolute_paths = {}
|
|
deduplicated_matches = []
|
|
for match in matches:
|
|
model_dict = match["model"]
|
|
absolute_path = model_dict.get("path", "")
|
|
if absolute_path:
|
|
absolute_path = os.path.normpath(absolute_path)
|
|
|
|
if absolute_path not in seen_absolute_paths:
|
|
seen_absolute_paths[absolute_path] = match
|
|
deduplicated_matches.append(match)
|
|
else:
|
|
existing_match = seen_absolute_paths[absolute_path]
|
|
if match["confidence"] > existing_match["confidence"]:
|
|
idx = deduplicated_matches.index(existing_match)
|
|
deduplicated_matches[idx] = match
|
|
seen_absolute_paths[absolute_path] = match
|
|
|
|
missing_with_matches.append({**missing, "matches": deduplicated_matches})
|
|
|
|
return {
|
|
"missing_models": missing_with_matches,
|
|
"total_missing": len(missing_with_matches),
|
|
"total_models_analyzed": len(all_model_refs),
|
|
}
|
|
|
|
|
|
def apply_resolution(
|
|
workflow_json: Dict[str, Any], resolutions: List[Dict[str, Any]]
|
|
) -> Dict[str, Any]:
|
|
"""
|
|
Apply model resolutions to workflow.
|
|
|
|
Args:
|
|
workflow_json: Workflow JSON dictionary (will be modified)
|
|
resolutions: List of resolution dictionaries:
|
|
{
|
|
'node_id': node ID,
|
|
'widget_index': widget index,
|
|
'resolved_path': absolute path to resolved model,
|
|
'category': model category (optional),
|
|
'resolved_model': model dict from scanner (optional),
|
|
'nested_key': nested key for dict-type widgets (optional)
|
|
}
|
|
|
|
Returns:
|
|
Updated workflow JSON dictionary
|
|
"""
|
|
# Prepare mappings for workflow_updater
|
|
mappings = []
|
|
for resolution in resolutions:
|
|
mapping = {
|
|
"node_id": resolution.get("node_id"),
|
|
"widget_index": resolution.get("widget_index"),
|
|
"resolved_path": resolution.get("resolved_path"),
|
|
"category": resolution.get("category"),
|
|
"resolved_model": resolution.get("resolved_model"),
|
|
"subgraph_id": resolution.get(
|
|
"subgraph_id"
|
|
), # Include subgraph_id for subgraph nodes
|
|
"is_top_level": resolution.get(
|
|
"is_top_level"
|
|
), # True for top-level nodes, False for nodes in subgraph definitions
|
|
"is_lora_v2": resolution.get("is_lora_v2"), # Flag for LoraManager nodes
|
|
"original_lora_name": resolution.get(
|
|
"original_lora_name"
|
|
), # Original lora name for LoraManager replacement
|
|
"nested_key": resolution.get(
|
|
"nested_key"
|
|
), # For dict-type widgets (e.g. Power Lora Loader)
|
|
}
|
|
|
|
# If resolved_model provided, extract path if needed
|
|
if "resolved_model" in resolution and resolution["resolved_model"]:
|
|
resolved_model = resolution["resolved_model"]
|
|
if "path" in resolved_model and not mapping.get("resolved_path"):
|
|
mapping["resolved_path"] = resolved_model["path"]
|
|
if "base_directory" in resolved_model:
|
|
mapping["base_directory"] = resolved_model["base_directory"]
|
|
|
|
mappings.append(mapping)
|
|
|
|
# Update workflow
|
|
updated_workflow = update_workflow_nodes(workflow_json, mappings)
|
|
|
|
return updated_workflow
|
|
|
|
|
|
def get_resolution_summary(workflow_json: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Get summary of missing models and matches without applying resolutions.
|
|
|
|
This is a convenience method that calls analyze_and_find_matches with defaults.
|
|
|
|
Args:
|
|
workflow_json: Complete workflow JSON dictionary
|
|
|
|
Returns:
|
|
Same format as analyze_and_find_matches
|
|
"""
|
|
return analyze_and_find_matches(workflow_json)
|