nodes/Azornes-Comfyui-Model-Resolver
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enhance Model Linker routes, downloader, and linker

Browse Source 
Refactor and harden the Model Linker extension: normalize string quoting/formatting and improve logging and error responses across API routes. Make route registration more robust to missing PromptServer and dependencies. Add/expand download-related features: safer request parsing (including is_urn handling), CivitAI URN support & direct download URL resolution, popular/model-list/HuggingFace/CivitAI search fallbacks, improved header handling for HF/CivitAI, and better category/directory mapping. Improve downloader reliability and UX: folder_paths fallback, pre-initialize download progress entries, smoother speed calculation and CLI logging, cancellation cleanup on Windows, stronger exception handling, and consistent JSON responses from endpoints.
This commit is contained in:
Dariusz L
2026-04-17 21:23:25 +02:00
parent e8939d4bf1
commit 3ab6a9cc55
4 changed files with 773 additions and 500 deletions
Download Patch File Download Diff File Expand all files Collapse all files
__init__.py
+350 -225
View File
@@ -20,39 +20,45 @@ __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)
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:
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
@@ -60,179 +66,221 @@ class ModelLinkerExtension:
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
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.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
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}")
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')
workflow_json = data.get("workflow")
if not workflow_json:
return web.json_response(
{'error': 'Workflow JSON is required'},
status=400
{"error": "Workflow JSON is required"}, status=400
)
# Analyze and find matches
result = analyze_and_find_matches(workflow_json)
# If download available, auto-search for download sources when no 100% local match
if download_available:
for missing in result.get('missing_models', []):
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)
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]
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', '')
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'
if "huggingface.co" in workflow_url:
source = "huggingface"
elif "civitai.com" in workflow_url:
source = "civitai"
else:
source = 'workflow'
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)
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))
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
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'
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)
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')
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)
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')
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)
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')
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
)
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', [])
workflow_json = data.get("workflow")
resolutions = data.get("resolutions", [])
if not workflow_json:
return web.json_response(
{'error': 'Workflow JSON is required'},
status=400
{"error": "Workflow JSON is required"}, status=400
)
if not resolutions:
return web.json_response(
{'error': 'Resolutions array is required'},
status=400
{"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
})
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
{"error": str(e), "success": False}, status=500
)
@routes.get("/model_linker/models")
async def get_models(request):
"""Get list of all available models."""
@@ -240,161 +288,224 @@ class ModelLinkerExtension:
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
self.logger.error(
f"Model Linker get_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', '')
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
{"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
"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["popular"] = {
"source": "popular",
"filename": filename,
**popular_info,
}
results['found'] = True
results["found"] = True
# 2. Search model-list.json (ComfyUI Manager database with fuzzy matching)
if not results['found']:
# 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:
results['model_list'] = model_list_result
results['found'] = True
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']:
if not results["found"]:
hf_result = search_huggingface_for_file(filename)
if hf_result:
results['huggingface'] = hf_result
results['found'] = True
# 4. Search CivitAI for exact file match
if not results['found']:
civitai_result = search_civitai_for_file(filename)
if civitai_result:
results['civitai'] = civitai_result
results['found'] = True
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
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', '')
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
{"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])
filename = unquote(parsed.path.split("/")[-1])
if not filename:
return web.json_response(
{'error': 'Could not determine filename'},
status=400
{"error": "Could not determine filename"}, status=400
)
# Build headers if needed
headers = {}
if 'huggingface.co' in url:
hf_token = data.get('hf_token', '')
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}"
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
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
{
"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']
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
{"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
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."""
@@ -402,59 +513,71 @@ class ModelLinkerExtension:
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
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']
download_id = request.match_info["download_id"]
cancel_download(download_id)
return web.json_response({'success': True})
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
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'
"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.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}")
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)
self.logger.error(
f"Model Linker: Error setting up routes: {e}", exc_info=True
)
return False
@@ -463,4 +586,6 @@ try:
extension = ModelLinkerExtension()
extension.initialize()
except Exception as e:
logging.error(f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True)
logging.error(
f"ComfyUI Model Linker extension initialization failed: {e}", exc_info=True
)
core/downloader.py
+229 -138
View File
@@ -36,7 +36,7 @@ def format_bytes(bytes_value: int) -> str:
if bytes_value == 0:
return "0 B"
k = 1024
sizes = ['B', 'KB', 'MB', 'GB', 'TB']
sizes = ["B", "KB", "MB", "GB", "TB"]
i = 0
while bytes_value >= k and i < len(sizes) - 1:
bytes_value /= k
@@ -47,54 +47,73 @@ def format_bytes(bytes_value: int) -> str:
def get_download_directory(category: str) -> Optional[str]:
"""
Get the appropriate download directory for a model category.
Args:
category: Model category (e.g., 'checkpoints', 'loras', 'vae')
Returns:
Absolute path to the download directory, or None if not found
"""
global folder_paths
if folder_paths is None:
return None
# Try to import again - ComfyUI might have initialized since last check
try:
import folder_paths as fp
folder_paths = fp
except ImportError:
return None
# Map common category names to folder_paths keys
category_map = {
'checkpoint': 'checkpoints',
'checkpoints': 'checkpoints',
'lora': 'loras',
'loras': 'loras',
'vae': 'vae',
'controlnet': 'controlnet',
'clip': 'clip',
'clip_vision': 'clip_vision',
'upscaler': 'upscale_models',
'upscale_models': 'upscale_models',
'embeddings': 'embeddings',
'embedding': 'embeddings',
'diffusion_models': 'diffusion_models',
'unet': 'diffusion_models',
'text_encoders': 'text_encoders',
'text_encoder': 'text_encoders',
'ipadapter': 'ipadapter',
'ip-adapter': 'ipadapter',
"checkpoint": "checkpoints",
"checkpoints": "checkpoints",
"lora": "loras",
"loras": "loras",
"vae": "vae",
"controlnet": "controlnet",
"clip": "clip",
"clip_vision": "clip_vision",
"upscaler": "upscale_models",
"upscale_models": "upscale_models",
"embeddings": "embeddings",
"embedding": "embeddings",
"diffusion_models": "diffusion_models",
"unet": "diffusion_models",
"text_encoders": "text_encoders",
"text_encoder": "text_encoders",
"ipadapter": "ipadapter",
"ip-adapter": "ipadapter",
# Legacy/unsupported categories map to checkpoints as fallback
"default": "upscale_models", # Model list uses "default" for upscalers
}
folder_key = category_map.get(category.lower(), category.lower())
try:
paths = folder_paths.get_folder_paths(folder_key)
if paths:
# Return the first path (usually the main models directory)
return paths[0]
# If category not found, try to get any models directory as fallback
all_names = folder_paths.get_folder_names()
if all_names:
# Fall back to first available directory
fallback_paths = folder_paths.get_folder_paths(all_names[0])
if fallback_paths:
return fallback_paths[0]
except Exception as e:
logger.debug(f"Could not get folder path for {folder_key}: {e}")
return None
def generate_download_id() -> str:
"""Generate a unique download ID."""
import uuid
return str(uuid.uuid4())[:8]
@@ -104,11 +123,11 @@ def download_file(
download_id: str,
headers: Optional[Dict[str, str]] = None,
chunk_size: int = None,
progress_callback: Optional[Callable[[int, int], None]] = None
progress_callback: Optional[Callable[[int, int], None]] = None,
) -> Dict[str, Any]:
"""
Download a file from URL with progress tracking and speed calculation.
Args:
url: URL to download from
dest_path: Destination file path
@@ -116,213 +135,244 @@ def download_file(
headers: Optional HTTP headers (for auth tokens)
chunk_size: Download chunk size in bytes (defaults to 1MB)
progress_callback: Optional callback(downloaded_bytes, total_bytes)
Returns:
Result dictionary with status and info
"""
global download_progress, cancelled_downloads
# Use default 1MB chunk size if not specified
if chunk_size is None:
chunk_size = CHUNK_SIZE
result = {
'success': False,
'download_id': download_id,
'path': dest_path,
'error': None,
'size': 0
"success": False,
"download_id": download_id,
"path": dest_path,
"error": None,
"size": 0,
}
# Initialize progress tracking with speed calculation
start_time = time.time()
speed_history: deque = deque(maxlen=SPEED_HISTORY_SIZE)
last_speed_update = start_time
last_downloaded = 0
last_cli_log = start_time # Track when we last logged to CLI
with download_lock:
download_progress[download_id] = {
'status': 'starting',
'progress': 0,
'total_size': 0,
'downloaded': 0,
'filename': os.path.basename(dest_path),
'url': url,
'error': None,
'speed': 0, # bytes per second
'start_time': start_time
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": os.path.basename(dest_path),
"url": url,
"error": None,
"speed": 0, # bytes per second
"start_time": start_time,
}
try:
# Ensure destination directory exists
os.makedirs(os.path.dirname(dest_path), exist_ok=True)
# Verbose logging - what model and from where
filename = os.path.basename(dest_path)
source = "HuggingFace" if "huggingface.co" in url else "CivitAI" if "civitai.com" in url else "URL"
source = (
"HuggingFace"
if "huggingface.co" in url
else "CivitAI"
if "civitai.com" in url
else "URL"
)
print(f"\n[Model Linker] Starting download: {filename}")
print(f"[Model Linker] Source: {source}")
print(f"[Model Linker] URL: {url}")
# Start download
response = requests.get(
url,
headers=headers,
stream=True,
timeout=30
)
response = requests.get(url, headers=headers, stream=True, timeout=30)
response.raise_for_status()
# Get total size
total_size = int(response.headers.get('content-length', 0))
total_size = int(response.headers.get("content-length", 0))
total_size_str = format_bytes(total_size) if total_size > 0 else "unknown"
print(f"[Model Linker] Size: {total_size_str}")
with download_lock:
download_progress[download_id]['total_size'] = total_size
download_progress[download_id]['status'] = 'downloading'
download_progress[download_id]["total_size"] = total_size
download_progress[download_id]["status"] = "downloading"
downloaded = 0
# Download with progress and speed calculation
cancelled = False
with open(dest_path, 'wb') as f:
with open(dest_path, "wb") as f:
for chunk in response.iter_content(chunk_size=chunk_size):
# Check for cancellation
if download_id in cancelled_downloads:
cancelled = True
break
if chunk:
f.write(chunk)
downloaded += len(chunk)
# Calculate speed with smoothing
current_time = time.time()
time_delta = current_time - last_speed_update
# Update speed every 0.5 seconds to avoid too frequent calculations
if time_delta >= 0.5:
bytes_delta = downloaded - last_downloaded
instant_speed = bytes_delta / time_delta if time_delta > 0 else 0
instant_speed = (
bytes_delta / time_delta if time_delta > 0 else 0
)
speed_history.append(instant_speed)
# Calculate smoothed speed (average of recent samples)
smoothed_speed = sum(speed_history) / len(speed_history) if speed_history else 0
smoothed_speed = (
sum(speed_history) / len(speed_history)
if speed_history
else 0
)
last_speed_update = current_time
last_downloaded = downloaded
# Update progress with speed
with download_lock:
download_progress[download_id]['downloaded'] = downloaded
download_progress[download_id]['speed'] = int(smoothed_speed)
download_progress[download_id]["downloaded"] = downloaded
download_progress[download_id]["speed"] = int(
smoothed_speed
)
if total_size > 0:
download_progress[download_id]['progress'] = int((downloaded / total_size) * 100)
download_progress[download_id]["progress"] = int(
(downloaded / total_size) * 100
)
# CLI progress logging (every CLI_LOG_INTERVAL seconds)
if current_time - last_cli_log >= CLI_LOG_INTERVAL:
last_cli_log = current_time
progress_pct = int((downloaded / total_size) * 100) if total_size > 0 else 0
progress_pct = (
int((downloaded / total_size) * 100)
if total_size > 0
else 0
)
downloaded_str = format_bytes(downloaded)
total_str = format_bytes(total_size) if total_size > 0 else "?"
total_str = (
format_bytes(total_size) if total_size > 0 else "?"
)
speed_str = format_bytes(int(smoothed_speed)) + "/s"
print(f"[Model Linker] Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}")
print(
f"[Model Linker] Progress: {downloaded_str} / {total_str} ({progress_pct}%) - {speed_str}"
)
else:
# Just update downloaded bytes without recalculating speed
with download_lock:
download_progress[download_id]['downloaded'] = downloaded
download_progress[download_id]["downloaded"] = downloaded
if total_size > 0:
download_progress[download_id]['progress'] = int((downloaded / total_size) * 100)
download_progress[download_id]["progress"] = int(
(downloaded / total_size) * 100
)
if progress_callback:
progress_callback(downloaded, total_size)
# Handle cancellation after file is closed (so we can delete it on Windows)
# Also check if cancellation was requested while we were finishing up
if cancelled or download_id in cancelled_downloads:
with download_lock:
download_progress[download_id]['status'] = 'cancelled'
download_progress[download_id]["status"] = "cancelled"
# Clean up partial/incomplete file
try:
if os.path.exists(dest_path):
os.remove(dest_path)
print(f"[Model Linker] Cancelled: {filename} - incomplete file deleted")
print(
f"[Model Linker] Cancelled: {filename} - incomplete file deleted"
)
else:
print(f"[Model Linker] Cancelled: {filename} - no file to delete")
except Exception as e:
print(f"[Model Linker] Warning: Could not delete incomplete file {dest_path}: {e}")
print(
f"[Model Linker] Warning: Could not delete incomplete file {dest_path}: {e}"
)
# Try harder on Windows - sometimes the file handle takes a moment to release
try:
time.sleep(0.5) # time is already imported at module level
if os.path.exists(dest_path):
os.remove(dest_path)
print(f"[Model Linker] Cancelled: {filename} - incomplete file deleted (delayed)")
print(
f"[Model Linker] Cancelled: {filename} - incomplete file deleted (delayed)"
)
except Exception:
pass
result['error'] = 'Download cancelled'
result["error"] = "Download cancelled"
cancelled_downloads.discard(download_id)
return result
# Success
with download_lock:
download_progress[download_id]['status'] = 'completed'
download_progress[download_id]['progress'] = 100
download_progress[download_id]['speed'] = 0 # Reset speed on completion
result['success'] = True
result['size'] = downloaded
download_progress[download_id]["status"] = "completed"
download_progress[download_id]["progress"] = 100
download_progress[download_id]["speed"] = 0 # Reset speed on completion
result["success"] = True
result["size"] = downloaded
# CLI completion log
elapsed = time.time() - start_time
avg_speed = downloaded / elapsed if elapsed > 0 else 0
print(f"[Model Linker] ✓ Download complete: {filename}")
print(f"[Model Linker] Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s")
print(
f"[Model Linker] Size: {format_bytes(downloaded)}, Time: {elapsed:.1f}s, Avg speed: {format_bytes(int(avg_speed))}/s"
)
except requests.exceptions.RequestException as e:
error_msg = str(e)
# Check for specific HTTP errors
if hasattr(e, 'response') and e.response is not None:
if hasattr(e, "response") and e.response is not None:
status_code = e.response.status_code
if status_code in [401, 403]:
if 'huggingface.co' in url:
if "huggingface.co" in url:
error_msg = f"Unauthorized (HTTP {status_code}): HuggingFace token may be required."
elif 'civitai.com' in url:
elif "civitai.com" in url:
error_msg = f"Unauthorized (HTTP {status_code}): CivitAI API key may be required."
else:
error_msg = f"Unauthorized (HTTP {status_code}): Authentication required."
error_msg = (
f"Unauthorized (HTTP {status_code}): Authentication required."
)
elif status_code == 404:
error_msg = "Model not found (HTTP 404): The file may have been moved or deleted."
with download_lock:
download_progress[download_id]['status'] = 'error'
download_progress[download_id]['error'] = error_msg
result['error'] = error_msg
download_progress[download_id]["status"] = "error"
download_progress[download_id]["error"] = error_msg
result["error"] = error_msg
# CLI error log
print(f"[Model Linker] ✗ Download failed: {os.path.basename(dest_path)}")
print(f"[Model Linker] Error: {error_msg}")
# Clean up partial file
try:
if os.path.exists(dest_path):
os.remove(dest_path)
except:
pass
except Exception as e:
error_msg = str(e)
with download_lock:
download_progress[download_id]['status'] = 'error'
download_progress[download_id]['error'] = error_msg
result['error'] = error_msg
download_progress[download_id]["status"] = "error"
download_progress[download_id]["error"] = error_msg
result["error"] = error_msg
# CLI error log
print(f"[Model Linker] ✗ Download failed: {os.path.basename(dest_path)}")
print(f"[Model Linker] Error: {error_msg}")
logger.error(f"Download error: {e}", exc_info=True)
return result
@@ -332,11 +382,11 @@ def download_model(
category: str,
download_id: Optional[str] = None,
headers: Optional[Dict[str, str]] = None,
subfolder: str = ""
subfolder: str = "",
) -> Dict[str, Any]:
"""
Download a model to the appropriate directory.
Args:
url: URL to download from
filename: Filename to save as
@@ -344,37 +394,37 @@ def download_model(
download_id: Optional download ID (generated if not provided)
headers: Optional HTTP headers
subfolder: Optional subfolder within category directory
Returns:
Result dictionary
"""
if download_id is None:
download_id = generate_download_id()
# Get destination directory
dest_dir = get_download_directory(category)
if not dest_dir:
return {
'success': False,
'download_id': download_id,
'error': f'Could not find directory for category: {category}'
"success": False,
"download_id": download_id,
"error": f"Could not find directory for category: {category}",
}
# Add subfolder if specified
if subfolder:
dest_dir = os.path.join(dest_dir, subfolder)
dest_path = os.path.join(dest_dir, filename)
# Check if file already exists
if os.path.exists(dest_path):
return {
'success': False,
'download_id': download_id,
'error': f'File already exists: {dest_path}',
'path': dest_path
"success": False,
"download_id": download_id,
"error": f"File already exists: {dest_path}",
"path": dest_path,
}
return download_file(url, dest_path, download_id, headers)
@@ -400,8 +450,9 @@ def clear_completed_downloads():
"""Clear completed/failed downloads from progress tracking."""
with download_lock:
to_remove = [
did for did, info in download_progress.items()
if info.get('status') in ('completed', 'error', 'cancelled')
did
for did, info in download_progress.items()
if info.get("status") in ("completed", "error", "cancelled")
]
for did in to_remove:
del download_progress[did]
@@ -413,20 +464,60 @@ def start_background_download(
filename: str,
category: str,
headers: Optional[Dict[str, str]] = None,
subfolder: str = ""
subfolder: str = "",
) -> str:
"""
Start a download in a background thread.
Returns:
download_id for tracking progress
"""
download_id = generate_download_id()
# Pre-initialize progress dict so it's always available for polling
# even if download fails before download_file is called
with download_lock:
download_progress[download_id] = {
"status": "starting",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": filename,
"url": url,
"error": None,
"speed": 0,
"start_time": time.time(),
}
def run_download():
download_model(url, filename, category, download_id, headers, subfolder)
try:
result = download_model(
url, filename, category, download_id, headers, subfolder
)
if not result.get("success"):
# Mark as error if download failed
with download_lock:
if download_id in download_progress:
download_progress[download_id]["status"] = "error"
download_progress[download_id]["error"] = result.get(
"error", "Download failed"
)
except Exception as e:
# Ensure any exception is captured and logged
with download_lock:
download_progress[download_id] = {
"status": "error",
"progress": 0,
"total_size": 0,
"downloaded": 0,
"filename": filename,
"url": url,
"error": str(e),
"speed": 0,
"start_time": time.time(),
}
thread = threading.Thread(target=run_download, daemon=True)
thread.start()
return download_id
core/linker.py
+142 -130
View File
@@ -23,74 +23,76 @@ logger = logging.getLogger(__name__)
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')
MODEL_EXTENSIONS = (".safetensors", ".ckpt", ".pt", ".pth", ".bin", ".onnx")
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', [])
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', [])
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', [])
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', '')
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'
"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()
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))', re.IGNORECASE)
model_pattern = re.compile(
r"([\w\-\.%]+\.(?:safetensors|ckpt|pt|pth|bin|onnx))", 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():
@@ -100,79 +102,81 @@ def extract_workflow_urls(workflow_json: Dict[str, Any]) -> Dict[str, Dict[str,
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'):
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'
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'
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'
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:
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)
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
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:
{
@@ -204,109 +208,114 @@ def analyze_and_find_matches(
# 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', '')
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', '')
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 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')
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']}")
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:
target_for_matching = missing.get('original_path', '')
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']
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':
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')
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])
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
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', '')
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']:
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
})
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)
"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]]
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:
@@ -317,7 +326,7 @@ def apply_resolution(
'category': model category (optional),
'resolved_model': model dict from scanner (optional)
}
Returns:
Updated workflow JSON dictionary
"""
@@ -325,42 +334,45 @@ def apply_resolution(
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
"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
}
# 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']
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)
web/linker.js
+52 -7
View File
@@ -1861,8 +1861,36 @@ class LinkerManagerDialog extends ComfyDialog {
* Search online for a model
*/
async searchOnline(missing) {
const filename = missing.original_path?.split('/').pop()?.split('\\').pop() || '';
const category = missing.category || '';
let filename = missing.original_path?.split('/').pop()?.split('\\').pop() || '';
let category = missing.category || '';
// For URNs, use the CivitAI model name for searching instead of the URN itself
// and pass the URN type as category (CivitAI expects specific type names)
if (missing.is_urn) {
if (missing.civitai_info?.model_name) {
filename = missing.civitai_info.model_name;
}
// Pass URN type directly - CivitAPI expects types like 'Upscaler', 'Checkpoint'
const urnType = missing.urn_type || '';
if (urnType) {
// Map URN types to CivitAI type names
const typeMap = {
'checkpoint': 'Checkpoint',
'lora': 'LORA',
'vae': 'VAE',
'upscaler': 'Upscaler',
'upscale_model': 'Upscaler',
'embedding': 'TextualInversion',
'controlnet': 'Controlnet'
};
const civitaiType = typeMap[urnType.toLowerCase()];
if (civitaiType) {
category = civitaiType;
}
}
}
const isUrn = missing.is_urn || false;
const resultsId = `search-results-${missing.node_id}-${missing.widget_index}`;
const resultsDiv = this.contentElement?.querySelector(`#${resultsId}`);
const searchBtn = this.contentElement?.querySelector(`#search-${missing.node_id}-${missing.widget_index}`);
@@ -1877,10 +1905,19 @@ class LinkerManagerDialog extends ComfyDialog {
resultsDiv.innerHTML = '<span style="color: #2196F3;">Searching HuggingFace and CivitAI...</span>';
}
// For URNs, include model_id and version_id for direct download
const searchData = { filename, category, is_urn: isUrn };
if (isUrn && missing.urn) {
searchData.model_id = missing.urn.model_id;
searchData.version_id = missing.urn.version_id;
}
console.log('Model Linker: Search request:', JSON.stringify(searchData));
const response = await api.fetchApi('/model_linker/search', {
method: 'POST',
headers: { 'Content-Type': 'application/json' },
body: JSON.stringify({ filename, category })
body: JSON.stringify(searchData)
});
if (!response.ok) {
@@ -1888,6 +1925,7 @@ class LinkerManagerDialog extends ComfyDialog {
}
const data = await response.json();
console.log('Model Linker: Search response:', JSON.stringify(data));
this.displaySearchResults(missing, data, resultsDiv);
} catch (error) {
@@ -1911,8 +1949,9 @@ class LinkerManagerDialog extends ComfyDialog {
const popular = data.popular;
const modelListResult = data.model_list;
const hfResult = data.huggingface?.[0];
const civitaiResult = data.civitai?.[0];
// Backend sends these as objects, not arrays - handle both
const hfResult = data.huggingface ? (Array.isArray(data.huggingface) ? data.huggingface[0] : data.huggingface) : null;
const civitaiResult = data.civitai ? (Array.isArray(data.civitai) ? data.civitai[0] : data.civitai) : null;
const hasResults = popular || modelListResult || hfResult || civitaiResult;
if (!hasResults) {
@@ -1975,14 +2014,20 @@ class LinkerManagerDialog extends ComfyDialog {
// CivitAI result
if (civitaiResult && civitaiResult.download_url) {
const modelUrl = civitaiResult.url || `https://civitai.com/models/${civitaiResult.model_id || ''}`;
// Use expected_filename from URN resolution, or fallback to civitaiResult filename
const downloadFilename = missing.civitai_info?.expected_filename || civitaiResult.filename || civitaiResult.name;
html += `<div class="ml-status ml-status-warning" style="flex-direction: column; align-items: flex-start;">`;
html += `<strong>Found on CivitAI</strong>`;
html += `<div style="margin-top: 4px; font-size: 12px;">`;
html += `<span class="ml-chip">${civitaiResult.filename || civitaiResult.name}</span> `;
html += `<span style="color: var(--ml-text-muted);">${civitaiResult.type || civitaiResult.name}</span>`;
html += `</div>`;
html += `<div style="margin-top: 8px;">`;
html += `<button class="search-download-btn ml-btn ml-btn-sm" style="background: #FF9800;" data-url="${civitaiResult.download_url}" data-filename="${civitaiResult.filename || civitaiResult.name + '.safetensors'}" data-category="${missing.category}">`;
html += `<div style="margin-top: 8px; display: flex; gap: 8px;">`;
// Link to CivitAI model page
html += `<a href="${modelUrl}" target="_blank" class="ml-btn ml-btn-sm" style="background: #46b5e5; text-decoration: none;">🔗 View</a>`;
// Download button - use expected_filename from URN resolution
html += `<button class="search-download-btn ml-btn ml-btn-sm" style="background: #FF9800;" data-url="${civitaiResult.download_url}" data-filename="${downloadFilename}" data-category="${missing.category}">`;
html += `<span class="ml-btn-icon">☁</span> Download`;
html += `</button>`;
html += `</div></div>`;