Dariusz L
040fa64584
Enhance model linking with CivitAI integration and UI improvements. - __init__.py: Expand model path resolution by building filename→path lookup (including no-ext keys) and using folder_paths to map common categories; add a /model_linker/civitai-search POST route that attempts metadata/hash-based lookup (or filename fallback) and returns CivitAI model/version info for the frontend. - core/sources/civitai.py: Add hashing and metadata handling utilities: _get_sha256_hash, metadata discovery/reading, helpers to extract images/trained-words, get_model_info_by_hash and get_model_info_for_file (which prefers local metadata and otherwise queries CivitAI by hash). Also tidy/normalize existing search/resolve functions and add caching for hash lookups. - web/linker.js: Add a right-click context menu for model chips with actions (Show Info, Open in CivitAI), implement dialog UI to display model info, and wire calls to the new civitai-search endpoint. Includes positioning/closing behavior and dialog population from returned data. These changes enable accurate model identification via local metadata or sha256 lookups on CivitAI and provide convenient frontend actions to view model details or open CivitAI pages.
940 lines
45 KiB
Python
940 lines
45 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}
|
|
# Create lookup for full paths by filename (with and without extension)
|
|
path_by_filename = {}
|
|
for m in available_models:
|
|
rel_path = m.get("relative_path", "")
|
|
if rel_path:
|
|
filename = rel_path.split("/")[-1].split("\\")[-1]
|
|
path_by_filename[filename] = m.get("path")
|
|
# Also add without extension for matching (simple approach)
|
|
if "." in filename:
|
|
filename_no_ext = filename.rsplit(".", 1)[0]
|
|
if filename_no_ext not in path_by_filename:
|
|
path_by_filename[filename_no_ext] = m.get("path")
|
|
# Add the full relative path as key too
|
|
path_by_filename[rel_path] = m.get("path")
|
|
|
|
# Also use folder_paths.get_full_path() to get paths
|
|
import folder_paths
|
|
|
|
for cat in [
|
|
"loras",
|
|
"checkpoints",
|
|
"vae",
|
|
"controlnet",
|
|
"upscale_models",
|
|
]:
|
|
try:
|
|
filenames = folder_paths.get_filename_list(cat)
|
|
for fn in filenames:
|
|
full_path = folder_paths.get_full_path(cat, fn)
|
|
if (
|
|
full_path
|
|
and full_path not in path_by_filename.values()
|
|
):
|
|
path_by_filename[fn] = full_path
|
|
fn_no_ext = (
|
|
fn.rsplit(".", 1)[0] if "." in fn else fn
|
|
)
|
|
if fn_no_ext not in path_by_filename:
|
|
path_by_filename[fn_no_ext] = full_path
|
|
except Exception:
|
|
pass
|
|
|
|
# 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),
|
|
"resolved_path": (
|
|
path_by_filename.get(model_name)
|
|
or path_by_filename.get(original_path)
|
|
),
|
|
}
|
|
)
|
|
|
|
# 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)
|
|
|
|
# ==================== CIVITAI SEARCH ROUTE ====================
|
|
|
|
@routes.post("/model_linker/civitai-search")
|
|
async def civitai_search(request):
|
|
"""Search CivitAI for a model using file hash."""
|
|
try:
|
|
data = await request.json()
|
|
filename = data.get("filename", "")
|
|
category = data.get("category", "")
|
|
resolved_path = data.get("resolved_path", "")
|
|
|
|
if not filename:
|
|
return web.json_response(
|
|
{"error": "Filename is required"}, status=400
|
|
)
|
|
|
|
# Clean filename for display
|
|
import os as _os
|
|
|
|
clean_name = _os.path.splitext(filename)[0]
|
|
|
|
# Get the file path to hash
|
|
file_path = resolved_path if resolved_path else None
|
|
|
|
if not file_path and category:
|
|
# Try to find the file in the model directories using folder_paths
|
|
try:
|
|
import folder_paths
|
|
|
|
# Map category to folder_paths type
|
|
category_map = {
|
|
"loras": "loras",
|
|
"checkpoints": "checkpoints",
|
|
"vae": "vae",
|
|
"controlnet": "controlnet",
|
|
"upscale_models": "upscale_models",
|
|
}
|
|
folder_type = category_map.get(
|
|
category.lower(), category.lower()
|
|
)
|
|
file_path = folder_paths.get_full_path(
|
|
folder_type, filename
|
|
)
|
|
except Exception:
|
|
pass
|
|
|
|
# If not found, try scanner
|
|
if not file_path:
|
|
try:
|
|
from .core.scanner import get_model_files
|
|
|
|
available_models = get_model_files()
|
|
for m in available_models:
|
|
if (
|
|
m.get("relative_path", "").endswith(filename)
|
|
or m.get("filename", "") == filename
|
|
):
|
|
file_path = m.get("path")
|
|
break
|
|
except Exception:
|
|
pass
|
|
|
|
# Search CivitAI for the model using hash
|
|
if download_available and file_path and _os.path.exists(file_path):
|
|
try:
|
|
from .core.sources.civitai import (
|
|
get_model_info_for_file,
|
|
)
|
|
|
|
result = get_model_info_for_file(file_path)
|
|
if result and result.get("url"):
|
|
return web.json_response(
|
|
{
|
|
"filename": filename,
|
|
"url": result.get("url"),
|
|
"version_url": result.get("version_url"),
|
|
"model_id": result.get("model_id"),
|
|
"model_name": result.get(
|
|
"model_name", clean_name
|
|
),
|
|
"model_type": result.get("model_type", ""),
|
|
"version_id": result.get("version_id"),
|
|
"version_name": result.get("version_name", ""),
|
|
"sha256": result.get("sha256"),
|
|
"base_model": result.get("base_model"),
|
|
"tags": result.get("tags", []),
|
|
"trained_words": result.get(
|
|
"trained_words", []
|
|
),
|
|
"images": result.get("images", []),
|
|
"clip_skip": result.get("clip_skip"),
|
|
"description": result.get("description", ""),
|
|
"model_description": result.get(
|
|
"model_description", ""
|
|
),
|
|
}
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"CivitAI search error: {e}")
|
|
|
|
# No result found - try fallback to filename search
|
|
if download_available:
|
|
try:
|
|
from .core.sources.civitai import (
|
|
search_civitai_for_file,
|
|
)
|
|
|
|
result = search_civitai_for_file(filename)
|
|
if result and result.get("url"):
|
|
return web.json_response(
|
|
{
|
|
"url": result["url"],
|
|
"model_name": result.get("name", clean_name),
|
|
"version_id": result.get("version_id"),
|
|
}
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"CivitAI fallback search error: {e}")
|
|
|
|
# No result found
|
|
return web.json_response({"url": None})
|
|
|
|
except Exception as e:
|
|
self.logger.error(
|
|
f"Model Linker civitai-search 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
|
|
)
|