Dariusz L
a8ba99800c
Support both Comfy-Org and Kijai author indexes as known HuggingFace fallbacks, aggregate their status in the backend, refresh the cached author index, and update the frontend UI and tests.
5332 lines
249 KiB
Python
5332 lines
249 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 os
|
|
import sys
|
|
import threading
|
|
import time
|
|
from typing import Any, Dict, Optional
|
|
|
|
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.file_manager import (
|
|
FileManagerError,
|
|
FileManagerUnavailableError,
|
|
UnsupportedFileManagerPlatformError,
|
|
normalize_file_manager_path,
|
|
open_in_file_manager,
|
|
)
|
|
from .core.log_system import LogLevel, create_module_logger
|
|
from .core.log_system import logger as backend_log_controller
|
|
from .core.log_system.config import LOG_LEVEL as BACKEND_DEFAULT_LOG_LEVEL
|
|
from .core.log_system.logger import parse_rotated_log_filename
|
|
from .core.path_utils import (
|
|
MODEL_RESOLVER_METADATA_SCHEMA,
|
|
MODEL_RESOLVER_METADATA_SCHEMA_VERSION,
|
|
get_filename_from_path,
|
|
get_model_resolver_sidecar_path,
|
|
get_safe_model_resolver_sidecar_path,
|
|
)
|
|
|
|
# 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 = (
|
|
"This passive node declares Model Resolver as a workflow dependency and opens "
|
|
"its tools. It does not process images, models, or prompts. To stop adding it "
|
|
"automatically, open Model Resolver, go to Options -> Defaults -> and disable "
|
|
"'Embed opener node'. We recommend keeping it in an unobtrusive part of the "
|
|
"canvas: it helps ComfyUI-Manager identify and offer Model Resolver on other "
|
|
"ComfyUI installations, making the workflow easier to use in the future."
|
|
)
|
|
|
|
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__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS", "WEB_DIRECTORY"]
|
|
|
|
|
|
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 update_from_payload(
|
|
self,
|
|
progress_id: Optional[str],
|
|
progress_payload: Dict[str, Any],
|
|
default_stage: str = "running",
|
|
) -> None:
|
|
if not progress_id or not isinstance(progress_payload, dict):
|
|
return
|
|
data = dict(progress_payload)
|
|
status = data.pop("status", None)
|
|
stage = data.pop("stage", None) or default_stage
|
|
if not status:
|
|
if stage in ("done", "completed"):
|
|
status = "completed"
|
|
elif stage == "cancelled":
|
|
status = "cancelled"
|
|
else:
|
|
status = "running"
|
|
self.update(progress_id, status=status, stage=stage, **data)
|
|
|
|
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 _update_analysis_progress(
|
|
self,
|
|
analysis_id: Optional[str],
|
|
payload: Dict[str, Any],
|
|
) -> None:
|
|
self.analysis_progress.update_from_payload(analysis_id, payload)
|
|
|
|
def _update_metadata_build_progress(
|
|
self,
|
|
progress_id: Optional[str],
|
|
progress_payload: Dict[str, Any]
|
|
) -> None:
|
|
self.metadata_builder_progress.update_from_payload(progress_id, progress_payload)
|
|
|
|
|
|
def _update_loaded_progress(
|
|
self,
|
|
loaded_id: Optional[str],
|
|
stage: str,
|
|
message: str,
|
|
percent: Optional[float] = None,
|
|
status: str = "running",
|
|
current: int = 0,
|
|
total: int = 0,
|
|
**payload,
|
|
) -> None:
|
|
if not loaded_id:
|
|
return
|
|
self.loaded_progress.update_from_payload(
|
|
loaded_id,
|
|
{
|
|
"stage": stage,
|
|
"message": message,
|
|
"percent": percent,
|
|
"status": status,
|
|
"current": current,
|
|
"total": total,
|
|
**payload,
|
|
},
|
|
)
|
|
|
|
def _update_workflow_analysis_progress(
|
|
self,
|
|
loaded_id: Optional[str],
|
|
workflow_node_count: int,
|
|
interpolate_percent_fn,
|
|
payload: Dict[str, Any],
|
|
start_percent: float = 35,
|
|
end_percent: float = 78,
|
|
) -> None:
|
|
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)
|
|
self._update_loaded_progress(
|
|
loaded_id,
|
|
stage,
|
|
message,
|
|
percent=interpolate_percent_fn(
|
|
start_percent,
|
|
end_percent,
|
|
current,
|
|
total,
|
|
),
|
|
current=current,
|
|
total=total,
|
|
**progress_payload,
|
|
)
|
|
|
|
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.custom_nodes import (
|
|
adapt_custom_node_loaded_model,
|
|
should_skip_existing_custom_node_reference,
|
|
)
|
|
from .core.metadata_audit import audit_metadata_sizes
|
|
from .core.metadata_builder import (
|
|
build_missing_local_metadata,
|
|
get_metadata_build_capabilities,
|
|
normalize_metadata_build_mode,
|
|
)
|
|
from .core.network_utils import (
|
|
UnsafeUrlError,
|
|
host_matches_domain,
|
|
request_public_url,
|
|
validate_public_http_url,
|
|
)
|
|
from .core.path_templates import infer_download_path_templates
|
|
from .core.path_utils import (
|
|
HashCalculationCancelled,
|
|
calculate_file_sha256,
|
|
dedupe_local_base_directories,
|
|
extract_safetensors_header_sha256,
|
|
find_external_metadata_sidecar_path,
|
|
get_comfy_root_path,
|
|
get_filename_from_path,
|
|
get_path_abs,
|
|
get_path_identity,
|
|
get_path_key,
|
|
is_path_in_configured_model_roots,
|
|
is_path_within,
|
|
prefer_local_base_directory,
|
|
read_json_safe,
|
|
split_path_segments,
|
|
write_json_atomic,
|
|
)
|
|
from .core.resolver import (
|
|
analyze_and_find_matches,
|
|
apply_resolution,
|
|
get_local_model_hash_metadata,
|
|
invalidate_local_hash_match_cache,
|
|
search_local_matches,
|
|
search_local_matches_by_hash,
|
|
)
|
|
from .core.scanner import find_local_file_path, get_model_files, invalidate_model_files_cache
|
|
from .core.settings import (
|
|
TEMPLATE_KEY_ALIASES,
|
|
get_default_root_for_category,
|
|
get_settings_schema,
|
|
resolve_download_subfolder,
|
|
)
|
|
from .core.settings import (
|
|
bool_setting as resolver_bool_setting,
|
|
)
|
|
from .core.settings import (
|
|
load_settings as load_resolver_settings,
|
|
)
|
|
from .core.settings import (
|
|
save_settings as save_resolver_settings,
|
|
)
|
|
from .core.type_utils import (
|
|
build_search_result,
|
|
extract_sha256_from_metadata,
|
|
fetch_remote_file_size_cached,
|
|
first_non_empty,
|
|
format_size_bytes,
|
|
get_category_folder_keys,
|
|
get_enabled_download_categories,
|
|
looks_like_model_file,
|
|
normalize_category_to_model_type,
|
|
normalize_sha256,
|
|
select_primary_model_file,
|
|
to_bool,
|
|
to_int,
|
|
)
|
|
from .core.workflow_analyzer import (
|
|
NODE_TYPE_MODEL_WIDGET_CATEGORIES,
|
|
URN_TYPE_MAP,
|
|
get_workflow_model_inventory,
|
|
)
|
|
except ImportError as e:
|
|
self.logger.error(f"Model Resolver: Could not import core modules: {e}")
|
|
return False
|
|
|
|
# Import download modules
|
|
try:
|
|
from .core.aria2_installer import Aria2InstallError, install_aria2_engine
|
|
from .core.downloader import (
|
|
cancel_download,
|
|
clear_completed_downloads,
|
|
get_all_progress,
|
|
get_aria2_status,
|
|
get_download_directory,
|
|
get_existing_model_preview_path,
|
|
get_progress,
|
|
is_allowed_model_download_filename,
|
|
normalize_download_category,
|
|
pause_download,
|
|
resume_download,
|
|
sanitize_download_filename,
|
|
start_aria2_daemon,
|
|
start_background_download,
|
|
stop_aria2_daemon,
|
|
write_model_resolver_metadata,
|
|
)
|
|
from .core.sources import clear_all_search_caches
|
|
from .core.sources.civarchive import (
|
|
CivArchiveSearchError,
|
|
build_civarchive_custom_result,
|
|
get_civarchive_model_details,
|
|
is_civarchive_available,
|
|
parse_civarchive_url,
|
|
resolve_civarchive_by_hash,
|
|
resolve_civarchive_model_version,
|
|
search_civarchive_for_file,
|
|
)
|
|
from .core.sources.civarchive import (
|
|
clear_search_cache as clear_civarchive_search_cache,
|
|
)
|
|
from .core.sources.civitai import (
|
|
build_civitai_custom_result,
|
|
check_civitai_api_key,
|
|
check_civitai_session_token,
|
|
get_civitai_download_url,
|
|
get_civitai_model_details,
|
|
get_model_info_by_hash,
|
|
get_model_info_for_file,
|
|
parse_civitai_url,
|
|
resolve_civitai_version_custom_result,
|
|
resolve_urn,
|
|
search_civitai,
|
|
search_civitai_for_file,
|
|
)
|
|
from .core.sources.civitai import (
|
|
clear_search_cache as clear_civitai_search_cache,
|
|
)
|
|
from .core.sources.huggingface import (
|
|
build_huggingface_custom_result,
|
|
check_brave_search_api_key,
|
|
check_huggingface_token,
|
|
get_huggingface_download_url,
|
|
get_huggingface_model_details,
|
|
get_known_author_fallback_indexes_status,
|
|
parse_huggingface_url,
|
|
refresh_known_author_fallback_indexes,
|
|
search_huggingface_for_file,
|
|
)
|
|
from .core.sources.huggingface import (
|
|
clear_search_cache as clear_huggingface_search_cache,
|
|
)
|
|
from .core.sources.lora_manager_archive import (
|
|
clear_search_cache as clear_lora_manager_archive_search_cache,
|
|
)
|
|
from .core.sources.lora_manager_archive import (
|
|
is_lora_manager_archive_available,
|
|
search_lora_manager_archive_for_file,
|
|
)
|
|
from .core.sources.model_list import (
|
|
get_model_list_update_status,
|
|
reload_model_list,
|
|
search_model_list,
|
|
search_model_list_multiple,
|
|
update_model_list_from_remote,
|
|
)
|
|
from .core.sources.popular import (
|
|
get_base_models_config,
|
|
get_base_models_status,
|
|
get_popular_model_url,
|
|
search_popular_models,
|
|
update_base_models_from_remote,
|
|
)
|
|
from .core.sources.popular import (
|
|
reload_databases as reload_popular_databases,
|
|
)
|
|
|
|
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,
|
|
})
|
|
|
|
def run_in_background_thread(
|
|
tracker,
|
|
progress_id,
|
|
task_func,
|
|
on_success,
|
|
on_cancel=None,
|
|
on_error=None,
|
|
error_log_msg="Background task failed",
|
|
):
|
|
def wrapper():
|
|
try:
|
|
result = task_func()
|
|
if tracker.is_cancelled(progress_id):
|
|
if on_cancel:
|
|
on_cancel(result)
|
|
else:
|
|
tracker.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Task cancelled",
|
|
)
|
|
return
|
|
on_success(result)
|
|
except (HashCalculationCancelled, asyncio.CancelledError):
|
|
if on_cancel:
|
|
on_cancel()
|
|
else:
|
|
tracker.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Task cancelled",
|
|
)
|
|
except Exception as exc:
|
|
self.logger.exception(f"{error_log_msg}: {exc}")
|
|
if on_error:
|
|
on_error(exc)
|
|
else:
|
|
tracker.update(
|
|
progress_id,
|
|
status="error",
|
|
stage="error",
|
|
message=str(exc) or "Task failed",
|
|
percent=100,
|
|
error=str(exc) or "Task failed",
|
|
)
|
|
|
|
import threading
|
|
threading.Thread(target=wrapper, daemon=True).start()
|
|
|
|
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 popular base models configuration dictionary.
|
|
|
|
Returns {base_models: [{name, aliases}]} so the frontend
|
|
dropdown can populate correctly via baseModels.base_models.
|
|
"""
|
|
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"
|
|
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."""
|
|
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._update_analysis_progress(
|
|
analysis_id,
|
|
{
|
|
"status": "starting",
|
|
"stage": "starting",
|
|
"message": "Starting analysis...",
|
|
"current": 0,
|
|
"total": 0,
|
|
},
|
|
)
|
|
|
|
def update_analysis_progress(payload):
|
|
self._update_analysis_progress(analysis_id, 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,
|
|
)
|
|
|
|
missing_models = result.get("missing_models", [])
|
|
filtered_missing = []
|
|
for missing in missing_models:
|
|
name = missing.get("name") or missing.get("original_path", "")
|
|
if should_skip_existing_custom_node_reference(missing):
|
|
self.logger.info(
|
|
"Filtered existing custom-node model "
|
|
f"reference: {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
|
|
)
|
|
|
|
normalized_path = os.path.realpath(
|
|
os.path.abspath(os.path.normpath(str(path)))
|
|
)
|
|
if not is_path_in_configured_model_roots(normalized_path):
|
|
return web.json_response(
|
|
{"error": "path is outside configured model directories"},
|
|
status=403,
|
|
)
|
|
|
|
return web.json_response(
|
|
get_local_model_hash_metadata(normalized_path, model=model)
|
|
)
|
|
|
|
@routes.get("/model_resolver/model-preview")
|
|
async def get_model_preview(request):
|
|
"""Serve an adjacent model preview from a configured model directory."""
|
|
model_path = str(request.query.get("path") or "").strip()
|
|
if not model_path:
|
|
return web.Response(text="path is required", status=400)
|
|
|
|
try:
|
|
normalized_path = os.path.realpath(
|
|
os.path.abspath(os.path.normpath(model_path))
|
|
)
|
|
except (OSError, TypeError, ValueError):
|
|
return web.Response(text="invalid model path", status=400)
|
|
if not os.path.isfile(normalized_path):
|
|
return web.Response(text="model file does not exist", status=404)
|
|
if not is_path_in_configured_model_roots(normalized_path):
|
|
return web.Response(
|
|
text="path is outside configured model directories",
|
|
status=403,
|
|
)
|
|
|
|
preview_path = get_existing_model_preview_path(normalized_path)
|
|
if not preview_path:
|
|
return web.Response(text="preview not found", status=404)
|
|
if not is_path_in_configured_model_roots(preview_path):
|
|
return web.Response(
|
|
text="preview is outside configured model directories",
|
|
status=403,
|
|
)
|
|
|
|
return web.FileResponse(
|
|
preview_path,
|
|
headers={"Cache-Control": "private, no-cache"},
|
|
)
|
|
|
|
@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,
|
|
}
|
|
)
|
|
|
|
inventory = await asyncio.to_thread(
|
|
get_workflow_model_inventory,
|
|
workflow_json,
|
|
)
|
|
refs = inventory["model_refs"]
|
|
|
|
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):
|
|
"""Reveal a file or open a directory in the host file manager."""
|
|
try:
|
|
data = await request.json()
|
|
except Exception as exc:
|
|
return web.json_response(
|
|
{"success": False, "error": f"invalid JSON body: {exc}"},
|
|
status=400,
|
|
)
|
|
if not isinstance(data, dict):
|
|
return web.json_response(
|
|
{"success": False, "error": "JSON body must be an object"},
|
|
status=400,
|
|
)
|
|
target_path = data.get("path", "")
|
|
|
|
try:
|
|
normalized_path = normalize_file_manager_path(target_path)
|
|
except ValueError as exc:
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)}, status=400
|
|
)
|
|
|
|
if not os.path.exists(normalized_path):
|
|
return web.json_response(
|
|
{"success": False, "error": "path does not exist"}, status=404
|
|
)
|
|
if not is_path_in_configured_model_roots(normalized_path):
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": "path is outside configured model directories",
|
|
},
|
|
status=403,
|
|
)
|
|
|
|
try:
|
|
result = await asyncio.to_thread(
|
|
open_in_file_manager, normalized_path
|
|
)
|
|
except FileNotFoundError:
|
|
return web.json_response(
|
|
{"success": False, "error": "path no longer exists"},
|
|
status=404,
|
|
)
|
|
except UnsupportedFileManagerPlatformError as exc:
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)}, status=501
|
|
)
|
|
except FileManagerUnavailableError as exc:
|
|
return web.json_response(
|
|
{
|
|
"success": False,
|
|
"error": (
|
|
f"{exc} The ComfyUI host may be running without "
|
|
"a graphical desktop session."
|
|
),
|
|
},
|
|
status=503,
|
|
)
|
|
except FileManagerError as exc:
|
|
return web.json_response(
|
|
{"success": False, "error": str(exc)}, status=500
|
|
)
|
|
|
|
return web.json_response({"success": True, **result})
|
|
|
|
# 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.realpath(
|
|
_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,
|
|
sha256,
|
|
sha256_source="file",
|
|
):
|
|
import os as _os
|
|
|
|
resolved_metadata_path = get_safe_model_resolver_sidecar_path(
|
|
normalized_path
|
|
)
|
|
|
|
metadata_updated = False
|
|
try:
|
|
metadata = read_json_safe(resolved_metadata_path, {})
|
|
if not isinstance(metadata, dict):
|
|
metadata = {}
|
|
|
|
filename = get_filename_from_path(normalized_path)
|
|
stem, _ext = _os.path.splitext(filename)
|
|
hashes = metadata.get("hashes")
|
|
if not isinstance(hashes, dict):
|
|
hashes = {}
|
|
hashes["SHA256"] = sha256
|
|
|
|
metadata["schema"] = MODEL_RESOLVER_METADATA_SCHEMA
|
|
metadata["schema_version"] = (
|
|
MODEL_RESOLVER_METADATA_SCHEMA_VERSION
|
|
)
|
|
metadata["managed_by"] = MODEL_RESOLVER_METADATA_SCHEMA
|
|
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()
|
|
|
|
hash_source = ["file"]
|
|
stage_transitioned = [False]
|
|
_bytes_read = [0]
|
|
_last_update = [0.0]
|
|
|
|
def on_progress(bytes_read, total_bytes):
|
|
_bytes_read[0] = bytes_read
|
|
if not stage_transitioned[0]:
|
|
stage_transitioned[0] = True
|
|
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,
|
|
)
|
|
|
|
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
|
|
|
|
def on_hash_source(source_name):
|
|
hash_source[0] = source_name
|
|
if source_name == "safetensors_header":
|
|
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_source="safetensors_header",
|
|
)
|
|
|
|
sha256 = _calculate_file_sha256_core(
|
|
normalized_path,
|
|
chunk_size=1024 * 1024 * 4,
|
|
on_progress=on_progress,
|
|
is_cancelled=is_cancelled,
|
|
use_safetensors_header=True,
|
|
on_hash_source=on_hash_source,
|
|
)
|
|
|
|
if is_cancelled():
|
|
raise HashCalculationCancelled()
|
|
|
|
return sha256, hash_source[0]
|
|
|
|
@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()
|
|
|
|
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,
|
|
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()
|
|
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():
|
|
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,
|
|
sha256,
|
|
sha256_source,
|
|
)
|
|
return sha256, sha256_source, resolved_metadata_path, metadata_updated
|
|
|
|
def on_success(res):
|
|
sha256, sha256_source, resolved_metadata_path, metadata_updated = res
|
|
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,
|
|
)
|
|
|
|
def on_cancel(res=None):
|
|
self.hash_tracker.update(
|
|
progress_id,
|
|
status="cancelled",
|
|
stage="cancelled",
|
|
message="Hash calculation cancelled",
|
|
)
|
|
|
|
run_in_background_thread(
|
|
self.hash_tracker,
|
|
progress_id,
|
|
run_hash_task,
|
|
on_success,
|
|
on_cancel,
|
|
error_log_msg=f"Hash calculation failed for {normalized_path}",
|
|
)
|
|
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)
|
|
metadata_mode = normalize_metadata_build_mode(
|
|
payload.get("metadata_mode")
|
|
)
|
|
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,
|
|
metadata_mode=metadata_mode,
|
|
)
|
|
|
|
def update_metadata_build_progress(progress_payload):
|
|
self._update_metadata_build_progress(
|
|
progress_id, progress_payload
|
|
)
|
|
|
|
def is_metadata_build_cancelled():
|
|
return self.metadata_builder_progress.is_cancelled(progress_id)
|
|
|
|
def run_metadata_build_task():
|
|
result = build_missing_local_metadata(
|
|
force_rescan=force_rescan,
|
|
worker_count=worker_count,
|
|
metadata_mode=metadata_mode,
|
|
progress_callback=update_metadata_build_progress,
|
|
is_cancelled=is_metadata_build_cancelled,
|
|
)
|
|
return result
|
|
|
|
def on_success(result):
|
|
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,
|
|
)
|
|
|
|
def on_cancel(result=None):
|
|
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 or {},
|
|
**(result or {}),
|
|
)
|
|
|
|
def on_error(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",
|
|
)
|
|
|
|
run_in_background_thread(
|
|
self.metadata_builder_progress,
|
|
progress_id,
|
|
run_metadata_build_task,
|
|
on_success,
|
|
on_cancel,
|
|
on_error,
|
|
error_log_msg="Metadata build failed",
|
|
)
|
|
return web.json_response(
|
|
{
|
|
"success": True,
|
|
"progress_id": progress_id,
|
|
"metadata_mode": metadata_mode,
|
|
}
|
|
)
|
|
|
|
@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(*args, **kwargs):
|
|
self._update_loaded_progress(loaded_id, *args, **kwargs)
|
|
|
|
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 (
|
|
URN_TYPE_MAP,
|
|
get_lora_model_strength,
|
|
)
|
|
|
|
update_loaded_progress(
|
|
"scanning",
|
|
"Scanning local model index...",
|
|
percent=5,
|
|
)
|
|
|
|
workflow_node_count = get_workflow_node_count()
|
|
|
|
def update_workflow_analysis_progress(payload):
|
|
self._update_workflow_analysis_progress(
|
|
loaded_id,
|
|
workflow_node_count,
|
|
interpolate_percent,
|
|
payload,
|
|
start_percent=5,
|
|
end_percent=45,
|
|
)
|
|
|
|
inventory = get_workflow_model_inventory(
|
|
workflow_json,
|
|
progress_callback=update_workflow_analysis_progress,
|
|
)
|
|
available_models = inventory["available_models"]
|
|
all_model_refs = inventory["model_refs"]
|
|
|
|
# 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=45,
|
|
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(
|
|
45, 55, 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(
|
|
55, 65, 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
|
|
|
|
# Also extract from node.properties.models
|
|
workflow_nodes = workflow_json.get("nodes", [])
|
|
nodes = list(workflow_nodes) if isinstance(workflow_nodes, list) else []
|
|
node_scopes = {
|
|
id(node): {
|
|
"subgraph_id": "",
|
|
"subgraph_name": "",
|
|
"is_top_level": True,
|
|
}
|
|
for node in nodes
|
|
if isinstance(node, dict)
|
|
}
|
|
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)
|
|
for node in subgraph_nodes:
|
|
if not isinstance(node, dict):
|
|
continue
|
|
node_scopes[id(node)] = {
|
|
"subgraph_id": subgraph.get("id") or "",
|
|
"subgraph_name": (
|
|
subgraph.get("name")
|
|
or subgraph.get("id")
|
|
or ""
|
|
),
|
|
"is_top_level": False,
|
|
}
|
|
|
|
def get_node_scope(node):
|
|
return node_scopes.get(
|
|
id(node),
|
|
{
|
|
"subgraph_id": "",
|
|
"subgraph_name": "",
|
|
"is_top_level": True,
|
|
},
|
|
)
|
|
|
|
def node_matches_ref(node, ref):
|
|
if not isinstance(node, dict):
|
|
return False
|
|
if str(node.get("id")) != str(ref.get("node_id")):
|
|
return False
|
|
scope = get_node_scope(node)
|
|
if scope["is_top_level"] != (
|
|
ref.get("is_top_level") is not False
|
|
):
|
|
return False
|
|
if not scope["is_top_level"]:
|
|
return str(scope["subgraph_id"]) == str(
|
|
ref.get("subgraph_id") or ""
|
|
)
|
|
return True
|
|
|
|
# Collect all loaded models with their values
|
|
loaded_models = []
|
|
total_refs = len(all_model_refs)
|
|
update_loaded_progress(
|
|
"building",
|
|
"Building loaded models list...",
|
|
percent=65,
|
|
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(65, 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
|
|
|
|
if (
|
|
ref.get("strength") is None
|
|
and category in {"lora", "loras"}
|
|
):
|
|
for node in nodes:
|
|
if node_matches_ref(node, ref):
|
|
strength = get_lora_model_strength(
|
|
node,
|
|
widget_index,
|
|
)
|
|
break
|
|
|
|
if ref.get("strength") is not None:
|
|
strength = ref.get("strength")
|
|
|
|
model_name, strength = adapt_custom_node_loaded_model(
|
|
ref,
|
|
model_name,
|
|
strength,
|
|
)
|
|
|
|
# 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,
|
|
"node_title": ref.get("node_title", ""),
|
|
"subgraph_id": ref.get("subgraph_id") or "",
|
|
"subgraph_name": ref.get("subgraph_name") or "",
|
|
"is_top_level": ref.get("is_top_level") is not False,
|
|
"locate_node_id": ref.get("locate_node_id"),
|
|
"locate_node_type": ref.get("locate_node_type", ""),
|
|
"locate_node_title": ref.get("locate_node_title", ""),
|
|
"locate_subgraph_id": (
|
|
ref.get("locate_subgraph_id") or ""
|
|
),
|
|
"locate_subgraph_name": (
|
|
ref.get("locate_subgraph_name") or ""
|
|
),
|
|
"locate_is_top_level": (
|
|
ref.get("locate_is_top_level") is not False
|
|
),
|
|
"locate_via_promoted_widget": ref.get(
|
|
"locate_via_promoted_widget", False
|
|
),
|
|
"exists": exists,
|
|
"strength": strength,
|
|
"original_path": original_path,
|
|
"is_urn": ref.get("is_urn", False),
|
|
"custom_node_adapter": ref.get(
|
|
"custom_node_adapter"
|
|
),
|
|
"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", "")
|
|
node_scope = get_node_scope(node)
|
|
|
|
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,
|
|
"node_title": node.get("title", ""),
|
|
"subgraph_id": node_scope["subgraph_id"],
|
|
"subgraph_name": node_scope[
|
|
"subgraph_name"
|
|
],
|
|
"is_top_level": node_scope[
|
|
"is_top_level"
|
|
],
|
|
"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:
|
|
file_path = _os.path.realpath(
|
|
_os.path.abspath(_os.path.normpath(str(file_path)))
|
|
)
|
|
if not is_path_in_configured_model_roots(file_path):
|
|
if resolved_path:
|
|
return web.json_response(
|
|
{
|
|
"error": (
|
|
"resolved_path is outside configured model directories"
|
|
)
|
|
},
|
|
status=403,
|
|
)
|
|
file_path = None
|
|
elif not _os.path.isfile(file_path):
|
|
if resolved_path:
|
|
return web.json_response(
|
|
{"error": "resolved_path is not an existing model file"},
|
|
status=404,
|
|
)
|
|
file_path = None
|
|
|
|
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 += "/"
|
|
|
|
external_model_description = ""
|
|
if file_path:
|
|
external_metadata_path = find_external_metadata_sidecar_path(
|
|
file_path
|
|
)
|
|
external_metadata = (
|
|
read_json_safe(external_metadata_path, {})
|
|
if external_metadata_path
|
|
else {}
|
|
)
|
|
if isinstance(external_metadata, dict):
|
|
external_model_description = str(
|
|
external_metadata.get("modelDescription")
|
|
or external_metadata.get("model_description")
|
|
or external_metadata.get("description")
|
|
or ""
|
|
).strip()
|
|
|
|
def infer_model_type_from_category(value):
|
|
return normalize_category_to_model_type(value)
|
|
|
|
def split_result_descriptions(result, preferred_model_description=""):
|
|
result = result if isinstance(result, dict) else {}
|
|
explicit_model_description = (
|
|
result.get("model_description")
|
|
or result.get("modelDescription")
|
|
or ""
|
|
)
|
|
generic_description = result.get("description") or ""
|
|
model_description = (
|
|
preferred_model_description
|
|
or explicit_model_description
|
|
or generic_description
|
|
)
|
|
version_description = (
|
|
result.get("version_description")
|
|
or result.get("versionDescription")
|
|
or ""
|
|
)
|
|
if (
|
|
not version_description
|
|
and explicit_model_description
|
|
and generic_description
|
|
and generic_description != explicit_model_description
|
|
):
|
|
version_description = generic_description
|
|
if version_description == model_description:
|
|
version_description = ""
|
|
return model_description, version_description
|
|
|
|
def build_info_response(
|
|
result=None,
|
|
*,
|
|
metadata_path="",
|
|
metadata_saved=False,
|
|
civitai_checked=False,
|
|
local_payload=False,
|
|
):
|
|
result = result or {}
|
|
model_description, version_description = (
|
|
split_result_descriptions(
|
|
result,
|
|
external_model_description,
|
|
)
|
|
)
|
|
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": model_description,
|
|
"model_description": model_description,
|
|
"version_description": version_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 = get_filename_from_path(filename).lower()
|
|
candidates = [
|
|
result.get("filename"),
|
|
result.get("path"),
|
|
result.get("file_path"),
|
|
]
|
|
for candidate in candidates:
|
|
basename = get_filename_from_path(candidate).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
|
|
return not require_filename or result_filename_matches(result)
|
|
|
|
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):
|
|
headers = {
|
|
"User-Agent": "ComfyUI-Model-Resolver/1.0",
|
|
"Accept": "*/*",
|
|
}
|
|
try:
|
|
response, _final_url, _final_headers = request_public_url(
|
|
"HEAD",
|
|
url,
|
|
headers=headers,
|
|
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, _final_url, _final_headers = request_public_url(
|
|
"GET",
|
|
url,
|
|
headers={**headers, "Range": "bytes=0-0"},
|
|
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:
|
|
model_description, version_description = (
|
|
split_result_descriptions(result)
|
|
)
|
|
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": model_description,
|
|
"model_description": model_description,
|
|
"version_description": version_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_model_resolver_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 "",
|
|
create_preview=True,
|
|
) 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_model_resolver_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:
|
|
preview_path = get_existing_model_preview_path(file_path)
|
|
if preview_path:
|
|
preview_url = preview_path.replace("\\", "/")
|
|
model_description, version_description = (
|
|
split_result_descriptions(result)
|
|
)
|
|
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": model_description,
|
|
"model_description": model_description,
|
|
"version_description": version_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_model_resolver_metadata(
|
|
file_path,
|
|
metadata_payload,
|
|
category,
|
|
result.get("version_url")
|
|
or result.get("url")
|
|
or result.get("download_url")
|
|
or "",
|
|
create_preview=True,
|
|
) 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")),
|
|
"description": response.get("model_description")
|
|
or response.get("description")
|
|
or "",
|
|
"model_description": response.get("model_description")
|
|
or response.get("description")
|
|
or "",
|
|
"version_description": response.get(
|
|
"version_description"
|
|
)
|
|
or "",
|
|
"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_model_resolver_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 _custom_result_timestamp():
|
|
return time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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 = (
|
|
get_filename_from_path(parsed.path)
|
|
or get_filename_from_path(expected_filename)
|
|
or "model"
|
|
)
|
|
if expected_filename and "." in get_filename_from_path(expected_filename):
|
|
filename = get_filename_from_path(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:
|
|
try:
|
|
normalized_url = await asyncio.to_thread(
|
|
validate_public_http_url,
|
|
raw_url,
|
|
)
|
|
except UnsafeUrlError as exc:
|
|
return web.json_response(
|
|
{"error": str(exc)},
|
|
status=400,
|
|
)
|
|
|
|
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.hostname
|
|
if (
|
|
host_matches_domain(direct_host, "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", "")
|
|
hf_token = data.get("hf_token", "")
|
|
file_path = data.get("file_path", "")
|
|
branch = data.get("branch", "")
|
|
|
|
if not download_available:
|
|
return web.json_response(
|
|
{"error": "Download providers are not available"}, status=503
|
|
)
|
|
|
|
if source == "lora_manager_archive":
|
|
source = "civitai"
|
|
|
|
if source not in {"civitai", "civarchive", "huggingface"}:
|
|
return web.json_response(
|
|
{"error": "Unsupported model details source"}, status=400
|
|
)
|
|
|
|
if source == "huggingface":
|
|
model_id = str(model_id or "").strip()
|
|
branch = str(branch or version_id or "main").strip() or "main"
|
|
version_id = branch
|
|
else:
|
|
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 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,
|
|
)
|
|
elif source == "civarchive":
|
|
details = await asyncio.to_thread(
|
|
get_civarchive_model_details,
|
|
model_id,
|
|
version_id,
|
|
)
|
|
else:
|
|
details = await asyncio.to_thread(
|
|
get_huggingface_model_details,
|
|
model_id,
|
|
file_path,
|
|
branch,
|
|
hf_token or None,
|
|
)
|
|
|
|
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 (
|
|
dedupe_local_base_directories,
|
|
get_comfy_root_path,
|
|
prefer_local_base_directory,
|
|
)
|
|
from .core.path_utils import (
|
|
get_path_abs as get_local_path_abs,
|
|
)
|
|
from .core.path_utils import (
|
|
get_path_identity as get_local_path_identity,
|
|
)
|
|
from .core.path_utils import (
|
|
get_path_key as get_local_path_key,
|
|
)
|
|
from .core.path_utils import (
|
|
is_path_within as is_local_path_within,
|
|
)
|
|
|
|
# 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) or 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"] = build_search_result(
|
|
"civitai",
|
|
model_id=model_id_val,
|
|
version_id=version_id_val,
|
|
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}",
|
|
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"] = build_search_result(
|
|
"civitai",
|
|
model_id=first_result.get("model_id"),
|
|
version_id=first_result.get("version_id"),
|
|
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_known_author_fallback_indexes_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_known_author_fallback_indexes, 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
|
|
)
|
|
try:
|
|
url = await asyncio.to_thread(validate_public_http_url, url)
|
|
except UnsafeUrlError as exc:
|
|
return web.json_response({"error": str(exc)}, status=400)
|
|
|
|
from urllib.parse import urlparse as _download_urlparse
|
|
|
|
download_host = _download_urlparse(url).hostname
|
|
|
|
if not filename:
|
|
# Extract filename from URL
|
|
from urllib.parse import unquote, urlparse
|
|
|
|
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
|
|
)
|
|
if not is_allowed_model_download_filename(filename):
|
|
return web.json_response(
|
|
{"error": "Unsupported model file extension"}, status=400
|
|
)
|
|
|
|
# Build headers if needed
|
|
headers = {}
|
|
if host_matches_domain(download_host, "huggingface.co"):
|
|
hf_token = data.get("hf_token", "")
|
|
if hf_token:
|
|
headers["Authorization"] = f"Bearer {hf_token}"
|
|
elif host_matches_domain(
|
|
download_host,
|
|
"civitai.com",
|
|
"civitai.red",
|
|
):
|
|
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 host_matches_domain(download_host, "civitai.com", "civitai.red"):
|
|
inferred_source = "civitai"
|
|
elif host_matches_domain(download_host, "huggingface.co"):
|
|
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, *split_path_segments(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.post("/model_resolver/clear_completed_downloads")
|
|
@json_api_endpoint("clear_completed_downloads", return_success_on_error=True)
|
|
async def clear_completed_downloads_route(request):
|
|
"""Clear completed and failed downloads from progress memory."""
|
|
clear_completed_downloads()
|
|
return web.json_response({"success": True})
|
|
|
|
@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."""
|
|
from .core.workflow_analyzer import NODE_TYPE_MODEL_WIDGET_CATEGORIES
|
|
return web.json_response(
|
|
{
|
|
"sources": {
|
|
"civarchive": is_civarchive_available(),
|
|
"lora_manager_archive": is_lora_manager_archive_available()
|
|
},
|
|
"node_rules": NODE_TYPE_MODEL_WIDGET_CATEGORIES
|
|
}
|
|
)
|
|
|
|
@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 = "/".join(split_path_segments(rel_path))
|
|
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 = split_path_segments(rel_path)
|
|
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)
|
|
parsed_name = parse_rotated_log_filename(name)
|
|
if parsed_name is None:
|
|
return (name.lower(), 999)
|
|
base_name, rotation = parsed_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 parse_rotated_log_filename(name) is None:
|
|
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, 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)
|
|
from .core.settings import get_settings_schema
|
|
return web.json_response({
|
|
"settings": data,
|
|
"schema": get_settings_schema()
|
|
})
|
|
|
|
@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
|
|
)
|