refactor: consolidate inline mappings and precondition checks into shared helpers

This commit is contained in:
Dariusz L
2026-07-04 00:04:22 +02:00
parent 4091dbb88d
commit 5de1500537
6 changed files with 129 additions and 120 deletions
+6 -77
View File
@@ -188,7 +188,7 @@ class ModelResolverExtension:
search_local_matches,
)
from .core.path_templates import infer_download_path_templates
from .core.scanner import get_model_files, invalidate_model_files_cache
from .core.scanner import get_model_files, invalidate_model_files_cache, find_local_file_path
from .core.path_utils import (
get_filename_from_path,
read_json_safe,
@@ -204,6 +204,7 @@ class ModelResolverExtension:
extract_sha256_from_metadata,
fetch_remote_file_size_cached,
looks_like_model_file,
normalize_category_to_model_type,
)
from .core.settings import (
TEMPLATE_KEY_ALIASES,
@@ -1515,33 +1516,8 @@ class ModelResolverExtension:
file_path = resolved_path if resolved_path else None
file_location = ""
if not file_path and category:
# Try to find the file in the model directories using folder_paths
try:
import folder_paths
folder_type = normalize_download_category(category)
file_path = folder_paths.get_full_path(
folder_type, filename
)
except Exception:
pass
# If not found, try scanner
if not file_path:
try:
from .core.scanner import get_model_files
available_models = get_model_files()
for m in available_models:
if (
m.get("relative_path", "").endswith(filename)
or m.get("filename", "") == filename
):
file_path = m.get("path")
break
except Exception:
pass
if not file_path:
file_path = find_local_file_path(filename, category)
if file_path and _os.path.exists(file_path):
file_location = _os.path.dirname(file_path).replace("\\", "/")
@@ -1549,21 +1525,7 @@ class ModelResolverExtension:
file_location += "/"
def infer_model_type_from_category(value):
if not str(value or "").strip():
return ""
normalized = normalize_download_category(value or "")
type_map = {
"checkpoints": "checkpoint",
"loras": "lora",
"vae": "vae",
"text_encoders": "text_encoder",
"clip": "clip",
"clip_vision": "clip_vision",
"controlnet": "controlnet",
"upscale_models": "upscale",
"diffusion_models": "diffusion_model",
}
return type_map.get(normalized, normalized.rstrip("s"))
return normalize_category_to_model_type(value)
def build_info_response(
result=None,
@@ -2668,39 +2630,6 @@ class ModelResolverExtension:
elif isinstance(result, dict):
yield result
def get_result_sha256(result):
if not isinstance(result, dict):
return ""
hashes = (
result.get("hashes")
if isinstance(result.get("hashes"), dict)
else {}
)
file_info = (
result.get("file_info")
if isinstance(result.get("file_info"), dict)
else {}
)
file_hashes = (
file_info.get("hashes")
if isinstance(file_info.get("hashes"), dict)
else {}
)
for candidate in (
result.get("sha256"),
result.get("hash"),
hashes.get("SHA256"),
hashes.get("sha256"),
file_info.get("sha256"),
file_info.get("hash"),
file_hashes.get("SHA256"),
file_hashes.get("sha256"),
):
normalized = normalize_sha256(candidate)
if normalized:
return normalized
return ""
def collect_local_hash_matches(payload):
matches = []
seen_match_paths = set()
@@ -2715,7 +2644,7 @@ class ModelResolverExtension:
):
for source_result in iter_result_items(payload.get(source_key)):
raise_if_search_cancelled(source_key)
sha256 = get_result_sha256(source_result)
sha256 = extract_sha256_from_metadata(source_result)
if not sha256:
continue
+4 -8
View File
@@ -25,7 +25,7 @@ log = create_module_logger(__name__)
from .resolver import invalidate_local_hash_match_cache, normalize_sha256
from .scanner import invalidate_model_files_cache
from .path_utils import is_path_within, get_path_identity, write_json_atomic, read_json_safe, get_comfy_root_path, calculate_file_sha256 as _calculate_file_sha256, get_filename_from_path
from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size
from .type_utils import as_dict, as_list, first_non_empty, format_size_bytes as format_bytes, DEFAULT_BROWSER_USER_AGENT, extract_response_file_size, normalize_category_to_model_type
try:
import folder_paths
@@ -259,13 +259,9 @@ def get_metadata_sidecar_path(file_path: str) -> str:
def _resolve_lora_manager_model_type(category: str, source_type: Any = "") -> str:
category_key = normalize_download_category(category)
if category_key == "diffusion_models":
return "diffusion_model"
if category_key == "checkpoints":
return "checkpoint"
if category_key == "embeddings":
return "embedding"
res = normalize_category_to_model_type(category)
if res in ("checkpoint", "diffusion_model", "embedding"):
return res
source_token = (
str(source_type or "")
+40
View File
@@ -306,3 +306,43 @@ def get_model_files(force_rescan: bool = False) -> List[Dict[str, str]]:
_MODEL_FILES_CACHE = models
_MODEL_FILES_CACHE_AT = now
return models
def find_local_file_path(filename: str, category: Optional[str] = None) -> Optional[str]:
"""
Tries to find the absolute path of a model file locally using folder_paths or the local scanner.
"""
if not filename:
return None
global folder_paths
if folder_paths is None:
try:
import folder_paths as fp
folder_paths = fp
except ImportError:
pass
file_path = None
if category and folder_paths is not None:
try:
from .type_utils import normalize_download_category
folder_type = normalize_download_category(category)
file_path = folder_paths.get_full_path(folder_type, filename)
except Exception:
pass
if not file_path:
try:
available_models = get_model_files()
for m in available_models:
if (
m.get("relative_path", "").endswith(filename)
or m.get("filename", "") == filename
):
file_path = m.get("path")
break
except Exception:
pass
return file_path
+10 -16
View File
@@ -76,15 +76,12 @@ def build_civitai_session_cookie(session_token: Optional[str]) -> str:
def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]:
"""Check whether a CivitAI browser session token is accepted by civitai.com."""
token = (session_token or "").strip()
if not token:
return {
"success": False,
"valid": False,
"status": "missing",
"message": "Paste a CivitAI session token first.",
}
from ..type_utils import check_credential_preconditions
precheck = check_credential_preconditions(session_token, "CivitAI session token")
if precheck:
return precheck
token = (session_token or "").strip()
headers = {
"accept": "application/json",
"Cookie": build_civitai_session_cookie(token),
@@ -105,15 +102,12 @@ def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]:
def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]:
"""Check whether a CivitAI API key is accepted by civitai.com."""
key = (api_key or "").strip()
if not key:
return {
"success": False,
"valid": False,
"status": "missing",
"message": "Paste a CivitAI API key first.",
}
from ..type_utils import check_credential_preconditions
precheck = check_credential_preconditions(api_key, "CivitAI API key")
if precheck:
return precheck
key = (api_key or "").strip()
headers = {
"accept": "application/json",
"Authorization": f"Bearer {key}",
+10 -16
View File
@@ -51,15 +51,12 @@ _report_progress = get_progress_reporter("HuggingFace progress callback")
def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]:
"""Check whether a HuggingFace access token is accepted."""
value = (token or "").strip()
if not value:
return {
"success": False,
"valid": False,
"status": "missing",
"message": "Paste a HuggingFace token first.",
}
from ..type_utils import check_credential_preconditions
precheck = check_credential_preconditions(token, "HuggingFace token")
if precheck:
return precheck
value = (token or "").strip()
headers = {"Authorization": f"Bearer {value}"}
def get_user(data):
@@ -85,15 +82,12 @@ def check_huggingface_token(token: Optional[str]) -> Dict[str, Any]:
def check_brave_search_api_key(api_key: Optional[str]) -> Dict[str, Any]:
"""Check whether a Brave Search API key is accepted."""
key = (api_key or "").strip()
if not key:
return {
"success": False,
"valid": False,
"status": "missing",
"message": "Paste a Brave Search API key first.",
}
from ..type_utils import check_credential_preconditions
precheck = check_credential_preconditions(api_key, "Brave Search API key")
if precheck:
return precheck
key = (api_key or "").strip()
headers = {"X-Subscription-Token": key}
params = {"q": "test", "count": 1}
+59 -3
View File
@@ -743,12 +743,30 @@ def extract_sha256_from_metadata(metadata: Any) -> str:
if normalized:
return normalized
# Obsługa słownika file_info (np. z popular/model_list)
file_info = metadata.get("file_info")
if isinstance(file_info, dict):
for key in ("sha256", "hash", "SHA256", "Sha256"):
val = file_info.get(key)
if val:
normalized = normalize_sha256(val)
if normalized:
return normalized
f_hashes = file_info.get("hashes")
if isinstance(f_hashes, dict):
for key in ("SHA256", "sha256", "Sha256", "hash"):
val = f_hashes.get(key)
if val:
normalized = normalize_sha256(val)
if normalized:
return normalized
# Obsługa zagnieżdżonej listy plików w metadanych Civitai
files = metadata.get("files")
if isinstance(files, list):
for file_info in files:
if isinstance(file_info, dict):
f_hashes = file_info.get("hashes")
for file_info_item in files:
if isinstance(file_info_item, dict):
f_hashes = file_info_item.get("hashes")
if isinstance(f_hashes, dict):
for key in ("SHA256", "sha256", "Sha256", "hash"):
val = f_hashes.get(key)
@@ -878,5 +896,43 @@ def extract_file_size(file_info: Dict[str, Any]) -> Optional[int]:
return None
def normalize_category_to_model_type(category: str) -> str:
"""
Mapuje nazwę kategorii pobierania/ComfyUI (np. 'checkpoints') na znormalizowany typ modelu (np. 'checkpoint').
"""
if not category:
return ""
normalized = normalize_download_category(category)
type_map = {
"checkpoints": "checkpoint",
"loras": "lora",
"vae": "vae",
"text_encoders": "text_encoder",
"clip": "clip",
"clip_vision": "clip_vision",
"controlnet": "controlnet",
"upscale_models": "upscale",
"diffusion_models": "diffusion_model",
"embeddings": "embedding",
}
return type_map.get(normalized, normalized.rstrip("s"))
def check_credential_preconditions(value: Optional[str], token_name: str) -> Optional[Dict[str, Any]]:
"""
Sprawdza, czy poświadczenie (np. token, klucz) nie jest puste.
Zwraca standardowy słownik błędu, jeśli poświadczenie jest puste, w przeciwnym razie None.
"""
clean_val = (value or "").strip()
if not clean_val:
return {
"success": False,
"valid": False,
"status": "missing",
"message": f"Paste a {token_name} first.",
}
return None