1869 lines
92 KiB
Python
1869 lines
92 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,
|
||
log_warn,
|
||
)
|
||
|
||
# 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 (
|
||
CivArchiveSearchError,
|
||
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
|
||
|
||
# ==================== BASE MODELS CONFIG ROUTE ====================
|
||
|
||
@routes.get("/model_resolver/base-models")
|
||
async def get_base_models(request):
|
||
"""Return the base models config from metadata/base-models.json.
|
||
|
||
Returns {base_models: [{name, aliases}]} so the frontend
|
||
dropdown can populate correctly via baseModels.base_models.
|
||
"""
|
||
try:
|
||
from .core.sources.popular import get_base_models_config
|
||
data = get_base_models_config()
|
||
return web.json_response(data)
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver base-models error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.get("/model_resolver/base-models/status")
|
||
async def get_base_models_status_route(request):
|
||
"""Get local and optional remote base models status."""
|
||
try:
|
||
check_remote = request.query.get("check_remote") == "1"
|
||
from .core.sources.popular import get_base_models_status
|
||
status = await asyncio.to_thread(get_base_models_status, check_remote)
|
||
return web.json_response(status)
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver base-models status error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
@routes.post("/model_resolver/base-models/update")
|
||
async def update_base_models_route(request):
|
||
"""Update base models list from CivitAI."""
|
||
try:
|
||
from .core.sources.popular import update_base_models_from_remote
|
||
status = await asyncio.to_thread(update_base_models_from_remote)
|
||
return web.json_response(status)
|
||
except Exception as e:
|
||
self.logger.error(
|
||
f"Model Resolver base-models update error: {e}", exc_info=True
|
||
)
|
||
return web.json_response({"error": str(e)}, status=500)
|
||
|
||
# ==================== 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", "")
|
||
force_rescan = data.get("force_rescan", False)
|
||
force_rescan = (
|
||
force_rescan
|
||
if isinstance(force_rescan, bool)
|
||
else str(force_rescan).lower() == "true"
|
||
)
|
||
|
||
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,
|
||
force_rescan=force_rescan,
|
||
)
|
||
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
|
||
)
|
||
force_search = data.get("force_search", False)
|
||
force_search = (
|
||
force_search
|
||
if isinstance(force_search, bool)
|
||
else str(force_search).lower() == "true"
|
||
)
|
||
|
||
if force_search:
|
||
if search_local:
|
||
reload_popular_databases()
|
||
reload_model_list()
|
||
if search_huggingface_source:
|
||
clear_huggingface_search_cache()
|
||
if search_civitai_source:
|
||
clear_civitai_search_cache()
|
||
if search_civarchive_source:
|
||
clear_civarchive_search_cache()
|
||
if search_lora_manager_archive_source:
|
||
clear_lora_manager_archive_search_cache()
|
||
log_info(
|
||
f"Force search enabled: cleared cache for sources={sorted(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)}, force_search={force_search}, 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),
|
||
"source_errors": {},
|
||
}
|
||
|
||
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 (
|
||
None
|
||
if force_search
|
||
else self.search_result_timestamps.get(signature)
|
||
)
|
||
)
|
||
if not timestamp:
|
||
timestamp = current_search_timestamp()
|
||
if force_search:
|
||
self.search_result_timestamps[signature] = timestamp
|
||
else:
|
||
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,
|
||
force_refresh=force_search,
|
||
)
|
||
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}"
|
||
)
|
||
try:
|
||
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
|
||
except CivArchiveSearchError as e:
|
||
error_message = f"CivArchive search failed: {e}"
|
||
log_warn(error_message)
|
||
source_results["source_errors"] = {
|
||
"civarchive": error_message
|
||
}
|
||
|
||
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_key == "source_errors":
|
||
results["source_errors"].update(source_result or {})
|
||
continue
|
||
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
|
||
)
|