Dariusz L
6f9aa9cc55
- Adopt ComfyUI's ComfyExtension/comfy_entrypoint API in __init__.py when available, falling back to legacy NODE_CLASS_MAPPINGS/DISPLAY_NAME_MAPPINGS on older runtimes so the node loads across ComfyUI versions. - Hook into beforeRegisterNodeDef, loadedGraphNode, and afterConfigureGraph in web/resolver.js so marker nodes are patched during registration, graph load, and initial render. - Teach model_resolver.js to recognise the marker node by every known class property shape, patch its type on creation, enforce a minimum size with setSize() when present, and dirty the canvas so changes are reflected.
5253 lines
247 KiB
Python
5253 lines
247 KiB
Python
"""
|
|
@author: Azornes
|
|
@title: ComfyUI Model Resolver
|
|
@version: 1.0.0
|
|
@description: Extension for resolving missing models and downloading from HuggingFace/CivitAI
|
|
"""
|
|
|
|
import asyncio
|
|
import threading
|
|
import time
|
|
import sys
|
|
import os
|
|
|
|
if not __package__ or __package__ == "":
|
|
this_dir = os.path.dirname(os.path.abspath(__file__))
|
|
parent_dir = os.path.dirname(this_dir)
|
|
if parent_dir not in sys.path:
|
|
sys.path.insert(0, parent_dir)
|
|
package_name = os.path.basename(this_dir)
|
|
__package__ = package_name
|
|
|
|
current_module = sys.modules.get(__name__)
|
|
if current_module:
|
|
sys.modules[package_name] = current_module
|
|
if not hasattr(current_module, "__path__"):
|
|
current_module.__path__ = [this_dir]
|
|
from .core.log_system import LogLevel, create_module_logger, logger as backend_log_controller
|
|
from .core.log_system.config import LOG_LEVEL as BACKEND_DEFAULT_LOG_LEVEL
|
|
|
|
# Web directory for JavaScript interface
|
|
WEB_DIRECTORY = "./web"
|
|
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_TYPE = "ModelResolverDependency"
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_DISPLAY_NAME = "Model Resolver Opener"
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_CATEGORY = "Model Resolver/Workflow"
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_DESCRIPTION = (
|
|
"A passive opener node for workflows that intentionally depend on "
|
|
"Model Resolver metadata. It does not process images, models, or prompts."
|
|
)
|
|
|
|
try:
|
|
from comfy_api.latest import ComfyExtension, io
|
|
except Exception:
|
|
ComfyExtension = None
|
|
io = None
|
|
|
|
|
|
if ComfyExtension is not None and io is not None:
|
|
|
|
class ModelResolverDependencyNode(io.ComfyNode):
|
|
"""Canvas node that declares Model Resolver and opens its workflow tools."""
|
|
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id=MODEL_RESOLVER_DEPENDENCY_NODE_TYPE,
|
|
display_name=MODEL_RESOLVER_DEPENDENCY_NODE_DISPLAY_NAME,
|
|
category=MODEL_RESOLVER_DEPENDENCY_NODE_CATEGORY,
|
|
description=MODEL_RESOLVER_DEPENDENCY_NODE_DESCRIPTION,
|
|
inputs=[],
|
|
outputs=[],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls) -> io.NodeOutput:
|
|
return io.NodeOutput()
|
|
|
|
class ModelResolverNodeExtension(ComfyExtension):
|
|
async def get_node_list(self) -> list[type[io.ComfyNode]]:
|
|
return [ModelResolverDependencyNode]
|
|
|
|
async def comfy_entrypoint() -> ComfyExtension:
|
|
return ModelResolverNodeExtension()
|
|
|
|
__all__ = ["WEB_DIRECTORY", "comfy_entrypoint"]
|
|
else:
|
|
|
|
class ModelResolverDependencyNode:
|
|
"""Legacy fallback for ComfyUI builds without the v3 custom-node API."""
|
|
|
|
DESCRIPTION = MODEL_RESOLVER_DEPENDENCY_NODE_DESCRIPTION
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "noop"
|
|
CATEGORY = MODEL_RESOLVER_DEPENDENCY_NODE_CATEGORY
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {"required": {}}
|
|
|
|
def noop(self):
|
|
return ()
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_TYPE: ModelResolverDependencyNode,
|
|
}
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
MODEL_RESOLVER_DEPENDENCY_NODE_TYPE: MODEL_RESOLVER_DEPENDENCY_NODE_DISPLAY_NAME,
|
|
}
|
|
|
|
__all__ = ["WEB_DIRECTORY", "NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]
|
|
|
|
|
|
class JobProgressTracker:
|
|
"""Helper class for thread-safe job progress tracking and cancellation management."""
|
|
|
|
def __init__(self, default_message="Processing..."):
|
|
self.lock = threading.Lock()
|
|
self.progress = {}
|
|
self.cancelled = set()
|
|
self.default_message = default_message
|
|
|
|
def cleanup(self, max_age_seconds=300):
|
|
cutoff = time.time() - max_age_seconds
|
|
with self.lock:
|
|
expired = [
|
|
pid
|
|
for pid, data in self.progress.items()
|
|
if data.get("updated_at", data.get("created_at", 0)) < cutoff
|
|
]
|
|
for pid in expired:
|
|
self.progress.pop(pid, None)
|
|
self.cancelled.difference_update(expired)
|
|
|
|
def update(self, progress_id, source=None, stage=None, message=None, percent=None, status=None, **payload):
|
|
if not progress_id:
|
|
return
|
|
if source is not None:
|
|
payload["source"] = source
|
|
if stage is not None:
|
|
payload["stage"] = stage
|
|
if message is not None:
|
|
payload["message"] = message
|
|
if percent is not None:
|
|
payload["percent"] = percent
|
|
if status is not None:
|
|
payload["status"] = status
|
|
|
|
now = time.time()
|
|
with self.lock:
|
|
current = self.progress.get(progress_id, {})
|
|
|
|
# If the job was cancelled, force it to remain cancelled.
|
|
if progress_id in self.cancelled and payload.get("status") != "cancelled":
|
|
payload["status"] = "cancelled"
|
|
payload["stage"] = "cancelled"
|
|
payload["message"] = "Cancelled"
|
|
payload["percent"] = 100
|
|
payload["cancelled"] = True
|
|
|
|
# Normalize percent if present
|
|
if "percent" in payload and payload["percent"] is not None:
|
|
try:
|
|
payload["percent"] = max(0.0, min(100.0, float(payload["percent"])))
|
|
except (TypeError, ValueError):
|
|
pass
|
|
|
|
self.progress[progress_id] = {
|
|
"created_at": now,
|
|
"message": self.default_message,
|
|
**current,
|
|
**payload,
|
|
"progress_id": progress_id,
|
|
"updated_at": now,
|
|
}
|
|
|
|
def is_cancelled(self, progress_id) -> bool:
|
|
if not progress_id:
|
|
return False
|
|
with self.lock:
|
|
return progress_id in self.cancelled
|
|
|
|
def mark_cancelled(self, progress_id, cancel_message="Cancelled") -> bool:
|
|
self.cleanup()
|
|
with self.lock:
|
|
current = self.progress.get(progress_id, {})
|
|
self.cancelled.add(progress_id)
|
|
self.progress[progress_id] = {
|
|
"created_at": time.time(),
|
|
**current,
|
|
"progress_id": progress_id,
|
|
"status": "cancelled",
|
|
"stage": "cancelled",
|
|
"message": cancel_message,
|
|
"percent": 100,
|
|
"cancelled": True,
|
|
"updated_at": time.time(),
|
|
}
|
|
return True
|
|
|
|
def get(self, progress_id):
|
|
with self.lock:
|
|
val = self.progress.get(progress_id)
|
|
return dict(val) if val else None
|
|
|
|
|
|
class ModelResolverExtension:
|
|
"""Main extension class for Model Resolver."""
|
|
|
|
def __init__(self):
|
|
self.routes_setup = False
|
|
self.logger = create_module_logger(__name__)
|
|
self.analysis_progress = JobProgressTracker("Analyzing...")
|
|
self.loaded_progress = JobProgressTracker("Loading loaded models...")
|
|
self.search_tracker = JobProgressTracker("Searching...")
|
|
self.hash_tracker = JobProgressTracker("Preparing hash calculation...")
|
|
self.metadata_builder_progress = JobProgressTracker("Building local metadata...")
|
|
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,
|
|
get_local_model_hash_metadata,
|
|
invalidate_local_hash_match_cache,
|
|
search_local_matches_by_hash,
|
|
search_local_matches,
|
|
)
|
|
from .core.workflow_analyzer import analyze_workflow_models
|
|
from .core.path_templates import infer_download_path_templates
|
|
from .core.metadata_audit import audit_metadata_sizes
|
|
from .core.metadata_builder import (
|
|
build_missing_local_metadata,
|
|
get_metadata_build_capabilities,
|
|
)
|
|
from .core.scanner import get_model_files, invalidate_model_files_cache, find_local_file_path
|
|
from .core.path_utils import (
|
|
extract_safetensors_header_sha256,
|
|
get_filename_from_path,
|
|
is_path_in_configured_model_roots,
|
|
read_json_safe,
|
|
write_json_atomic,
|
|
)
|
|
from .core.type_utils import (
|
|
first_non_empty,
|
|
to_int,
|
|
to_bool,
|
|
MODEL_EXTENSIONS,
|
|
format_size_bytes,
|
|
normalize_sha256,
|
|
extract_sha256_from_metadata,
|
|
fetch_remote_file_size_cached,
|
|
looks_like_model_file,
|
|
normalize_category_to_model_type,
|
|
get_category_folder_keys,
|
|
get_enabled_download_categories,
|
|
)
|
|
from .core.settings import (
|
|
TEMPLATE_KEY_ALIASES,
|
|
bool_setting as resolver_bool_setting,
|
|
get_default_root_for_category,
|
|
load_settings as load_resolver_settings,
|
|
resolve_download_subfolder,
|
|
save_settings as save_resolver_settings,
|
|
)
|
|
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,
|
|
pause_download,
|
|
resume_download,
|
|
get_aria2_status,
|
|
start_aria2_daemon,
|
|
stop_aria2_daemon,
|
|
get_download_directory,
|
|
get_metadata_sidecar_path,
|
|
normalize_download_category,
|
|
write_lora_manager_metadata,
|
|
sanitize_download_filename,
|
|
)
|
|
from .core.aria2_installer import Aria2InstallError, install_aria2_engine
|
|
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,
|
|
parse_huggingface_url,
|
|
get_huggingface_download_url,
|
|
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,
|
|
parse_civitai_url,
|
|
get_civitai_download_url,
|
|
get_civitai_model_details,
|
|
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,
|
|
parse_civarchive_url,
|
|
resolve_civarchive_by_hash,
|
|
resolve_civarchive_model_version,
|
|
get_civarchive_model_details,
|
|
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,
|
|
)
|
|
from .core.sources import clear_all_search_caches
|
|
|
|
download_available = True
|
|
except ImportError as e:
|
|
self.logger.warning(
|
|
f"Model Resolver: Download features not available: {e}"
|
|
)
|
|
download_available = False
|
|
|
|
def json_api_endpoint(error_prefix, return_success_on_error=False):
|
|
def decorator(func):
|
|
from functools import wraps
|
|
@wraps(func)
|
|
async def wrapper(request, *args, **kwargs):
|
|
try:
|
|
return await func(request, *args, **kwargs)
|
|
except Exception as e:
|
|
self.logger.error(
|
|
f"Model Resolver {error_prefix} error: {e}", exc_info=True
|
|
)
|
|
response_data = {"error": str(e)}
|
|
if return_success_on_error:
|
|
response_data["success"] = False
|
|
return web.json_response(response_data, status=500)
|
|
return wrapper
|
|
return decorator
|
|
|
|
def get_progress_response(tracker, request, param_name="progress_id", not_found_payload=None, not_found_status=200, found_wrapper=None):
|
|
job_id = request.match_info.get(param_name, "").strip()
|
|
if not job_id:
|
|
return web.json_response({"error": f"{param_name} is required"}, status=400)
|
|
tracker.cleanup()
|
|
progress = tracker.get(job_id)
|
|
if not progress:
|
|
payload = not_found_payload or {"error": "progress not found"}
|
|
return web.json_response(payload, status=not_found_status)
|
|
if found_wrapper:
|
|
return web.json_response(found_wrapper(progress))
|
|
return web.json_response(progress)
|
|
|
|
def cancel_progress_response(tracker, request, param_name="progress_id", cancel_message="Cancelled"):
|
|
job_id = request.match_info.get(param_name, "").strip()
|
|
if not job_id:
|
|
return web.json_response({"error": f"{param_name} is required"}, status=400)
|
|
cancelled = tracker.mark_cancelled(job_id, cancel_message)
|
|
return web.json_response({
|
|
"success": True,
|
|
"cancelled": cancelled,
|
|
"progress_id": job_id,
|
|
})
|
|
|
|
async def get_override_settings_from_request(request):
|
|
settings = await asyncio.to_thread(load_resolver_settings)
|
|
if request.method == "POST":
|
|
try:
|
|
payload = await request.json()
|
|
if isinstance(payload, dict) and "aria2c_path" in payload:
|
|
settings = dict(settings)
|
|
settings["aria2c_path"] = payload.get("aria2c_path", "")
|
|
except Exception:
|
|
pass
|
|
return settings
|
|
|
|
# ==================== BASE MODELS CONFIG ROUTE ====================
|
|
|
|
@routes.get("/model_resolver/base-models")
|
|
@json_api_endpoint("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.
|
|
"""
|
|
from .core.sources.popular import get_base_models_config
|
|
data = get_base_models_config()
|
|
return web.json_response(data)
|
|
|
|
@routes.get("/model_resolver/base-models/status")
|
|
@json_api_endpoint("base-models status")
|
|
async def get_base_models_status_route(request):
|
|
"""Get local and optional remote base models status."""
|
|
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)
|
|
|
|
@routes.post("/model_resolver/base-models/update")
|
|
@json_api_endpoint("base-models update")
|
|
async def update_base_models_route(request):
|
|
"""Update base models list from CivitAI."""
|
|
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)
|
|
|
|
# ==================== 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()
|
|
force_rescan = to_bool(data.get("force_rescan"), False)
|
|
if force_rescan:
|
|
invalidate_local_hash_match_cache()
|
|
|
|
if workflow_json is None:
|
|
return web.json_response(
|
|
{"error": "Workflow JSON is required"}, status=400
|
|
)
|
|
if not isinstance(workflow_json, dict):
|
|
return web.json_response(
|
|
{"error": "Workflow JSON must be an object"}, status=400
|
|
)
|
|
|
|
if analysis_id:
|
|
self.analysis_progress.update(
|
|
analysis_id,
|
|
status="starting",
|
|
stage="starting",
|
|
message="Starting analysis...",
|
|
current=0,
|
|
total=0,
|
|
)
|
|
|
|
def update_analysis_progress(payload):
|
|
if not analysis_id:
|
|
return
|
|
status = "running" if payload.get("stage") != "completed" else "completed"
|
|
self.analysis_progress.update(
|
|
analysis_id,
|
|
status=status,
|
|
**payload
|
|
)
|
|
|
|
# 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,
|
|
force_rescan=force_rescan,
|
|
)
|
|
|
|
# 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 = get_filename_from_path(
|
|
missing.get("original_path", "")
|
|
)
|
|
|
|
# 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 using cached remote helper
|
|
file_size = fetch_remote_file_size_cached(workflow_url, timeout=5)
|
|
|
|
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.update(
|
|
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.update(
|
|
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}")
|
|
@json_api_endpoint("analyze-progress")
|
|
async def get_analyze_progress(request):
|
|
"""Get workflow analysis progress."""
|
|
return get_progress_response(
|
|
self.analysis_progress,
|
|
request,
|
|
param_name="analysis_id",
|
|
not_found_payload={
|
|
"status": "unknown",
|
|
"stage": "unknown",
|
|
"message": "No analysis progress available",
|
|
"current": 0,
|
|
"total": 0,
|
|
}
|
|
)
|
|
|
|
@routes.post("/model_resolver/resolve")
|
|
@json_api_endpoint("resolve", return_success_on_error=True)
|
|
async def resolve_models(request):
|
|
"""Apply model resolution and return updated workflow."""
|
|
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}
|
|
)
|
|
|
|
@routes.post("/model_resolver/local-matches")
|
|
@json_api_endpoint("local-matches")
|
|
async def local_matches(request):
|
|
"""Search local model files by filename/path."""
|
|
data = await request.json()
|
|
filename = data.get("filename", "")
|
|
category = data.get("category", "")
|
|
force_rescan = to_bool(data.get("force_rescan"), False)
|
|
if force_rescan:
|
|
invalidate_local_hash_match_cache()
|
|
|
|
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})
|
|
|
|
@routes.post("/model_resolver/local-model-hashes")
|
|
@json_api_endpoint("local-model-hashes")
|
|
async def local_model_hashes(request):
|
|
"""Return SHA256 hashes already stored in local sidecar metadata."""
|
|
data = await request.json()
|
|
model = data.get("model") if isinstance(data.get("model"), dict) else {}
|
|
path = (
|
|
data.get("path")
|
|
or data.get("file_path")
|
|
or data.get("resolved_path")
|
|
or model.get("path")
|
|
or model.get("resolved_path")
|
|
or ""
|
|
)
|
|
|
|
if not path:
|
|
return web.json_response(
|
|
{"error": "path is required"}, status=400
|
|
)
|
|
|
|
return web.json_response(
|
|
get_local_model_hash_metadata(path, model=model)
|
|
)
|
|
|
|
@routes.post("/model_resolver/workflow-model-hashes")
|
|
@json_api_endpoint("workflow-model-hashes")
|
|
async def workflow_model_hashes(request):
|
|
"""Return hash metadata for existing local models used by a workflow."""
|
|
data = await request.json()
|
|
workflow_json = data.get("workflow")
|
|
if not isinstance(workflow_json, dict):
|
|
return web.json_response(
|
|
{"error": "Workflow JSON must be an object"}, status=400
|
|
)
|
|
|
|
settings = await asyncio.to_thread(load_resolver_settings)
|
|
if not resolver_bool_setting(
|
|
settings.get("workflow_hash_metadata_enabled"), True
|
|
):
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"enabled": False,
|
|
"models": [],
|
|
"by_node": {},
|
|
"by_path": {},
|
|
"count": 0,
|
|
}
|
|
)
|
|
|
|
available_models = await asyncio.to_thread(get_model_files, False)
|
|
refs = await asyncio.to_thread(
|
|
analyze_workflow_models,
|
|
workflow_json,
|
|
available_models,
|
|
)
|
|
|
|
by_node = {}
|
|
by_path = {}
|
|
models = []
|
|
seen = set()
|
|
|
|
for ref in refs:
|
|
if not isinstance(ref, dict) or not ref.get("exists"):
|
|
continue
|
|
full_path = str(ref.get("full_path") or "").strip()
|
|
if not full_path:
|
|
continue
|
|
model_info = {
|
|
"path": full_path,
|
|
"filename": get_filename_from_path(full_path),
|
|
"relative_path": ref.get("original_path") or "",
|
|
"category": ref.get("category") or "",
|
|
}
|
|
metadata = get_local_model_hash_metadata(
|
|
full_path, model=model_info
|
|
)
|
|
sha256 = normalize_sha256(metadata.get("sha256"))
|
|
if not sha256:
|
|
continue
|
|
|
|
entry = {
|
|
"node_id": ref.get("node_id"),
|
|
"node_type": ref.get("node_type") or "",
|
|
"widget_index": ref.get("widget_index"),
|
|
"widget_name": ref.get("widget_name") or "",
|
|
"path": ref.get("original_path") or "",
|
|
"filename": get_filename_from_path(
|
|
ref.get("original_path") or full_path
|
|
),
|
|
"category": ref.get("category") or "",
|
|
"sha256": sha256,
|
|
"size": metadata.get("size") or 0,
|
|
}
|
|
entry_key = (
|
|
str(entry.get("node_id")),
|
|
str(entry.get("widget_index")),
|
|
entry.get("path"),
|
|
sha256,
|
|
)
|
|
if entry_key in seen:
|
|
continue
|
|
seen.add(entry_key)
|
|
models.append(entry)
|
|
|
|
path_key = str(entry.get("path") or entry.get("filename") or "")
|
|
if path_key:
|
|
by_path[path_key] = entry
|
|
by_path[get_filename_from_path(path_key)] = entry
|
|
node_key = f"{entry.get('node_id')}:{entry.get('widget_index')}"
|
|
by_node[node_key] = entry
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"enabled": True,
|
|
"models": models,
|
|
"by_node": by_node,
|
|
"by_path": by_path,
|
|
"count": len(models),
|
|
}
|
|
)
|
|
|
|
@routes.post("/model_resolver/local-matches-by-hash")
|
|
@json_api_endpoint("local-matches-by-hash")
|
|
async def local_matches_by_hash(request):
|
|
"""Search local model metadata sidecars for a remote SHA256."""
|
|
data = await request.json()
|
|
sha256 = normalize_sha256(
|
|
data.get("sha256")
|
|
or data.get("hash")
|
|
or data.get("SHA256")
|
|
or ""
|
|
)
|
|
if not sha256:
|
|
return web.json_response(
|
|
{"error": "sha256 is required"}, status=400
|
|
)
|
|
|
|
category = data.get("category", "")
|
|
source = str(
|
|
data.get("source")
|
|
or data.get("hash_lookup_source")
|
|
or "download_source"
|
|
).strip().lower().replace("-", "_")
|
|
filename = (
|
|
data.get("filename")
|
|
or data.get("path")
|
|
or data.get("model_name")
|
|
or ""
|
|
)
|
|
max_matches = to_int(data.get("max_matches"), 20)
|
|
force_rescan = to_bool(data.get("force_rescan"), False)
|
|
|
|
matches = search_local_matches_by_hash(
|
|
sha256,
|
|
category=category or None,
|
|
max_matches=max_matches,
|
|
force_rescan=force_rescan,
|
|
)
|
|
enriched_matches = [
|
|
{
|
|
**match,
|
|
"hash_lookup_source": source or "download_source",
|
|
"hash_lookup_filename": filename,
|
|
"hash_lookup_sha256": sha256,
|
|
}
|
|
for match in matches
|
|
]
|
|
return web.json_response(
|
|
{
|
|
"sha256": sha256,
|
|
"local_hash_matches": enriched_matches,
|
|
"matches": enriched_matches,
|
|
}
|
|
)
|
|
|
|
@routes.post("/model_resolver/open-containing-folder")
|
|
@json_api_endpoint("open-containing-folder")
|
|
async def open_containing_folder(request):
|
|
"""Open Explorer at the folder containing the selected model."""
|
|
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.abspath(os.path.normpath(target_path))
|
|
if not os.path.exists(normalized_path):
|
|
return web.json_response(
|
|
{"error": "path does not exist"}, status=404
|
|
)
|
|
if not is_path_in_configured_model_roots(normalized_path):
|
|
return web.json_response(
|
|
{"error": "path is outside configured model directories"},
|
|
status=403,
|
|
)
|
|
|
|
if os.path.isfile(normalized_path):
|
|
subprocess.Popen(
|
|
["explorer.exe", "/select,", normalized_path],
|
|
shell=False,
|
|
)
|
|
else:
|
|
os.startfile(normalized_path)
|
|
|
|
return web.json_response({"success": True})
|
|
|
|
# cleanup_hash_progress, update_hash_progress, is_hash_progress_cancelled, and mark_hash_progress_cancelled removed
|
|
|
|
from .core.path_utils import HashCalculationCancelled
|
|
|
|
def resolve_hash_file_request(data):
|
|
import os as _os
|
|
|
|
file_path = (
|
|
data.get("file_path")
|
|
or data.get("resolved_path")
|
|
or data.get("path")
|
|
or ""
|
|
)
|
|
if not file_path:
|
|
return "", "file_path is required"
|
|
|
|
normalized_path = _os.path.abspath(_os.path.normpath(file_path))
|
|
if not _os.path.exists(normalized_path) or not _os.path.isfile(normalized_path):
|
|
return "", "file does not exist"
|
|
if not is_path_in_configured_model_roots(normalized_path):
|
|
return "", "file is outside configured model directories"
|
|
return normalized_path, ""
|
|
|
|
def write_calculated_hash_metadata(
|
|
normalized_path,
|
|
metadata_path,
|
|
sha256,
|
|
sha256_source="file",
|
|
):
|
|
import os as _os
|
|
|
|
model_dir = _os.path.abspath(_os.path.dirname(normalized_path))
|
|
resolved_metadata_path = ""
|
|
if metadata_path:
|
|
candidate_path = _os.path.abspath(_os.path.normpath(metadata_path))
|
|
try:
|
|
same_dir = _os.path.commonpath([model_dir, candidate_path]) == model_dir
|
|
except ValueError:
|
|
same_dir = False
|
|
if same_dir:
|
|
resolved_metadata_path = candidate_path
|
|
|
|
if not resolved_metadata_path:
|
|
resolved_metadata_path = get_metadata_sidecar_path(normalized_path)
|
|
|
|
metadata_updated = False
|
|
try:
|
|
metadata = read_json_safe(resolved_metadata_path, {})
|
|
if not isinstance(metadata, dict):
|
|
metadata = {}
|
|
|
|
filename = _os.path.basename(normalized_path)
|
|
stem, _ext = _os.path.splitext(filename)
|
|
hashes = metadata.get("hashes")
|
|
if not isinstance(hashes, dict):
|
|
hashes = {}
|
|
hashes["SHA256"] = sha256
|
|
|
|
metadata["sha256"] = sha256
|
|
metadata["hashes"] = hashes
|
|
metadata["hash_status"] = "completed"
|
|
metadata["sha256_source"] = sha256_source or "file"
|
|
metadata["last_checked_at"] = time.time()
|
|
metadata.setdefault("file_name", stem)
|
|
metadata.setdefault("model_name", stem)
|
|
metadata.setdefault("file_path", normalized_path.replace("\\", "/"))
|
|
try:
|
|
metadata.setdefault("size", _os.path.getsize(normalized_path))
|
|
except Exception:
|
|
pass
|
|
|
|
write_json_atomic(resolved_metadata_path, metadata, indent=2)
|
|
metadata_updated = True
|
|
self.logger.info(
|
|
f"Stored SHA256 and updated metadata: {resolved_metadata_path}"
|
|
)
|
|
except Exception as metadata_error:
|
|
self.logger.warning(
|
|
f"Could not update metadata with calculated SHA256 for {normalized_path}: {metadata_error}"
|
|
)
|
|
|
|
return resolved_metadata_path, metadata_updated
|
|
|
|
def calculate_sha256_with_progress(normalized_path, progress_id=""):
|
|
import os as _os
|
|
from .core.path_utils import calculate_file_sha256 as _calculate_file_sha256_core
|
|
|
|
total_bytes = max(0, _os.path.getsize(normalized_path))
|
|
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="running",
|
|
stage="header",
|
|
message="Checking safetensors header...",
|
|
percent=0,
|
|
bytes_read=0,
|
|
total_bytes=total_bytes,
|
|
)
|
|
|
|
if self.hash_tracker.is_cancelled(progress_id):
|
|
raise HashCalculationCancelled()
|
|
|
|
is_safetensors = normalized_path.lower().endswith(".safetensors")
|
|
if is_safetensors:
|
|
self.logger.info(
|
|
f"Checking safetensors header for SHA256: {normalized_path}"
|
|
)
|
|
else:
|
|
self.logger.info(
|
|
"Skipping safetensors header lookup for non-safetensors file; "
|
|
f"calculating full SHA256: {normalized_path}"
|
|
)
|
|
|
|
header_sha256 = extract_safetensors_header_sha256(normalized_path)
|
|
if header_sha256:
|
|
self.logger.info(
|
|
f"Found SHA256 in safetensors header: {normalized_path}"
|
|
)
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="running",
|
|
stage="header",
|
|
message="Found SHA256 in safetensors header",
|
|
percent=98,
|
|
bytes_read=0,
|
|
total_bytes=total_bytes,
|
|
sha256=header_sha256,
|
|
hash=header_sha256,
|
|
sha256_source="safetensors_header",
|
|
)
|
|
return header_sha256, "safetensors_header"
|
|
|
|
if is_safetensors:
|
|
self.logger.info(
|
|
"No SHA256 found in safetensors header; calculating full "
|
|
f"file SHA256: {normalized_path}"
|
|
)
|
|
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="running",
|
|
stage="hashing",
|
|
message="Calculating SHA256...",
|
|
percent=0,
|
|
bytes_read=0,
|
|
total_bytes=total_bytes,
|
|
)
|
|
|
|
_bytes_read = [0]
|
|
_last_update = [0.0]
|
|
|
|
def on_progress(bytes_read, total_bytes):
|
|
_bytes_read[0] = bytes_read
|
|
now = time.time()
|
|
if now - _last_update[0] >= 0.15 or bytes_read >= total_bytes:
|
|
percent = 98 if total_bytes <= 0 else min(
|
|
98,
|
|
(bytes_read / total_bytes) * 98,
|
|
)
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="running",
|
|
stage="hashing",
|
|
message="Calculating SHA256...",
|
|
percent=percent,
|
|
bytes_read=bytes_read,
|
|
total_bytes=total_bytes,
|
|
)
|
|
_last_update[0] = now
|
|
|
|
def is_cancelled():
|
|
if self.hash_tracker.is_cancelled(progress_id):
|
|
percent = 0 if total_bytes <= 0 else min(
|
|
98,
|
|
(_bytes_read[0] / total_bytes) * 98,
|
|
)
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Hash calculation cancelled",
|
|
percent=percent,
|
|
bytes_read=_bytes_read[0],
|
|
total_bytes=total_bytes,
|
|
)
|
|
return True
|
|
return False
|
|
|
|
sha256 = _calculate_file_sha256_core(
|
|
normalized_path,
|
|
chunk_size=1024 * 1024 * 4,
|
|
on_progress=on_progress,
|
|
is_cancelled=is_cancelled,
|
|
use_safetensors_header=False,
|
|
)
|
|
return sha256, "file"
|
|
|
|
@routes.post("/model_resolver/calculate-file-hash")
|
|
@json_api_endpoint("calculate-file-hash")
|
|
async def calculate_file_hash_route(request):
|
|
"""Calculate SHA256 for a local model and persist it to sidecar metadata."""
|
|
data = await request.json()
|
|
metadata_path = data.get("metadata_path") or ""
|
|
|
|
normalized_path, error = resolve_hash_file_request(data)
|
|
if error == "file_path is required":
|
|
return web.json_response(
|
|
{"error": "file_path is required"}, status=400
|
|
)
|
|
if error:
|
|
status = 403 if error == "file is outside configured model directories" else 404
|
|
return web.json_response(
|
|
{"error": error}, status=status
|
|
)
|
|
|
|
sha256, sha256_source = calculate_sha256_with_progress(normalized_path)
|
|
if not sha256:
|
|
return web.json_response(
|
|
{"error": "could not calculate hash"}, status=500
|
|
)
|
|
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
|
normalized_path,
|
|
metadata_path,
|
|
sha256,
|
|
sha256_source,
|
|
)
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"sha256": sha256,
|
|
"hash": sha256,
|
|
"sha256_source": sha256_source,
|
|
"file_path": normalized_path,
|
|
"metadata_path": resolved_metadata_path,
|
|
"metadata_updated": metadata_updated,
|
|
}
|
|
)
|
|
|
|
@routes.post("/model_resolver/calculate-file-hash/start")
|
|
@json_api_endpoint("calculate-file-hash-start")
|
|
async def calculate_file_hash_start_route(request):
|
|
"""Start SHA256 calculation in a background thread."""
|
|
import uuid
|
|
|
|
data = await request.json()
|
|
metadata_path = data.get("metadata_path") or ""
|
|
normalized_path, error = resolve_hash_file_request(data)
|
|
if error == "file_path is required":
|
|
return web.json_response(
|
|
{"error": "file_path is required"}, status=400
|
|
)
|
|
if error:
|
|
status = 403 if error == "file is outside configured model directories" else 404
|
|
return web.json_response(
|
|
{"error": error}, status=status
|
|
)
|
|
|
|
self.hash_tracker.cleanup()
|
|
progress_id = f"hash_{uuid.uuid4().hex}"
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="queued",
|
|
stage="queued",
|
|
message="Preparing hash calculation...",
|
|
percent=0,
|
|
file_path=normalized_path,
|
|
)
|
|
|
|
def run_hash_task():
|
|
try:
|
|
sha256, sha256_source = calculate_sha256_with_progress(
|
|
normalized_path,
|
|
progress_id=progress_id,
|
|
)
|
|
if self.hash_tracker.is_cancelled(progress_id):
|
|
raise HashCalculationCancelled()
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="running",
|
|
stage="metadata",
|
|
message="Saving metadata...",
|
|
percent=99,
|
|
)
|
|
resolved_metadata_path, metadata_updated = write_calculated_hash_metadata(
|
|
normalized_path,
|
|
metadata_path,
|
|
sha256,
|
|
sha256_source,
|
|
)
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="done",
|
|
stage="done",
|
|
message="Hash calculated",
|
|
percent=100,
|
|
sha256=sha256,
|
|
hash=sha256,
|
|
sha256_source=sha256_source,
|
|
file_path=normalized_path,
|
|
metadata_path=resolved_metadata_path,
|
|
metadata_updated=metadata_updated,
|
|
)
|
|
except HashCalculationCancelled:
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Hash calculation cancelled",
|
|
)
|
|
except Exception as exc:
|
|
self.logger.exception(f"Hash calculation failed for {normalized_path}: {exc}")
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="error",
|
|
stage="error",
|
|
message=str(exc) or "Hash calculation failed",
|
|
percent=100,
|
|
error=str(exc) or "Hash calculation failed",
|
|
)
|
|
|
|
threading.Thread(target=run_hash_task, daemon=True).start()
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"progress_id": progress_id,
|
|
}
|
|
)
|
|
|
|
@routes.get("/model_resolver/calculate-file-hash/progress/{progress_id}")
|
|
@json_api_endpoint("calculate-file-hash-progress")
|
|
async def calculate_file_hash_progress_route(request):
|
|
"""Return progress for a background SHA256 calculation."""
|
|
return get_progress_response(
|
|
self.hash_tracker,
|
|
request,
|
|
not_found_status=404
|
|
)
|
|
|
|
@routes.post("/model_resolver/calculate-file-hash/cancel/{progress_id}")
|
|
@json_api_endpoint("calculate-file-hash-cancel")
|
|
async def calculate_file_hash_cancel_route(request):
|
|
"""Cancel a background SHA256 calculation."""
|
|
return cancel_progress_response(
|
|
self.hash_tracker,
|
|
request,
|
|
cancel_message="Stopping hash calculation..."
|
|
)
|
|
|
|
@routes.get("/model_resolver/models")
|
|
@json_api_endpoint("get_models")
|
|
async def get_models(request):
|
|
"""Get list of all available models."""
|
|
force_rescan = to_bool(
|
|
request.query.get("force") or request.query.get("force_rescan"),
|
|
False
|
|
)
|
|
if force_rescan:
|
|
invalidate_local_hash_match_cache()
|
|
models = get_model_files(force_rescan=force_rescan)
|
|
return web.json_response(models)
|
|
|
|
@routes.post("/model_resolver/metadata-size-audit")
|
|
@json_api_endpoint("metadata-size-audit")
|
|
async def metadata_size_audit_route(request):
|
|
"""Check local model .metadata.json sidecars for stale size values."""
|
|
payload = {}
|
|
try:
|
|
if request.can_read_body:
|
|
payload = await request.json()
|
|
except Exception:
|
|
payload = {}
|
|
if not isinstance(payload, dict):
|
|
payload = {}
|
|
|
|
force_rescan = to_bool(payload.get("force_rescan"), True)
|
|
result = await asyncio.to_thread(
|
|
audit_metadata_sizes,
|
|
force_rescan=force_rescan,
|
|
worker_count=payload.get("worker_count"),
|
|
batch_size=payload.get("batch_size"),
|
|
)
|
|
return web.json_response(result)
|
|
|
|
@routes.get("/model_resolver/metadata-build/capabilities")
|
|
@json_api_endpoint("metadata-build-capabilities")
|
|
async def metadata_build_capabilities_route(request):
|
|
"""Return local CPU and worker limits for the metadata builder."""
|
|
return web.json_response(get_metadata_build_capabilities())
|
|
|
|
@routes.post("/model_resolver/metadata-build/start")
|
|
@json_api_endpoint("metadata-build-start")
|
|
async def metadata_build_start_route(request):
|
|
"""Create missing local metadata sidecars and fill missing SHA256 values."""
|
|
import uuid
|
|
|
|
payload = {}
|
|
try:
|
|
if request.can_read_body:
|
|
payload = await request.json()
|
|
except Exception:
|
|
payload = {}
|
|
if not isinstance(payload, dict):
|
|
payload = {}
|
|
|
|
force_rescan = to_bool(payload.get("force_rescan"), True)
|
|
worker_count = to_int(payload.get("worker_count"), 0)
|
|
self.metadata_builder_progress.cleanup()
|
|
progress_id = f"metadata_build_{uuid.uuid4().hex}"
|
|
self.metadata_builder_progress.update(
|
|
progress_id,
|
|
status="queued",
|
|
stage="queued",
|
|
message="Preparing local metadata build...",
|
|
percent=0,
|
|
current=0,
|
|
total=0,
|
|
requested_worker_count=worker_count,
|
|
)
|
|
|
|
def update_metadata_build_progress(progress_payload):
|
|
if not isinstance(progress_payload, dict):
|
|
return
|
|
data = dict(progress_payload)
|
|
status = data.pop("status", None)
|
|
stage = data.get("stage") or ""
|
|
if not status:
|
|
if stage == "done":
|
|
status = "done"
|
|
elif stage == "cancelled":
|
|
status = "cancelled"
|
|
else:
|
|
status = "running"
|
|
self.metadata_builder_progress.update(
|
|
progress_id,
|
|
status=status,
|
|
**data,
|
|
)
|
|
|
|
def is_metadata_build_cancelled():
|
|
return self.metadata_builder_progress.is_cancelled(progress_id)
|
|
|
|
def run_metadata_build_task():
|
|
try:
|
|
result = build_missing_local_metadata(
|
|
force_rescan=force_rescan,
|
|
worker_count=worker_count,
|
|
progress_callback=update_metadata_build_progress,
|
|
is_cancelled=is_metadata_build_cancelled,
|
|
)
|
|
if result.get("cancelled") or is_metadata_build_cancelled():
|
|
self.metadata_builder_progress.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Metadata build cancelled",
|
|
percent=100,
|
|
active_models=[],
|
|
active_worker_count=0,
|
|
current_model="",
|
|
current_path="",
|
|
bytes_read=0,
|
|
total_bytes=0,
|
|
result=result,
|
|
**result,
|
|
)
|
|
return
|
|
|
|
self.metadata_builder_progress.update(
|
|
progress_id,
|
|
status="done",
|
|
stage="done",
|
|
message="Local metadata build completed.",
|
|
percent=100,
|
|
active_models=[],
|
|
active_worker_count=0,
|
|
current_model="",
|
|
current_path="",
|
|
bytes_read=0,
|
|
total_bytes=0,
|
|
result=result,
|
|
**result,
|
|
)
|
|
except Exception as exc:
|
|
self.logger.exception(f"Metadata build failed: {exc}")
|
|
self.metadata_builder_progress.update(
|
|
progress_id,
|
|
status="error",
|
|
stage="error",
|
|
message=str(exc) or "Metadata build failed",
|
|
percent=100,
|
|
active_models=[],
|
|
active_worker_count=0,
|
|
current_model="",
|
|
current_path="",
|
|
bytes_read=0,
|
|
total_bytes=0,
|
|
error=str(exc) or "Metadata build failed",
|
|
)
|
|
|
|
threading.Thread(target=run_metadata_build_task, daemon=True).start()
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"progress_id": progress_id,
|
|
}
|
|
)
|
|
|
|
@routes.get("/model_resolver/metadata-build/progress/{progress_id}")
|
|
@json_api_endpoint("metadata-build-progress")
|
|
async def metadata_build_progress_route(request):
|
|
"""Return progress for local metadata sidecar creation."""
|
|
return get_progress_response(
|
|
self.metadata_builder_progress,
|
|
request,
|
|
not_found_status=404,
|
|
)
|
|
|
|
@routes.post("/model_resolver/metadata-build/cancel/{progress_id}")
|
|
@json_api_endpoint("metadata-build-cancel")
|
|
async def metadata_build_cancel_route(request):
|
|
"""Cancel local metadata sidecar creation."""
|
|
return cancel_progress_response(
|
|
self.metadata_builder_progress,
|
|
request,
|
|
cancel_message="Stopping metadata build...",
|
|
)
|
|
|
|
@routes.post("/model_resolver/loaded")
|
|
@json_api_endpoint("get_loaded_models")
|
|
async def get_loaded_models(request):
|
|
"""Get all currently loaded models in the workflow."""
|
|
data = await request.json()
|
|
workflow_json = data.get("workflow")
|
|
loaded_id = str(
|
|
data.get("loaded_id") or data.get("progress_id") or ""
|
|
).strip()
|
|
|
|
if not workflow_json:
|
|
return web.json_response(
|
|
{"error": "Workflow JSON is required"}, status=400
|
|
)
|
|
if not isinstance(workflow_json, dict):
|
|
return web.json_response(
|
|
{"error": "Workflow JSON must be an object"}, status=400
|
|
)
|
|
|
|
def update_loaded_progress(
|
|
stage,
|
|
message,
|
|
percent=None,
|
|
status="running",
|
|
current=0,
|
|
total=0,
|
|
**payload,
|
|
):
|
|
if not loaded_id:
|
|
return
|
|
self.loaded_progress.update(
|
|
loaded_id,
|
|
status=status,
|
|
stage=stage,
|
|
message=message,
|
|
percent=percent,
|
|
current=current,
|
|
total=total,
|
|
**payload,
|
|
)
|
|
|
|
def get_workflow_node_count():
|
|
node_count = 0
|
|
nodes = workflow_json.get("nodes", [])
|
|
if isinstance(nodes, list):
|
|
node_count += len(nodes)
|
|
|
|
definitions = workflow_json.get("definitions", {})
|
|
subgraphs = (
|
|
definitions.get("subgraphs", [])
|
|
if isinstance(definitions, dict)
|
|
else []
|
|
)
|
|
if isinstance(subgraphs, list):
|
|
for subgraph in subgraphs:
|
|
if not isinstance(subgraph, dict):
|
|
continue
|
|
subgraph_nodes = subgraph.get("nodes", [])
|
|
if isinstance(subgraph_nodes, list):
|
|
node_count += len(subgraph_nodes)
|
|
return node_count
|
|
|
|
def interpolate_percent(start, end, current, total):
|
|
if not total:
|
|
return start
|
|
try:
|
|
ratio = max(0.0, min(1.0, float(current) / float(total)))
|
|
except (TypeError, ValueError, ZeroDivisionError):
|
|
ratio = 0.0
|
|
return start + ((end - start) * ratio)
|
|
|
|
def build_loaded_models_response():
|
|
from .core.workflow_analyzer import (
|
|
analyze_workflow_models,
|
|
URN_TYPE_MAP,
|
|
)
|
|
|
|
update_loaded_progress(
|
|
"scanning",
|
|
"Scanning local model index...",
|
|
percent=5,
|
|
)
|
|
|
|
# Get available models for existence checking
|
|
available_models = get_model_files()
|
|
|
|
# Create lookup for full paths by filename (with and without extension)
|
|
path_by_filename = {}
|
|
total_local_models = len(available_models)
|
|
update_loaded_progress(
|
|
"indexing",
|
|
f"Indexing {total_local_models} local models...",
|
|
percent=15,
|
|
current=0,
|
|
total=total_local_models,
|
|
)
|
|
for index, model_info in enumerate(available_models, start=1):
|
|
rel_path = model_info.get("relative_path", "")
|
|
if rel_path:
|
|
filename = get_filename_from_path(rel_path)
|
|
path_by_filename[filename] = model_info.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] = model_info.get("path")
|
|
# Add the full relative path as key too
|
|
path_by_filename[rel_path] = model_info.get("path")
|
|
|
|
if (
|
|
total_local_models
|
|
and (index == total_local_models or index % 250 == 0)
|
|
):
|
|
update_loaded_progress(
|
|
"indexing",
|
|
f"Indexing local model {index} of {total_local_models}",
|
|
percent=interpolate_percent(
|
|
15, 25, index, total_local_models
|
|
),
|
|
current=index,
|
|
total=total_local_models,
|
|
)
|
|
|
|
# Also use folder_paths.get_full_path() to get paths
|
|
try:
|
|
import folder_paths
|
|
except Exception:
|
|
folder_paths = None
|
|
|
|
folder_categories = [
|
|
"loras",
|
|
"checkpoints",
|
|
"vae",
|
|
"controlnet",
|
|
"upscale_models",
|
|
]
|
|
if folder_paths:
|
|
for index, category_name in enumerate(folder_categories, start=1):
|
|
update_loaded_progress(
|
|
"indexing",
|
|
f"Reading {category_name} model list...",
|
|
percent=interpolate_percent(
|
|
25, 35, index - 1, len(folder_categories)
|
|
),
|
|
current=index - 1,
|
|
total=len(folder_categories),
|
|
)
|
|
try:
|
|
filenames = folder_paths.get_filename_list(category_name)
|
|
for filename in filenames:
|
|
full_path = folder_paths.get_full_path(
|
|
category_name, filename
|
|
)
|
|
if (
|
|
full_path
|
|
and full_path not in path_by_filename.values()
|
|
):
|
|
path_by_filename[filename] = full_path
|
|
filename_no_ext = (
|
|
filename.rsplit(".", 1)[0]
|
|
if "." in filename
|
|
else filename
|
|
)
|
|
if filename_no_ext not in path_by_filename:
|
|
path_by_filename[filename_no_ext] = full_path
|
|
except Exception:
|
|
pass
|
|
|
|
workflow_node_count = get_workflow_node_count()
|
|
update_loaded_progress(
|
|
"analyzing",
|
|
"Analyzing workflow nodes...",
|
|
percent=35,
|
|
current=0,
|
|
total=workflow_node_count,
|
|
)
|
|
|
|
def update_workflow_analysis_progress(payload):
|
|
progress_payload = dict(payload or {})
|
|
current = progress_payload.pop("current", 0)
|
|
total = progress_payload.pop("total", workflow_node_count)
|
|
stage = progress_payload.pop("stage", "analyzing")
|
|
message = progress_payload.pop(
|
|
"message", "Analyzing workflow nodes..."
|
|
)
|
|
progress_payload.pop("percent", None)
|
|
update_loaded_progress(
|
|
stage,
|
|
message,
|
|
percent=interpolate_percent(35, 78, current, total),
|
|
current=current,
|
|
total=total,
|
|
**progress_payload,
|
|
)
|
|
|
|
# Analyze workflow to get all model references
|
|
all_model_refs = analyze_workflow_models(
|
|
workflow_json,
|
|
available_models=available_models,
|
|
progress_callback=update_workflow_analysis_progress,
|
|
)
|
|
|
|
# Also extract from node.properties.models
|
|
workflow_nodes = workflow_json.get("nodes", [])
|
|
nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else []
|
|
definitions = workflow_json.get("definitions", {})
|
|
subgraphs = (
|
|
definitions.get("subgraphs", [])
|
|
if isinstance(definitions, dict)
|
|
else []
|
|
)
|
|
for subgraph in subgraphs:
|
|
if isinstance(subgraph, dict):
|
|
subgraph_nodes = subgraph.get("nodes", [])
|
|
if isinstance(subgraph_nodes, list):
|
|
nodes.extend(subgraph_nodes)
|
|
|
|
# Collect all loaded models with their values
|
|
loaded_models = []
|
|
total_refs = len(all_model_refs)
|
|
update_loaded_progress(
|
|
"building",
|
|
"Building loaded models list...",
|
|
percent=78,
|
|
current=0,
|
|
total=total_refs,
|
|
)
|
|
|
|
# Process each model reference from analyze_workflow_models
|
|
for index, ref in enumerate(all_model_refs, start=1):
|
|
if total_refs and (index == total_refs or index % 50 == 0):
|
|
update_loaded_progress(
|
|
"building",
|
|
f"Building loaded model {index} of {total_refs}",
|
|
percent=interpolate_percent(78, 94, index, total_refs),
|
|
current=index,
|
|
total=total_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 = get_filename_from_path(original_path)
|
|
strength = None
|
|
|
|
# For standard LoraLoader nodes, strength is in next widget_value
|
|
if (
|
|
node_type in ["LoraLoader", "LoraLoaderModelOnly"]
|
|
and isinstance(widget_index, int)
|
|
):
|
|
# 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 (
|
|
isinstance(widgets_values, list)
|
|
and len(widgets_values) > widget_index + 1
|
|
):
|
|
try:
|
|
strength = float(
|
|
widgets_values[widget_index + 1]
|
|
)
|
|
except (ValueError, TypeError):
|
|
strength = 1.0
|
|
break
|
|
|
|
if ref.get("strength") is not None:
|
|
strength = ref.get("strength")
|
|
|
|
# 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
|
|
total_nodes = len(nodes)
|
|
for node_index, node in enumerate(nodes, start=1):
|
|
if total_nodes and (
|
|
node_index == total_nodes or node_index % 50 == 0
|
|
):
|
|
update_loaded_progress(
|
|
"embedded",
|
|
"Checking embedded model metadata...",
|
|
percent=interpolate_percent(
|
|
94, 98, node_index, total_nodes
|
|
),
|
|
current=node_index,
|
|
total=total_nodes,
|
|
)
|
|
|
|
node_type = node.get("type", "")
|
|
properties = node.get("properties", {})
|
|
if not isinstance(properties, dict):
|
|
properties = {}
|
|
models_list = properties.get("models", [])
|
|
if not isinstance(models_list, list):
|
|
models_list = []
|
|
|
|
for model_info in models_list:
|
|
if isinstance(model_info, dict):
|
|
name = model_info.get("name", "")
|
|
directory = model_info.get("directory", "")
|
|
|
|
if name:
|
|
# Check if this model is already in loaded_models
|
|
existing = next(
|
|
(
|
|
model
|
|
for model in loaded_models
|
|
if model.get("original_path") == name
|
|
),
|
|
None,
|
|
)
|
|
if not existing:
|
|
loaded_models.append(
|
|
{
|
|
"name": get_filename_from_path(name),
|
|
"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,
|
|
}
|
|
)
|
|
|
|
update_loaded_progress(
|
|
"finalizing",
|
|
"Loaded models ready",
|
|
percent=99,
|
|
current=len(loaded_models),
|
|
total=len(loaded_models),
|
|
)
|
|
return {
|
|
"loaded_models": loaded_models,
|
|
"total": len(loaded_models),
|
|
}
|
|
|
|
update_loaded_progress(
|
|
"starting",
|
|
"Preparing loaded model scan...",
|
|
percent=0,
|
|
status="starting",
|
|
current=0,
|
|
total=0,
|
|
)
|
|
|
|
try:
|
|
result = await asyncio.to_thread(build_loaded_models_response)
|
|
except Exception as e:
|
|
update_loaded_progress(
|
|
"error",
|
|
str(e),
|
|
percent=100,
|
|
status="error",
|
|
current=0,
|
|
total=0,
|
|
)
|
|
self.logger.error(
|
|
f"Model Resolver loaded models error: {e}", exc_info=True
|
|
)
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
update_loaded_progress(
|
|
"completed",
|
|
"Loaded models ready",
|
|
percent=100,
|
|
status="completed",
|
|
current=result.get("total", 0),
|
|
total=result.get("total", 0),
|
|
)
|
|
return web.json_response(result)
|
|
|
|
@routes.get("/model_resolver/loaded-progress/{loaded_id}")
|
|
@json_api_endpoint("loaded-progress")
|
|
async def get_loaded_models_progress(request):
|
|
"""Get loaded models inspection progress."""
|
|
return get_progress_response(
|
|
self.loaded_progress,
|
|
request,
|
|
param_name="loaded_id",
|
|
not_found_payload={
|
|
"status": "unknown",
|
|
"stage": "unknown",
|
|
"message": "No loaded models progress available",
|
|
"percent": 0,
|
|
"current": 0,
|
|
"total": 0,
|
|
}
|
|
)
|
|
|
|
# ==================== MODEL METADATA LOOKUP ROUTE ====================
|
|
|
|
@routes.post("/model_resolver/civitai-search")
|
|
@json_api_endpoint("civitai-search")
|
|
async def civitai_search(request):
|
|
"""Fetch model metadata from trusted exact-match sources."""
|
|
data = await request.json()
|
|
filename = data.get("filename", "")
|
|
category = data.get("category", "")
|
|
resolved_path = data.get("resolved_path", "")
|
|
local_only = to_bool(data.get("local_only"), False)
|
|
force_refresh = to_bool(
|
|
data.get("force_refresh") or data.get("force"), False
|
|
)
|
|
provided_hash = normalize_sha256(
|
|
data.get("sha256")
|
|
or data.get("hash")
|
|
or data.get("file_hash")
|
|
or ""
|
|
)
|
|
hf_token = data.get("hf_token", "")
|
|
brave_search_api_key = data.get("brave_search_api_key", "")
|
|
hf_use_brave_fallback = to_bool(
|
|
data.get("hf_use_brave_fallback", True),
|
|
True,
|
|
)
|
|
|
|
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]
|
|
|
|
local_metadata_fields = (
|
|
"author",
|
|
"creator",
|
|
"license",
|
|
"usage_hint",
|
|
"usage_tips",
|
|
"preview_url",
|
|
"metadata_source",
|
|
"from_safetensors_header",
|
|
"local_metadata_available",
|
|
"header_metadata_keys",
|
|
"metadata_summary",
|
|
"base_model_raw",
|
|
"sha256_source",
|
|
"hashes",
|
|
"hash_status",
|
|
)
|
|
|
|
def add_local_metadata_fields(payload, result):
|
|
if not isinstance(payload, dict) or not isinstance(result, dict):
|
|
return payload
|
|
for key in local_metadata_fields:
|
|
value = result.get(key)
|
|
if value in (None, "", [], {}):
|
|
continue
|
|
payload[key] = value
|
|
return payload
|
|
|
|
# Get the file path to hash
|
|
file_path = resolved_path if resolved_path else None
|
|
file_location = ""
|
|
|
|
if not file_path:
|
|
file_path = find_local_file_path(filename, category)
|
|
|
|
if file_path and _os.path.exists(file_path):
|
|
file_location = _os.path.dirname(file_path).replace("\\", "/")
|
|
if file_location and not file_location.endswith("/"):
|
|
file_location += "/"
|
|
|
|
def infer_model_type_from_category(value):
|
|
return normalize_category_to_model_type(value)
|
|
|
|
def build_info_response(
|
|
result=None,
|
|
*,
|
|
metadata_path="",
|
|
metadata_saved=False,
|
|
civitai_checked=False,
|
|
local_payload=False,
|
|
):
|
|
result = result or {}
|
|
size_value = result.get("size")
|
|
if not size_value and file_path and _os.path.exists(file_path):
|
|
try:
|
|
size_value = _os.path.getsize(file_path)
|
|
except Exception:
|
|
size_value = None
|
|
|
|
model_type = (
|
|
result.get("model_type")
|
|
or result.get("type")
|
|
or infer_model_type_from_category(category)
|
|
)
|
|
|
|
response = {
|
|
"filename": result.get("filename") or filename,
|
|
"category": category,
|
|
"file_path": result.get("file_path") or file_path or "",
|
|
"resolved_path": result.get("resolved_path") or file_path or "",
|
|
"metadata_path": metadata_path
|
|
or result.get("metadata_path")
|
|
or "",
|
|
"metadata_saved": bool(metadata_saved),
|
|
"location": result.get("location") or file_location,
|
|
"source": result.get("source") or "",
|
|
"details_source": result.get("details_source")
|
|
or result.get("source")
|
|
or "",
|
|
"url": result.get("url"),
|
|
"version_url": result.get("version_url"),
|
|
"download_url": result.get("download_url"),
|
|
"platform_url": result.get("platform_url"),
|
|
"repo_id": result.get("repo_id"),
|
|
"path": result.get("path"),
|
|
"model_id": result.get("model_id"),
|
|
"model_name": result.get("model_name") or clean_name,
|
|
"model_type": model_type,
|
|
"version_id": result.get("version_id"),
|
|
"version_name": result.get("version_name", ""),
|
|
"sha256": result.get("sha256") or provided_hash,
|
|
"size": size_value,
|
|
"base_model": result.get("base_model"),
|
|
"base_model_source": result.get("base_model_source"),
|
|
"base_model_inferred": bool(result.get("base_model_inferred")),
|
|
"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", ""),
|
|
"from_metadata": bool(result.get("from_metadata")),
|
|
"local_only": bool(local_payload),
|
|
"metadata_checked": bool(civitai_checked),
|
|
"civitai_checked": bool(civitai_checked),
|
|
}
|
|
return add_local_metadata_fields(response, result)
|
|
def extract_result_sha256(result):
|
|
return extract_sha256_from_metadata(result)
|
|
|
|
def result_filename_matches(result):
|
|
if not isinstance(result, dict):
|
|
return False
|
|
expected = _os.path.basename(filename or "").lower()
|
|
candidates = [
|
|
result.get("filename"),
|
|
result.get("path"),
|
|
result.get("file_path"),
|
|
]
|
|
for candidate in candidates:
|
|
basename = _os.path.basename(str(candidate or "")).lower()
|
|
if basename and basename == expected:
|
|
return True
|
|
return False
|
|
|
|
def result_hash_matches(result, *, require_filename=False):
|
|
if not provided_hash or not isinstance(result, dict):
|
|
return False
|
|
result_hash = extract_result_sha256(result)
|
|
if result_hash != provided_hash:
|
|
return False
|
|
if require_filename and not result_filename_matches(result):
|
|
return False
|
|
return True
|
|
|
|
def huggingface_page_url(result):
|
|
try:
|
|
from urllib.parse import quote as _quote
|
|
except Exception:
|
|
_quote = None
|
|
repo_id = str(
|
|
result.get("repo_id") or result.get("repo") or ""
|
|
).strip()
|
|
hf_path = str(
|
|
result.get("path") or result.get("filename") or ""
|
|
).strip()
|
|
if not repo_id or not hf_path:
|
|
return ""
|
|
if _quote:
|
|
hf_path = _quote(hf_path.replace("\\", "/"), safe="/")
|
|
return f"https://huggingface.co/{repo_id}/blob/main/{hf_path}"
|
|
|
|
def prepare_remote_result(result, source_name):
|
|
result = dict(result or {})
|
|
source_name = str(source_name or result.get("source") or "").lower()
|
|
if source_name == "huggingface":
|
|
download_url = result.get("download_url") or result.get("url")
|
|
page_url = result.get("page_url") or huggingface_page_url(result)
|
|
if page_url:
|
|
result["url"] = page_url
|
|
result["version_url"] = page_url
|
|
result["download_url"] = download_url
|
|
result["model_name"] = (
|
|
result.get("model_name")
|
|
or result.get("name")
|
|
or result.get("repo_id")
|
|
or clean_name
|
|
)
|
|
result["model_type"] = (
|
|
result.get("model_type")
|
|
or infer_model_type_from_category(category)
|
|
)
|
|
else:
|
|
result["model_name"] = (
|
|
result.get("model_name")
|
|
or result.get("name")
|
|
or clean_name
|
|
)
|
|
if result.get("url") and not result.get("version_url"):
|
|
result["version_url"] = result.get("url")
|
|
|
|
result["source"] = source_name or result.get("source") or ""
|
|
result["details_source"] = (
|
|
result.get("details_source")
|
|
or source_name
|
|
or result.get("source")
|
|
or ""
|
|
)
|
|
result["file_path"] = file_path or result.get("file_path") or ""
|
|
result["resolved_path"] = (
|
|
file_path or result.get("resolved_path") or ""
|
|
)
|
|
result["location"] = result.get("location") or file_location
|
|
result["filename"] = result.get("filename") or filename
|
|
result["sha256"] = (
|
|
extract_result_sha256(result)
|
|
or provided_hash
|
|
or result.get("sha256")
|
|
)
|
|
if not result.get("size") and file_path and _os.path.exists(file_path):
|
|
try:
|
|
result["size"] = _os.path.getsize(file_path)
|
|
except Exception:
|
|
pass
|
|
return result
|
|
|
|
def result_has_remote_identity(result):
|
|
if not isinstance(result, dict):
|
|
return False
|
|
return bool(
|
|
result.get("url")
|
|
or result.get("version_url")
|
|
or result.get("download_url")
|
|
or result.get("platform_url")
|
|
or result.get("repo_id")
|
|
or result.get("model_id")
|
|
or result.get("version_id")
|
|
or result.get("from_metadata")
|
|
)
|
|
|
|
def remote_link_is_marked_dead(item):
|
|
if not isinstance(item, dict):
|
|
return False
|
|
status = str(item.get("status") or "").lower()
|
|
return bool(
|
|
item.get("deletedAt")
|
|
or item.get("deleted_at")
|
|
or item.get("is_dead")
|
|
or item.get("isDead")
|
|
or item.get("likelyDead")
|
|
or item.get("likely_dead")
|
|
or item.get("dead")
|
|
or status in {"dead", "deleted", "unavailable", "missing"}
|
|
)
|
|
|
|
def collect_result_download_urls(result):
|
|
urls = []
|
|
if remote_link_is_marked_dead(result):
|
|
return urls
|
|
expected_filename = result.get("filename") or filename
|
|
dead_urls = set()
|
|
mirrors = result.get("mirrors") or []
|
|
if not isinstance(mirrors, list):
|
|
mirrors = [mirrors]
|
|
for mirror in mirrors:
|
|
if not isinstance(mirror, dict):
|
|
continue
|
|
if remote_link_is_marked_dead(mirror):
|
|
dead_url = str(mirror.get("url") or "").strip()
|
|
if dead_url.startswith(("http://", "https://")):
|
|
dead_urls.add(dead_url)
|
|
continue
|
|
url = str(mirror.get("url") or "").strip()
|
|
mirror_filename = (
|
|
mirror.get("filename")
|
|
or mirror.get("name")
|
|
or expected_filename
|
|
)
|
|
if (
|
|
looks_like_model_file(url, mirror_filename)
|
|
and url not in urls
|
|
):
|
|
urls.append(url)
|
|
raw_urls = result.get("download_urls") or []
|
|
if not isinstance(raw_urls, list):
|
|
raw_urls = [raw_urls]
|
|
for raw_url in raw_urls:
|
|
url = str(raw_url or "").strip()
|
|
if (
|
|
looks_like_model_file(url, expected_filename)
|
|
and url not in dead_urls
|
|
and url not in urls
|
|
):
|
|
urls.append(url)
|
|
for key in ("download_url", "downloadUrl"):
|
|
url = str(result.get(key) or "").strip()
|
|
if (
|
|
looks_like_model_file(url, expected_filename)
|
|
and url not in dead_urls
|
|
and url not in urls
|
|
):
|
|
urls.append(url)
|
|
return urls
|
|
|
|
def remote_download_url_is_alive(url):
|
|
try:
|
|
import requests
|
|
except Exception:
|
|
return True
|
|
|
|
headers = {
|
|
"User-Agent": "ComfyUI-Model-Resolver/1.0",
|
|
"Accept": "*/*",
|
|
}
|
|
try:
|
|
response = requests.head(
|
|
url,
|
|
headers=headers,
|
|
allow_redirects=True,
|
|
timeout=8,
|
|
)
|
|
try:
|
|
if response.status_code < 400:
|
|
return True
|
|
if response.status_code in {401, 403, 404, 410}:
|
|
return False
|
|
finally:
|
|
response.close()
|
|
except Exception:
|
|
pass
|
|
|
|
try:
|
|
response = requests.get(
|
|
url,
|
|
headers={**headers, "Range": "bytes=0-0"},
|
|
allow_redirects=True,
|
|
stream=True,
|
|
timeout=8,
|
|
)
|
|
try:
|
|
return response.status_code < 400
|
|
finally:
|
|
response.close()
|
|
except Exception:
|
|
return False
|
|
|
|
def civarchive_result_has_live_download(result):
|
|
urls = collect_result_download_urls(result)
|
|
if not urls:
|
|
return False
|
|
for url in urls[:3]:
|
|
if remote_download_url_is_alive(url):
|
|
result["download_url"] = url
|
|
result["download_urls"] = [
|
|
url,
|
|
*[other for other in urls if other != url],
|
|
]
|
|
return True
|
|
return False
|
|
|
|
def save_remote_metadata(result, source_name):
|
|
metadata_path = result.get("metadata_path") or ""
|
|
metadata_saved = False
|
|
if (
|
|
result.get("from_metadata")
|
|
or not file_path
|
|
or not _os.path.exists(file_path)
|
|
):
|
|
return metadata_path, metadata_saved
|
|
try:
|
|
metadata_payload = {
|
|
"source": source_name,
|
|
"details_source": result.get("details_source")
|
|
or source_name,
|
|
"filename": filename,
|
|
"category": category,
|
|
"model_name": result.get("model_name", clean_name),
|
|
"name": result.get("model_name", clean_name),
|
|
"model_type": result.get("model_type", "")
|
|
or result.get("type", ""),
|
|
"type": result.get("model_type", "")
|
|
or result.get("type", ""),
|
|
"model_id": result.get("model_id"),
|
|
"version_id": result.get("version_id"),
|
|
"version_name": result.get("version_name", ""),
|
|
"sha256": result.get("sha256") or provided_hash,
|
|
"size": result.get("size"),
|
|
"base_model": result.get("base_model"),
|
|
"base_model_source": result.get("base_model_source"),
|
|
"base_model_inferred": bool(result.get("base_model_inferred")),
|
|
"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", ""),
|
|
"download_url": result.get("download_url"),
|
|
"source_url": result.get("version_url") or result.get("url"),
|
|
"version_url": result.get("version_url") or result.get("url"),
|
|
"model_url": result.get("url"),
|
|
"url": result.get("version_url") or result.get("url"),
|
|
"platform_url": result.get("platform_url"),
|
|
"repo_id": result.get("repo_id"),
|
|
"path": result.get("path"),
|
|
"path_metadata": {
|
|
"filename": filename,
|
|
"category": category,
|
|
"source": source_name,
|
|
"model_id": result.get("model_id"),
|
|
"version_id": result.get("version_id"),
|
|
"repo_id": result.get("repo_id"),
|
|
"path": result.get("path"),
|
|
},
|
|
}
|
|
add_local_metadata_fields(metadata_payload, result)
|
|
metadata_path = write_lora_manager_metadata(
|
|
file_path,
|
|
metadata_payload,
|
|
category,
|
|
result.get("version_url")
|
|
or result.get("url")
|
|
or result.get("platform_url")
|
|
or result.get("download_url")
|
|
or "",
|
|
) or ""
|
|
metadata_saved = bool(metadata_path)
|
|
if metadata_path:
|
|
result["metadata_path"] = metadata_path
|
|
except Exception as metadata_error:
|
|
self.logger.warning(
|
|
f"{source_name} metadata sidecar save failed: {metadata_error}"
|
|
)
|
|
return metadata_path, metadata_saved
|
|
|
|
if local_only:
|
|
result = None
|
|
metadata_path = ""
|
|
metadata_saved = False
|
|
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,
|
|
local_only=True,
|
|
)
|
|
source_metadata_path = (
|
|
result.get("metadata_path")
|
|
if isinstance(result, dict)
|
|
else ""
|
|
)
|
|
canonical_metadata_path = get_metadata_sidecar_path(file_path)
|
|
if (
|
|
isinstance(result, dict)
|
|
and result.get("from_metadata")
|
|
and source_metadata_path
|
|
and canonical_metadata_path
|
|
and _os.path.abspath(source_metadata_path)
|
|
!= _os.path.abspath(canonical_metadata_path)
|
|
and not _os.path.exists(canonical_metadata_path)
|
|
):
|
|
preview_url = result.get("preview_url") or ""
|
|
if not preview_url:
|
|
model_base_path, _model_ext = _os.path.splitext(file_path)
|
|
for preview_ext in (
|
|
".preview.png",
|
|
".preview.jpg",
|
|
".preview.jpeg",
|
|
".preview.webp",
|
|
):
|
|
preview_path = f"{model_base_path}{preview_ext}"
|
|
if _os.path.exists(preview_path):
|
|
preview_url = preview_path.replace("\\", "/")
|
|
break
|
|
metadata_payload = {
|
|
"source": "metadata_import",
|
|
"details_source": result.get("source") or "metadata",
|
|
"filename": filename,
|
|
"category": category,
|
|
"model_name": result.get("model_name", clean_name),
|
|
"name": result.get("model_name", clean_name),
|
|
"model_type": result.get("model_type", ""),
|
|
"type": result.get("model_type", ""),
|
|
"model_id": result.get("model_id"),
|
|
"version_id": result.get("version_id"),
|
|
"version_name": result.get("version_name", ""),
|
|
"sha256": result.get("sha256"),
|
|
"size": result.get("size"),
|
|
"base_model": result.get("base_model"),
|
|
"base_model_source": result.get("base_model_source"),
|
|
"base_model_inferred": bool(result.get("base_model_inferred")),
|
|
"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", ""),
|
|
"download_url": result.get("download_url"),
|
|
"preview_url": preview_url,
|
|
"source_url": result.get("version_url")
|
|
or result.get("url"),
|
|
"version_url": result.get("version_url")
|
|
or result.get("url"),
|
|
"model_url": result.get("url"),
|
|
"url": result.get("version_url") or result.get("url"),
|
|
"path_metadata": {
|
|
"filename": filename,
|
|
"category": category,
|
|
"source": "metadata_import",
|
|
"model_id": result.get("model_id"),
|
|
"version_id": result.get("version_id"),
|
|
"imported_from": source_metadata_path,
|
|
},
|
|
}
|
|
add_local_metadata_fields(metadata_payload, result)
|
|
metadata_path = write_lora_manager_metadata(
|
|
file_path,
|
|
metadata_payload,
|
|
category,
|
|
result.get("version_url")
|
|
or result.get("url")
|
|
or result.get("download_url")
|
|
or "",
|
|
) or ""
|
|
metadata_saved = bool(metadata_path)
|
|
if metadata_path:
|
|
result["metadata_path"] = metadata_path
|
|
result["metadata_imported_from"] = source_metadata_path
|
|
except Exception as e:
|
|
self.logger.warning(f"Local model info error: {e}")
|
|
return web.json_response(
|
|
build_info_response(
|
|
result,
|
|
metadata_path=metadata_path,
|
|
metadata_saved=metadata_saved,
|
|
local_payload=True,
|
|
civitai_checked=False,
|
|
)
|
|
)
|
|
|
|
# Search remote metadata. CivitAI and CivArchive are hash-only;
|
|
# HuggingFace name results must still confirm the exact SHA.
|
|
if download_available and file_path and _os.path.exists(file_path):
|
|
try:
|
|
from .core.sources.civitai import (
|
|
get_model_info_by_hash,
|
|
get_model_info_for_file,
|
|
)
|
|
|
|
if provided_hash:
|
|
result = get_model_info_by_hash(
|
|
provided_hash,
|
|
use_cache=not force_refresh,
|
|
)
|
|
if result:
|
|
result["file_path"] = file_path
|
|
result["resolved_path"] = file_path
|
|
result["location"] = file_location
|
|
if not result.get("size"):
|
|
try:
|
|
result["size"] = _os.path.getsize(file_path)
|
|
except Exception:
|
|
pass
|
|
if not extract_result_sha256(result):
|
|
result["sha256"] = provided_hash
|
|
elif not force_refresh:
|
|
result = get_model_info_for_file(file_path)
|
|
else:
|
|
result = None
|
|
|
|
if result and (
|
|
result.get("url")
|
|
or result.get("version_url")
|
|
or result.get("from_metadata")
|
|
or result.get("trained_words")
|
|
or result.get("base_model_inferred")
|
|
or result.get("local_metadata_available")
|
|
or result.get("from_safetensors_header")
|
|
) and (
|
|
result.get("from_metadata")
|
|
or not provided_hash
|
|
or result_hash_matches(result)
|
|
):
|
|
result_source = (
|
|
"local"
|
|
if result.get("source") == "local"
|
|
and not result_has_remote_identity(result)
|
|
else "civitai"
|
|
)
|
|
result = prepare_remote_result(result, result_source)
|
|
metadata_path, metadata_saved = save_remote_metadata(
|
|
result,
|
|
result_source,
|
|
)
|
|
|
|
return web.json_response(
|
|
build_info_response(
|
|
result,
|
|
metadata_path=metadata_path,
|
|
metadata_saved=metadata_saved,
|
|
civitai_checked=True,
|
|
)
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"CivitAI search error: {e}")
|
|
|
|
if provided_hash:
|
|
try:
|
|
result = resolve_civarchive_by_hash(
|
|
provided_hash,
|
|
query=filename,
|
|
exact_only=False,
|
|
model_type=infer_model_type_from_category(category),
|
|
)
|
|
if result:
|
|
if not extract_result_sha256(result):
|
|
result["sha256"] = provided_hash
|
|
result = prepare_remote_result(result, "civarchive")
|
|
if result_hash_matches(result) and civarchive_result_has_live_download(result):
|
|
metadata_path, metadata_saved = save_remote_metadata(
|
|
result,
|
|
"civarchive",
|
|
)
|
|
return web.json_response(
|
|
build_info_response(
|
|
result,
|
|
metadata_path=metadata_path,
|
|
metadata_saved=metadata_saved,
|
|
civitai_checked=True,
|
|
)
|
|
)
|
|
if result_hash_matches(result):
|
|
self.logger.info(
|
|
f"CivArchive metadata candidate rejected: "
|
|
f"no live download link for {filename}"
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(f"CivArchive hash lookup error: {e}")
|
|
|
|
if provided_hash:
|
|
hf_attempts = [
|
|
{
|
|
"label": "HuggingFace",
|
|
"use_api_search": True,
|
|
"use_comfy_org_fallback": True,
|
|
"use_brave_fallback": False,
|
|
}
|
|
]
|
|
if hf_use_brave_fallback and brave_search_api_key:
|
|
hf_attempts.append(
|
|
{
|
|
"label": "HuggingFace Brave",
|
|
"use_api_search": False,
|
|
"use_comfy_org_fallback": False,
|
|
"use_brave_fallback": True,
|
|
}
|
|
)
|
|
|
|
for hf_attempt in hf_attempts:
|
|
try:
|
|
result = search_huggingface_for_file(
|
|
filename,
|
|
token=hf_token or None,
|
|
exact_only=True,
|
|
brave_api_key=brave_search_api_key or None,
|
|
use_api_search=hf_attempt["use_api_search"],
|
|
use_comfy_org_fallback=hf_attempt[
|
|
"use_comfy_org_fallback"
|
|
],
|
|
use_brave_fallback=hf_attempt[
|
|
"use_brave_fallback"
|
|
],
|
|
force_refresh=force_refresh,
|
|
)
|
|
if result and result_hash_matches(
|
|
result,
|
|
require_filename=True,
|
|
):
|
|
result = prepare_remote_result(
|
|
result,
|
|
"huggingface",
|
|
)
|
|
metadata_path, metadata_saved = save_remote_metadata(
|
|
result,
|
|
"huggingface",
|
|
)
|
|
return web.json_response(
|
|
build_info_response(
|
|
result,
|
|
metadata_path=metadata_path,
|
|
metadata_saved=metadata_saved,
|
|
civitai_checked=True,
|
|
)
|
|
)
|
|
if result:
|
|
self.logger.info(
|
|
f"{hf_attempt['label']} metadata candidate rejected: "
|
|
f"filename/hash mismatch for {filename}"
|
|
)
|
|
except Exception as e:
|
|
self.logger.warning(
|
|
f"{hf_attempt['label']} metadata lookup error: {e}"
|
|
)
|
|
|
|
# No remote result found. Keep any useful local safetensors
|
|
# header metadata, including base-model fingerprints, model
|
|
# titles, descriptions, tags, trigger words, and author data.
|
|
local_fallback_result = None
|
|
if file_path and _os.path.exists(file_path):
|
|
try:
|
|
from .core.sources.civitai import get_model_info_for_file
|
|
|
|
local_candidate = get_model_info_for_file(
|
|
file_path,
|
|
local_only=True,
|
|
)
|
|
if (
|
|
isinstance(local_candidate, dict)
|
|
and (
|
|
local_candidate.get("base_model")
|
|
or local_candidate.get("local_metadata_available")
|
|
or local_candidate.get("from_safetensors_header")
|
|
)
|
|
):
|
|
local_fallback_result = local_candidate
|
|
except Exception as local_fallback_error:
|
|
self.logger.debug(
|
|
f"Local base model inference fallback failed: {local_fallback_error}"
|
|
)
|
|
|
|
response = build_info_response(
|
|
local_fallback_result,
|
|
civitai_checked=True,
|
|
local_payload=bool(local_fallback_result),
|
|
)
|
|
response["url"] = None
|
|
if force_refresh and file_path and _os.path.exists(file_path):
|
|
try:
|
|
metadata_payload = {
|
|
"filename": filename,
|
|
"category": category,
|
|
"model_name": response.get("model_name") or clean_name,
|
|
"name": response.get("model_name") or clean_name,
|
|
"model_type": response.get("model_type")
|
|
or infer_model_type_from_category(category),
|
|
"sha256": response.get("sha256") or provided_hash,
|
|
"size": response.get("size") or _os.path.getsize(file_path),
|
|
"base_model": response.get("base_model"),
|
|
"base_model_source": response.get("base_model_source"),
|
|
"base_model_inferred": bool(response.get("base_model_inferred")),
|
|
"civitai_deleted": True,
|
|
"civitai_checked": True,
|
|
"remote_metadata_missing": True,
|
|
"source": "local",
|
|
"details_source": response.get("details_source") or "",
|
|
}
|
|
add_local_metadata_fields(metadata_payload, response)
|
|
metadata_path = write_lora_manager_metadata(
|
|
file_path,
|
|
metadata_payload,
|
|
category,
|
|
"",
|
|
) or ""
|
|
response["metadata_path"] = metadata_path
|
|
response["metadata_saved"] = bool(metadata_path)
|
|
except Exception as metadata_error:
|
|
self.logger.warning(
|
|
f"Remote metadata no-match sidecar save failed: {metadata_error}"
|
|
)
|
|
return web.json_response(response)
|
|
|
|
def _select_custom_url_file(files, expected_filename=""):
|
|
"""Pick the most likely downloadable model file from provider details."""
|
|
if not isinstance(files, list):
|
|
files = [files] if files else []
|
|
|
|
candidates = [item for item in files if isinstance(item, dict)]
|
|
if not candidates:
|
|
return {}
|
|
|
|
expected = os.path.basename(str(expected_filename or "")).lower()
|
|
|
|
def file_name(file_info):
|
|
return str(
|
|
file_info.get("name")
|
|
or file_info.get("filename")
|
|
or file_info.get("path")
|
|
or ""
|
|
).strip()
|
|
|
|
def has_download(file_info):
|
|
return bool(
|
|
file_info.get("download_url")
|
|
or file_info.get("downloadUrl")
|
|
or file_info.get("url")
|
|
)
|
|
|
|
if expected:
|
|
for file_info in candidates:
|
|
if os.path.basename(file_name(file_info)).lower() == expected:
|
|
return file_info
|
|
|
|
for file_info in candidates:
|
|
if file_info.get("primary") and has_download(file_info):
|
|
return file_info
|
|
|
|
for file_info in candidates:
|
|
if str(file_info.get("type") or "").lower() == "model" and has_download(file_info):
|
|
return file_info
|
|
|
|
for file_info in candidates:
|
|
if has_download(file_info):
|
|
return file_info
|
|
|
|
return candidates[0]
|
|
|
|
def _custom_result_timestamp():
|
|
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
|
|
def _quote_url_path(value):
|
|
from urllib.parse import quote as _quote
|
|
|
|
return _quote(str(value or "").replace("\\", "/"), safe="/")
|
|
|
|
def _build_civitai_custom_result(details, expected_filename="", api_key=None):
|
|
if not isinstance(details, dict):
|
|
return None
|
|
|
|
selected_version = (
|
|
details.get("selected_version")
|
|
if isinstance(details.get("selected_version"), dict)
|
|
else {}
|
|
)
|
|
version_id = details.get("version_id") or selected_version.get("id")
|
|
model_id = details.get("model_id")
|
|
file_info = _select_custom_url_file(
|
|
selected_version.get("files") or [],
|
|
expected_filename,
|
|
)
|
|
filename = (
|
|
file_info.get("name")
|
|
or file_info.get("filename")
|
|
or os.path.basename(str(expected_filename or ""))
|
|
or details.get("name")
|
|
or f"civitai-{version_id or model_id}"
|
|
)
|
|
download_url = (
|
|
file_info.get("download_url")
|
|
or file_info.get("downloadUrl")
|
|
or (get_civitai_download_url(version_id, api_key) if version_id else "")
|
|
)
|
|
if not download_url:
|
|
return None
|
|
|
|
hashes = (
|
|
file_info.get("hashes")
|
|
if isinstance(file_info.get("hashes"), dict)
|
|
else {}
|
|
)
|
|
sha256 = (
|
|
file_info.get("sha256")
|
|
or file_info.get("hash")
|
|
or hashes.get("SHA256")
|
|
or hashes.get("sha256")
|
|
)
|
|
version_url = (
|
|
details.get("version_url")
|
|
or selected_version.get("url")
|
|
or (
|
|
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
|
|
if model_id and version_id
|
|
else details.get("url")
|
|
)
|
|
)
|
|
return {
|
|
"source": "civitai",
|
|
"details_source": "civitai",
|
|
"model_id": model_id,
|
|
"version_id": version_id,
|
|
"name": details.get("name") or filename,
|
|
"version_name": selected_version.get("name") or "",
|
|
"type": details.get("type") or file_info.get("type") or "",
|
|
"filename": filename,
|
|
"url": version_url or details.get("url"),
|
|
"version_url": version_url or details.get("url"),
|
|
"download_url": download_url,
|
|
"size": file_info.get("size"),
|
|
"base_model": selected_version.get("base_model"),
|
|
"tags": details.get("tags") or [],
|
|
"trained_words": selected_version.get("trained_words") or [],
|
|
"images": details.get("images") or selected_version.get("images") or [],
|
|
"description": selected_version.get("description")
|
|
or details.get("description")
|
|
or "",
|
|
"model_description": details.get("description") or "",
|
|
"sha256": sha256,
|
|
"hashes": hashes,
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
|
|
def _resolve_civitai_version_custom_result(version_id, expected_filename="", api_key=None):
|
|
try:
|
|
import requests
|
|
except Exception:
|
|
requests = None
|
|
|
|
if not requests or not version_id:
|
|
return None
|
|
|
|
headers = {}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
try:
|
|
response = requests.get(
|
|
f"https://civitai.com/api/v1/model-versions/{version_id}",
|
|
headers=headers,
|
|
timeout=20,
|
|
)
|
|
try:
|
|
if response.status_code != 200:
|
|
return None
|
|
data = response.json()
|
|
finally:
|
|
response.close()
|
|
except Exception as e:
|
|
self.logger.warning(
|
|
f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}"
|
|
)
|
|
return None
|
|
|
|
model_block = data.get("model") if isinstance(data.get("model"), dict) else {}
|
|
model_id = data.get("modelId") or data.get("model_id") or model_block.get("id")
|
|
if model_id:
|
|
try:
|
|
details = get_civitai_model_details(model_id, version_id, api_key)
|
|
result = _build_civitai_custom_result(
|
|
details,
|
|
expected_filename=expected_filename,
|
|
api_key=api_key,
|
|
)
|
|
if result:
|
|
return result
|
|
except Exception as e:
|
|
self.logger.warning(
|
|
f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}"
|
|
)
|
|
|
|
file_info = _select_custom_url_file(data.get("files") or [], expected_filename)
|
|
hashes = (
|
|
file_info.get("hashes")
|
|
if isinstance(file_info.get("hashes"), dict)
|
|
else {}
|
|
)
|
|
sha256 = (
|
|
file_info.get("sha256")
|
|
or file_info.get("hash")
|
|
or hashes.get("SHA256")
|
|
or hashes.get("sha256")
|
|
)
|
|
filename = (
|
|
file_info.get("name")
|
|
or file_info.get("filename")
|
|
or os.path.basename(str(expected_filename or ""))
|
|
or f"civitai-{version_id}"
|
|
)
|
|
download_url = (
|
|
file_info.get("downloadUrl")
|
|
or file_info.get("download_url")
|
|
or get_civitai_download_url(version_id, api_key)
|
|
)
|
|
model_url = (
|
|
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
|
|
if model_id
|
|
else f"https://civitai.com/api/download/models/{version_id}"
|
|
)
|
|
return {
|
|
"source": "civitai",
|
|
"details_source": "civitai",
|
|
"model_id": model_id,
|
|
"version_id": version_id,
|
|
"name": model_block.get("name") or data.get("modelName") or filename,
|
|
"version_name": data.get("name") or "",
|
|
"type": model_block.get("type") or data.get("modelType") or file_info.get("type") or "",
|
|
"filename": filename,
|
|
"url": model_url,
|
|
"version_url": model_url,
|
|
"download_url": download_url,
|
|
"size": file_info.get("size")
|
|
or ((file_info.get("sizeKB") or 0) * 1024 if file_info.get("sizeKB") else None),
|
|
"base_model": data.get("baseModel"),
|
|
"tags": model_block.get("tags") or [],
|
|
"trained_words": data.get("trainedWords") or data.get("trained_words") or [],
|
|
"images": data.get("images") or [],
|
|
"description": data.get("description") or "",
|
|
"sha256": sha256,
|
|
"hashes": hashes,
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
|
|
def _build_civarchive_custom_result(details, expected_filename=""):
|
|
if not isinstance(details, dict):
|
|
return None
|
|
|
|
selected_version = (
|
|
details.get("selected_version")
|
|
if isinstance(details.get("selected_version"), dict)
|
|
else {}
|
|
)
|
|
file_info = _select_custom_url_file(
|
|
selected_version.get("files") or [],
|
|
expected_filename,
|
|
)
|
|
download_urls = file_info.get("download_urls") or []
|
|
if not isinstance(download_urls, list):
|
|
download_urls = [download_urls]
|
|
download_url = (
|
|
file_info.get("download_url")
|
|
or file_info.get("downloadUrl")
|
|
or (download_urls[0] if download_urls else "")
|
|
)
|
|
if not download_url:
|
|
return None
|
|
|
|
filename = (
|
|
file_info.get("name")
|
|
or file_info.get("filename")
|
|
or os.path.basename(str(expected_filename or ""))
|
|
or details.get("name")
|
|
or "model"
|
|
)
|
|
hashes = (
|
|
file_info.get("hashes")
|
|
if isinstance(file_info.get("hashes"), dict)
|
|
else {}
|
|
)
|
|
sha256 = (
|
|
file_info.get("sha256")
|
|
or file_info.get("hash")
|
|
or hashes.get("SHA256")
|
|
or hashes.get("sha256")
|
|
)
|
|
version_id = details.get("version_id") or selected_version.get("id")
|
|
model_id = details.get("model_id")
|
|
version_url = (
|
|
details.get("version_url")
|
|
or selected_version.get("url")
|
|
or (
|
|
f"https://civarchive.com/models/{model_id}?modelVersionId={version_id}"
|
|
if model_id and version_id
|
|
else details.get("url")
|
|
)
|
|
)
|
|
return {
|
|
"source": "civarchive",
|
|
"details_source": "civarchive",
|
|
"model_id": model_id,
|
|
"version_id": version_id,
|
|
"name": details.get("name") or filename,
|
|
"version_name": selected_version.get("name") or "",
|
|
"type": details.get("type") or file_info.get("type") or "",
|
|
"filename": filename,
|
|
"url": version_url or details.get("url"),
|
|
"version_url": version_url or details.get("url"),
|
|
"platform_url": details.get("platform_url")
|
|
or selected_version.get("platform_url"),
|
|
"download_url": download_url,
|
|
"download_urls": [url for url in download_urls if url],
|
|
"size": file_info.get("size"),
|
|
"base_model": selected_version.get("base_model"),
|
|
"tags": details.get("tags") or [],
|
|
"trained_words": selected_version.get("trained_words") or [],
|
|
"images": details.get("images") or selected_version.get("images") or [],
|
|
"description": selected_version.get("description")
|
|
or details.get("description")
|
|
or "",
|
|
"sha256": sha256,
|
|
"hash": sha256,
|
|
"hashes": hashes,
|
|
"platform": details.get("platform"),
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
|
|
def _build_huggingface_custom_result(url, expected_filename="", token=None):
|
|
from urllib.parse import unquote
|
|
|
|
parsed = parse_huggingface_url(url)
|
|
if not parsed:
|
|
return None
|
|
|
|
repo_id = parsed.get("repo") or ""
|
|
branch = parsed.get("branch") or "main"
|
|
file_path = unquote(parsed.get("filename") or "")
|
|
if not repo_id or not file_path:
|
|
return None
|
|
|
|
filename = os.path.basename(file_path.replace("\\", "/"))
|
|
download_url = get_huggingface_download_url(repo_id, file_path, branch)
|
|
if not looks_like_model_file(download_url, expected_filename or filename):
|
|
return None
|
|
|
|
headers = {}
|
|
if token:
|
|
headers["Authorization"] = f"Bearer {token}"
|
|
size = fetch_remote_file_size_cached(
|
|
download_url,
|
|
headers=headers,
|
|
timeout=10,
|
|
)
|
|
page_url = (
|
|
f"https://huggingface.co/{repo_id}/blob/{branch}/{_quote_url_path(file_path)}"
|
|
)
|
|
return {
|
|
"source": "huggingface",
|
|
"details_source": "huggingface",
|
|
"repo_id": repo_id,
|
|
"repo": repo_id,
|
|
"path": file_path,
|
|
"filename": filename,
|
|
"name": repo_id,
|
|
"url": download_url,
|
|
"download_url": download_url,
|
|
"page_url": page_url,
|
|
"version_url": page_url,
|
|
"size": size,
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
|
|
def _build_direct_custom_download_result(url, source, expected_filename=""):
|
|
from urllib.parse import urlparse
|
|
|
|
if not looks_like_model_file(url, expected_filename):
|
|
return None
|
|
|
|
parsed = urlparse(url)
|
|
filename = (
|
|
os.path.basename(parsed.path)
|
|
or os.path.basename(str(expected_filename or ""))
|
|
or "model"
|
|
)
|
|
if expected_filename and "." in os.path.basename(str(expected_filename)):
|
|
filename = os.path.basename(str(expected_filename))
|
|
source = str(source or "custom").strip().lower()
|
|
source_label = {
|
|
"civarchive": "CivArchive",
|
|
"civitai": "CivitAI",
|
|
"huggingface": "HuggingFace",
|
|
}.get(source, "Custom URL")
|
|
return {
|
|
"source": source,
|
|
"details_source": source,
|
|
"name": source_label,
|
|
"filename": filename,
|
|
"url": url,
|
|
"version_url": url,
|
|
"download_url": url,
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
|
|
def _collect_custom_url_local_hash_matches(result, category):
|
|
source_key = str(result.get("source") or "custom").strip().lower()
|
|
sha256 = normalize_sha256(extract_sha256_from_metadata(result))
|
|
if not sha256:
|
|
return []
|
|
|
|
try:
|
|
matches = search_local_matches_by_hash(
|
|
sha256,
|
|
category=category or None,
|
|
max_matches=20,
|
|
)
|
|
except Exception as hash_error:
|
|
self.logger.warning(
|
|
f"Custom URL local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}"
|
|
)
|
|
return []
|
|
|
|
return [
|
|
{
|
|
**match,
|
|
"hash_lookup_source": source_key,
|
|
"hash_lookup_filename": result.get("filename")
|
|
or result.get("path")
|
|
or "",
|
|
"hash_lookup_sha256": sha256,
|
|
}
|
|
for match in matches
|
|
]
|
|
|
|
@routes.post("/model_resolver/custom-url")
|
|
@json_api_endpoint("custom-url")
|
|
async def custom_url(request):
|
|
"""Resolve a user-provided provider URL into a normal search result."""
|
|
data = await request.json()
|
|
raw_url = str(data.get("url") or data.get("custom_url") or "").strip()
|
|
if not raw_url:
|
|
return web.json_response(
|
|
{"error": "URL is required"}, status=400
|
|
)
|
|
|
|
if raw_url.startswith("hf://"):
|
|
normalized_url = raw_url
|
|
else:
|
|
from urllib.parse import urlparse
|
|
|
|
parsed_input = urlparse(raw_url)
|
|
if parsed_input.scheme not in {"http", "https"}:
|
|
return web.json_response(
|
|
{"error": "Only http, https, and hf:// URLs are supported"},
|
|
status=400,
|
|
)
|
|
normalized_url = raw_url
|
|
|
|
category = data.get("category") or ""
|
|
expected_filename = (
|
|
data.get("filename")
|
|
or get_filename_from_path(data.get("original_path") or "")
|
|
or ""
|
|
)
|
|
civitai_key = data.get("civitai_key", "")
|
|
hf_token = data.get("hf_token", "")
|
|
|
|
result = None
|
|
source = ""
|
|
try:
|
|
civitai_parsed = parse_civitai_url(normalized_url)
|
|
except Exception:
|
|
civitai_parsed = None
|
|
try:
|
|
civarchive_parsed = parse_civarchive_url(normalized_url)
|
|
except Exception:
|
|
civarchive_parsed = None
|
|
|
|
if civitai_parsed:
|
|
source = "civitai"
|
|
model_id = civitai_parsed.get("model_id")
|
|
version_id = civitai_parsed.get("version_id")
|
|
if model_id:
|
|
details = await asyncio.to_thread(
|
|
get_civitai_model_details,
|
|
model_id,
|
|
version_id,
|
|
civitai_key or None,
|
|
)
|
|
result = _build_civitai_custom_result(
|
|
details,
|
|
expected_filename=expected_filename,
|
|
api_key=civitai_key or None,
|
|
)
|
|
elif version_id:
|
|
result = await asyncio.to_thread(
|
|
_resolve_civitai_version_custom_result,
|
|
version_id,
|
|
expected_filename,
|
|
civitai_key or None,
|
|
)
|
|
if not result and version_id:
|
|
result = {
|
|
"source": "civitai",
|
|
"details_source": "civitai",
|
|
"version_id": version_id,
|
|
"name": expected_filename or f"CivitAI version {version_id}",
|
|
"filename": expected_filename or f"civitai-{version_id}",
|
|
"url": normalized_url,
|
|
"version_url": normalized_url,
|
|
"download_url": get_civitai_download_url(
|
|
version_id,
|
|
civitai_key or None,
|
|
),
|
|
"match_type": "custom_url",
|
|
"custom_url": True,
|
|
}
|
|
elif civarchive_parsed:
|
|
source = "civarchive"
|
|
if civarchive_parsed.get("sha256"):
|
|
result = await asyncio.to_thread(
|
|
resolve_civarchive_by_hash,
|
|
civarchive_parsed.get("sha256"),
|
|
expected_filename,
|
|
False,
|
|
normalize_category_to_model_type(category),
|
|
)
|
|
else:
|
|
model_id = civarchive_parsed.get("model_id")
|
|
version_id = civarchive_parsed.get("version_id")
|
|
if model_id:
|
|
result = await asyncio.to_thread(
|
|
resolve_civarchive_model_version,
|
|
model_id,
|
|
version_id,
|
|
expected_filename or str(model_id),
|
|
False,
|
|
True,
|
|
)
|
|
if not result:
|
|
details = await asyncio.to_thread(
|
|
get_civarchive_model_details,
|
|
model_id,
|
|
version_id,
|
|
True,
|
|
)
|
|
result = _build_civarchive_custom_result(
|
|
details,
|
|
expected_filename=expected_filename,
|
|
)
|
|
if result:
|
|
result = dict(result)
|
|
result["source"] = "civarchive"
|
|
result["details_source"] = "civarchive"
|
|
result["match_type"] = "custom_url"
|
|
result["custom_url"] = True
|
|
else:
|
|
from urllib.parse import urlparse
|
|
|
|
parsed_direct = urlparse(normalized_url)
|
|
direct_host = parsed_direct.netloc.lower()
|
|
if (
|
|
direct_host.endswith("civarchive.com")
|
|
and "/api/download/" in parsed_direct.path
|
|
):
|
|
source = "civarchive"
|
|
result = _build_direct_custom_download_result(
|
|
normalized_url,
|
|
"civarchive",
|
|
expected_filename,
|
|
)
|
|
else:
|
|
result = await asyncio.to_thread(
|
|
_build_huggingface_custom_result,
|
|
normalized_url,
|
|
expected_filename,
|
|
hf_token or None,
|
|
)
|
|
if result:
|
|
source = "huggingface"
|
|
|
|
if not result:
|
|
return web.json_response(
|
|
{
|
|
"error": (
|
|
"Unsupported or unresolved URL. Use a HuggingFace file URL, "
|
|
"CivitAI model/download URL, or CivArchive model/hash URL."
|
|
)
|
|
},
|
|
status=400,
|
|
)
|
|
|
|
result = dict(result)
|
|
result["provided_url"] = normalized_url
|
|
result["url_source"] = "custom"
|
|
result["searched_at"] = result.get("searched_at") or _custom_result_timestamp()
|
|
result.setdefault("category", category)
|
|
if expected_filename and not result.get("filename"):
|
|
result["filename"] = expected_filename
|
|
|
|
if not (result.get("download_url") or result.get("url")):
|
|
return web.json_response(
|
|
{"error": "The URL resolved, but no download URL was found"},
|
|
status=400,
|
|
)
|
|
|
|
source = source or result.get("source") or "custom"
|
|
local_hash_matches = _collect_custom_url_local_hash_matches(
|
|
result,
|
|
category,
|
|
)
|
|
response = {
|
|
"success": True,
|
|
"source": source,
|
|
"result": result,
|
|
"custom": [result],
|
|
"searched_sources": ["custom"],
|
|
"local_hash_matches": local_hash_matches,
|
|
}
|
|
response[source] = result
|
|
return web.json_response(response)
|
|
|
|
@routes.post("/model_resolver/model-details")
|
|
@json_api_endpoint("model-details")
|
|
async def model_details(request):
|
|
"""Return normalized full model details for sources that expose model pages."""
|
|
data = await request.json()
|
|
source = str(data.get("source", "")).strip().lower()
|
|
model_id = data.get("model_id")
|
|
version_id = data.get("version_id")
|
|
civitai_key = data.get("civitai_key", "")
|
|
|
|
if not download_available:
|
|
return web.json_response(
|
|
{"error": "Download providers are not available"}, status=503
|
|
)
|
|
|
|
try:
|
|
model_id = (
|
|
int(model_id)
|
|
if model_id is not None and str(model_id).strip()
|
|
else None
|
|
)
|
|
except (TypeError, ValueError):
|
|
model_id = None
|
|
|
|
try:
|
|
version_id = (
|
|
int(version_id)
|
|
if version_id is not None and str(version_id).strip()
|
|
else None
|
|
)
|
|
except (TypeError, ValueError):
|
|
version_id = None
|
|
|
|
if source == "lora_manager_archive":
|
|
source = "civitai"
|
|
|
|
if source not in {"civitai", "civarchive"}:
|
|
return web.json_response(
|
|
{"error": "Unsupported model details source"}, status=400
|
|
)
|
|
if not model_id:
|
|
return web.json_response(
|
|
{"error": "model_id is required"}, status=400
|
|
)
|
|
|
|
if source == "civitai":
|
|
details = await asyncio.to_thread(
|
|
get_civitai_model_details,
|
|
model_id,
|
|
version_id,
|
|
civitai_key or None,
|
|
)
|
|
else:
|
|
details = await asyncio.to_thread(
|
|
get_civarchive_model_details,
|
|
model_id,
|
|
version_id,
|
|
)
|
|
|
|
if not details:
|
|
return web.json_response(
|
|
{"error": "Model details not found"}, status=404
|
|
)
|
|
|
|
return web.json_response(details)
|
|
|
|
# ==================== DOWNLOAD ROUTES ====================
|
|
|
|
if download_available:
|
|
|
|
from .core.path_utils import (
|
|
get_path_abs as get_local_path_abs,
|
|
get_path_key as get_local_path_key,
|
|
get_path_identity as get_local_path_identity,
|
|
get_comfy_root_path,
|
|
is_path_within as is_local_path_within,
|
|
prefer_local_base_directory,
|
|
dedupe_local_base_directories,
|
|
)
|
|
|
|
# cleanup_search_progress, is_search_progress_cancelled, mark_search_progress_cancelled, and update_search_progress removed
|
|
|
|
@routes.get("/model_resolver/search-progress/{progress_id}")
|
|
@json_api_endpoint("search-progress")
|
|
async def get_search_progress_route(request):
|
|
"""Return live progress for an in-flight source search."""
|
|
return get_progress_response(
|
|
self.search_tracker,
|
|
request,
|
|
not_found_payload={"exists": False},
|
|
found_wrapper=lambda prog: {"exists": True, **prog}
|
|
)
|
|
|
|
@routes.post("/model_resolver/search-cancel/{progress_id}")
|
|
@json_api_endpoint("search-cancel", return_success_on_error=True)
|
|
async def cancel_search_progress_route(request):
|
|
"""Mark an in-flight source search as cancelled."""
|
|
return cancel_progress_response(
|
|
self.search_tracker,
|
|
request,
|
|
cancel_message="Cancelled"
|
|
)
|
|
|
|
@routes.post("/model_resolver/search")
|
|
async def search_sources(request):
|
|
"""Search for model download sources."""
|
|
try:
|
|
class SearchCancelled(BaseException):
|
|
# Progress helpers swallow ordinary callback errors; cancellation
|
|
# must bubble out of source loops and stop follow-up requests.
|
|
pass
|
|
|
|
def raise_if_search_cancelled(source=""):
|
|
if self.search_tracker.is_cancelled(progress_id):
|
|
raise SearchCancelled("Search cancelled")
|
|
|
|
def format_log_value(value):
|
|
if value is None or value == "":
|
|
return None
|
|
if isinstance(value, bool):
|
|
return "yes" if value else "no"
|
|
if isinstance(value, (list, tuple, set)):
|
|
return ",".join(str(item) for item in value)
|
|
|
|
text = str(value)
|
|
if any(char.isspace() for char in text):
|
|
return '"' + text.replace("\\", "\\\\").replace('"', '\\"') + '"'
|
|
return text
|
|
|
|
def format_log_size(value):
|
|
return format_size_bytes(value, include_space=False)
|
|
|
|
def format_log_fields(**fields):
|
|
parts = []
|
|
for key, value in fields.items():
|
|
formatted = format_log_value(value)
|
|
if formatted is not None:
|
|
parts.append(f"{key}={formatted}")
|
|
return " ".join(parts)
|
|
|
|
def normalize_result_extra(extra):
|
|
if not extra:
|
|
return {}
|
|
normalized = dict(extra)
|
|
model_id = normalized.pop("model_id", None)
|
|
version_id = normalized.pop("version_id", None)
|
|
if model_id or version_id:
|
|
normalized["ids"] = (
|
|
f"{model_id}@{version_id}"
|
|
if model_id and version_id
|
|
else model_id or version_id
|
|
)
|
|
if "files_count" in normalized:
|
|
normalized["files"] = normalized.pop("files_count")
|
|
return normalized
|
|
|
|
def format_result_details(result, extra=None):
|
|
if isinstance(result, list):
|
|
return format_log_fields(count=len(result))
|
|
if not isinstance(result, dict):
|
|
fields = {"result": "none" if result is None else result}
|
|
fields.update(normalize_result_extra(extra))
|
|
return format_log_fields(**fields)
|
|
|
|
model_id = result.get("model_id")
|
|
version_id = result.get("version_id")
|
|
ids = (
|
|
f"{model_id}@{version_id}"
|
|
if model_id and version_id
|
|
else model_id or version_id
|
|
)
|
|
fields = {
|
|
"name": result.get("name"),
|
|
"file": result.get("filename") or result.get("path"),
|
|
"match": result.get("match_type"),
|
|
"repo": result.get("repo_id") or result.get("repo"),
|
|
"ids": ids,
|
|
"size": format_log_size(result.get("size")),
|
|
"files": result.get("files_count"),
|
|
}
|
|
fields.update(normalize_result_extra(extra))
|
|
return format_log_fields(**fields)
|
|
|
|
def log_search_result(source_name, result, extra=None):
|
|
details = format_result_details(result, extra)
|
|
if result and (
|
|
not isinstance(result, list) or len(result) > 0
|
|
):
|
|
self.logger.info(f"Search [{source_name}] found {details}")
|
|
else:
|
|
self.logger.info(f"Search [{source_name}] miss {details}")
|
|
|
|
data = await request.json()
|
|
filename = data.get("filename", "")
|
|
category = data.get("category", "")
|
|
base_model_context = data.get("base_model_context", "")
|
|
progress_id = str(data.get("progress_id") or "").strip()
|
|
progress_source = str(data.get("progress_source") or "").strip()
|
|
civitai_candidate_limit = to_int(data.get("civitai_candidate_limit"), 5)
|
|
civitai_candidate_limit = max(
|
|
1, min(civitai_candidate_limit, 20)
|
|
)
|
|
civarchive_candidate_limit = to_int(data.get("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_key = data.get("civitai_key", "")
|
|
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_api_search = data.get(
|
|
"civitai_use_api_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 = to_bool(is_urn_raw, False)
|
|
hf_use_api_search = to_bool(hf_use_api_search, True)
|
|
civitai_use_trpc_search = to_bool(civitai_use_trpc_search, True)
|
|
civitai_use_api_search = to_bool(civitai_use_api_search, True)
|
|
civitai_use_html_fallback = to_bool(civitai_use_html_fallback, True)
|
|
hf_use_comfy_org_fallback = to_bool(hf_use_comfy_org_fallback, True)
|
|
hf_use_brave_fallback = to_bool(hf_use_brave_fallback, 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
|
|
)
|
|
if not progress_source:
|
|
progress_source = (
|
|
next(iter(normalized_sources))
|
|
if len(normalized_sources) == 1
|
|
else "all"
|
|
)
|
|
force_search = to_bool(data.get("force_search"), False)
|
|
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
progress_source,
|
|
"starting",
|
|
"Preparing search",
|
|
8,
|
|
)
|
|
raise_if_search_cancelled(progress_source)
|
|
|
|
if force_search:
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
progress_source,
|
|
"cache",
|
|
"Refreshing search caches",
|
|
12,
|
|
)
|
|
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()
|
|
self.logger.debug(
|
|
"Force search enabled: cleared cache "
|
|
+ format_log_fields(sources=sorted(normalized_sources))
|
|
)
|
|
raise_if_search_cancelled(progress_source)
|
|
|
|
self.logger.info(
|
|
f"Search [{','.join(sorted(normalized_sources))}] request "
|
|
+ format_log_fields(
|
|
file=filename,
|
|
cat=category,
|
|
urn=is_urn,
|
|
ids=(
|
|
f"{data.get('model_id')}@{data.get('version_id')}"
|
|
if data.get("model_id") and data.get("version_id")
|
|
else data.get("model_id")
|
|
or data.get("version_id")
|
|
),
|
|
base=base_model_context,
|
|
force=force_search,
|
|
)
|
|
)
|
|
|
|
results = {
|
|
"popular": None,
|
|
"model_list": None,
|
|
"huggingface": None,
|
|
"civitai": None,
|
|
"civarchive": None,
|
|
"lora_manager_archive": None,
|
|
"local_hash_matches": [],
|
|
"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 mark_any_model_fallback(result):
|
|
if isinstance(result, list):
|
|
return [
|
|
mark_any_model_fallback(item)
|
|
for item in result
|
|
]
|
|
if not isinstance(result, dict):
|
|
return result
|
|
|
|
marked = dict(result)
|
|
marked["any_model_match"] = True
|
|
marked["base_model_fallback"] = True
|
|
marked["requested_base_model"] = base_model_context
|
|
return marked
|
|
|
|
def iter_result_items(result):
|
|
if isinstance(result, list):
|
|
for item in result:
|
|
if isinstance(item, dict):
|
|
yield item
|
|
elif isinstance(result, dict):
|
|
yield result
|
|
|
|
def collect_local_hash_matches(payload):
|
|
matches = []
|
|
seen_match_paths = set()
|
|
seen_hashes = set()
|
|
for source_key in (
|
|
"huggingface",
|
|
"civitai",
|
|
"civarchive",
|
|
"lora_manager_archive",
|
|
"popular",
|
|
"model_list",
|
|
):
|
|
for source_result in iter_result_items(payload.get(source_key)):
|
|
raise_if_search_cancelled(source_key)
|
|
sha256 = extract_sha256_from_metadata(source_result)
|
|
if not sha256:
|
|
continue
|
|
|
|
if sha256 in seen_hashes:
|
|
continue
|
|
seen_hashes.add(sha256)
|
|
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
progress_source,
|
|
"local_hash",
|
|
"Checking local metadata hashes",
|
|
94,
|
|
)
|
|
try:
|
|
hash_matches = search_local_matches_by_hash(
|
|
sha256,
|
|
category=category or None,
|
|
max_matches=20,
|
|
force_rescan=force_search,
|
|
)
|
|
except Exception as hash_error:
|
|
self.logger.warning(
|
|
f"Local metadata hash lookup failed for {source_key}:{sha256}: {hash_error}"
|
|
)
|
|
continue
|
|
|
|
for match in hash_matches:
|
|
model_path = (
|
|
match.get("model", {}).get("path")
|
|
or match.get("path")
|
|
or ""
|
|
)
|
|
path_key = model_path.lower()
|
|
if path_key and path_key in seen_match_paths:
|
|
continue
|
|
if path_key:
|
|
seen_match_paths.add(path_key)
|
|
enriched = {
|
|
**match,
|
|
"hash_lookup_source": source_key,
|
|
"hash_lookup_filename": source_result.get("filename")
|
|
or source_result.get("path")
|
|
or filename,
|
|
"hash_lookup_sha256": sha256,
|
|
}
|
|
matches.append(enriched)
|
|
|
|
if matches:
|
|
self.logger.info(
|
|
"Search local hash matches "
|
|
+ format_log_fields(count=len(matches))
|
|
)
|
|
return matches
|
|
|
|
def make_source_progress_callback(
|
|
source_key,
|
|
percent_min=None,
|
|
percent_max=None,
|
|
):
|
|
def source_progress_callback(payload):
|
|
raise_if_search_cancelled(source_key)
|
|
if not isinstance(payload, dict):
|
|
return
|
|
|
|
progress_payload = dict(payload)
|
|
stage = progress_payload.pop("stage", "running")
|
|
message = progress_payload.pop(
|
|
"message", "Searching..."
|
|
)
|
|
percent = progress_payload.pop("percent", None)
|
|
status = progress_payload.pop("status", "running")
|
|
progress_payload.pop("source", None)
|
|
|
|
if (
|
|
percent is not None
|
|
and percent_min is not None
|
|
and percent_max is not None
|
|
):
|
|
try:
|
|
normalized_percent = max(
|
|
0.0, min(100.0, float(percent))
|
|
)
|
|
percent = percent_min + (
|
|
normalized_percent / 100.0
|
|
) * (percent_max - percent_min)
|
|
except (TypeError, ValueError):
|
|
percent = None
|
|
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
source_key,
|
|
stage,
|
|
message,
|
|
percent,
|
|
status=status,
|
|
**progress_payload,
|
|
)
|
|
|
|
return source_progress_callback
|
|
|
|
def run_source_search(
|
|
source_key,
|
|
search_task_fn,
|
|
initial_stage="query",
|
|
initial_message=None,
|
|
initial_percent=30,
|
|
log_start_fields=None,
|
|
error_handlers=None,
|
|
):
|
|
if initial_message is None:
|
|
initial_message = f"Querying {source_key.capitalize()}"
|
|
raise_if_search_cancelled(source_key)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
source_key,
|
|
initial_stage,
|
|
initial_message,
|
|
initial_percent,
|
|
)
|
|
start_fields = log_start_fields or {"file": filename}
|
|
self.logger.info(
|
|
f"Search [{source_key}] start "
|
|
+ format_log_fields(**start_fields)
|
|
)
|
|
try:
|
|
raise_if_search_cancelled(source_key)
|
|
source_results, source_found = search_task_fn()
|
|
raise_if_search_cancelled(source_key)
|
|
done_messages = {
|
|
"local": "Local database checked",
|
|
"huggingface": "HuggingFace checked",
|
|
"civitai": "CivitAI checked",
|
|
"civarchive": "CivArchive checked",
|
|
"lora_manager_archive": "LoRA Manager archive checked",
|
|
}
|
|
done_msg = done_messages.get(source_key, f"{source_key} checked")
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
source_key,
|
|
"done",
|
|
done_msg,
|
|
92,
|
|
)
|
|
return source_results, source_found
|
|
except Exception as e:
|
|
if error_handlers:
|
|
for exc_type, handler_fn in error_handlers.items():
|
|
if isinstance(e, exc_type):
|
|
return handler_fn(e)
|
|
raise e
|
|
|
|
def execute_search_with_fallback(
|
|
source_key,
|
|
search_fn,
|
|
any_model_label,
|
|
):
|
|
raise_if_search_cancelled(source_key)
|
|
res = search_fn(
|
|
base_model_context or None,
|
|
make_source_progress_callback(source_key),
|
|
)
|
|
raise_if_search_cancelled(source_key)
|
|
log_search_result(source_key, res)
|
|
|
|
if not res and base_model_context:
|
|
raise_if_search_cancelled(source_key)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
source_key,
|
|
"any_model",
|
|
f"Retrying {any_model_label} any model",
|
|
72,
|
|
)
|
|
self.logger.info(
|
|
f"Search [{source_key}] retry any model "
|
|
+ format_log_fields(
|
|
file=filename,
|
|
cat=category,
|
|
base=base_model_context,
|
|
)
|
|
)
|
|
res = search_fn(
|
|
None,
|
|
make_source_progress_callback(source_key, 72, 92),
|
|
)
|
|
raise_if_search_cancelled(source_key)
|
|
log_search_result(f"{source_key}/any_model", res)
|
|
if res:
|
|
res = mark_any_model_fallback(res)
|
|
return res
|
|
|
|
def search_local_sources():
|
|
def task():
|
|
source_results = {"popular": None, "model_list": None}
|
|
source_found = False
|
|
|
|
popular_info = get_popular_model_url(filename)
|
|
log_search_result("popular", popular_info)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"local",
|
|
"model_list",
|
|
"Checking local model database",
|
|
58,
|
|
)
|
|
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
|
|
|
|
return run_source_search(
|
|
"local",
|
|
task,
|
|
initial_stage="popular",
|
|
initial_message="Checking popular models",
|
|
initial_percent=28,
|
|
log_start_fields={"file": filename, "cat": category},
|
|
)
|
|
|
|
def search_huggingface_source_task():
|
|
def task():
|
|
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,
|
|
progress_callback=make_source_progress_callback(
|
|
"huggingface"
|
|
),
|
|
)
|
|
log_search_result("huggingface", hf_result)
|
|
return {"huggingface": hf_result}, bool(hf_result)
|
|
|
|
return run_source_search(
|
|
"huggingface",
|
|
task,
|
|
initial_stage="query",
|
|
initial_message="Querying HuggingFace",
|
|
initial_percent=32,
|
|
log_start_fields={"file": filename},
|
|
)
|
|
|
|
def search_civitai_source_task():
|
|
def task():
|
|
source_results = {"civitai": None}
|
|
source_found = False
|
|
|
|
if is_urn:
|
|
model_id_val = data.get("model_id")
|
|
version_id_val = data.get("version_id")
|
|
|
|
if model_id_val and version_id_val:
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"civitai",
|
|
"urn",
|
|
"Resolving CivitAI URN",
|
|
46,
|
|
)
|
|
model_info = resolve_urn(model_id_val, version_id_val)
|
|
if model_info:
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"civitai",
|
|
"file",
|
|
"Selecting CivitAI file",
|
|
76,
|
|
)
|
|
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_val
|
|
)
|
|
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_val}?modelVersionId={version_id_val}",
|
|
"model_id": model_id_val,
|
|
"version_id": version_id_val,
|
|
"size": primary_file.get("size"),
|
|
"base_model": model_info.get("base_model"),
|
|
"tags": model_info.get("tags", []),
|
|
"sha256": primary_file.get("sha256")
|
|
or (primary_file.get("hashes") or {}).get("SHA256")
|
|
or (primary_file.get("hashes") or {}).get("sha256"),
|
|
"hashes": primary_file.get("hashes") or {},
|
|
"match_type": "exact",
|
|
"confidence": 100.0,
|
|
}
|
|
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_val,
|
|
"version_id": version_id_val,
|
|
},
|
|
)
|
|
elif category:
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"civitai",
|
|
"fallback",
|
|
"Searching CivitAI fallback",
|
|
58,
|
|
)
|
|
self.logger.info(
|
|
"Search [civitai] URN ids missing; falling back"
|
|
)
|
|
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 = execute_search_with_fallback(
|
|
"civitai",
|
|
lambda base_ctx, cb: search_civitai_for_file(
|
|
filename,
|
|
api_key=civitai_key or None,
|
|
model_type=category,
|
|
base_model_context=base_ctx,
|
|
session_token=civitai_session_token or None,
|
|
candidate_limit=civitai_candidate_limit,
|
|
use_trpc_search=civitai_use_trpc_search,
|
|
use_api_search=civitai_use_api_search,
|
|
use_html_fallback=civitai_use_html_fallback,
|
|
progress_callback=cb,
|
|
),
|
|
"CivitAI"
|
|
)
|
|
if civitai_result:
|
|
source_results["civitai"] = civitai_result
|
|
source_found = True
|
|
|
|
return source_results, source_found
|
|
|
|
return run_source_search(
|
|
"civitai",
|
|
task,
|
|
initial_stage="query",
|
|
initial_message="Querying CivitAI",
|
|
initial_percent=30,
|
|
log_start_fields={
|
|
"file": filename,
|
|
"cat": category,
|
|
"urn": is_urn,
|
|
},
|
|
)
|
|
|
|
def search_civarchive_source_task():
|
|
def task():
|
|
source_results = {"civarchive": None}
|
|
source_found = False
|
|
|
|
if is_urn:
|
|
model_id_val = data.get("model_id")
|
|
version_id_val = data.get("version_id")
|
|
if model_id_val and version_id_val:
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"civarchive",
|
|
"urn",
|
|
"Resolving CivArchive version",
|
|
50,
|
|
)
|
|
civarchive_result = resolve_civarchive_model_version(
|
|
model_id_val,
|
|
version_id_val,
|
|
query=filename,
|
|
)
|
|
log_search_result(
|
|
"civarchive/urn",
|
|
civarchive_result,
|
|
{
|
|
"model_id": model_id_val,
|
|
"version_id": version_id_val,
|
|
},
|
|
)
|
|
if civarchive_result:
|
|
source_results["civarchive"] = civarchive_result
|
|
source_found = True
|
|
else:
|
|
log_search_result(
|
|
"civarchive/urn",
|
|
None,
|
|
{
|
|
"model_id": model_id_val,
|
|
"version_id": version_id_val,
|
|
},
|
|
)
|
|
else:
|
|
civarchive_result = execute_search_with_fallback(
|
|
"civarchive",
|
|
lambda base_ctx, cb: search_civarchive_for_file(
|
|
filename,
|
|
model_type=category,
|
|
base_model_context=base_ctx,
|
|
limit=civarchive_candidate_limit,
|
|
progress_callback=cb,
|
|
),
|
|
"CivArchive"
|
|
)
|
|
if civarchive_result:
|
|
source_results["civarchive"] = civarchive_result
|
|
source_found = True
|
|
|
|
return source_results, source_found
|
|
|
|
def handle_civarchive_error(e):
|
|
error_message = f"CivArchive search failed: {e}"
|
|
self.logger.warning(error_message)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
"civarchive",
|
|
"error",
|
|
error_message,
|
|
100,
|
|
status="error",
|
|
)
|
|
return {
|
|
"civarchive": None,
|
|
"source_errors": {"civarchive": error_message},
|
|
}, False
|
|
|
|
return run_source_search(
|
|
"civarchive",
|
|
task,
|
|
initial_stage="query",
|
|
initial_message="Querying CivArchive",
|
|
initial_percent=30,
|
|
log_start_fields={
|
|
"file": filename,
|
|
"cat": category,
|
|
"urn": is_urn,
|
|
},
|
|
error_handlers={CivArchiveSearchError: handle_civarchive_error},
|
|
)
|
|
|
|
def search_lora_manager_archive_source_task():
|
|
def task():
|
|
lora_manager_archive_result = execute_search_with_fallback(
|
|
"lora_manager_archive",
|
|
lambda base_ctx, cb: search_lora_manager_archive_for_file(
|
|
filename,
|
|
model_type=category,
|
|
base_model_context=base_ctx,
|
|
progress_callback=cb,
|
|
),
|
|
"LoRA archive"
|
|
)
|
|
return (
|
|
{
|
|
"lora_manager_archive": lora_manager_archive_result
|
|
},
|
|
bool(lora_manager_archive_result),
|
|
)
|
|
|
|
return run_source_search(
|
|
"lora_manager_archive",
|
|
task,
|
|
initial_stage="query",
|
|
initial_message="Searching LoRA Manager archive",
|
|
initial_percent=36,
|
|
log_start_fields={"file": filename, "cat": category},
|
|
)
|
|
|
|
|
|
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:
|
|
self.logger.debug(
|
|
"Search sources async "
|
|
+ format_log_fields(count=len(search_tasks))
|
|
)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
progress_source,
|
|
"running",
|
|
"Waiting for search sources",
|
|
18,
|
|
)
|
|
raise_if_search_cancelled(progress_source)
|
|
|
|
for source_results, source_found in await asyncio.gather(
|
|
*search_tasks
|
|
):
|
|
raise_if_search_cancelled(progress_source)
|
|
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
|
|
|
|
raise_if_search_cancelled(progress_source)
|
|
results["local_hash_matches"] = collect_local_hash_matches(results)
|
|
raise_if_search_cancelled(progress_source)
|
|
|
|
self.logger.info(
|
|
f"Search [{','.join(results['searched_sources'])}] done "
|
|
+ format_log_fields(
|
|
found=results["found"],
|
|
)
|
|
)
|
|
self.search_tracker.update(
|
|
progress_id,
|
|
progress_source,
|
|
"completed",
|
|
"Search complete",
|
|
100,
|
|
status="completed",
|
|
)
|
|
stamp_search_results(results)
|
|
return web.json_response(results)
|
|
|
|
except SearchCancelled:
|
|
self.search_tracker.update(
|
|
progress_id if "progress_id" in locals() else "",
|
|
progress_source if "progress_source" in locals() else "",
|
|
"cancelled",
|
|
"Cancelled",
|
|
100,
|
|
status="cancelled",
|
|
cancelled=True,
|
|
)
|
|
self.logger.info(
|
|
"Search cancelled "
|
|
+ format_log_fields(
|
|
source=progress_source if "progress_source" in locals() else "",
|
|
progress_id=progress_id if "progress_id" in locals() else "",
|
|
)
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"cancelled": True,
|
|
"found": False,
|
|
"searched_sources": sorted(normalized_sources)
|
|
if "normalized_sources" in locals()
|
|
else [],
|
|
"source_errors": {},
|
|
}
|
|
)
|
|
|
|
except Exception as e:
|
|
self.search_tracker.update(
|
|
progress_id if "progress_id" in locals() else "",
|
|
progress_source if "progress_source" in locals() else "",
|
|
"error",
|
|
str(e),
|
|
100,
|
|
status="error",
|
|
)
|
|
self.logger.exception(f"Model Resolver search error: {e}")
|
|
return web.json_response({"error": str(e)}, status=500)
|
|
|
|
@routes.post("/model_resolver/clear-search-cache")
|
|
@json_api_endpoint("Clear search cache")
|
|
async def clear_search_cache_route(request):
|
|
"""Clear backend search caches after token/settings changes."""
|
|
clear_all_search_caches()
|
|
reload_popular_databases()
|
|
reload_model_list()
|
|
invalidate_model_files_cache()
|
|
invalidate_local_hash_match_cache()
|
|
self.search_result_timestamps.clear()
|
|
self.logger.info("Cleared backend search caches")
|
|
return web.json_response({"success": True, "cleared": "all"})
|
|
|
|
async def _check_credential_helper(request, payload_key, check_func, log_name):
|
|
try:
|
|
data = await request.json()
|
|
val = data.get(payload_key, "")
|
|
result = await asyncio.to_thread(check_func, val)
|
|
return web.json_response(result)
|
|
except Exception as e:
|
|
self.logger.exception(f"{log_name} check 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."""
|
|
return await _check_credential_helper(
|
|
request, "civitai_session_token", check_civitai_session_token, "CivitAI session token"
|
|
)
|
|
|
|
@routes.post("/model_resolver/civitai/api-key/check")
|
|
async def civitai_api_key_check_route(request):
|
|
"""Check whether a CivitAI API key is valid."""
|
|
return await _check_credential_helper(
|
|
request, "civitai_key", check_civitai_api_key, "CivitAI API key"
|
|
)
|
|
|
|
@routes.post("/model_resolver/huggingface/token/check")
|
|
async def huggingface_token_check_route(request):
|
|
"""Check whether a HuggingFace token is valid."""
|
|
return await _check_credential_helper(
|
|
request, "hf_token", check_huggingface_token, "HuggingFace token"
|
|
)
|
|
|
|
@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."""
|
|
return await _check_credential_helper(
|
|
request, "brave_search_api_key", check_brave_search_api_key, "Brave Search API key"
|
|
)
|
|
|
|
@routes.get("/model_resolver/huggingface/author-index/status")
|
|
@json_api_endpoint("HuggingFace author index status")
|
|
async def huggingface_author_index_status_route(request):
|
|
"""Return local HuggingFace author fallback index status."""
|
|
return web.json_response(get_author_fallback_index_status())
|
|
|
|
@routes.post("/model_resolver/huggingface/author-index/refresh")
|
|
@json_api_endpoint("HuggingFace author index refresh")
|
|
async def huggingface_author_index_refresh_route(request):
|
|
"""Refresh HuggingFace author fallback index."""
|
|
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)
|
|
|
|
@routes.get("/model_resolver/model-list/status")
|
|
@json_api_endpoint("Model list status")
|
|
async def model_list_status_route(request):
|
|
"""Return local model-list status and optionally compare with GitHub."""
|
|
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)
|
|
)
|
|
|
|
@routes.post("/model_resolver/model-list/update")
|
|
@json_api_endpoint("Model list update")
|
|
async def model_list_update_route(request):
|
|
"""Download latest ComfyUI-Manager model-list.json."""
|
|
result = await asyncio.to_thread(update_model_list_from_remote)
|
|
clear_all_search_caches()
|
|
self.search_result_timestamps.clear()
|
|
return web.json_response(result)
|
|
|
|
@routes.post("/model_resolver/download")
|
|
@json_api_endpoint("download", return_success_on_error=True)
|
|
async def download_model(request):
|
|
"""Start downloading a model."""
|
|
data = await request.json()
|
|
url = data.get("url", "")
|
|
filename = data.get("filename", "")
|
|
category = data.get("category", "checkpoints")
|
|
category = normalize_download_category(category)
|
|
subfolder = data.get("subfolder", "")
|
|
base_directory = data.get("base_directory", "")
|
|
path_metadata = data.get("path_metadata", {})
|
|
if not isinstance(path_metadata, dict):
|
|
path_metadata = {}
|
|
download_metadata = data.get("download_metadata") or data.get(
|
|
"metadata", {}
|
|
)
|
|
if not isinstance(download_metadata, dict):
|
|
download_metadata = {}
|
|
download_metadata = dict(download_metadata)
|
|
settings = load_resolver_settings()
|
|
if not base_directory:
|
|
base_directory = get_default_root_for_category(category, settings)
|
|
subfolder = resolve_download_subfolder(
|
|
category,
|
|
subfolder,
|
|
path_metadata,
|
|
settings,
|
|
)
|
|
|
|
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])
|
|
filename = sanitize_download_filename(filename)
|
|
|
|
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 or "civitai.red" 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}"
|
|
)
|
|
civitai_session_token = str(
|
|
data.get("civitai_session_token", "") or ""
|
|
).strip()
|
|
if civitai_session_token:
|
|
headers["Cookie"] = (
|
|
f"__Secure-civitai-token={civitai_session_token}"
|
|
)
|
|
|
|
def _first_metadata_value(*values):
|
|
return first_non_empty(*values)
|
|
|
|
def _metadata_int(value):
|
|
return to_int(value)
|
|
|
|
inferred_source = ""
|
|
if "civitai.com" in url:
|
|
inferred_source = "civitai"
|
|
elif "huggingface.co" in url:
|
|
inferred_source = "huggingface"
|
|
|
|
download_metadata.setdefault("filename", filename)
|
|
download_metadata.setdefault("category", category)
|
|
download_metadata.setdefault("download_url", url)
|
|
download_metadata.setdefault("source_url", url)
|
|
download_metadata.setdefault("path_metadata", path_metadata)
|
|
download_metadata.setdefault(
|
|
"source",
|
|
_first_metadata_value(
|
|
download_metadata.get("details_source"),
|
|
path_metadata.get("source"),
|
|
inferred_source,
|
|
),
|
|
)
|
|
|
|
model_id = _metadata_int(
|
|
_first_metadata_value(
|
|
download_metadata.get("model_id"),
|
|
download_metadata.get("modelId"),
|
|
path_metadata.get("model_id"),
|
|
)
|
|
)
|
|
version_id = _metadata_int(
|
|
_first_metadata_value(
|
|
download_metadata.get("version_id"),
|
|
download_metadata.get("versionId"),
|
|
path_metadata.get("version_id"),
|
|
)
|
|
)
|
|
source_name = str(
|
|
_first_metadata_value(
|
|
download_metadata.get("details_source"),
|
|
download_metadata.get("source"),
|
|
)
|
|
).lower()
|
|
try:
|
|
if (
|
|
source_name == "civitai"
|
|
and model_id
|
|
and not download_metadata.get("civitai_details")
|
|
):
|
|
details = await asyncio.to_thread(
|
|
get_civitai_model_details,
|
|
model_id,
|
|
version_id,
|
|
data.get("civitai_key", ""),
|
|
)
|
|
if details:
|
|
download_metadata["civitai_details"] = details
|
|
elif (
|
|
source_name == "civarchive"
|
|
and model_id
|
|
and not download_metadata.get("civitai_details")
|
|
):
|
|
details = await asyncio.to_thread(
|
|
get_civarchive_model_details,
|
|
model_id,
|
|
version_id,
|
|
)
|
|
if details:
|
|
download_metadata["civitai_details"] = details
|
|
except Exception as metadata_error:
|
|
self.logger.warning(
|
|
f"Model metadata lookup failed: {metadata_error}"
|
|
)
|
|
|
|
target_directory = ""
|
|
target_path = ""
|
|
try:
|
|
import os as _download_os
|
|
|
|
target_directory = (
|
|
get_download_directory(category, base_directory) or ""
|
|
)
|
|
if target_directory and subfolder:
|
|
target_directory = _download_os.path.join(
|
|
target_directory, subfolder
|
|
)
|
|
if target_directory:
|
|
target_path = _download_os.path.join(
|
|
target_directory, filename
|
|
)
|
|
except Exception:
|
|
target_directory = ""
|
|
target_path = ""
|
|
|
|
# Start background download
|
|
download_id = start_background_download(
|
|
url=url,
|
|
filename=filename,
|
|
category=category,
|
|
headers=headers if headers else None,
|
|
subfolder=subfolder,
|
|
base_directory=base_directory,
|
|
metadata=download_metadata,
|
|
)
|
|
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"download_id": download_id,
|
|
"filename": filename,
|
|
"category": category,
|
|
"path": target_path,
|
|
"directory": target_directory,
|
|
"download_backend": settings.get("download_backend", "python"),
|
|
}
|
|
)
|
|
|
|
@routes.get("/model_resolver/progress/{download_id}")
|
|
@json_api_endpoint("progress")
|
|
async def get_download_progress(request):
|
|
"""Get progress for a specific download."""
|
|
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
|
|
)
|
|
|
|
@routes.get("/model_resolver/progress")
|
|
@json_api_endpoint("progress")
|
|
async def get_all_downloads_progress(request):
|
|
"""Get progress for all downloads."""
|
|
progress = get_all_progress()
|
|
return web.json_response(progress)
|
|
|
|
@routes.post("/model_resolver/cancel/{download_id}")
|
|
@json_api_endpoint("cancel", return_success_on_error=True)
|
|
async def cancel_download_route(request):
|
|
"""Cancel a download in progress."""
|
|
download_id = request.match_info["download_id"]
|
|
cancel_download(download_id)
|
|
return web.json_response({"success": True})
|
|
|
|
@routes.post("/model_resolver/pause/{download_id}")
|
|
@json_api_endpoint("pause", return_success_on_error=True)
|
|
async def pause_download_route(request):
|
|
"""Pause an aria2 download."""
|
|
download_id = request.match_info["download_id"]
|
|
result = pause_download(download_id)
|
|
status = 200 if result.get("success") else 400
|
|
return web.json_response(result, status=status)
|
|
|
|
@routes.post("/model_resolver/resume/{download_id}")
|
|
@json_api_endpoint("resume", return_success_on_error=True)
|
|
async def resume_download_route(request):
|
|
"""Resume an aria2 download."""
|
|
download_id = request.match_info["download_id"]
|
|
result = resume_download(download_id)
|
|
status = 200 if result.get("success") else 400
|
|
return web.json_response(result, status=status)
|
|
|
|
@routes.get("/model_resolver/aria2/status")
|
|
@routes.post("/model_resolver/aria2/status")
|
|
@json_api_endpoint("aria2 status")
|
|
async def aria2_status_route(request):
|
|
"""Report aria2 availability using saved settings, with optional override via POST."""
|
|
settings = await get_override_settings_from_request(request)
|
|
return web.json_response(await asyncio.to_thread(get_aria2_status, settings))
|
|
|
|
|
|
@routes.post("/model_resolver/aria2/start")
|
|
@json_api_endpoint("aria2 start", return_success_on_error=True)
|
|
async def aria2_start_route(request):
|
|
"""Start the aria2 daemon without starting a download."""
|
|
settings = await get_override_settings_from_request(request)
|
|
result = await asyncio.to_thread(start_aria2_daemon, settings)
|
|
status = 200 if result.get("success") else 400
|
|
return web.json_response(result, status=status)
|
|
|
|
@routes.get("/model_resolver/aria2/stop")
|
|
@routes.post("/model_resolver/aria2/stop")
|
|
@json_api_endpoint("aria2 stop", return_success_on_error=True)
|
|
async def aria2_stop_route(request):
|
|
"""Stop the aria2 daemon started by Model Resolver."""
|
|
result = await asyncio.to_thread(stop_aria2_daemon)
|
|
status = 200 if result.get("success") else 400
|
|
return web.json_response(result, status=status)
|
|
|
|
@routes.post("/model_resolver/aria2/install")
|
|
@json_api_endpoint("aria2 install")
|
|
async def aria2_install_route(request):
|
|
"""Download official aria2c binary and save its path."""
|
|
try:
|
|
payload = await request.json()
|
|
except Exception:
|
|
payload = {}
|
|
force = False
|
|
if isinstance(payload, dict):
|
|
force = to_bool(payload.get("force"), False)
|
|
try:
|
|
install_result = await asyncio.to_thread(install_aria2_engine, force)
|
|
except Aria2InstallError as exc:
|
|
self.logger.warning(f"Model Resolver aria2 install failed: {exc}")
|
|
return web.json_response({"success": False, "error": str(exc)}, status=500)
|
|
settings = await asyncio.to_thread(
|
|
save_resolver_settings,
|
|
{
|
|
"aria2c_path": install_result.get("aria2c_path", ""),
|
|
"download_backend": "aria2",
|
|
},
|
|
)
|
|
install_result["settings"] = settings
|
|
return web.json_response(install_result)
|
|
|
|
@routes.get("/model_resolver/directories")
|
|
@json_api_endpoint("directories")
|
|
async def get_directories(request):
|
|
"""Get available model directories."""
|
|
import folder_paths
|
|
|
|
categories = get_enabled_download_categories(list(folder_paths.folder_names_and_paths.keys()))
|
|
directories = {}
|
|
for cat in categories:
|
|
path = get_download_directory(cat)
|
|
if path:
|
|
directories[cat] = path
|
|
|
|
return web.json_response(directories)
|
|
|
|
@routes.get("/model_resolver/root-directories")
|
|
@json_api_endpoint("root directories")
|
|
async def get_root_directories(request):
|
|
"""Get configured ComfyUI root directories for path settings."""
|
|
import os
|
|
import folder_paths
|
|
|
|
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
|
skip_categories = {"custom_nodes", "configs"}
|
|
categories = get_enabled_download_categories(list(known_categories))
|
|
category_source_keys = {}
|
|
for cat in [*known_categories]:
|
|
folder_key = normalize_download_category(cat)
|
|
if folder_key:
|
|
category_source_keys.setdefault(folder_key, []).append(cat)
|
|
|
|
roots = {}
|
|
settings = load_resolver_settings()
|
|
comfy_root = get_comfy_root_path(folder_paths)
|
|
for cat in categories:
|
|
folder_key = normalize_download_category(cat)
|
|
if folder_key in known_categories:
|
|
raw_candidate_keys = [folder_key]
|
|
else:
|
|
raw_candidate_keys = [
|
|
folder_key,
|
|
*TEMPLATE_KEY_ALIASES.get(folder_key, ()),
|
|
*category_source_keys.get(folder_key, []),
|
|
]
|
|
if folder_key == "ultralytics":
|
|
raw_candidate_keys = [
|
|
candidate_key
|
|
for candidate_key in raw_candidate_keys
|
|
if str(candidate_key or "").strip().lower() != "yolo"
|
|
]
|
|
|
|
candidate_keys = []
|
|
for candidate_key in raw_candidate_keys:
|
|
if (
|
|
candidate_key
|
|
and candidate_key in known_categories
|
|
and candidate_key not in candidate_keys
|
|
):
|
|
candidate_keys.append(candidate_key)
|
|
paths = []
|
|
for candidate_key in candidate_keys:
|
|
paths.extend(folder_paths.get_folder_paths(candidate_key) or [])
|
|
if folder_key == "ultralytics":
|
|
normalized_ultralytics_paths = []
|
|
for path in paths:
|
|
normalized_path = os.path.normpath(str(path or ""))
|
|
basename = os.path.basename(normalized_path).lower()
|
|
if basename == "yolo":
|
|
continue
|
|
parent_dir = os.path.dirname(normalized_path)
|
|
if (
|
|
basename in {"bbox", "segm"}
|
|
and os.path.basename(parent_dir).lower()
|
|
== "ultralytics"
|
|
):
|
|
normalized_path = parent_dir
|
|
normalized_ultralytics_paths.append(normalized_path)
|
|
paths = normalized_ultralytics_paths
|
|
preferred_directory = get_default_root_for_category(
|
|
folder_key, settings
|
|
)
|
|
normalized_paths = dedupe_local_base_directories(
|
|
paths,
|
|
preferred_directory=preferred_directory,
|
|
comfy_root=comfy_root,
|
|
)
|
|
roots[folder_key] = normalized_paths
|
|
|
|
for raw_key in known_categories:
|
|
if (
|
|
not raw_key
|
|
or raw_key in skip_categories
|
|
or raw_key in roots
|
|
):
|
|
continue
|
|
raw_paths = folder_paths.get_folder_paths(raw_key) or []
|
|
roots[raw_key] = dedupe_local_base_directories(
|
|
raw_paths,
|
|
comfy_root=comfy_root,
|
|
)
|
|
|
|
return web.json_response(roots)
|
|
|
|
@routes.get("/model_resolver/path-template-suggestions")
|
|
@json_api_endpoint("path template suggestions")
|
|
async def get_path_template_suggestions(request):
|
|
"""Infer path template presets from existing local model folders."""
|
|
from .core.sources.popular import get_base_models_config
|
|
|
|
force_rescan = request.query.get("force") == "1"
|
|
models = await asyncio.to_thread(get_model_files, force_rescan)
|
|
base_models_config = get_base_models_config()
|
|
suggestions = await asyncio.to_thread(
|
|
infer_download_path_templates,
|
|
models,
|
|
base_models_config,
|
|
)
|
|
return web.json_response(suggestions)
|
|
|
|
@routes.get("/model_resolver/capabilities")
|
|
@json_api_endpoint("capabilities")
|
|
async def get_capabilities(request):
|
|
"""Get optional source capabilities available in this install."""
|
|
return web.json_response(
|
|
{
|
|
"sources": {
|
|
"civarchive": is_civarchive_available(),
|
|
"lora_manager_archive": is_lora_manager_archive_available()
|
|
}
|
|
}
|
|
)
|
|
|
|
@routes.get("/model_resolver/subfolders/{category}")
|
|
@json_api_endpoint("subfolders")
|
|
async def get_subfolders(request):
|
|
"""Get known subfolders for a category using ComfyUI folder_paths."""
|
|
import os
|
|
import folder_paths
|
|
|
|
raw_category = (request.match_info.get("category") or "").strip()
|
|
category = normalize_download_category(raw_category)
|
|
|
|
if not category or category == "unknown":
|
|
return web.json_response([])
|
|
|
|
known_categories = set(folder_paths.folder_names_and_paths.keys())
|
|
folder_keys = get_category_folder_keys(category)
|
|
available_folder_keys = [
|
|
folder_key for folder_key in folder_keys if folder_key in known_categories
|
|
]
|
|
if not available_folder_keys:
|
|
self.logger.debug(
|
|
f"Model Resolver: skipping subfolder lookup for unknown category '{raw_category}' -> '{category}'"
|
|
)
|
|
return web.json_response([])
|
|
|
|
subfolders = {}
|
|
settings = load_resolver_settings()
|
|
comfy_root = get_comfy_root_path(folder_paths)
|
|
preferred_directory = (
|
|
get_default_root_for_category(category, settings)
|
|
or get_download_directory(category)
|
|
or ""
|
|
)
|
|
|
|
def add_subfolder(rel_path, base_dir=""):
|
|
rel_path = os.path.normpath(str(rel_path or "")).replace(
|
|
os.sep, "\\"
|
|
)
|
|
if not rel_path or rel_path == ".":
|
|
return
|
|
base_dir = os.path.abspath(base_dir) if base_dir else ""
|
|
base_identity = (
|
|
get_local_path_identity(base_dir) if base_dir else ""
|
|
)
|
|
key = (rel_path.lower(), base_identity)
|
|
base_label = (
|
|
os.path.basename(os.path.normpath(base_dir))
|
|
if base_dir
|
|
else ""
|
|
)
|
|
current = subfolders.get(key)
|
|
if current and not prefer_local_base_directory(
|
|
base_dir,
|
|
current.get("base_directory", ""),
|
|
preferred_directory,
|
|
comfy_root,
|
|
):
|
|
return
|
|
subfolders[key] = {
|
|
"value": rel_path,
|
|
"label": rel_path,
|
|
"base_label": base_label,
|
|
"base_directory": base_dir,
|
|
}
|
|
|
|
raw_base_dirs = []
|
|
for folder_key in available_folder_keys:
|
|
for base_dir in folder_paths.get_folder_paths(folder_key) or []:
|
|
if not base_dir or not os.path.isdir(base_dir):
|
|
continue
|
|
raw_base_dirs.append(base_dir)
|
|
base_dirs = dedupe_local_base_directories(
|
|
raw_base_dirs,
|
|
preferred_directory=preferred_directory,
|
|
comfy_root=comfy_root,
|
|
)
|
|
|
|
def find_base_dir(full_path):
|
|
if not full_path:
|
|
return ""
|
|
full_path_identity = get_local_path_identity(full_path)
|
|
for base_dir in base_dirs:
|
|
base_identity = get_local_path_identity(base_dir)
|
|
try:
|
|
if (
|
|
os.path.commonpath(
|
|
[full_path_identity, base_identity]
|
|
)
|
|
== base_identity
|
|
):
|
|
return base_dir
|
|
except Exception:
|
|
continue
|
|
return ""
|
|
|
|
for folder_key in available_folder_keys:
|
|
filenames = folder_paths.get_filename_list(folder_key) or []
|
|
for rel_path in filenames:
|
|
if not isinstance(rel_path, str):
|
|
continue
|
|
base_dir = ""
|
|
try:
|
|
base_dir = find_base_dir(
|
|
folder_paths.get_full_path(folder_key, rel_path)
|
|
)
|
|
except Exception:
|
|
base_dir = ""
|
|
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
|
|
add_subfolder(current, base_dir)
|
|
|
|
for base_dir in base_dirs:
|
|
for root, dirs, _files in os.walk(base_dir):
|
|
rel_root = os.path.relpath(root, base_dir)
|
|
for dirname in dirs:
|
|
rel_path = (
|
|
dirname
|
|
if rel_root in ("", ".")
|
|
else os.path.join(rel_root, dirname)
|
|
)
|
|
add_subfolder(rel_path, base_dir)
|
|
|
|
value_counts = {}
|
|
for item in subfolders.values():
|
|
value_key = item.get("value", "").lower()
|
|
value_counts[value_key] = value_counts.get(value_key, 0) + 1
|
|
|
|
response_items = []
|
|
for item in subfolders.values():
|
|
value = item.get("value", "")
|
|
base_label = item.get("base_label", "")
|
|
label = (
|
|
f"{value} ({base_label})"
|
|
if base_label and value_counts.get(value.lower(), 0) > 1
|
|
else value
|
|
)
|
|
response_items.append(
|
|
{
|
|
"value": value,
|
|
"label": label,
|
|
"base_directory": item.get("base_directory", ""),
|
|
}
|
|
)
|
|
|
|
return web.json_response(
|
|
sorted(
|
|
response_items,
|
|
key=lambda item: (
|
|
item.get("value", "").lower(),
|
|
item.get("base_directory", "").lower(),
|
|
),
|
|
)
|
|
)
|
|
|
|
# ==================== SETTINGS (server-side persistence) ====================
|
|
|
|
def _flush_backend_log_handlers():
|
|
for handler in getattr(backend_log_controller, "file_handlers", {}).values():
|
|
try:
|
|
handler.flush()
|
|
except Exception:
|
|
pass
|
|
|
|
def _backend_log_sort_key(path):
|
|
name = os.path.basename(path)
|
|
rotation = 0
|
|
if ".log." in name:
|
|
base_name, rotation_text = name.rsplit(".log.", 1)
|
|
try:
|
|
rotation = int(rotation_text)
|
|
except ValueError:
|
|
rotation = 999
|
|
elif name.endswith(".log"):
|
|
base_name = name[:-4]
|
|
else:
|
|
base_name = name
|
|
return (base_name.lower(), rotation)
|
|
|
|
def _collect_backend_log_files(log_dir):
|
|
if not log_dir or not os.path.isdir(log_dir):
|
|
return []
|
|
|
|
files = []
|
|
for entry in os.listdir(log_dir):
|
|
name = str(entry or "")
|
|
if not name.startswith("azlogs_"):
|
|
continue
|
|
if not (name.endswith(".log") or ".log." in name):
|
|
continue
|
|
|
|
path = os.path.abspath(os.path.join(log_dir, name))
|
|
try:
|
|
if os.path.commonpath([log_dir, path]) != log_dir:
|
|
continue
|
|
except ValueError:
|
|
continue
|
|
if os.path.isfile(path):
|
|
files.append(path)
|
|
return sorted(files, key=_backend_log_sort_key)
|
|
|
|
def _build_backend_log_export():
|
|
_flush_backend_log_handlers()
|
|
raw_log_dir = str(backend_log_controller.config.get("log_dir") or "")
|
|
log_dir = os.path.abspath(raw_log_dir) if raw_log_dir else ""
|
|
exported_at = time.strftime("%Y-%m-%d %H:%M:%S")
|
|
lines = [
|
|
"Model Resolver Backend Logs",
|
|
f"Exported: {exported_at}",
|
|
f"Log directory: {log_dir or 'not configured'}",
|
|
f"File logging: {bool(backend_log_controller.config.get('log_to_file'))}",
|
|
"",
|
|
]
|
|
|
|
log_files = _collect_backend_log_files(log_dir)
|
|
if not log_files:
|
|
lines.append("No backend log files found.")
|
|
lines.append("")
|
|
return "\n".join(lines)
|
|
|
|
for path in log_files:
|
|
name = os.path.basename(path)
|
|
try:
|
|
stat = os.stat(path)
|
|
modified_at = time.strftime(
|
|
"%Y-%m-%d %H:%M:%S",
|
|
time.localtime(stat.st_mtime),
|
|
)
|
|
lines.append(
|
|
f"===== {name} ({stat.st_size} bytes, modified {modified_at}) ====="
|
|
)
|
|
with open(path, "r", encoding="utf-8", errors="replace") as log_file:
|
|
lines.append(log_file.read().rstrip())
|
|
except OSError as exc:
|
|
lines.append(f"===== {name} =====")
|
|
lines.append(f"Could not read log file: {exc}")
|
|
lines.append("")
|
|
|
|
return "\n".join(lines)
|
|
|
|
@routes.get("/model_resolver/logs/backend/export")
|
|
@json_api_endpoint("backend log export")
|
|
async def export_backend_logs_route(request):
|
|
"""Download Model Resolver backend logs as a text file."""
|
|
content = await asyncio.to_thread(_build_backend_log_export)
|
|
filename = f"model_resolver_backend_logs_{time.strftime('%Y%m%d_%H%M%S')}.txt"
|
|
return web.Response(
|
|
text=content,
|
|
content_type="text/plain",
|
|
headers={
|
|
"Content-Disposition": f'attachment; filename="{filename}"',
|
|
"Cache-Control": "no-store",
|
|
},
|
|
)
|
|
|
|
def _log_level_setting(value, default: str = BACKEND_DEFAULT_LOG_LEVEL) -> LogLevel:
|
|
normalized = str(value or default or "INFO").strip().upper()
|
|
if hasattr(LogLevel, normalized):
|
|
return getattr(LogLevel, normalized)
|
|
fallback = str(default or "INFO").strip().upper()
|
|
return getattr(LogLevel, fallback, LogLevel.INFO)
|
|
|
|
def _apply_backend_logging_settings(settings: dict) -> None:
|
|
enabled = resolver_bool_setting(settings.get("backend_logs_enabled"), True)
|
|
level = _log_level_setting(settings.get("backend_log_level"))
|
|
backend_log_controller.set_enabled(enabled)
|
|
backend_log_controller.set_global_level(level)
|
|
|
|
_apply_backend_logging_settings(load_resolver_settings())
|
|
|
|
@routes.get("/model_resolver/settings")
|
|
@json_api_endpoint("settings GET")
|
|
async def get_settings_route(request):
|
|
"""Return persisted settings (API keys, preferences)."""
|
|
data = await asyncio.to_thread(load_resolver_settings)
|
|
return web.json_response(data)
|
|
|
|
@routes.post("/model_resolver/settings")
|
|
@json_api_endpoint("settings POST")
|
|
async def save_settings_route(request):
|
|
"""Persist settings (API keys, preferences) to disk."""
|
|
payload = await request.json()
|
|
if not isinstance(payload, dict):
|
|
return web.json_response({"error": "Expected JSON object"}, status=400)
|
|
settings = await asyncio.to_thread(save_resolver_settings, payload)
|
|
_apply_backend_logging_settings(settings)
|
|
return web.json_response({"success": True})
|
|
|
|
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
|
|
_module_log = create_module_logger(__name__)
|
|
try:
|
|
extension = ModelResolverExtension()
|
|
extension.initialize()
|
|
except Exception as e:
|
|
_module_log.error(
|
|
f"ComfyUI Model Resolver extension initialization failed: {e}", exc_info=True
|
|
)
|