refactor: consolidate inline mappings and precondition checks into shared helpers
This commit is contained in:
+6
-77
@@ -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
@@ -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 "")
|
||||
|
||||
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user