Files
Azornes-Comfyui-Model-Resolver/__init__.py
T
Dariusz L 259e491c96 Handle LoraManager lora detection and updates
Filter out LoraManager (LoraLoaderV2) entries that already exist locally so they don't appear as missing, and skip matching for those entries. Integrate scanner data to detect local lora files (full_path/exists) in workflow analysis. Add special handling in update_model_path/update_workflow_nodes to update LoraManager lora lists and the formatted lora text widget (widgets_values[2] and [1]) when a new lora is resolved, using mapping fields is_lora_v2 and original_lora_name. Improve convert_to_relative_path and get_base_directory_for_model to prefer folder_paths lookup and preserve exact ComfyUI filename formats. In the web UI, make resolve button IDs unique by including a missingIndex, avoid duplicate listeners, forward is_lora_v2/original_lora_name to the resolve API, and add extra logging for debugging.
2026-04-18 02:51:06 +02:00

767 lines
36 KiB
Python

"""
@author: Model Linker Team
@title: ComfyUI Model Linker
@nickname: Model Linker
@version: 1.1.0
@description: Extension for relinking missing models and downloading from HuggingFace/CivitAI
"""
import logging
# Web directory for JavaScript interface
WEB_DIRECTORY = "./web"
# Empty NODE_CLASS_MAPPINGS - we don't provide custom nodes, only web extension
# This prevents ComfyUI from showing "IMPORT FAILED" message
NODE_CLASS_MAPPINGS = {}
__all__ = ["WEB_DIRECTORY"]
class ModelLinkerExtension:
"""Main extension class for Model Linker."""
def __init__(self):
self.routes_setup = False
self.logger = logging.getLogger(__name__)
def initialize(self):
"""Initialize the extension and set up API routes."""
try:
self.setup_routes()
self.logger.info("Model Linker: Extension initialized successfully")
except Exception as e:
self.logger.error(
f"Model Linker: Extension initialization failed: {e}", exc_info=True
)
def setup_routes(self):
"""Register API routes for the Model Linker extension."""
if self.routes_setup:
return # Already set up
try:
from aiohttp import web
# Try to get routes from PromptServer
try:
from server import PromptServer
if (
not hasattr(PromptServer, "instance")
or PromptServer.instance is None
):
self.logger.debug("Model Linker: PromptServer not available yet")
return False
routes = PromptServer.instance.routes
except (ImportError, AttributeError) as e:
self.logger.debug(f"Model Linker: Could not access PromptServer: {e}")
return False
# Import linker modules
try:
from .core.linker import analyze_and_find_matches, apply_resolution
from .core.scanner import get_model_files
except ImportError as e:
self.logger.error(f"Model Linker: Could not import core modules: {e}")
return False
# Import download modules
try:
from .core.downloader import (
start_background_download,
get_progress,
get_all_progress,
cancel_download,
get_download_directory,
)
from .core.sources.popular import (
get_popular_model_url,
search_popular_models,
)
from .core.sources.model_list import (
search_model_list,
search_model_list_multiple,
)
from .core.sources.huggingface import search_huggingface_for_file
from .core.sources.civitai import (
search_civitai_for_file,
search_civitai,
get_civitai_download_url,
)
download_available = True
except ImportError as e:
self.logger.warning(
f"Model Linker: Download features not available: {e}"
)
download_available = False
# ==================== ANALYZE ROUTES ====================
@routes.post("/model_linker/analyze")
async def analyze_workflow(request):
"""Analyze workflow and return missing models with matches."""
try:
data = await request.json()
workflow_json = data.get("workflow")
if not workflow_json:
return web.json_response(
{"error": "Workflow JSON is required"}, status=400
)
# Analyze and find matches
result = analyze_and_find_matches(workflow_json)
# Filter out LoraManager lorAs that already exist locally (exists=True)
# These should not appear in missing models at all
missing_models = result.get("missing_models", [])
filtered_missing = []
for missing in missing_models:
is_lora = missing.get("is_lora_v2")
exists = missing.get("exists")
name = missing.get("name") or missing.get("original_path", "")
# Log for debugging
import logging
logging.getLogger(__name__).debug(
f"Filtering: {name} is_lora_v2={is_lora} exists={exists}"
)
# Skip LoraManager lorAs that already exist locally
if is_lora and exists:
logging.getLogger(__name__).info(
f"Filtered out LoraManager lora: {name}"
)
continue
filtered_missing.append(missing)
result["missing_models"] = filtered_missing
result["total_missing"] = len(filtered_missing)
# If download available, auto-search for download sources when no 100% local match
if download_available:
for missing in result.get("missing_models", []):
# Check if there's a 100% local match
matches = missing.get("matches", [])
has_perfect_match = any(
m.get("confidence", 0) == 100 for m in matches
)
if not has_perfect_match:
filename = (
missing.get("original_path", "")
.split("/")[-1]
.split("\\")[-1]
)
# 0. Check workflow URL first (highest priority - directly from workflow)
workflow_url = missing.get("workflow_url", "")
if workflow_url:
# Determine source from URL
if "huggingface.co" in workflow_url:
source = "huggingface"
elif "civitai.com" in workflow_url:
source = "civitai"
else:
source = "workflow"
# Try to get file size with HEAD request (non-blocking, timeout quickly)
file_size = None
try:
import requests
head_response = requests.head(
workflow_url,
allow_redirects=True,
timeout=5,
)
if head_response.status_code == 200:
file_size = int(
head_response.headers.get(
"content-length", 0
)
)
except Exception:
pass # Size unknown is fine
missing["download_source"] = {
"source": source,
"url": workflow_url,
"filename": filename,
"directory": missing.get(
"workflow_directory", ""
)
or missing.get("category", "checkpoints"),
"match_type": "exact",
"url_source": "workflow",
"size": file_size,
}
continue
# 1. Check popular models (always exact match)
popular_info = get_popular_model_url(filename)
if popular_info:
missing["download_source"] = {
"source": "popular",
"url": popular_info.get("url"),
"filename": filename,
"type": popular_info.get("type"),
"directory": popular_info.get("directory"),
"match_type": "exact",
}
continue
# 2. Check model list (ComfyUI Manager database)
# Use exact_only=True to avoid confusing fuzzy matches for downloads
model_list_result = search_model_list(
filename, exact_only=True
)
if model_list_result:
missing["download_source"] = {
"source": "model_list",
"url": model_list_result.get("url"),
"filename": model_list_result.get("filename"),
"name": model_list_result.get("name"),
"type": model_list_result.get("type"),
"directory": model_list_result.get("directory"),
"size": model_list_result.get("size"),
"match_type": model_list_result.get(
"match_type"
),
"confidence": model_list_result.get(
"confidence"
),
}
continue
# 3. Search HuggingFace (exact_only=True for downloads)
hf_result = search_huggingface_for_file(
filename, exact_only=True
)
if hf_result:
missing["download_source"] = {
"source": "huggingface",
"url": hf_result.get("url"),
"filename": hf_result.get("filename"),
"name": hf_result.get("repo_id", ""),
"size": hf_result.get("size", ""),
"match_type": hf_result.get(
"match_type", "exact"
),
}
continue
# 4. Search CivitAI (exact_only=True for downloads)
civitai_result = search_civitai_for_file(
filename, exact_only=True
)
if civitai_result:
missing["download_source"] = {
"source": "civitai",
"url": civitai_result.get(
"download_url"
), # CivitAI uses download_url
"filename": civitai_result.get("filename"),
"name": civitai_result.get("name", ""),
"size": civitai_result.get("size", ""),
"match_type": civitai_result.get(
"match_type", "exact"
),
}
return web.json_response(result)
except Exception as e:
self.logger.error(f"Model Linker analyze error: {e}", exc_info=True)
return web.json_response({"error": str(e)}, status=500)
@routes.post("/model_linker/resolve")
async def resolve_models(request):
"""Apply model resolution and return updated workflow."""
try:
data = await request.json()
workflow_json = data.get("workflow")
resolutions = data.get("resolutions", [])
if not workflow_json:
return web.json_response(
{"error": "Workflow JSON is required"}, status=400
)
if not resolutions:
return web.json_response(
{"error": "Resolutions array is required"}, status=400
)
# Apply resolutions
updated_workflow = apply_resolution(workflow_json, resolutions)
return web.json_response(
{"workflow": updated_workflow, "success": True}
)
except Exception as e:
self.logger.error(f"Model Linker resolve error: {e}", exc_info=True)
return web.json_response(
{"error": str(e), "success": False}, status=500
)
@routes.get("/model_linker/models")
async def get_models(request):
"""Get list of all available models."""
try:
models = get_model_files()
return web.json_response(models)
except Exception as e:
self.logger.error(
f"Model Linker get_models error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
@routes.post("/model_linker/loaded")
async def get_loaded_models(request):
"""Get all currently loaded models in the workflow."""
try:
data = await request.json()
workflow_json = data.get("workflow")
if not workflow_json:
return web.json_response(
{"error": "Workflow JSON is required"}, status=400
)
# Import workflow analyzer to extract models
from .core.workflow_analyzer import (
analyze_workflow_models,
try_resolve_model_path,
is_model_filename,
URN_REGEX,
URN_TYPE_MAP,
)
# Get available models for existence checking
available_models = get_model_files()
available_paths = {m.get("path") for m in available_models}
# Analyze workflow to get all model references
all_model_refs = analyze_workflow_models(workflow_json)
# Also extract from node.properties.models
nodes = list(workflow_json.get("nodes", []))
definitions = workflow_json.get("definitions", {})
subgraphs = definitions.get("subgraphs", [])
for subgraph in subgraphs:
nodes.extend(subgraph.get("nodes", []))
# Collect all loaded models with their values
loaded_models = []
# Process each model reference from analyze_workflow_models
for ref in all_model_refs:
original_path = ref.get("original_path", "")
node_id = ref.get("node_id")
widget_index = ref.get("widget_index")
node_type = ref.get("node_type", "")
category = ref.get("category", "unknown")
# Determine model name and strength
model_name = original_path.split("/")[-1].split("\\")[-1]
strength = None
# For standard LoraLoader nodes, strength is in next widget_value
if node_type in ["LoraLoader", "LoraLoaderModelOnly"]:
# Find the node in workflow to get strength value
for node in nodes:
if str(node.get("id")) == str(node_id):
widgets_values = node.get("widgets_values", [])
if len(widgets_values) > widget_index + 1:
try:
strength = float(
widgets_values[widget_index + 1]
)
except (ValueError, TypeError):
strength = 1.0
break
# For text-based lora loaders (LoraLoaderV2, LoraManager), get strength from ref
if ref.get("is_lora_v2"):
strength = ref.get("strength")
model_name = ref.get("name", model_name)
# Check if model exists locally
exists = ref.get("exists", False)
# If URN, resolve to display name
if ref.get("is_urn"):
urn = ref.get("urn", {})
# Use model name from URN as display name
model_name = (
f"urn:{urn.get('type', 'model')}:{urn.get('model_id')}"
)
category = urn.get("type", category)
if category in URN_TYPE_MAP:
category = URN_TYPE_MAP[category]
loaded_models.append(
{
"name": model_name,
"category": category,
"node_id": node_id,
"widget_index": widget_index,
"node_type": node_type,
"exists": exists,
"strength": strength,
"original_path": original_path,
"is_urn": ref.get("is_urn", False),
"is_lora_v2": ref.get("is_lora_v2", False),
"active": ref.get("active"),
"connected": ref.get("connected", True),
}
)
# Also check node.properties.models for embedded models
for node in 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:
# Check if this model is already in loaded_models
existing = next(
(
m
for m in loaded_models
if m.get("original_path") == name
),
None,
)
if not existing:
loaded_models.append(
{
"name": name.split("/")[-1].split("\\")[
-1
],
"category": directory or "checkpoints",
"node_id": node.get("id"),
"widget_index": None,
"node_type": node_type,
"exists": True, # Embedded models are loaded
"strength": None,
"original_path": name,
"is_urn": False,
}
)
return web.json_response(
{"loaded_models": loaded_models, "total": len(loaded_models)}
)
except Exception as e:
self.logger.error(
f"Model Linker get_loaded_models error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
# ==================== DOWNLOAD ROUTES ====================
if download_available:
@routes.post("/model_linker/search")
async def search_sources(request):
"""Search for model download sources."""
try:
data = await request.json()
filename = data.get("filename", "")
category = data.get("category", "")
# Handle both boolean and string forms
is_urn_raw = data.get("is_urn", False)
is_urn = (
is_urn_raw
if isinstance(is_urn_raw, bool)
else (str(is_urn_raw).lower() == "true")
)
if not filename:
return web.json_response(
{"error": "Filename is required"}, status=400
)
# Debug logging
self.logger.info(
f"Search request: filename={filename}, category={category}, is_urn={is_urn}, model_id={data.get('model_id')}, version_id={data.get('version_id')}"
)
results = {
"popular": None,
"model_list": None,
"huggingface": None,
"civitai": None,
"found": False,
}
# 1. Check popular models first (curated database)
popular_info = get_popular_model_url(filename)
if popular_info:
results["popular"] = {
"source": "popular",
"filename": filename,
**popular_info,
}
results["found"] = True
# 2. Search model-list.json (ComfyUI Manager database with fuzzy matching)
# Only accept if confidence >= 70% (or if not URN, allow any)
if not results["found"]:
model_list_result = search_model_list(filename)
if model_list_result:
confidence = model_list_result.get("confidence", 0)
# For URNs, only accept high-confidence matches (>70%)
# Otherwise continue to CivitAI search for exact model
if is_urn and confidence >= 70:
results["model_list"] = model_list_result
results["found"] = True
elif not is_urn:
results["model_list"] = model_list_result
results["found"] = True
# 3. Search HuggingFace for exact file match
if not results["found"]:
hf_result = search_huggingface_for_file(filename)
if hf_result:
results["huggingface"] = hf_result
results["found"] = True
# 4. Search CivitAI - use direct download for URNs
if not results["found"]:
# For URNs, use direct model_id/version_id to get download URL
if is_urn:
# Get model_id and version_id from request data
model_id = data.get("model_id")
version_id = data.get("version_id")
if model_id and version_id:
# Use direct download URL - no search needed
self.logger.info(
f"URN: Using direct download for model_id={model_id}, version_id={version_id}"
)
download_url = get_civitai_download_url(version_id)
results["civitai"] = {
"source": "civitai",
"name": filename,
"filename": filename,
"type": category,
"download_url": download_url,
"url": f"https://civitai.com/models/{model_id}",
}
results["found"] = True
elif category:
# Fallback to search if no IDs
self.logger.info(
f"URN: No model_id/version_id, falling back to CivitAI search"
)
civitai_results = search_civitai(
filename, model_type=category
)
if civitai_results:
first_result = civitai_results[0]
results["civitai"] = {
"source": "civitai",
"name": first_result.get("name"),
"filename": first_result.get("filename"),
"type": first_result.get("type"),
"download_url": first_result.get(
"download_url"
),
"url": first_result.get("url"),
"size": first_result.get("size"),
}
results["found"] = True
else:
civitai_result = search_civitai_for_file(filename)
if civitai_result:
results["civitai"] = civitai_result
results["found"] = True
return web.json_response(results)
except Exception as e:
self.logger.error(
f"Model Linker search error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
@routes.post("/model_linker/download")
async def download_model(request):
"""Start downloading a model."""
try:
data = await request.json()
url = data.get("url", "")
filename = data.get("filename", "")
category = data.get("category", "checkpoints")
subfolder = data.get("subfolder", "")
if not url:
return web.json_response(
{"error": "URL is required"}, status=400
)
if not filename:
# Extract filename from URL
from urllib.parse import urlparse, unquote
parsed = urlparse(url)
filename = unquote(parsed.path.split("/")[-1])
if not filename:
return web.json_response(
{"error": "Could not determine filename"}, status=400
)
# Build headers if needed
headers = {}
if "huggingface.co" in url:
hf_token = data.get("hf_token", "")
if hf_token:
headers["Authorization"] = f"Bearer {hf_token}"
elif "civitai.com" in url:
civitai_key = data.get("civitai_key", "")
if civitai_key and "token=" not in url:
url += (
f"{'&' if '?' in url else '?'}token={civitai_key}"
)
# Start background download
download_id = start_background_download(
url=url,
filename=filename,
category=category,
headers=headers if headers else None,
subfolder=subfolder,
)
return web.json_response(
{
"success": True,
"download_id": download_id,
"filename": filename,
"category": category,
}
)
except Exception as e:
self.logger.error(
f"Model Linker download error: {e}", exc_info=True
)
return web.json_response(
{"error": str(e), "success": False}, status=500
)
@routes.get("/model_linker/progress/{download_id}")
async def get_download_progress(request):
"""Get progress for a specific download."""
try:
download_id = request.match_info["download_id"]
progress = get_progress(download_id)
if progress:
return web.json_response(progress)
else:
return web.json_response(
{"error": "Download not found"}, status=404
)
except Exception as e:
self.logger.error(
f"Model Linker progress error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
@routes.get("/model_linker/progress")
async def get_all_downloads_progress(request):
"""Get progress for all downloads."""
try:
progress = get_all_progress()
return web.json_response(progress)
except Exception as e:
self.logger.error(
f"Model Linker progress error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
@routes.post("/model_linker/cancel/{download_id}")
async def cancel_download_route(request):
"""Cancel a download in progress."""
try:
download_id = request.match_info["download_id"]
cancel_download(download_id)
return web.json_response({"success": True})
except Exception as e:
self.logger.error(
f"Model Linker cancel error: {e}", exc_info=True
)
return web.json_response(
{"error": str(e), "success": False}, status=500
)
@routes.get("/model_linker/directories")
async def get_directories(request):
"""Get available model directories."""
try:
categories = [
"checkpoints",
"loras",
"vae",
"controlnet",
"clip",
"clip_vision",
"embeddings",
"upscale_models",
"diffusion_models",
"text_encoders",
"ipadapter",
"sams",
]
directories = {}
for cat in categories:
path = get_download_directory(cat)
if path:
directories[cat] = path
return web.json_response(directories)
except Exception as e:
self.logger.error(
f"Model Linker directories error: {e}", exc_info=True
)
return web.json_response({"error": str(e)}, status=500)
self.routes_setup = True
self.logger.info("Model Linker: API routes registered successfully")
return True
except ImportError as e:
self.logger.warning(
f"Model Linker: Could not register routes (missing dependency): {e}"
)
return False
except Exception as e:
self.logger.error(
f"Model Linker: Error setting up routes: {e}", exc_info=True
)
return False
# Initialize the extension
try:
extension = ModelLinkerExtension()
extension.initialize()
except Exception as e:
logging.error(
f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True
)