Dariusz L
57018b4e0f
Add searched_at timestamps to search results so cached entries are accurately tracked and preserved across merges. Expand the clear-search-cache command to also invalidate popular databases, model list, and the new result timestamp store. Frontend gains new cache-clearing methods and a full resolver cache reset for maintenance. The search panel and missing browser now show source result dates and toggle the search button between 'Search' and 'Search Again' when results already exist.
1773 lines
86 KiB
Python
1773 lines
86 KiB
Python
"""
|
||
@author: Model Resolver Team
|
||
@title: ComfyUI Model Resolver
|
||
@nickname: Model Resolver
|
||
@version: 1.1.0
|
||
@description: Extension for resolving missing models and downloading from HuggingFace/CivitAI
|
||
"""
|
||
|
||
import asyncio
|
||
from .core.log_system.log_funcs import (
|
||
create_module_logger,
|
||
log_info,
|
||
log_error,
|
||
log_exception,
|
||
)
|
||
|
||
# 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 ModelResolverExtension:
|
||
"""Main extension class for Model Resolver."""
|
||
|
||
def __init__(self):
|
||
self.routes_setup = False
|
||
self.logger = create_module_logger(__name__)
|
||
self.analysis_progress = {}
|
||
self.search_result_timestamps = {}
|
||
|
||
def initialize(self):
|
||
"""Initialize the extension and set up API routes."""
|
||
try:
|
||
self.setup_routes()
|
||
self.logger.info("Model Resolver: Extension initialized successfully")
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver: Extension initialization failed: {e}", exc_info=True
|
||
)
|
||
|
||
def setup_routes(self):
|
||
"""Register API routes for the Model Resolver 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 Resolver: PromptServer not available yet")
|
||
return False
|
||
|
||
routes = PromptServer.instance.routes
|
||
except (ImportError, AttributeError) as e:
|
||
self.logger.debug(f"Model Resolver: Could not access PromptServer: {e}")
|
||
return False
|
||
|
||
# Import resolver modules
|
||
try:
|
||
from .core.resolver import (
|
||
analyze_and_find_matches,
|
||
apply_resolution,
|
||
search_local_matches,
|
||
)
|
||
from .core.scanner import get_model_files, invalidate_model_files_cache
|
||
except ImportError as e:
|
||
self.logger.error(f"Model Resolver: 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,
|
||
reload_databases as reload_popular_databases,
|
||
)
|
||
from .core.sources.model_list import (
|
||
search_model_list,
|
||
search_model_list_multiple,
|
||
get_model_list_update_status,
|
||
update_model_list_from_remote,
|
||
reload_model_list,
|
||
)
|
||
from .core.sources.huggingface import (
|
||
search_huggingface_for_file,
|
||
get_author_fallback_index_status,
|
||
refresh_author_fallback_index,
|
||
check_huggingface_token,
|
||
check_brave_search_api_key,
|
||
clear_search_cache as clear_huggingface_search_cache,
|
||
)
|
||
from .core.sources.civitai import (
|
||
search_civitai_for_file,
|
||
search_civitai,
|
||
get_civitai_download_url,
|
||
resolve_urn,
|
||
check_civitai_api_key,
|
||
check_civitai_session_token,
|
||
clear_search_cache as clear_civitai_search_cache,
|
||
)
|
||
from .core.sources.civarchive import (
|
||
is_civarchive_available,
|
||
search_civarchive_for_file,
|
||
resolve_civarchive_model_version,
|
||
clear_search_cache as clear_civarchive_search_cache,
|
||
)
|
||
from .core.sources.lora_manager_archive import (
|
||
is_lora_manager_archive_available,
|
||
search_lora_manager_archive_for_file,
|
||
clear_search_cache as clear_lora_manager_archive_search_cache,
|
||
)
|
||
|
||
download_available = True
|
||
except ImportError as e:
|
||
self.logger.warning(
|
||
f"Model Resolver: Download features not available: {e}"
|
||
)
|
||
download_available = False
|
||
|
||
# ==================== ANALYZE ROUTES ====================
|
||
|
||
@routes.post("/model_resolver/analyze")
|
||
async def analyze_workflow(request):
|
||
"""Analyze workflow and return missing models with matches."""
|
||
try:
|
||
data = await request.json()
|
||
workflow_json = data.get("workflow")
|
||
analysis_id = str(data.get("analysis_id") or "").strip()
|
||
|
||
if not workflow_json:
|
||
return web.json_response(
|
||
{"error": "Workflow JSON is required"}, status=400
|
||
)
|
||
|
||
if analysis_id:
|
||
self.analysis_progress[analysis_id] = {
|
||
"status": "starting",
|
||
"stage": "starting",
|
||
"message": "Starting analysis...",
|
||
"current": 0,
|
||
"total": 0,
|
||
}
|
||
|
||
def update_analysis_progress(payload):
|
||
if not analysis_id:
|
||
return
|
||
self.analysis_progress[analysis_id] = {
|
||
**self.analysis_progress.get(analysis_id, {}),
|
||
**payload,
|
||
"status": "running"
|
||
if payload.get("stage") != "completed"
|
||
else "completed",
|
||
}
|
||
|
||
# Analyze and find matches
|
||
result = await asyncio.to_thread(
|
||
analyze_and_find_matches,
|
||
workflow_json,
|
||
0.0,
|
||
10,
|
||
update_analysis_progress if analysis_id else None,
|
||
)
|
||
|
||
# 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", "")
|
||
self.logger.debug(
|
||
f"Filtering: {name} is_lora_v2={is_lora} exists={exists}"
|
||
)
|
||
|
||
# Skip LoraManager lorAs that already exist locally
|
||
if is_lora and exists:
|
||
self.logger.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, check for download sources only from LOCAL sources
|
||
# (workflow_url, popular, model-list.json) - skip automatic online search
|
||
# Online search is now only triggered on-demand via search button
|
||
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,
|
||
"model_url": missing.get(
|
||
"workflow_model_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:
|
||
popular_model_list_result = search_model_list(
|
||
filename, exact_only=True
|
||
)
|
||
missing["download_source"] = {
|
||
"source": "popular",
|
||
"url": popular_info.get("url"),
|
||
"filename": filename,
|
||
"type": popular_info.get("type"),
|
||
"directory": popular_info.get("directory"),
|
||
"size": (
|
||
popular_model_list_result.get("size")
|
||
if popular_model_list_result
|
||
else None
|
||
)
|
||
or popular_info.get("size"),
|
||
"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
|
||
|
||
# NOTE: Search for online sources (HuggingFace, CivitAI) is
|
||
# now done on-demand via /model_resolver/search endpoint
|
||
# when user clicks "Search Online" button, not automatically
|
||
|
||
if analysis_id:
|
||
self.analysis_progress[analysis_id] = {
|
||
**self.analysis_progress.get(analysis_id, {}),
|
||
"status": "completed",
|
||
"stage": "completed",
|
||
"message": "Analysis complete",
|
||
"current": result.get("total_missing", 0),
|
||
"total": result.get("total_missing", 0),
|
||
}
|
||
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
if "analysis_id" in locals() and analysis_id:
|
||
self.analysis_progress[analysis_id] = {
|
||
"status": "error",
|
||
"stage": "error",
|
||
"message": str(e),
|
||
"current": 0,
|
||
"total": 0,
|
||
}
|
||
self.logger.error(f"Model Resolver analyze error: {e}", exc_info=True)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/analyze-progress/{analysis_id}")
|
||
async def get_analyze_progress(request):
|
||
"""Get workflow analysis progress."""
|
||
analysis_id = request.match_info.get("analysis_id", "").strip()
|
||
if not analysis_id:
|
||
return web.json_response(
|
||
{"error": "Analysis ID is required"}, status=400
|
||
)
|
||
|
||
progress = self.analysis_progress.get(analysis_id)
|
||
if not progress:
|
||
return web.json_response(
|
||
{
|
||
"status": "unknown",
|
||
"stage": "unknown",
|
||
"message": "No analysis progress available",
|
||
"current": 0,
|
||
"total": 0,
|
||
}
|
||
)
|
||
|
||
return web.json_response(progress)
|
||
|
||
@routes.post("/model_resolver/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 Resolver resolve error: {e}", exc_info=True)
|
||
return web.json_response(
|
||
{"error": str(e), "success": False}, status=500
|
||
)
|
||
|
||
@routes.post("/model_resolver/local-matches")
|
||
async def local_matches(request):
|
||
"""Search local model files by filename/path."""
|
||
try:
|
||
data = await request.json()
|
||
filename = data.get("filename", "")
|
||
category = data.get("category", "")
|
||
|
||
if not filename:
|
||
return web.json_response(
|
||
{"error": "filename is required"}, status=400
|
||
)
|
||
|
||
matches = search_local_matches(
|
||
filename,
|
||
category=category or None,
|
||
similarity_threshold=0.0,
|
||
max_matches_per_model=10,
|
||
)
|
||
return web.json_response({"matches": matches})
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver local-matches error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/open-containing-folder")
|
||
async def open_containing_folder(request):
|
||
"""Open Explorer at the folder containing the selected model."""
|
||
try:
|
||
import os
|
||
import subprocess
|
||
|
||
data = await request.json()
|
||
target_path = data.get("path", "")
|
||
|
||
if not target_path:
|
||
return web.json_response(
|
||
{"error": "path is required"}, status=400
|
||
)
|
||
|
||
normalized_path = os.path.normpath(target_path)
|
||
if not os.path.exists(normalized_path):
|
||
return web.json_response(
|
||
{"error": "path does not exist"}, status=404
|
||
)
|
||
|
||
if os.path.isfile(normalized_path):
|
||
absolute_path = os.path.abspath(normalized_path)
|
||
subprocess.Popen(
|
||
["explorer.exe", "/select,", absolute_path],
|
||
shell=False,
|
||
)
|
||
else:
|
||
os.startfile(normalized_path)
|
||
|
||
return web.json_response({"success": True})
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver open-containing-folder error: {e}",
|
||
exc_info=True,
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/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 Resolver get_models error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/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 Resolver get_loaded_models error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
# ==================== CIVITAI SEARCH ROUTE ====================
|
||
|
||
@routes.post("/model_resolver/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, model_type=category
|
||
)
|
||
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 Resolver civitai-search error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
# ==================== DOWNLOAD ROUTES ====================
|
||
|
||
if download_available:
|
||
|
||
@routes.post("/model_resolver/search")
|
||
async def search_sources(request):
|
||
"""Search for model download sources."""
|
||
try:
|
||
def summarize_result(source_name, result, extra=None):
|
||
if isinstance(result, dict):
|
||
summary = {
|
||
"source": source_name,
|
||
"name": result.get("name"),
|
||
"filename": result.get("filename"),
|
||
"match_type": result.get("match_type"),
|
||
"repo_id": result.get("repo_id"),
|
||
"model_id": result.get("model_id"),
|
||
"version_id": result.get("version_id"),
|
||
"url": result.get("url"),
|
||
"download_url": result.get("download_url"),
|
||
"size": result.get("size"),
|
||
}
|
||
else:
|
||
summary = {"source": source_name, "result": result}
|
||
|
||
if extra:
|
||
summary.update(extra)
|
||
|
||
return summary
|
||
|
||
def log_search_result(source_name, result, extra=None):
|
||
details = summarize_result(source_name, result, extra)
|
||
if result and (
|
||
not isinstance(result, list) or len(result) > 0
|
||
):
|
||
log_info(f"Search source [{source_name}] FOUND: {details}")
|
||
else:
|
||
log_info(f"Search source [{source_name}] no result: {details}")
|
||
|
||
data = await request.json()
|
||
filename = data.get("filename", "")
|
||
category = data.get("category", "")
|
||
civitai_candidate_limit_raw = data.get(
|
||
"civitai_candidate_limit", 5
|
||
)
|
||
try:
|
||
civitai_candidate_limit = int(
|
||
civitai_candidate_limit_raw
|
||
)
|
||
except (TypeError, ValueError):
|
||
civitai_candidate_limit = 5
|
||
civitai_candidate_limit = max(
|
||
1, min(civitai_candidate_limit, 20)
|
||
)
|
||
civarchive_candidate_limit_raw = data.get(
|
||
"civarchive_candidate_limit", 10
|
||
)
|
||
try:
|
||
civarchive_candidate_limit = int(
|
||
civarchive_candidate_limit_raw
|
||
)
|
||
except (TypeError, ValueError):
|
||
civarchive_candidate_limit = 10
|
||
civarchive_candidate_limit = max(
|
||
1, min(civarchive_candidate_limit, 30)
|
||
)
|
||
# Handle both boolean and string forms
|
||
is_urn_raw = data.get("is_urn", False)
|
||
civitai_session_token = data.get("civitai_session_token", "")
|
||
hf_token = data.get("hf_token", "")
|
||
brave_search_api_key = data.get("brave_search_api_key", "")
|
||
civitai_use_trpc_search = data.get(
|
||
"civitai_use_trpc_search", True
|
||
)
|
||
civitai_use_html_fallback = data.get(
|
||
"civitai_use_html_fallback", True
|
||
)
|
||
hf_use_api_search = data.get("hf_use_api_search", True)
|
||
hf_use_comfy_org_fallback = data.get(
|
||
"hf_use_comfy_org_fallback", True
|
||
)
|
||
hf_use_brave_fallback = data.get(
|
||
"hf_use_brave_fallback", True
|
||
)
|
||
is_urn = (
|
||
is_urn_raw
|
||
if isinstance(is_urn_raw, bool)
|
||
else (str(is_urn_raw).lower() == "true")
|
||
)
|
||
hf_use_api_search = (
|
||
hf_use_api_search
|
||
if isinstance(hf_use_api_search, bool)
|
||
else str(hf_use_api_search).lower() == "true"
|
||
)
|
||
civitai_use_trpc_search = (
|
||
civitai_use_trpc_search
|
||
if isinstance(civitai_use_trpc_search, bool)
|
||
else str(civitai_use_trpc_search).lower() == "true"
|
||
)
|
||
civitai_use_html_fallback = (
|
||
civitai_use_html_fallback
|
||
if isinstance(civitai_use_html_fallback, bool)
|
||
else str(civitai_use_html_fallback).lower() == "true"
|
||
)
|
||
hf_use_comfy_org_fallback = (
|
||
hf_use_comfy_org_fallback
|
||
if isinstance(hf_use_comfy_org_fallback, bool)
|
||
else str(hf_use_comfy_org_fallback).lower() == "true"
|
||
)
|
||
hf_use_brave_fallback = (
|
||
hf_use_brave_fallback
|
||
if isinstance(hf_use_brave_fallback, bool)
|
||
else str(hf_use_brave_fallback).lower() == "true"
|
||
)
|
||
|
||
# For URN-only requests, model_id and version_id are required instead of filename
|
||
model_id = data.get("model_id")
|
||
version_id = data.get("version_id")
|
||
if not filename and not (is_urn and model_id and version_id):
|
||
return web.json_response(
|
||
{
|
||
"error": "Filename is required for non-URN, or model_id+version_id for URN"
|
||
},
|
||
status=400,
|
||
)
|
||
|
||
raw_sources = data.get("sources", ["all"])
|
||
if isinstance(raw_sources, str):
|
||
raw_sources = [raw_sources]
|
||
elif not isinstance(raw_sources, list):
|
||
raw_sources = ["all"]
|
||
|
||
normalized_sources = {
|
||
str(source).strip().lower()
|
||
for source in raw_sources
|
||
if str(source).strip()
|
||
}
|
||
if not normalized_sources:
|
||
normalized_sources = {"all"}
|
||
|
||
if "all" in normalized_sources:
|
||
normalized_sources = {
|
||
"local",
|
||
"huggingface",
|
||
"civitai",
|
||
"civarchive",
|
||
"lora_manager_archive",
|
||
}
|
||
|
||
search_local = "local" in normalized_sources
|
||
search_huggingface_source = "huggingface" in normalized_sources
|
||
search_civitai_source = "civitai" in normalized_sources
|
||
search_civarchive_source = "civarchive" in normalized_sources
|
||
search_lora_manager_archive_source = (
|
||
"lora_manager_archive" in normalized_sources
|
||
)
|
||
|
||
log_info(
|
||
f"Search request: filename={filename}, category={category}, is_urn={is_urn}, model_id={data.get('model_id')}, version_id={data.get('version_id')}, sources={sorted(normalized_sources)}, civitai_candidate_limit={civitai_candidate_limit}, civarchive_candidate_limit={civarchive_candidate_limit}"
|
||
)
|
||
|
||
results = {
|
||
"popular": None,
|
||
"model_list": None,
|
||
"huggingface": None,
|
||
"civitai": None,
|
||
"civarchive": None,
|
||
"lora_manager_archive": None,
|
||
"found": False,
|
||
"searched_sources": sorted(normalized_sources),
|
||
}
|
||
|
||
def current_search_timestamp():
|
||
from datetime import datetime, timezone
|
||
|
||
return (
|
||
datetime.now(timezone.utc)
|
||
.replace(microsecond=0)
|
||
.isoformat()
|
||
)
|
||
|
||
def get_search_result_signature(source_key, result):
|
||
if isinstance(result, list):
|
||
return "|".join(
|
||
get_search_result_signature(source_key, item)
|
||
for item in result
|
||
)
|
||
if not isinstance(result, dict):
|
||
return ""
|
||
|
||
parts = [
|
||
source_key,
|
||
result.get("download_url")
|
||
or result.get("url")
|
||
or result.get("model_url")
|
||
or "",
|
||
result.get("filename") or result.get("path") or "",
|
||
result.get("repo_id") or result.get("repo") or "",
|
||
result.get("model_id") or "",
|
||
result.get("version_id") or "",
|
||
result.get("name") or "",
|
||
]
|
||
return "::".join(str(part).strip() for part in parts)
|
||
|
||
def stamp_search_result(source_key, result):
|
||
if isinstance(result, list):
|
||
return [
|
||
stamp_search_result(source_key, item)
|
||
for item in result
|
||
]
|
||
if not isinstance(result, dict):
|
||
return result
|
||
|
||
signature = get_search_result_signature(source_key, result)
|
||
if not signature:
|
||
return result
|
||
|
||
timestamp = (
|
||
result.get("searched_at")
|
||
or result.get("searchedAt")
|
||
or self.search_result_timestamps.get(signature)
|
||
)
|
||
if not timestamp:
|
||
timestamp = current_search_timestamp()
|
||
self.search_result_timestamps.setdefault(
|
||
signature, timestamp
|
||
)
|
||
result["searched_at"] = timestamp
|
||
return result
|
||
|
||
def stamp_search_results(payload):
|
||
for source_key in (
|
||
"popular",
|
||
"model_list",
|
||
"huggingface",
|
||
"civitai",
|
||
"civarchive",
|
||
"lora_manager_archive",
|
||
):
|
||
if payload.get(source_key):
|
||
payload[source_key] = stamp_search_result(
|
||
source_key, payload[source_key]
|
||
)
|
||
return payload
|
||
|
||
def search_local_sources():
|
||
source_results = {"popular": None, "model_list": None}
|
||
source_found = False
|
||
|
||
log_info(
|
||
f"Search source [local] start: filename={filename}, category={category}"
|
||
)
|
||
popular_info = get_popular_model_url(filename)
|
||
log_search_result("popular", popular_info)
|
||
model_list_result = search_model_list(filename)
|
||
log_search_result(
|
||
"model_list",
|
||
model_list_result,
|
||
{
|
||
"confidence": model_list_result.get("confidence")
|
||
if model_list_result
|
||
else None
|
||
},
|
||
)
|
||
if popular_info:
|
||
popular_result = {
|
||
"source": "popular",
|
||
"filename": filename,
|
||
**popular_info,
|
||
}
|
||
if (
|
||
model_list_result
|
||
and model_list_result.get("filename", "").lower()
|
||
== filename.lower()
|
||
and model_list_result.get("size")
|
||
):
|
||
popular_result["size"] = model_list_result.get(
|
||
"size"
|
||
)
|
||
source_results["popular"] = popular_result
|
||
source_found = True
|
||
|
||
if model_list_result:
|
||
confidence = model_list_result.get("confidence", 0)
|
||
if is_urn and confidence >= 70:
|
||
source_results["model_list"] = model_list_result
|
||
source_found = True
|
||
elif not is_urn:
|
||
source_results["model_list"] = model_list_result
|
||
source_found = True
|
||
|
||
return source_results, source_found
|
||
|
||
def search_huggingface_source_task():
|
||
log_info(f"Search source [huggingface] start: filename={filename}")
|
||
hf_result = search_huggingface_for_file(
|
||
filename,
|
||
token=hf_token or None,
|
||
brave_api_key=brave_search_api_key or None,
|
||
use_api_search=hf_use_api_search,
|
||
use_comfy_org_fallback=hf_use_comfy_org_fallback,
|
||
use_brave_fallback=hf_use_brave_fallback,
|
||
)
|
||
log_search_result("huggingface", hf_result)
|
||
return {"huggingface": hf_result}, bool(hf_result)
|
||
|
||
def search_civitai_source_task():
|
||
source_results = {"civitai": None}
|
||
source_found = False
|
||
|
||
log_info(
|
||
f"Search source [civitai] start: filename={filename}, category={category}, is_urn={is_urn}"
|
||
)
|
||
# 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 resolve_urn to get model info (cached)
|
||
model_info = resolve_urn(model_id, version_id)
|
||
if model_info:
|
||
primary_file = None
|
||
for file_info in model_info.get("files", []):
|
||
if (
|
||
file_info.get("name")
|
||
== model_info.get("expected_filename")
|
||
):
|
||
primary_file = file_info
|
||
break
|
||
if primary_file is None:
|
||
primary_file = (
|
||
model_info.get("files") or [{}]
|
||
)[0]
|
||
|
||
download_url = get_civitai_download_url(
|
||
version_id
|
||
)
|
||
source_results["civitai"] = {
|
||
"source": "civitai",
|
||
"name": model_info.get("model_name"),
|
||
"version_name": model_info.get(
|
||
"version_name"
|
||
),
|
||
"filename": model_info.get(
|
||
"expected_filename"
|
||
),
|
||
"type": category,
|
||
"download_url": download_url,
|
||
"url": f"https://civitai.com/models/{model_id}?modelVersionId={version_id}",
|
||
"model_id": model_id,
|
||
"version_id": version_id,
|
||
"size": primary_file.get("size"),
|
||
"base_model": model_info.get("base_model"),
|
||
"tags": model_info.get("tags", []),
|
||
}
|
||
log_search_result(
|
||
"civitai/urn",
|
||
source_results["civitai"],
|
||
{
|
||
"files_count": len(
|
||
model_info.get("files", [])
|
||
)
|
||
},
|
||
)
|
||
source_found = True
|
||
else:
|
||
log_search_result(
|
||
"civitai/urn",
|
||
None,
|
||
{
|
||
"model_id": model_id,
|
||
"version_id": version_id,
|
||
},
|
||
)
|
||
elif category:
|
||
# Fallback to search if no IDs
|
||
log_info(
|
||
f"URN: No model_id/version_id, falling back to CivitAI search"
|
||
)
|
||
civitai_results = search_civitai(
|
||
filename,
|
||
model_type=category,
|
||
)
|
||
log_search_result(
|
||
"civitai/fallback",
|
||
civitai_results[0] if civitai_results else None,
|
||
{
|
||
"results_count": len(civitai_results),
|
||
},
|
||
)
|
||
if civitai_results:
|
||
first_result = civitai_results[0]
|
||
source_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"),
|
||
"base_model": first_result.get("base_model"),
|
||
"tags": first_result.get("tags", []),
|
||
}
|
||
source_found = True
|
||
else:
|
||
civitai_result = search_civitai_for_file(
|
||
filename,
|
||
model_type=category,
|
||
session_token=civitai_session_token or None,
|
||
candidate_limit=civitai_candidate_limit,
|
||
use_trpc_search=civitai_use_trpc_search,
|
||
use_html_fallback=civitai_use_html_fallback,
|
||
)
|
||
log_search_result("civitai", civitai_result)
|
||
if civitai_result:
|
||
source_results["civitai"] = civitai_result
|
||
source_found = True
|
||
|
||
return source_results, source_found
|
||
|
||
def search_civarchive_source_task():
|
||
source_results = {"civarchive": None}
|
||
source_found = False
|
||
|
||
log_info(
|
||
f"Search source [civarchive] start: filename={filename}, category={category}, is_urn={is_urn}"
|
||
)
|
||
if is_urn:
|
||
model_id = data.get("model_id")
|
||
version_id = data.get("version_id")
|
||
if model_id and version_id:
|
||
civarchive_result = resolve_civarchive_model_version(
|
||
model_id,
|
||
version_id,
|
||
query=filename,
|
||
)
|
||
log_search_result(
|
||
"civarchive/urn",
|
||
civarchive_result,
|
||
{
|
||
"model_id": model_id,
|
||
"version_id": version_id,
|
||
},
|
||
)
|
||
if civarchive_result:
|
||
source_results["civarchive"] = civarchive_result
|
||
source_found = True
|
||
else:
|
||
log_search_result(
|
||
"civarchive/urn",
|
||
None,
|
||
{
|
||
"model_id": model_id,
|
||
"version_id": version_id,
|
||
},
|
||
)
|
||
else:
|
||
civarchive_result = search_civarchive_for_file(
|
||
filename,
|
||
model_type=category,
|
||
limit=civarchive_candidate_limit,
|
||
)
|
||
log_search_result("civarchive", civarchive_result)
|
||
if civarchive_result:
|
||
source_results["civarchive"] = civarchive_result
|
||
source_found = True
|
||
|
||
return source_results, source_found
|
||
|
||
def search_lora_manager_archive_source_task():
|
||
log_info(
|
||
f"Search source [lora_manager_archive] start: filename={filename}, category={category}"
|
||
)
|
||
lora_manager_archive_result = (
|
||
search_lora_manager_archive_for_file(
|
||
filename,
|
||
model_type=category,
|
||
)
|
||
)
|
||
log_search_result(
|
||
"lora_manager_archive",
|
||
lora_manager_archive_result,
|
||
)
|
||
return (
|
||
{
|
||
"lora_manager_archive": lora_manager_archive_result
|
||
},
|
||
bool(lora_manager_archive_result),
|
||
)
|
||
|
||
search_tasks = []
|
||
if search_local:
|
||
search_tasks.append(
|
||
asyncio.to_thread(search_local_sources)
|
||
)
|
||
if search_huggingface_source:
|
||
search_tasks.append(
|
||
asyncio.to_thread(search_huggingface_source_task)
|
||
)
|
||
if search_civitai_source:
|
||
search_tasks.append(
|
||
asyncio.to_thread(search_civitai_source_task)
|
||
)
|
||
if search_civarchive_source and (
|
||
filename or (is_urn and model_id and version_id)
|
||
):
|
||
search_tasks.append(
|
||
asyncio.to_thread(search_civarchive_source_task)
|
||
)
|
||
if search_lora_manager_archive_source and filename:
|
||
search_tasks.append(
|
||
asyncio.to_thread(
|
||
search_lora_manager_archive_source_task
|
||
)
|
||
)
|
||
|
||
if len(search_tasks) > 1:
|
||
log_info(
|
||
f"Search sources running asynchronously: count={len(search_tasks)}"
|
||
)
|
||
|
||
for source_results, source_found in await asyncio.gather(
|
||
*search_tasks
|
||
):
|
||
for source_key, source_result in source_results.items():
|
||
if source_result:
|
||
results[source_key] = source_result
|
||
if source_found:
|
||
results["found"] = True
|
||
|
||
log_info(
|
||
f"Search summary: found={results['found']}, searched_sources={results['searched_sources']}"
|
||
)
|
||
stamp_search_results(results)
|
||
return web.json_response(results)
|
||
|
||
except Exception as e:
|
||
log_exception(f"Model Resolver search error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/clear-search-cache")
|
||
async def clear_search_cache_route(request):
|
||
"""Clear backend search caches after token/settings changes."""
|
||
try:
|
||
clear_huggingface_search_cache()
|
||
clear_civitai_search_cache()
|
||
clear_civarchive_search_cache()
|
||
clear_lora_manager_archive_search_cache()
|
||
reload_popular_databases()
|
||
reload_model_list()
|
||
invalidate_model_files_cache()
|
||
self.search_result_timestamps.clear()
|
||
log_info("Cleared backend search caches")
|
||
return web.json_response({"success": True, "cleared": "all"})
|
||
except Exception as e:
|
||
log_exception(f"Clear search cache error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/civitai/session-token/check")
|
||
async def civitai_session_token_check_route(request):
|
||
"""Check whether a CivitAI browser session token is valid."""
|
||
try:
|
||
data = await request.json()
|
||
token = data.get("civitai_session_token", "")
|
||
result = await asyncio.to_thread(
|
||
check_civitai_session_token, token
|
||
)
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"CivitAI session token check error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/civitai/api-key/check")
|
||
async def civitai_api_key_check_route(request):
|
||
"""Check whether a CivitAI API key is valid."""
|
||
try:
|
||
data = await request.json()
|
||
api_key = data.get("civitai_key", "")
|
||
result = await asyncio.to_thread(
|
||
check_civitai_api_key, api_key
|
||
)
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"CivitAI API key check error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/huggingface/token/check")
|
||
async def huggingface_token_check_route(request):
|
||
"""Check whether a HuggingFace token is valid."""
|
||
try:
|
||
data = await request.json()
|
||
token = data.get("hf_token", "")
|
||
result = await asyncio.to_thread(check_huggingface_token, token)
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"HuggingFace token check error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/brave/api-key/check")
|
||
async def brave_api_key_check_route(request):
|
||
"""Check whether a Brave Search API key is valid."""
|
||
try:
|
||
data = await request.json()
|
||
api_key = data.get("brave_search_api_key", "")
|
||
result = await asyncio.to_thread(
|
||
check_brave_search_api_key, api_key
|
||
)
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"Brave Search API key check error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/huggingface/author-index/status")
|
||
async def huggingface_author_index_status_route(request):
|
||
"""Return local HuggingFace author fallback index status."""
|
||
try:
|
||
return web.json_response(get_author_fallback_index_status())
|
||
except Exception as e:
|
||
log_exception(f"HuggingFace author index status error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/huggingface/author-index/refresh")
|
||
async def huggingface_author_index_refresh_route(request):
|
||
"""Refresh HuggingFace author fallback index."""
|
||
try:
|
||
data = await request.json()
|
||
hf_token = data.get("hf_token", "")
|
||
result = await asyncio.to_thread(
|
||
refresh_author_fallback_index, hf_token or None
|
||
)
|
||
clear_huggingface_search_cache()
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"HuggingFace author index refresh error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/model-list/status")
|
||
async def model_list_status_route(request):
|
||
"""Return local model-list status and optionally compare with GitHub."""
|
||
try:
|
||
check_remote = (
|
||
str(request.query.get("check_remote", "")).lower()
|
||
in {"1", "true", "yes"}
|
||
)
|
||
return web.json_response(
|
||
get_model_list_update_status(check_remote=check_remote)
|
||
)
|
||
except Exception as e:
|
||
log_exception(f"Model list status error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/model-list/update")
|
||
async def model_list_update_route(request):
|
||
"""Download latest ComfyUI-Manager model-list.json."""
|
||
try:
|
||
result = await asyncio.to_thread(update_model_list_from_remote)
|
||
clear_huggingface_search_cache()
|
||
clear_civitai_search_cache()
|
||
clear_civarchive_search_cache()
|
||
clear_lora_manager_archive_search_cache()
|
||
self.search_result_timestamps.clear()
|
||
return web.json_response(result)
|
||
except Exception as e:
|
||
log_exception(f"Model list update error: {e}")
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/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 Resolver download error: {e}", exc_info=True
|
||
)
|
||
return web.json_response(
|
||
{"error": str(e), "success": False}, status=500
|
||
)
|
||
|
||
@routes.get("/model_resolver/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 Resolver progress error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/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 Resolver progress error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/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 Resolver cancel error: {e}", exc_info=True
|
||
)
|
||
return web.json_response(
|
||
{"error": str(e), "success": False}, status=500
|
||
)
|
||
|
||
@routes.get("/model_resolver/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 Resolver directories error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/capabilities")
|
||
async def get_capabilities(request):
|
||
"""Get optional source capabilities available in this install."""
|
||
try:
|
||
return web.json_response(
|
||
{
|
||
"sources": {
|
||
"civarchive": is_civarchive_available(),
|
||
"lora_manager_archive": is_lora_manager_archive_available()
|
||
}
|
||
}
|
||
)
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver capabilities error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/subfolders/{category}")
|
||
async def get_subfolders(request):
|
||
"""Get known subfolders for a category using ComfyUI folder_paths."""
|
||
try:
|
||
import folder_paths
|
||
|
||
raw_category = (request.match_info.get("category") or "").strip()
|
||
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",
|
||
"default": "upscale_models",
|
||
}
|
||
category = category_map.get(raw_category.lower(), raw_category.lower())
|
||
|
||
if not category or category == "unknown":
|
||
return web.json_response([])
|
||
|
||
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
||
if category not in known_categories:
|
||
self.logger.debug(
|
||
f"Model Resolver: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'"
|
||
)
|
||
return web.json_response([])
|
||
|
||
subfolders = set()
|
||
filenames = folder_paths.get_filename_list(category) or []
|
||
for rel_path in filenames:
|
||
if not isinstance(rel_path, str):
|
||
continue
|
||
parts = [p for p in rel_path.replace("/", "\\").split("\\") if p]
|
||
if len(parts) <= 1:
|
||
continue
|
||
current = ""
|
||
for part in parts[:-1]:
|
||
current = f"{current}\\{part}" if current else part
|
||
subfolders.add(current)
|
||
|
||
return web.json_response(sorted(subfolders))
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver subfolders error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
# ==================== SETTINGS (server-side persistence) ====================
|
||
|
||
import os as _os
|
||
import json as _json
|
||
|
||
_SETTINGS_FILE = _os.path.join(
|
||
_os.path.dirname(_os.path.abspath(__file__)),
|
||
"model_resolver_settings.json",
|
||
)
|
||
|
||
# Keys that contain sensitive credentials – stored in the file but
|
||
# never logged.
|
||
_CREDENTIAL_KEYS = {
|
||
"civitai_key",
|
||
"civitai_session_token",
|
||
"hf_token",
|
||
"brave_search_api_key",
|
||
}
|
||
|
||
def _load_settings() -> dict:
|
||
try:
|
||
if _os.path.isfile(_SETTINGS_FILE):
|
||
with open(_SETTINGS_FILE, "r", encoding="utf-8") as f:
|
||
data = _json.load(f)
|
||
if isinstance(data, dict):
|
||
return data
|
||
except Exception as exc:
|
||
self.logger.warning(f"Model Resolver: could not read settings file: {exc}")
|
||
return {}
|
||
|
||
def _save_settings(payload: dict) -> None:
|
||
# Merge with existing settings so we never lose keys not sent
|
||
# by the current request.
|
||
current = _load_settings()
|
||
current.update({k: v for k, v in payload.items() if k})
|
||
with open(_SETTINGS_FILE, "w", encoding="utf-8") as f:
|
||
_json.dump(current, f, indent=2, ensure_ascii=False)
|
||
|
||
@routes.get("/model_resolver/settings")
|
||
async def get_settings_route(request):
|
||
"""Return persisted settings (API keys, preferences)."""
|
||
try:
|
||
data = await asyncio.to_thread(_load_settings)
|
||
return web.json_response(data)
|
||
except Exception as e:
|
||
self.logger.error(f"Model Resolver settings GET error: {e}", exc_info=True)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/settings")
|
||
async def save_settings_route(request):
|
||
"""Persist settings (API keys, preferences) to disk."""
|
||
try:
|
||
payload = await request.json()
|
||
if not isinstance(payload, dict):
|
||
return web.json_response({"error": "Expected JSON object"}, status=400)
|
||
await asyncio.to_thread(_save_settings, payload)
|
||
return web.json_response({"success": True})
|
||
except Exception as e:
|
||
self.logger.error(f"Model Resolver settings POST error: {e}", exc_info=True)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
self.routes_setup = True
|
||
self.logger.info("Model Resolver: API routes registered successfully")
|
||
return True
|
||
|
||
except ImportError as e:
|
||
self.logger.warning(
|
||
f"Model Resolver: Could not register routes (missing dependency): {e}"
|
||
)
|
||
return False
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver: Error setting up routes: {e}", exc_info=True
|
||
)
|
||
return False
|
||
|
||
|
||
# Initialize the extension
|
||
try:
|
||
extension = ModelResolverExtension()
|
||
extension.initialize()
|
||
except Exception as e:
|
||
log_error(
|
||
f"ComfyUI Model Resolver extension initialization failed: {e}", exc_info=True
|
||
)
|