406 lines
16 KiB
Python
406 lines
16 KiB
Python
"""Shared settings helpers for Model Resolver."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import re
|
|
from pathlib import Path
|
|
from typing import Any, Dict, List, Mapping, Optional
|
|
|
|
from .path_utils import read_json_safe, write_json_atomic
|
|
|
|
SETTINGS_FILE = Path(__file__).resolve().parents[1] / "model_resolver_settings.json"
|
|
|
|
from .type_utils import normalize_download_category, to_bool
|
|
|
|
DOWNLOAD_PATH_MODES = {"suggested", "template", "manual"}
|
|
DOWNLOAD_BACKENDS = {"python", "aria2"}
|
|
|
|
SETTINGS_SCHEMA: List[Dict[str, Any]] = [
|
|
{ "serverKey": "civitai_key", "localKey": "ModelResolver.civitaiApiKey", "type": "string", "default": "" },
|
|
{ "serverKey": "civitai_session_token", "localKey": "ModelResolver.civitaiSessionToken", "type": "string", "default": "" },
|
|
{ "serverKey": "civitai_use_trpc_search", "localKey": "ModelResolver.civitaiUseTrpcSearch", "type": "boolean", "default": True },
|
|
{ "serverKey": "civitai_use_api_search", "localKey": "ModelResolver.civitaiUseApiSearch", "type": "boolean", "default": True },
|
|
{ "serverKey": "civitai_use_html_fallback", "localKey": "ModelResolver.civitaiUseHtmlFallback", "type": "boolean", "default": True },
|
|
{ "serverKey": "hf_token", "localKey": "ModelResolver.huggingFaceToken", "type": "string", "default": "" },
|
|
{ "serverKey": "brave_search_api_key", "localKey": "ModelResolver.braveSearchApiKey", "type": "string", "default": "" },
|
|
{ "serverKey": "hf_use_api_search", "localKey": "ModelResolver.hfUseApiSearch", "type": "boolean", "default": True },
|
|
{ "serverKey": "hf_use_comfy_org_fallback", "localKey": "ModelResolver.hfUseComfyOrgFallback", "type": "boolean", "default": True },
|
|
{ "serverKey": "hf_use_brave_fallback", "localKey": "ModelResolver.hfUseBraveFallback", "type": "boolean", "default": True },
|
|
{ "serverKey": "auto_fill_base_model", "localKey": "ModelResolver.autoFillBaseModel", "type": "boolean", "default": True },
|
|
{ "serverKey": "auto_fill_subfolder", "localKey": "ModelResolver.autoFillSubfolder", "type": "boolean", "default": True },
|
|
{ "serverKey": "auto_refresh_comfy_models_after_apply", "localKey": "ModelResolver.autoRefreshComfyModelsAfterApply", "type": "boolean", "default": True },
|
|
{ "serverKey": "auto_open_on_missing", "localKey": "ModelResolver.autoOpenOnMissing", "type": "boolean", "default": False },
|
|
{ "serverKey": "workflow_hash_metadata_enabled", "localKey": "ModelResolver.workflowHashMetadataEnabled", "type": "boolean", "default": True },
|
|
{ "serverKey": "workflow_dependency_marker_enabled", "localKey": "ModelResolver.workflowDependencyMarkerEnabled", "type": "boolean", "default": False },
|
|
{ "serverKey": "download_backend", "localKey": "ModelResolver.downloadBackend", "type": "backend", "default": "python" },
|
|
{ "serverKey": "aria2c_path", "localKey": "ModelResolver.aria2cPath", "type": "string", "default": "" },
|
|
{ "serverKey": "aria2_auto_stop_daemon", "localKey": "ModelResolver.aria2AutoStopDaemon", "type": "boolean", "default": True },
|
|
{ "serverKey": "download_path_mode", "localKey": "ModelResolver.downloadPathMode", "type": "pathMode", "default": "suggested" },
|
|
{ "serverKey": "download_path_templates", "localKey": "ModelResolver.downloadPathTemplates", "type": "json", "default": {} },
|
|
{ "serverKey": "base_model_path_mappings", "localKey": "ModelResolver.baseModelPathMappings", "type": "json", "default": {} },
|
|
{ "serverKey": "frontend_logs_enabled", "localKey": "ModelResolver.frontendLogsEnabled", "type": "frontendLogsEnabled", "default": True },
|
|
{ "serverKey": "backend_logs_enabled", "localKey": "ModelResolver.backendLogsEnabled", "type": "backendLogsEnabled", "default": True },
|
|
{ "serverKey": "frontend_log_level", "localKey": "ModelResolver.frontendLogLevel", "type": "string", "default": "ERROR" },
|
|
{ "serverKey": "backend_log_level", "localKey": "ModelResolver.backendLogLevel", "type": "string", "default": "ERROR" },
|
|
{ "serverKey": "civitai_candidate_limit", "localKey": "ModelResolver.civitaiCandidateLimit", "type": "candidateLimit", "default": 5 },
|
|
]
|
|
|
|
DEFAULT_DOWNLOAD_PATH_TEMPLATES: Dict[str, str] = {
|
|
"loras": "{base_model}/{first_tag}",
|
|
"checkpoints": "{base_model}",
|
|
"embeddings": "{base_model}",
|
|
"diffusion_models": "{base_model}",
|
|
"text_encoders": "",
|
|
"controlnet": "{base_model}",
|
|
"vae": "",
|
|
"upscale_models": "",
|
|
"sams": "",
|
|
"ultralytics": "",
|
|
}
|
|
|
|
DEFAULT_ROOT_KEYS = {
|
|
"loras": "default_lora_root",
|
|
"checkpoints": "default_checkpoint_root",
|
|
"diffusion_models": "default_unet_root",
|
|
"embeddings": "default_embedding_root",
|
|
"text_encoders": "default_text_encoder_root",
|
|
"vae": "default_vae_root",
|
|
"upscale_models": "default_upscale_model_root",
|
|
}
|
|
|
|
TEMPLATE_KEY_ALIASES = {
|
|
"loras": ("loras", "lora"),
|
|
"checkpoints": ("checkpoints", "checkpoint"),
|
|
"embeddings": ("embeddings", "embedding"),
|
|
"diffusion_models": (
|
|
"diffusion_models",
|
|
"diffusion_model",
|
|
"unet",
|
|
"unet_gguf",
|
|
"model_gguf",
|
|
),
|
|
"text_encoders": ("text_encoders", "text_encoder", "clip", "clips", "clip_gguf"),
|
|
"controlnet": ("controlnet", "control_net"),
|
|
"vae": ("vae",),
|
|
"upscale_models": ("upscale_models", "upscale_model", "upscaler"),
|
|
"sams": ("sams", "sam", "sam_model", "sam_models"),
|
|
"ultralytics": ("ultralytics", "ultralytics_bbox", "ultralytics_segm", "yolo"),
|
|
}
|
|
|
|
PRIORITY_TAGS = (
|
|
"concept",
|
|
"style",
|
|
"character",
|
|
"clothing",
|
|
"pose",
|
|
"object",
|
|
"vehicle",
|
|
"artist",
|
|
"celebrity",
|
|
)
|
|
|
|
_INVALID_SEGMENT_CHARS_RE = re.compile(r'[<>:"|?*\x00-\x1f]+')
|
|
_UNKNOWN_PLACEHOLDER_RE = re.compile(r"\{[^{}]+\}")
|
|
|
|
|
|
def bool_setting(value: Any, default: bool = True) -> bool:
|
|
return to_bool(value, default)
|
|
|
|
|
|
def sanitize_folder_name(value: Any, fallback: str = "") -> str:
|
|
text = str(value or "").strip()
|
|
if not text:
|
|
text = fallback
|
|
text = text.replace("/", "_").replace("\\", "_")
|
|
text = _INVALID_SEGMENT_CHARS_RE.sub("_", text)
|
|
text = re.sub(r"\s+", " ", text).strip(" .")
|
|
if text in {"", ".", ".."}:
|
|
text = fallback
|
|
return text.strip(" .")
|
|
|
|
|
|
def normalize_relative_subfolder(path_value: Any) -> str:
|
|
from .path_utils import split_path_segments
|
|
parts = [sanitize_folder_name(part) for part in split_path_segments(path_value)]
|
|
return "/".join([part for part in parts if part])
|
|
|
|
|
|
def normalize_download_path_template(template: Any) -> str:
|
|
from .path_utils import split_path_segments
|
|
return "/".join(split_path_segments(template))
|
|
|
|
|
|
def _read_settings_file() -> Dict[str, Any]:
|
|
data = read_json_safe(str(SETTINGS_FILE), {})
|
|
return data if isinstance(data, dict) else {}
|
|
|
|
|
|
def _coerce_dict(value: Any) -> Dict[str, Any]:
|
|
if isinstance(value, dict):
|
|
return dict(value)
|
|
if isinstance(value, str) and value.strip():
|
|
try:
|
|
parsed = json.loads(value)
|
|
except json.JSONDecodeError:
|
|
return {}
|
|
if isinstance(parsed, dict):
|
|
return dict(parsed)
|
|
return {}
|
|
|
|
|
|
def normalize_download_path_templates(value: Any) -> Dict[str, str]:
|
|
raw_templates = _coerce_dict(value)
|
|
normalized = dict(DEFAULT_DOWNLOAD_PATH_TEMPLATES)
|
|
for raw_key, raw_template in raw_templates.items():
|
|
key = normalize_download_category(str(raw_key))
|
|
normalized[key] = normalize_download_path_template(raw_template)
|
|
return normalized
|
|
|
|
|
|
def normalize_base_model_path_mappings(value: Any) -> Dict[str, str]:
|
|
raw_mappings = _coerce_dict(value)
|
|
mappings: Dict[str, str] = {}
|
|
for raw_key, raw_value in raw_mappings.items():
|
|
key = str(raw_key or "").strip()
|
|
value = str(raw_value or "").strip()
|
|
if key and value:
|
|
mappings[key] = value
|
|
return mappings
|
|
|
|
|
|
def normalize_download_path_mode(value: Any, auto_fill_subfolder: Any = None) -> str:
|
|
mode = str(value or "").strip().lower()
|
|
if mode in DOWNLOAD_PATH_MODES:
|
|
return mode
|
|
if auto_fill_subfolder is False:
|
|
return "manual"
|
|
return "suggested"
|
|
|
|
|
|
def normalize_download_backend(value: Any) -> str:
|
|
backend = str(value or "").strip().lower()
|
|
return backend if backend in DOWNLOAD_BACKENDS else "python"
|
|
|
|
|
|
def normalize_settings(settings: Optional[Mapping[str, Any]]) -> Dict[str, Any]:
|
|
data = dict(settings or {})
|
|
data["workflow_hash_metadata_enabled"] = bool_setting(
|
|
data.get("workflow_hash_metadata_enabled"), True
|
|
)
|
|
data["workflow_dependency_marker_enabled"] = bool_setting(
|
|
data.get("workflow_dependency_marker_enabled"), False
|
|
)
|
|
data["download_backend"] = normalize_download_backend(data.get("download_backend"))
|
|
data["aria2c_path"] = str(data.get("aria2c_path") or "").strip()
|
|
data["aria2_auto_stop_daemon"] = bool_setting(
|
|
data.get("aria2_auto_stop_daemon"), True
|
|
)
|
|
data["auto_refresh_comfy_models_after_apply"] = bool_setting(
|
|
data.get("auto_refresh_comfy_models_after_apply"), True
|
|
)
|
|
data["auto_open_on_missing"] = bool_setting(
|
|
data.get("auto_open_on_missing"), False
|
|
)
|
|
data["download_path_mode"] = normalize_download_path_mode(
|
|
data.get("download_path_mode"), data.get("auto_fill_subfolder")
|
|
)
|
|
data["download_path_templates"] = normalize_download_path_templates(
|
|
data.get("download_path_templates")
|
|
)
|
|
data["base_model_path_mappings"] = normalize_base_model_path_mappings(
|
|
data.get("base_model_path_mappings")
|
|
)
|
|
for key in DEFAULT_ROOT_KEYS.values():
|
|
data[key] = str(data.get(key) or "").strip()
|
|
return data
|
|
|
|
|
|
def load_settings() -> Dict[str, Any]:
|
|
return normalize_settings(_read_settings_file())
|
|
|
|
|
|
def save_settings(payload: Mapping[str, Any]) -> Dict[str, Any]:
|
|
current = _read_settings_file()
|
|
current.update({str(key): value for key, value in payload.items() if key})
|
|
normalized = normalize_settings(current)
|
|
write_json_atomic(str(SETTINGS_FILE), normalized, indent=2)
|
|
return normalized
|
|
|
|
|
|
def _template_for_category(settings: Mapping[str, Any], category: str) -> str:
|
|
templates = normalize_download_path_templates(settings.get("download_path_templates"))
|
|
category_key = normalize_download_category(category)
|
|
aliases = TEMPLATE_KEY_ALIASES.get(category_key, (category_key,))
|
|
for alias in aliases:
|
|
if alias in templates:
|
|
return templates.get(alias, "")
|
|
return templates.get(category_key, "")
|
|
|
|
|
|
def _listify_tags(value: Any) -> List[str]:
|
|
from .type_utils import as_list
|
|
return as_list(value)
|
|
|
|
|
|
def _normalize_tag(value: str) -> str:
|
|
from .type_utils import normalize_alphanumeric_key
|
|
return normalize_alphanumeric_key(value)
|
|
|
|
|
|
|
|
def _resolve_base_model_mapping(mappings: Mapping[str, str], base_model: Any) -> str:
|
|
base_text = str(base_model or "")
|
|
if base_text in mappings:
|
|
return mappings[base_text]
|
|
|
|
base_token = _normalize_tag(base_text)
|
|
if not base_token:
|
|
return base_text
|
|
|
|
normalized_mappings = {
|
|
_normalize_tag(key): value
|
|
for key, value in mappings.items()
|
|
if _normalize_tag(key)
|
|
}
|
|
if base_token in normalized_mappings:
|
|
return normalized_mappings[base_token]
|
|
|
|
for key_token, value in sorted(
|
|
normalized_mappings.items(),
|
|
key=lambda item: len(item[0]),
|
|
reverse=True,
|
|
):
|
|
if len(key_token) < 4:
|
|
continue
|
|
if base_token.startswith(key_token) or key_token in base_token or base_token in key_token:
|
|
return value
|
|
|
|
return base_text
|
|
|
|
|
|
def _first_tag(tags: Any) -> str:
|
|
from .type_utils import as_list
|
|
tag_list = list(as_list(tags))
|
|
if not tag_list:
|
|
return "no tags"
|
|
normalized_tags = {tag: _normalize_tag(tag) for tag in tag_list}
|
|
for priority in PRIORITY_TAGS:
|
|
priority_token = _normalize_tag(priority)
|
|
for tag, normalized in normalized_tags.items():
|
|
if normalized == priority_token:
|
|
return tag
|
|
return tag_list[0]
|
|
|
|
|
|
def _creator_name(metadata: Mapping[str, Any]) -> str:
|
|
creator = metadata.get("creator")
|
|
if isinstance(creator, Mapping):
|
|
name = creator.get("username") or creator.get("name")
|
|
if name:
|
|
return str(name)
|
|
if isinstance(creator, str) and creator.strip():
|
|
return creator
|
|
for key in ("author", "username", "creator_username", "creator_name"):
|
|
value = metadata.get(key)
|
|
if value:
|
|
return str(value)
|
|
repo_id = metadata.get("repo_id") or metadata.get("repo")
|
|
if isinstance(repo_id, str) and "/" in repo_id:
|
|
return repo_id.split("/", 1)[0]
|
|
return "Anonymous"
|
|
|
|
|
|
def _filename_stem(metadata: Mapping[str, Any]) -> str:
|
|
filename = str(metadata.get("filename") or metadata.get("file_name") or "").strip()
|
|
if not filename:
|
|
return ""
|
|
return Path(filename).stem
|
|
|
|
|
|
def calculate_template_subfolder(
|
|
category: str,
|
|
metadata: Optional[Mapping[str, Any]] = None,
|
|
settings: Optional[Mapping[str, Any]] = None,
|
|
) -> str:
|
|
active_settings = normalize_settings(settings if settings is not None else load_settings())
|
|
template = _template_for_category(active_settings, category)
|
|
if not template:
|
|
return ""
|
|
|
|
model_metadata = dict(metadata or {})
|
|
base_model = (
|
|
model_metadata.get("base_model")
|
|
or model_metadata.get("baseModel")
|
|
or model_metadata.get("base_model_context")
|
|
or "Unknown Base Model"
|
|
)
|
|
mappings = normalize_base_model_path_mappings(
|
|
active_settings.get("base_model_path_mappings")
|
|
)
|
|
mapped_base_model = _resolve_base_model_mapping(mappings, base_model)
|
|
model_name = (
|
|
model_metadata.get("model_name")
|
|
or model_metadata.get("name")
|
|
or _filename_stem(model_metadata)
|
|
or "Model"
|
|
)
|
|
version_name = (
|
|
model_metadata.get("version_name")
|
|
or model_metadata.get("versionName")
|
|
or model_metadata.get("version")
|
|
or ""
|
|
)
|
|
|
|
replacements = {
|
|
"{base_model}": (
|
|
normalize_relative_subfolder(mapped_base_model)
|
|
or sanitize_folder_name(mapped_base_model, "Unknown Base Model")
|
|
),
|
|
"{author}": sanitize_folder_name(_creator_name(model_metadata), "Anonymous"),
|
|
"{first_tag}": sanitize_folder_name(
|
|
_first_tag(model_metadata.get("tags")), "no tags"
|
|
),
|
|
"{model_name}": sanitize_folder_name(model_name, "Model"),
|
|
"{version_name}": sanitize_folder_name(version_name, ""),
|
|
}
|
|
|
|
formatted = template
|
|
for token, replacement in replacements.items():
|
|
formatted = formatted.replace(token, replacement)
|
|
formatted = _UNKNOWN_PLACEHOLDER_RE.sub("", formatted)
|
|
return normalize_relative_subfolder(formatted)
|
|
|
|
|
|
def resolve_download_subfolder(
|
|
category: str,
|
|
requested_subfolder: Any = "",
|
|
metadata: Optional[Mapping[str, Any]] = None,
|
|
settings: Optional[Mapping[str, Any]] = None,
|
|
) -> str:
|
|
requested = str(requested_subfolder or "").strip()
|
|
if requested:
|
|
return requested
|
|
|
|
active_settings = normalize_settings(settings if settings is not None else load_settings())
|
|
if active_settings.get("download_path_mode") != "template":
|
|
return ""
|
|
|
|
return calculate_template_subfolder(category, metadata, active_settings)
|
|
|
|
|
|
def get_default_root_for_category(
|
|
category: str,
|
|
settings: Optional[Mapping[str, Any]] = None,
|
|
) -> str:
|
|
active_settings = normalize_settings(settings if settings is not None else load_settings())
|
|
key = DEFAULT_ROOT_KEYS.get(normalize_download_category(category))
|
|
if not key:
|
|
return ""
|
|
return str(active_settings.get(key) or "").strip()
|
|
|
|
|
|
def get_settings_schema() -> List[Dict[str, Any]]:
|
|
"""Return the schema dictionary for frontend setting options synchronization."""
|
|
return SETTINGS_SCHEMA
|