2347 lines
79 KiB
Python
2347 lines
79 KiB
Python
"""
|
|
CivitAI Source Module
|
|
|
|
Search and download models from CivitAI.
|
|
"""
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
from typing import Any, Callable, Dict, List, Optional
|
|
from urllib.parse import quote, urlencode
|
|
|
|
import requests
|
|
|
|
from ..log_system import create_module_logger
|
|
from ..matcher import (
|
|
base_model_matches as _base_model_matches,
|
|
)
|
|
from ..matcher import (
|
|
base_model_score as _base_model_score,
|
|
)
|
|
from ..matcher import (
|
|
calculate_filename_confidence,
|
|
)
|
|
from ..matcher import (
|
|
has_known_model_extension as _has_known_model_extension,
|
|
)
|
|
from ..matcher import (
|
|
normalize_base_model as _normalize_base_model,
|
|
)
|
|
from ..network_utils import execute_provider_json_request, request_source_response
|
|
from ..path_utils import (
|
|
calculate_file_sha256,
|
|
extract_safetensors_header_metadata,
|
|
find_metadata_sidecar_path,
|
|
get_filename_from_path,
|
|
infer_safetensors_base_model,
|
|
read_json_safe,
|
|
read_merged_model_metadata,
|
|
)
|
|
from ..progress import get_progress_reporter
|
|
from ..type_utils import (
|
|
CIVITAI_API_TYPE_MAP,
|
|
DEFAULT_BROWSER_USER_AGENT,
|
|
as_dict,
|
|
as_list,
|
|
check_credential_http,
|
|
extract_trained_words,
|
|
first_non_empty,
|
|
normalize_model_file_info,
|
|
normalize_model_image,
|
|
parse_size_to_bytes,
|
|
resolve_model_category,
|
|
select_primary_model_file,
|
|
)
|
|
|
|
log = create_module_logger(__name__)
|
|
|
|
|
|
CIVITAI_API_URL = "https://civitai.com/api/v1"
|
|
|
|
|
|
# Cache for search results and URN resolutions
|
|
_search_cache: Dict[str, Any] = {}
|
|
_urn_cache: Dict[tuple[int, int], Dict[str, Any]] = {}
|
|
_hash_cache: Dict[str, Dict[str, Any]] = {}
|
|
DEFAULT_CIVITAI_CANDIDATE_LIMIT = 5
|
|
MAX_CIVITAI_CANDIDATE_LIMIT = 20
|
|
MODEL_TITLE_MATCH_THRESHOLD = 82.0
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def clear_search_cache():
|
|
"""Clear cached CivitAI search, URN, and hash results."""
|
|
global _search_cache, _urn_cache, _hash_cache
|
|
_search_cache.clear()
|
|
_urn_cache.clear()
|
|
_hash_cache.clear()
|
|
|
|
|
|
_report_progress = get_progress_reporter("CivitAI progress callback")
|
|
|
|
|
|
def build_civitai_session_cookie(session_token: Optional[str]) -> str:
|
|
"""Build a CivitAI session cookie header value accepted by current and older endpoints."""
|
|
token = str(session_token or "").strip()
|
|
if not token:
|
|
return ""
|
|
return f"__Secure-civ-token={token}; __Secure-civitai-token={token}"
|
|
|
|
|
|
def _extract_civitai_username(data: Dict[str, Any]) -> str:
|
|
return data.get("username") or data.get("name") or data.get("email") or ""
|
|
|
|
|
|
def check_civitai_session_token(session_token: Optional[str]) -> Dict[str, Any]:
|
|
"""Check whether a CivitAI browser session token is accepted by civitai.com."""
|
|
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),
|
|
"user-agent": DEFAULT_BROWSER_USER_AGENT,
|
|
}
|
|
|
|
return check_credential_http(
|
|
"https://civitai.com/api/v1/me",
|
|
headers=headers,
|
|
success_message="Session token is valid.",
|
|
get_username=_extract_civitai_username,
|
|
error_msg_401_403="Session token is not accepted by CivitAI.",
|
|
)
|
|
|
|
|
|
def check_civitai_api_key(api_key: Optional[str]) -> Dict[str, Any]:
|
|
"""Check whether a CivitAI API key is accepted by civitai.com."""
|
|
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}",
|
|
"user-agent": DEFAULT_BROWSER_USER_AGENT,
|
|
}
|
|
|
|
return check_credential_http(
|
|
"https://civitai.com/api/v1/me",
|
|
headers=headers,
|
|
success_message="CivitAI API key is valid.",
|
|
get_username=_extract_civitai_username,
|
|
error_msg_401_403="CivitAI API key is not accepted.",
|
|
)
|
|
|
|
|
|
from .common import build_custom_url_result, build_unified_search_result
|
|
|
|
|
|
def _build_civitai_result_from_version(
|
|
model_id: int,
|
|
model_name: str,
|
|
model_type: str,
|
|
version: Dict[str, Any],
|
|
file_info: Dict[str, Any],
|
|
tags: Optional[List[str]] = None,
|
|
match_type: str = "exact",
|
|
) -> Dict[str, Any]:
|
|
"""Normalize CivitAI model/version/file data into the search result format."""
|
|
version_id = version.get("id")
|
|
hashes = file_info.get("hashes") if isinstance(file_info.get("hashes"), dict) else {}
|
|
sha256 = file_info.get("sha256") or hashes.get("SHA256") or hashes.get("sha256")
|
|
size = file_info.get("sizeKB", 0) * 1024 if file_info.get("sizeKB") else file_info.get("size")
|
|
download_url = file_info.get("downloadUrl") or get_civitai_download_url(version_id)
|
|
return build_unified_search_result(
|
|
source="civitai",
|
|
model_id=model_id,
|
|
version_id=version_id,
|
|
name=model_name,
|
|
type=model_type,
|
|
filename=file_info.get("name", ""),
|
|
url=f"https://civitai.com/models/{model_id}?modelVersionId={version_id}",
|
|
download_url=download_url,
|
|
size=size,
|
|
base_model=version.get("baseModel"),
|
|
tags=tags or [],
|
|
match_type=match_type,
|
|
sha256=sha256,
|
|
hashes=hashes,
|
|
)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _find_model_title_match_in_model(
|
|
model_id: int,
|
|
model_data: Dict[str, Any],
|
|
title_query: str,
|
|
api_key: Optional[str] = None,
|
|
base_model_context: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""For extensionless workflow values, resolve by CivitAI model page title."""
|
|
from ..matcher import match_model_by_title_generic
|
|
|
|
model_name = model_data.get("name", "")
|
|
versions = [
|
|
version
|
|
for version in (model_data.get("modelVersions") or [])
|
|
if isinstance(version, dict)
|
|
]
|
|
|
|
def get_base_model(v):
|
|
return v.get("baseModel")
|
|
|
|
def select_file(v):
|
|
return select_primary_model_file(v.get("files") or [])
|
|
|
|
def build_result(v, file_info, confidence):
|
|
res = _build_civitai_result_from_version(
|
|
model_id=model_id,
|
|
model_name=model_name,
|
|
model_type=model_data.get("type", ""),
|
|
version=v,
|
|
file_info=file_info,
|
|
tags=model_data.get("tags", []),
|
|
match_type="model_title",
|
|
)
|
|
if res:
|
|
res["version_name"] = v.get("name", "")
|
|
return res
|
|
|
|
return match_model_by_title_generic(
|
|
model_id=model_id,
|
|
title_query=title_query,
|
|
model_name=model_name,
|
|
versions=versions,
|
|
base_model_context=base_model_context,
|
|
get_base_model_fn=get_base_model,
|
|
select_file_fn=select_file,
|
|
build_result_fn=build_result,
|
|
log_prefix="CivitAI",
|
|
)
|
|
|
|
|
|
def _filename_base_partial_match(target_base: str, candidate_base: str) -> bool:
|
|
"""Return True for meaningful filename-base containment matches."""
|
|
target_base = str(target_base or "").strip().lower()
|
|
candidate_base = str(candidate_base or "").strip().lower()
|
|
if not target_base or not candidate_base:
|
|
return False
|
|
shorter = min(target_base, candidate_base, key=len)
|
|
if len(shorter) < 4:
|
|
return False
|
|
return target_base in candidate_base or candidate_base in target_base
|
|
|
|
|
|
|
|
def _find_matching_file_in_versions(
|
|
versions: List[Dict[str, Any]],
|
|
filename: str,
|
|
exact_only: bool = False,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Search all versions/files for an exact or partial filename match."""
|
|
filename_lower = filename.lower()
|
|
filename_base = os.path.splitext(filename_lower)[0]
|
|
best_match = None
|
|
best_confidence = 0.0
|
|
|
|
for version in versions:
|
|
version_id = version.get("id")
|
|
files = version.get("files", [])
|
|
log.debug(
|
|
f"CivitAI checking version_id={version_id} with {len(files)} files"
|
|
)
|
|
|
|
for file_info in files:
|
|
file_name = file_info.get("name", "")
|
|
if not file_name:
|
|
continue
|
|
|
|
file_name_lower = file_name.lower()
|
|
file_base = os.path.splitext(file_name_lower)[0]
|
|
|
|
if file_name_lower == filename_lower:
|
|
return {"version": version, "file_info": file_info, "match_type": "exact"}
|
|
|
|
if not exact_only and _filename_base_partial_match(filename_base, file_base):
|
|
return {
|
|
"version": version,
|
|
"file_info": file_info,
|
|
"match_type": "partial",
|
|
}
|
|
|
|
if not exact_only:
|
|
confidence = calculate_filename_confidence(filename, file_name)
|
|
if confidence > best_confidence:
|
|
best_confidence = confidence
|
|
best_match = {
|
|
"version": version,
|
|
"file_info": file_info,
|
|
"match_type": "similar",
|
|
"confidence": confidence,
|
|
}
|
|
|
|
if best_match and best_confidence >= 50.0:
|
|
return best_match
|
|
|
|
return None
|
|
|
|
|
|
def _extract_civitai_red_candidates(html: str, limit: int = 5) -> List[Dict[str, int]]:
|
|
"""Extract model/version candidates from civitai.red search results HTML."""
|
|
candidates: List[Dict[str, int]] = []
|
|
seen = set()
|
|
|
|
for match in re.finditer(r"/models/(\d+)(?:\?modelVersionId=(\d+))?", html):
|
|
model_id = int(match.group(1))
|
|
version_id = int(match.group(2)) if match.group(2) else None
|
|
key = (model_id, version_id)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
candidates.append({"model_id": model_id, "version_id": version_id})
|
|
if len(candidates) >= limit:
|
|
break
|
|
|
|
return candidates
|
|
|
|
|
|
_DEVALUE_UNDEFINED = object()
|
|
|
|
|
|
def _decode_devalue_trpc_data(value: str) -> Optional[Dict[str, Any]]:
|
|
"""Decode CivitAI's alternate devalue-serialized tRPC data payload."""
|
|
try:
|
|
parsed = json.loads(value)
|
|
except (TypeError, ValueError):
|
|
return None
|
|
|
|
if isinstance(parsed, dict):
|
|
return parsed
|
|
if not isinstance(parsed, list) or not parsed:
|
|
return None
|
|
|
|
memo: Dict[int, Any] = {}
|
|
|
|
def hydrate_reference(reference: Any) -> Any:
|
|
if isinstance(reference, bool) or reference is None:
|
|
return reference
|
|
if not isinstance(reference, int):
|
|
return reference
|
|
if reference < 0:
|
|
return _DEVALUE_UNDEFINED
|
|
if reference >= len(parsed):
|
|
return _DEVALUE_UNDEFINED
|
|
return hydrate_index(reference)
|
|
|
|
def hydrate_index(index: int) -> Any:
|
|
if index in memo:
|
|
return memo[index]
|
|
|
|
raw_value = parsed[index]
|
|
if isinstance(raw_value, dict):
|
|
hydrated: Any = {}
|
|
memo[index] = hydrated
|
|
for key, reference in raw_value.items():
|
|
resolved = hydrate_reference(reference)
|
|
if resolved is not _DEVALUE_UNDEFINED:
|
|
hydrated[key] = resolved
|
|
return hydrated
|
|
|
|
if isinstance(raw_value, list):
|
|
hydrated = []
|
|
memo[index] = hydrated
|
|
for reference in raw_value:
|
|
resolved = hydrate_reference(reference)
|
|
if resolved is not _DEVALUE_UNDEFINED:
|
|
hydrated.append(resolved)
|
|
return hydrated
|
|
|
|
memo[index] = raw_value
|
|
return raw_value
|
|
|
|
decoded = hydrate_index(0)
|
|
return decoded if isinstance(decoded, dict) else None
|
|
|
|
|
|
def _extract_trpc_model_candidates(
|
|
payload: Any, limit: int = 5
|
|
) -> List[Dict[str, Optional[int]]]:
|
|
"""Extract model/version candidates from a tRPC JSON response."""
|
|
candidates: List[Dict[str, Optional[int]]] = []
|
|
seen = set()
|
|
|
|
if not isinstance(payload, dict):
|
|
return candidates
|
|
|
|
result = payload.get("result")
|
|
if not isinstance(result, dict):
|
|
return candidates
|
|
|
|
data = result.get("data")
|
|
if isinstance(data, str):
|
|
data = _decode_devalue_trpc_data(data)
|
|
if not isinstance(data, dict):
|
|
return candidates
|
|
|
|
json_data = data.get("json", data)
|
|
if isinstance(json_data, str):
|
|
json_data = _decode_devalue_trpc_data(json_data)
|
|
if not isinstance(json_data, dict):
|
|
return candidates
|
|
|
|
items = json_data.get("items", [])
|
|
|
|
if not isinstance(items, list):
|
|
return candidates
|
|
|
|
for item in items:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
|
|
model_id = item.get("id")
|
|
version = item.get("version", {})
|
|
version_id = version.get("id") if isinstance(version, dict) else None
|
|
|
|
if not isinstance(model_id, int):
|
|
continue
|
|
|
|
key = (model_id, version_id if isinstance(version_id, int) else None)
|
|
if key in seen:
|
|
continue
|
|
|
|
seen.add(key)
|
|
candidates.append(
|
|
{
|
|
"model_id": model_id,
|
|
"version_id": version_id if isinstance(version_id, int) else None,
|
|
}
|
|
)
|
|
|
|
if len(candidates) >= limit:
|
|
break
|
|
|
|
return candidates
|
|
|
|
|
|
def _search_civitai_trpc_candidates(
|
|
filename: str,
|
|
model_type: Optional[str] = None,
|
|
session_token: Optional[str] = None,
|
|
timeout: int = 15,
|
|
limit: int = 5,
|
|
) -> List[Dict[str, Optional[int]]]:
|
|
"""Try CivitAI.red tRPC search endpoint and log the raw outcome for diagnostics."""
|
|
civitai_type = CIVITAI_API_TYPE_MAP.get(str(model_type).lower()) if model_type else None
|
|
search_query = get_filename_from_path(filename)
|
|
if _has_known_model_extension(search_query):
|
|
search_query = os.path.splitext(search_query)[0]
|
|
|
|
headers = {
|
|
"accept": "application/json",
|
|
"content-type": "application/json",
|
|
"referer": f"https://civitai.red/models?query={quote(search_query)}",
|
|
"user-agent": DEFAULT_BROWSER_USER_AGENT,
|
|
"x-client": "web",
|
|
"x-client-version": "5.0.1657",
|
|
}
|
|
if session_token:
|
|
headers["Cookie"] = build_civitai_session_cookie(session_token)
|
|
|
|
def build_url(type_filter: Optional[str]) -> str:
|
|
input_payload = {
|
|
"json": {
|
|
"period": "Month",
|
|
"periodMode": "stats",
|
|
"sort": "Highest Rated",
|
|
"query": search_query,
|
|
"pending": False,
|
|
"browsingLevel": 28,
|
|
"excludedTagIds": [
|
|
415792,
|
|
426772,
|
|
5351,
|
|
5161,
|
|
5162,
|
|
5188,
|
|
5249,
|
|
306619,
|
|
5351,
|
|
154326,
|
|
161829,
|
|
163032,
|
|
130818,
|
|
130820,
|
|
133182,
|
|
],
|
|
"disablePoi": True,
|
|
"disableMinor": True,
|
|
"limit": limit,
|
|
"authed": True,
|
|
},
|
|
"meta": {"values": {"cursor": ["undefined"]}},
|
|
}
|
|
if type_filter:
|
|
input_payload["json"]["types"] = [type_filter]
|
|
return (
|
|
"https://civitai.red/api/trpc/model.getAll?input="
|
|
+ quote(json.dumps(input_payload, separators=(",", ":")))
|
|
)
|
|
|
|
url = build_url(civitai_type)
|
|
log.info(
|
|
f"CivitAI tRPC search start: filename={filename}, model_type={model_type}, type_filter={civitai_type or 'none'}, session_token={'yes' if session_token else 'no'}, url={url}"
|
|
)
|
|
|
|
response = request_source_response(url, headers=headers, timeout=timeout, log_name="CivitAI tRPC")
|
|
if response is None:
|
|
log.warning(f"CivitAI tRPC request failed for filename={filename}: no response")
|
|
return []
|
|
|
|
text_preview = response.text[:800].replace("\n", " ").replace("\r", " ")
|
|
log.info(
|
|
f"CivitAI tRPC response: status={response.status_code}, content_type={response.headers.get('content-type')}, text_preview={text_preview}"
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
return []
|
|
|
|
try:
|
|
payload = response.json()
|
|
except Exception as e:
|
|
log.warning(f"CivitAI tRPC JSON parse failed for filename={filename}: {e}")
|
|
return []
|
|
|
|
candidates = _extract_trpc_model_candidates(payload, limit=limit)
|
|
log.info(
|
|
f"CivitAI tRPC extracted {len(candidates)} candidates for filename={filename}: {candidates}"
|
|
)
|
|
return candidates
|
|
|
|
|
|
def _search_civitai_red_candidates(
|
|
filename: str,
|
|
model_type: Optional[str] = None,
|
|
session_token: Optional[str] = None,
|
|
timeout: int = 15,
|
|
limit: int = 5,
|
|
) -> List[Dict[str, int]]:
|
|
"""Search civitai.red by full filename and return model/version candidates."""
|
|
search_query = get_filename_from_path(filename)
|
|
if _has_known_model_extension(search_query):
|
|
search_query = os.path.splitext(search_query)[0]
|
|
params = {
|
|
"sortBy": "models_v9",
|
|
"query": search_query,
|
|
}
|
|
civitai_type = CIVITAI_API_TYPE_MAP.get(str(model_type).lower()) if model_type else None
|
|
if civitai_type:
|
|
params["modelType"] = civitai_type
|
|
|
|
search_url = "https://civitai.red/search/models?" + urlencode(params)
|
|
headers = {"user-agent": DEFAULT_BROWSER_USER_AGENT}
|
|
if session_token:
|
|
headers["Cookie"] = build_civitai_session_cookie(session_token)
|
|
log.info(f"CivitAI.red search start: filename={filename}, url={search_url}")
|
|
|
|
response = request_source_response(search_url, headers=headers, timeout=timeout, log_name="CivitAI.red")
|
|
if response is None or response.status_code != 200:
|
|
if response is not None and response.status_code in {401, 403}:
|
|
return []
|
|
status = response.status_code if response is not None else "no-response"
|
|
log.warning(
|
|
f"CivitAI.red search returned {status} for filename={filename}"
|
|
)
|
|
return []
|
|
|
|
candidates = _extract_civitai_red_candidates(response.text, limit=limit)
|
|
log.info(
|
|
f"CivitAI.red search extracted {len(candidates)} candidates for filename={filename}"
|
|
)
|
|
return candidates
|
|
|
|
|
|
def _extract_public_api_model_candidates(
|
|
payload: Any, limit: int = 5
|
|
) -> List[Dict[str, Optional[int]]]:
|
|
"""Extract model/version candidates from the public CivitAI /models API."""
|
|
candidates: List[Dict[str, Optional[int]]] = []
|
|
seen = set()
|
|
items = payload.get("items", []) if isinstance(payload, dict) else []
|
|
if not isinstance(items, list):
|
|
return candidates
|
|
|
|
for item in items:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
|
|
model_id = item.get("id")
|
|
if not isinstance(model_id, int):
|
|
continue
|
|
|
|
version_id = None
|
|
versions = item.get("modelVersions") or []
|
|
if isinstance(versions, list):
|
|
for version in versions:
|
|
if isinstance(version, dict) and isinstance(version.get("id"), int):
|
|
version_id = version.get("id")
|
|
break
|
|
|
|
key = (model_id, version_id)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
candidates.append({"model_id": model_id, "version_id": version_id})
|
|
if len(candidates) >= limit:
|
|
break
|
|
|
|
return candidates
|
|
|
|
|
|
def _search_civitai_public_api_candidates(
|
|
filename: str,
|
|
model_type: Optional[str] = None,
|
|
api_key: Optional[str] = None,
|
|
session_token: Optional[str] = None,
|
|
timeout: int = 15,
|
|
limit: int = 5,
|
|
) -> List[Dict[str, Optional[int]]]:
|
|
"""Search the public CivitAI API by model name and return model/version candidates."""
|
|
search_query = get_filename_from_path(filename)
|
|
if _has_known_model_extension(search_query):
|
|
search_query = os.path.splitext(search_query)[0]
|
|
params = {
|
|
"query": search_query,
|
|
"limit": max(1, min(int(limit), MAX_CIVITAI_CANDIDATE_LIMIT)),
|
|
"sort": "Highest Rated",
|
|
"period": "AllTime",
|
|
}
|
|
civitai_type = resolve_model_category(model_type, target_format="civitai") if model_type else None
|
|
if civitai_type:
|
|
params["types"] = civitai_type
|
|
|
|
headers = {
|
|
"accept": "application/json",
|
|
"user-agent": DEFAULT_BROWSER_USER_AGENT,
|
|
}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
if session_token:
|
|
headers["Cookie"] = build_civitai_session_cookie(session_token)
|
|
|
|
log.info(
|
|
f"CivitAI public API search start: filename={filename}, model_type={model_type}, api_key={'yes' if api_key else 'no'}, session_token={'yes' if session_token else 'no'}"
|
|
)
|
|
|
|
response = request_source_response(
|
|
f"{CIVITAI_API_URL}/models",
|
|
params=params,
|
|
headers=headers,
|
|
timeout=timeout,
|
|
log_name="CivitAI public API"
|
|
)
|
|
if response is None:
|
|
log.warning(f"CivitAI public API request failed for filename={filename}: no response")
|
|
return []
|
|
|
|
text_preview = response.text[:800].replace("\n", " ").replace("\r", " ")
|
|
log.info(
|
|
f"CivitAI public API response: status={response.status_code}, content_type={response.headers.get('content-type')}, text_preview={text_preview}"
|
|
)
|
|
|
|
if response.status_code != 200:
|
|
return []
|
|
|
|
try:
|
|
payload = response.json()
|
|
except Exception as e:
|
|
log.warning(f"CivitAI public API JSON parse failed for filename={filename}: {e}")
|
|
return []
|
|
|
|
candidates = _extract_public_api_model_candidates(payload, limit=limit)
|
|
log.info(
|
|
f"CivitAI public API extracted {len(candidates)} candidates for filename={filename}: {candidates}"
|
|
)
|
|
return candidates
|
|
|
|
|
|
def _find_civitai_file_in_model(
|
|
model_id: int,
|
|
filename: str,
|
|
api_key: Optional[str] = None,
|
|
exact_only: bool = False,
|
|
preferred_version_id: Optional[int] = None,
|
|
base_model_context: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Load one CivitAI model and search all its versions for the requested file."""
|
|
filename_lower = filename.lower()
|
|
filename_base = os.path.splitext(filename_lower)[0]
|
|
allow_model_title_match = not exact_only and not _has_known_model_extension(
|
|
get_filename_from_path(filename)
|
|
)
|
|
|
|
def build_result_from_resolved_version(
|
|
resolved: Dict[str, Any], version_id: int
|
|
) -> Dict[str, Any]:
|
|
expected_filename = resolved.get("expected_filename", "")
|
|
primary_file = None
|
|
for file_info in resolved.get("files", []):
|
|
if file_info.get("name") == expected_filename:
|
|
primary_file = file_info
|
|
break
|
|
if primary_file is None:
|
|
primary_file = (resolved.get("files") or [{}])[0]
|
|
|
|
hashes = primary_file.get("hashes") if isinstance(primary_file.get("hashes"), dict) else {}
|
|
sha256 = primary_file.get("sha256") or hashes.get("SHA256") or hashes.get("sha256")
|
|
|
|
return {
|
|
"source": "civitai",
|
|
"model_id": model_id,
|
|
"version_id": version_id,
|
|
"name": resolved.get("model_name", ""),
|
|
"type": "",
|
|
"filename": expected_filename,
|
|
"url": f"https://civitai.com/models/{model_id}?modelVersionId={version_id}",
|
|
"download_url": get_civitai_download_url(version_id, api_key),
|
|
"size": primary_file.get("size"),
|
|
"base_model": resolved.get("base_model"),
|
|
"tags": resolved.get("tags", []),
|
|
"match_type": "exact",
|
|
"confidence": calculate_filename_confidence(
|
|
filename, expected_filename
|
|
),
|
|
"sha256": sha256,
|
|
"hashes": hashes,
|
|
}
|
|
|
|
def resolved_version_matches(resolved: Dict[str, Any]) -> bool:
|
|
expected_filename = str(resolved.get("expected_filename", "")).lower()
|
|
expected_base = os.path.splitext(expected_filename)[0]
|
|
log.debug(
|
|
f"CivitAI resolved version filename check: expected_filename={resolved.get('expected_filename')}, target_filename={filename}"
|
|
)
|
|
if expected_filename == filename_lower:
|
|
return True
|
|
return not exact_only and _filename_base_partial_match(
|
|
filename_base, expected_base
|
|
)
|
|
|
|
if preferred_version_id is not None:
|
|
resolved = resolve_urn(model_id, preferred_version_id, api_key)
|
|
if (
|
|
resolved
|
|
and resolved_version_matches(resolved)
|
|
and _base_model_matches(resolved.get("base_model"), base_model_context)
|
|
):
|
|
return build_result_from_resolved_version(resolved, preferred_version_id)
|
|
|
|
best_resolved_result = None
|
|
best_resolved_confidence = 0.0
|
|
data = execute_provider_json_request(
|
|
"CivitAI model lookup",
|
|
f"{CIVITAI_API_URL}/models/{model_id}",
|
|
api_key=api_key,
|
|
timeout=15,
|
|
)
|
|
if data is None:
|
|
return None
|
|
versions = data.get("modelVersions", [])
|
|
|
|
if allow_model_title_match:
|
|
title_match = _find_model_title_match_in_model(
|
|
model_id=model_id,
|
|
model_data=data,
|
|
title_query=filename,
|
|
api_key=api_key,
|
|
base_model_context=base_model_context,
|
|
)
|
|
if title_match:
|
|
return title_match
|
|
|
|
if preferred_version_id is not None:
|
|
preferred = [v for v in versions if v.get("id") == preferred_version_id]
|
|
others = [v for v in versions if v.get("id") != preferred_version_id]
|
|
versions = preferred + others
|
|
|
|
log.debug(
|
|
f"CivitAI model lookup model_id={model_id} returned {len(versions)} versions"
|
|
)
|
|
|
|
for version in versions:
|
|
version_id = version.get("id")
|
|
if not version_id:
|
|
continue
|
|
|
|
resolved = resolve_urn(model_id, version_id, api_key)
|
|
if resolved:
|
|
if resolved_version_matches(resolved) and _base_model_matches(
|
|
resolved.get("base_model"), base_model_context
|
|
):
|
|
return build_result_from_resolved_version(resolved, version_id)
|
|
|
|
if not exact_only:
|
|
expected_filename = resolved.get("expected_filename", "")
|
|
confidence = calculate_filename_confidence(filename, expected_filename)
|
|
ranking_score = confidence + _base_model_score(
|
|
resolved.get("base_model"), base_model_context
|
|
)
|
|
log.debug(
|
|
f"CivitAI resolved version confidence: model_id={model_id}, version_id={version_id}, expected_filename={expected_filename}, confidence={confidence}"
|
|
)
|
|
if ranking_score > best_resolved_confidence:
|
|
best_resolved_confidence = ranking_score
|
|
best_resolved_result = build_result_from_resolved_version(
|
|
resolved, version_id
|
|
)
|
|
best_resolved_result["match_type"] = "similar"
|
|
best_resolved_result["confidence"] = confidence
|
|
|
|
files = version.get("files", [])
|
|
file_names = [f.get("name", "") for f in files if isinstance(f, dict)]
|
|
log.debug(
|
|
f"CivitAI version file names: model_id={model_id}, version_id={version_id}, files={file_names}"
|
|
)
|
|
|
|
if best_resolved_result and best_resolved_confidence >= 50.0:
|
|
log.info(
|
|
f"CivitAI best probable match: model_id={model_id}, version_id={best_resolved_result.get('version_id')}, filename={best_resolved_result.get('filename')}, confidence={best_resolved_confidence}"
|
|
)
|
|
return best_resolved_result
|
|
|
|
match = _find_matching_file_in_versions(versions, filename, exact_only=exact_only)
|
|
if match:
|
|
version = match["version"]
|
|
file_info = match["file_info"]
|
|
result = _build_civitai_result_from_version(
|
|
model_id=model_id,
|
|
model_name=data.get("name", ""),
|
|
model_type=data.get("type", ""),
|
|
version=version,
|
|
file_info=file_info,
|
|
tags=data.get("tags", []),
|
|
match_type=match["match_type"],
|
|
)
|
|
result["confidence"] = match.get(
|
|
"confidence",
|
|
calculate_filename_confidence(filename, result.get("filename", "")),
|
|
)
|
|
if match["match_type"] == "similar":
|
|
log.info(
|
|
f"CivitAI version-list probable match: model_id={model_id}, version_id={result.get('version_id')}, filename={result.get('filename')}, confidence={result['confidence']}"
|
|
)
|
|
if _base_model_matches(result.get("base_model"), base_model_context):
|
|
return result
|
|
if not base_model_context:
|
|
return result
|
|
|
|
return None
|
|
|
|
|
|
|
|
def parse_civitai_url(url: str) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Parse a CivitAI URL to extract model/version info.
|
|
"""
|
|
from ..type_utils import parse_provider_model_url
|
|
return parse_provider_model_url(url, ["civitai.com", "civitai.red"])
|
|
|
|
|
|
def get_civitai_download_url(version_id: int, api_key: Optional[str] = None) -> str:
|
|
"""Get download URL for a CivitAI model version."""
|
|
url = f"https://civitai.com/api/download/models/{version_id}"
|
|
if api_key:
|
|
url += f"?token={api_key}"
|
|
return url
|
|
|
|
|
|
def search_civitai_for_file(
|
|
filename: str,
|
|
api_key: Optional[str] = None,
|
|
exact_only: bool = False,
|
|
model_type: Optional[str] = None,
|
|
base_model_context: Optional[str] = None,
|
|
session_token: Optional[str] = None,
|
|
candidate_limit: int = DEFAULT_CIVITAI_CANDIDATE_LIMIT,
|
|
use_trpc_search: bool = True,
|
|
use_api_search: bool = True,
|
|
use_html_fallback: bool = True,
|
|
progress_callback: Optional[Callable[[Dict[str, Any]], None]] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Search CivitAI for a specific model file.
|
|
Returns the first model that actually has this exact filename.
|
|
|
|
Args:
|
|
filename: Exact filename to search for
|
|
api_key: Optional API key
|
|
exact_only: If True, only return exact filename matches (for downloads).
|
|
If False, also try partial matching (for local file resolution).
|
|
|
|
Returns:
|
|
Dict with download info if found, None otherwise
|
|
"""
|
|
global _search_cache
|
|
|
|
candidate_limit = max(1, min(int(candidate_limit), MAX_CIVITAI_CANDIDATE_LIMIT))
|
|
session_key = "session" if session_token else "anon"
|
|
model_type_key = str(model_type or "").lower()
|
|
base_model_key = _normalize_base_model(base_model_context or "")
|
|
methods_key = (
|
|
f"trpc{int(bool(use_trpc_search))}"
|
|
f"_api{int(bool(use_api_search))}"
|
|
f"_html{int(bool(use_html_fallback))}"
|
|
)
|
|
cache_key = (
|
|
f"civit_{filename}_exact{exact_only}_type{model_type_key}_base{base_model_key}_session{session_key}_limit{candidate_limit}_{methods_key}"
|
|
)
|
|
if cache_key in _search_cache:
|
|
log.debug(f"CivitAI search cache hit for {filename} (exact_only={exact_only})")
|
|
_report_progress(
|
|
progress_callback,
|
|
"cache",
|
|
"Using CivitAI cache",
|
|
86,
|
|
)
|
|
return _search_cache[cache_key]
|
|
|
|
try:
|
|
# Prefer the CivitAI.red tRPC search because it is much closer to what
|
|
# the browser search UI returns than the broad public /models API.
|
|
if use_trpc_search:
|
|
_report_progress(
|
|
progress_callback,
|
|
"trpc",
|
|
"Searching CivitAI tRPC",
|
|
32,
|
|
)
|
|
trpc_candidates = (
|
|
_search_civitai_trpc_candidates(
|
|
filename,
|
|
model_type=model_type,
|
|
session_token=session_token,
|
|
limit=candidate_limit,
|
|
)
|
|
if use_trpc_search
|
|
else []
|
|
)
|
|
_report_progress(
|
|
progress_callback,
|
|
"trpc_candidates",
|
|
f"CivitAI tRPC candidates: {len(trpc_candidates)}",
|
|
42,
|
|
candidate_count=len(trpc_candidates),
|
|
)
|
|
candidates_to_check: List[Dict[str, Optional[int]]] = []
|
|
seen_candidates = set()
|
|
|
|
def add_candidates(candidates: List[Dict[str, Optional[int]]]) -> None:
|
|
for candidate in candidates:
|
|
key = (candidate.get("model_id"), candidate.get("version_id"))
|
|
if key in seen_candidates:
|
|
continue
|
|
seen_candidates.add(key)
|
|
candidates_to_check.append(candidate)
|
|
if len(candidates_to_check) >= candidate_limit:
|
|
break
|
|
|
|
add_candidates(trpc_candidates)
|
|
|
|
if use_html_fallback and not candidates_to_check:
|
|
_report_progress(
|
|
progress_callback,
|
|
"html",
|
|
"Searching CivitAI HTML fallback",
|
|
50,
|
|
)
|
|
html_candidates = _search_civitai_red_candidates(
|
|
filename,
|
|
model_type=model_type,
|
|
session_token=session_token,
|
|
limit=candidate_limit,
|
|
)
|
|
add_candidates(html_candidates)
|
|
else:
|
|
html_candidates = []
|
|
_report_progress(
|
|
progress_callback,
|
|
"html_candidates",
|
|
f"CivitAI HTML candidates: {len(html_candidates)}",
|
|
58,
|
|
candidate_count=len(html_candidates),
|
|
)
|
|
if use_api_search and not candidates_to_check:
|
|
_report_progress(
|
|
progress_callback,
|
|
"api",
|
|
"Searching CivitAI API",
|
|
58,
|
|
)
|
|
api_candidates = _search_civitai_public_api_candidates(
|
|
filename,
|
|
model_type=model_type,
|
|
api_key=api_key,
|
|
session_token=session_token,
|
|
limit=candidate_limit,
|
|
)
|
|
add_candidates(api_candidates)
|
|
else:
|
|
api_candidates = []
|
|
_report_progress(
|
|
progress_callback,
|
|
"api_candidates",
|
|
f"CivitAI API candidates: {len(api_candidates)}",
|
|
60,
|
|
candidate_count=len(api_candidates),
|
|
)
|
|
_report_progress(
|
|
progress_callback,
|
|
"candidates",
|
|
f"CivitAI candidates to check: {len(candidates_to_check)}",
|
|
62,
|
|
candidate_count=len(candidates_to_check),
|
|
candidate_limit=candidate_limit,
|
|
)
|
|
|
|
best_result = None
|
|
best_confidence = 0.0
|
|
|
|
total_candidates = len(candidates_to_check)
|
|
for candidate_index, candidate in enumerate(candidates_to_check, start=1):
|
|
model_id = candidate["model_id"]
|
|
version_id = candidate.get("version_id")
|
|
_report_progress(
|
|
progress_callback,
|
|
"candidate",
|
|
f"Checking CivitAI candidate {candidate_index}/{total_candidates}",
|
|
62 + (candidate_index / max(1, total_candidates)) * 24,
|
|
model_id=model_id,
|
|
version_id=version_id,
|
|
candidate_index=candidate_index,
|
|
candidate_count=total_candidates,
|
|
)
|
|
log.info(
|
|
f"CivitAI candidate check: model_id={model_id}, preferred_version_id={version_id}, filename={filename}"
|
|
)
|
|
result = _find_civitai_file_in_model(
|
|
model_id=model_id,
|
|
filename=filename,
|
|
api_key=api_key,
|
|
exact_only=exact_only,
|
|
preferred_version_id=version_id,
|
|
base_model_context=base_model_context,
|
|
)
|
|
if result:
|
|
confidence = float(result.get("confidence") or 0.0)
|
|
if confidence >= 100.0:
|
|
_search_cache[cache_key] = result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found exact CivitAI match",
|
|
92,
|
|
model_id=result.get("model_id"),
|
|
version_id=result.get("version_id"),
|
|
confidence=confidence,
|
|
)
|
|
log.info(
|
|
f"Found exact CivitAI match for {filename}: model_id={result.get('model_id')}, version_id={result.get('version_id')}, candidate_limit={candidate_limit}"
|
|
)
|
|
return result
|
|
if confidence > best_confidence:
|
|
best_result = result
|
|
best_confidence = confidence
|
|
|
|
if best_result:
|
|
_search_cache[cache_key] = best_result
|
|
_report_progress(
|
|
progress_callback,
|
|
"found",
|
|
"Found CivitAI match",
|
|
92,
|
|
model_id=best_result.get("model_id"),
|
|
version_id=best_result.get("version_id"),
|
|
confidence=best_confidence,
|
|
)
|
|
log.info(
|
|
f"Found best CivitAI match for {filename}: model_id={best_result.get('model_id')}, version_id={best_result.get('version_id')}, confidence={best_confidence}, candidate_limit={candidate_limit}"
|
|
)
|
|
return best_result
|
|
|
|
# Not found
|
|
_search_cache[cache_key] = None
|
|
_report_progress(
|
|
progress_callback,
|
|
"done",
|
|
"CivitAI checked",
|
|
92,
|
|
trpc_candidate_count=len(trpc_candidates),
|
|
html_candidate_count=len(html_candidates),
|
|
api_candidate_count=len(api_candidates),
|
|
checked_candidate_count=len(candidates_to_check),
|
|
candidate_limit=candidate_limit,
|
|
)
|
|
log.info(
|
|
f"CivitAI search no result: filename={filename}, trpc_candidates={len(trpc_candidates)}, html_candidates={len(html_candidates)}, api_candidates={len(api_candidates)}, checked_candidates={len(candidates_to_check)}, candidate_limit={candidate_limit}"
|
|
)
|
|
return None
|
|
|
|
except Exception as e:
|
|
log.exception(f"CivitAI search error for {filename}: {e}")
|
|
_report_progress(
|
|
progress_callback,
|
|
"error",
|
|
f"CivitAI search error: {e}",
|
|
100,
|
|
status="error",
|
|
)
|
|
return None
|
|
|
|
|
|
def search_civitai(
|
|
query: str,
|
|
model_type: Optional[str] = None,
|
|
limit: int = 10,
|
|
api_key: Optional[str] = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Search CivitAI for models (general search).
|
|
Returns models that might be relevant.
|
|
"""
|
|
results = []
|
|
|
|
try:
|
|
params = {"query": query, "limit": limit, "nsfw": "false"}
|
|
|
|
if model_type:
|
|
civitai_type = CIVITAI_API_TYPE_MAP.get(model_type.lower())
|
|
if civitai_type:
|
|
params["types"] = civitai_type
|
|
|
|
data = execute_provider_json_request(
|
|
"CivitAI search models",
|
|
f"{CIVITAI_API_URL}/models",
|
|
params=params,
|
|
api_key=api_key,
|
|
timeout=15,
|
|
)
|
|
|
|
if data:
|
|
|
|
for model in data.get("items", []):
|
|
model_id = model.get("id")
|
|
model_name = model.get("name", "")
|
|
model_type = model.get("type", "")
|
|
|
|
versions = model.get("modelVersions", [])
|
|
if versions:
|
|
latest = versions[0]
|
|
version_id = latest.get("id")
|
|
|
|
files = latest.get("files", [])
|
|
primary_file = None
|
|
for f in files:
|
|
if f.get("primary", False) or f.get("type") == "Model":
|
|
primary_file = f
|
|
break
|
|
|
|
if not primary_file and files:
|
|
primary_file = files[0]
|
|
|
|
result = {
|
|
"source": "civitai",
|
|
"model_id": model_id,
|
|
"version_id": version_id,
|
|
"name": model_name,
|
|
"type": model_type,
|
|
"url": f"https://civitai.com/models/{model_id}",
|
|
"download_url": get_civitai_download_url(version_id, api_key),
|
|
"downloads": model.get("stats", {}).get("downloadCount", 0),
|
|
"base_model": latest.get("baseModel"),
|
|
"tags": model.get("tags", []),
|
|
}
|
|
|
|
if primary_file:
|
|
result["filename"] = primary_file.get("name", "")
|
|
result["size"] = primary_file.get("sizeKB", 0) * 1024
|
|
|
|
results.append(result)
|
|
|
|
except Exception as e:
|
|
log.error(f"CivitAI search error: {e}")
|
|
|
|
return results
|
|
|
|
|
|
def _normalize_civitai_file(
|
|
file_info: Dict[str, Any],
|
|
model_id: Optional[int],
|
|
version_id: Optional[int],
|
|
api_key: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
fallback_url = get_civitai_download_url(version_id, api_key) if version_id else None
|
|
return normalize_model_file_info(
|
|
file_info,
|
|
model_id=model_id,
|
|
version_id=version_id,
|
|
fallback_download_url=fallback_url
|
|
)
|
|
|
|
|
|
def _build_civitai_sidecar_payload(
|
|
model_data: Dict[str, Any],
|
|
version_data: Dict[str, Any],
|
|
model_id: int,
|
|
) -> Dict[str, Any]:
|
|
"""Preserve the selected CivitAI version for LoRA Manager sidecars."""
|
|
if not isinstance(version_data, dict):
|
|
return {}
|
|
|
|
payload = dict(version_data)
|
|
payload.setdefault("modelId", model_id)
|
|
|
|
model_payload = {
|
|
key: value
|
|
for key, value in model_data.items()
|
|
if key not in {"id", "modelVersions", "creator", "stats"}
|
|
}
|
|
existing_model = payload.get("model")
|
|
if isinstance(existing_model, dict):
|
|
model_payload.update(existing_model)
|
|
if model_payload:
|
|
payload["model"] = model_payload
|
|
|
|
creator = model_data.get("creator")
|
|
if isinstance(creator, dict) and "creator" not in payload:
|
|
payload["creator"] = creator
|
|
|
|
return payload
|
|
|
|
|
|
def get_civitai_model_details(
|
|
model_id: int,
|
|
version_id: Optional[int] = None,
|
|
api_key: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch and normalize full CivitAI model details for the resolver UI."""
|
|
if not model_id:
|
|
return None
|
|
|
|
data = execute_provider_json_request(
|
|
"CivitAI model details",
|
|
f"{CIVITAI_API_URL}/models/{model_id}",
|
|
api_key=api_key,
|
|
timeout=20,
|
|
)
|
|
if not data:
|
|
return None
|
|
versions = data.get("modelVersions") or []
|
|
normalized_versions = []
|
|
selected_version = None
|
|
selected_raw_version = None
|
|
raw_versions = []
|
|
|
|
for version in versions:
|
|
if not isinstance(version, dict):
|
|
continue
|
|
|
|
raw_versions.append(version)
|
|
current_version_id = version.get("id")
|
|
files = [
|
|
_normalize_civitai_file(file_info, model_id, current_version_id, api_key)
|
|
for file_info in (version.get("files") or [])
|
|
if isinstance(file_info, dict)
|
|
]
|
|
normalized = {
|
|
"id": current_version_id,
|
|
"name": version.get("name") or "",
|
|
"base_model": version.get("baseModel"),
|
|
"published_at": version.get("publishedAt"),
|
|
"updated_at": version.get("updatedAt"),
|
|
"description": version.get("description") or "",
|
|
"trained_words": extract_trained_words(version),
|
|
"clip_skip": version.get("clipSkip"),
|
|
"stats": version.get("stats") or {},
|
|
"files": files,
|
|
"images": _extract_model_images(
|
|
version,
|
|
model_info=data,
|
|
model_id=model_id,
|
|
),
|
|
"url": f"https://civitai.com/models/{model_id}?modelVersionId={current_version_id}",
|
|
}
|
|
normalized_versions.append(normalized)
|
|
if version_id and str(current_version_id) == str(version_id):
|
|
selected_version = normalized
|
|
selected_raw_version = version
|
|
|
|
if not selected_version and normalized_versions:
|
|
selected_version = normalized_versions[0]
|
|
selected_raw_version = raw_versions[0]
|
|
|
|
selected_images = selected_version.get("images", []) if selected_version else []
|
|
civitai_payload = _build_civitai_sidecar_payload(
|
|
data,
|
|
selected_raw_version or {},
|
|
model_id,
|
|
)
|
|
if selected_images:
|
|
normalized_by_url = {
|
|
str(image.get("url")): image
|
|
for image in selected_images
|
|
if isinstance(image, dict) and image.get("url")
|
|
}
|
|
raw_images = civitai_payload.get("images")
|
|
if isinstance(raw_images, list):
|
|
enriched_raw_images = []
|
|
for raw_image in raw_images:
|
|
if not isinstance(raw_image, dict):
|
|
enriched_raw_images.append(raw_image)
|
|
continue
|
|
enriched_image = dict(raw_image)
|
|
normalized_image = normalized_by_url.get(str(raw_image.get("url") or ""))
|
|
if normalized_image and normalized_image.get("resources"):
|
|
enriched_image["resources"] = normalized_image["resources"]
|
|
enriched_raw_images.append(enriched_image)
|
|
civitai_payload["images"] = enriched_raw_images
|
|
else:
|
|
civitai_payload["images"] = selected_images
|
|
return {
|
|
"source": "civitai",
|
|
"model_id": model_id,
|
|
"version_id": selected_version.get("id") if selected_version else version_id,
|
|
"name": data.get("name") or "",
|
|
"type": data.get("type") or "",
|
|
"description": data.get("description") or "",
|
|
"tags": data.get("tags") or [],
|
|
"stats": data.get("stats") or {},
|
|
"creator": data.get("creator") or {},
|
|
"url": f"https://civitai.com/models/{model_id}",
|
|
"version_url": selected_version.get("url") if selected_version else f"https://civitai.com/models/{model_id}",
|
|
"versions": normalized_versions,
|
|
"selected_version": selected_version,
|
|
"images": selected_images,
|
|
"civitai": civitai_payload,
|
|
}
|
|
|
|
|
|
def _enrich_model_info_with_details(
|
|
result: Dict[str, Any],
|
|
model_id: Optional[int],
|
|
version_id: Optional[int],
|
|
api_key: Optional[str] = None,
|
|
) -> Dict[str, Any]:
|
|
"""Fill hash lookup results with the fuller /models/{id} payload when useful."""
|
|
if not result or not model_id:
|
|
return result
|
|
|
|
needs_details = (
|
|
not result.get("images")
|
|
or not result.get("description")
|
|
or not result.get("model_description")
|
|
or not result.get("tags")
|
|
)
|
|
if not needs_details:
|
|
return result
|
|
|
|
try:
|
|
details = get_civitai_model_details(model_id, version_id, api_key)
|
|
except Exception as exc:
|
|
log.debug(f"CivitAI details enrichment failed for model_id={model_id}: {exc}")
|
|
return result
|
|
|
|
if not details:
|
|
return result
|
|
|
|
selected_version = as_dict(details.get("selected_version"))
|
|
result["model_name"] = result.get("model_name") or details.get("name")
|
|
result["model_type"] = result.get("model_type") or details.get("type")
|
|
result["description"] = (
|
|
result.get("description")
|
|
or selected_version.get("description")
|
|
or details.get("description")
|
|
or ""
|
|
)
|
|
result["model_description"] = (
|
|
result.get("model_description")
|
|
or details.get("description")
|
|
or ""
|
|
)
|
|
result["tags"] = result.get("tags") or details.get("tags") or []
|
|
result["images"] = result.get("images") or details.get("images") or []
|
|
result["version_url"] = result.get("version_url") or details.get("version_url")
|
|
result["url"] = result.get("url") or details.get("url")
|
|
return result
|
|
|
|
|
|
|
|
def resolve_urn(
|
|
model_id: int, version_id: int, api_key: Optional[str] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Resolve URN model_id/version_id to model info and expected filename.
|
|
|
|
Args:
|
|
model_id: CivitAI model ID
|
|
version_id: CivitAI version ID
|
|
api_key: Optional API key
|
|
|
|
Returns:
|
|
Dict with model name and primary filename, or None
|
|
"""
|
|
cache_key = (model_id, version_id)
|
|
if cache_key in _urn_cache:
|
|
return _urn_cache[cache_key]
|
|
|
|
try:
|
|
data = execute_provider_json_request(
|
|
"CivitAI URN resolve",
|
|
f"{CIVITAI_API_URL}/models/{model_id}",
|
|
params={"modelVersionId": version_id},
|
|
api_key=api_key,
|
|
timeout=15,
|
|
)
|
|
if not data:
|
|
_urn_cache[cache_key] = None
|
|
return None
|
|
versions = data.get("modelVersions", [])
|
|
|
|
if not versions:
|
|
log.warning(f"No versions found for model {model_id}/version {version_id}")
|
|
_urn_cache[cache_key] = None
|
|
return None
|
|
|
|
# Get specific version
|
|
target_version = None
|
|
for v in versions:
|
|
if v.get("id") == version_id:
|
|
target_version = v
|
|
break
|
|
|
|
if not target_version:
|
|
log.warning(f"Version {version_id} not found in model {model_id}")
|
|
_urn_cache[cache_key] = None
|
|
return None
|
|
|
|
files = target_version.get("files", [])
|
|
primary_file = None
|
|
|
|
# Prefer primary file or type=='Model'
|
|
for f in files:
|
|
if f.get("primary") or f.get("type") == "Model":
|
|
primary_file = f
|
|
break
|
|
|
|
if not primary_file and files:
|
|
primary_file = files[0] # Fallback to first
|
|
|
|
if not primary_file:
|
|
log.warning(f"No files found for version {version_id}")
|
|
_urn_cache[cache_key] = None
|
|
return None
|
|
|
|
result = {
|
|
"model_name": data.get("name", "Unknown"),
|
|
"version_name": target_version.get("name", "Unknown"),
|
|
"expected_filename": primary_file.get("name", "Unknown"),
|
|
"base_model": target_version.get("baseModel"),
|
|
"tags": data.get("tags", []),
|
|
"files": [
|
|
{
|
|
"name": f.get("name"),
|
|
"size": f.get("sizeKB", 0) * 1024,
|
|
"sha256": (
|
|
f.get("sha256")
|
|
or (f.get("hashes") or {}).get("SHA256")
|
|
or (f.get("hashes") or {}).get("sha256")
|
|
),
|
|
"hashes": f.get("hashes") or {},
|
|
}
|
|
for f in files
|
|
],
|
|
}
|
|
|
|
_urn_cache[cache_key] = result
|
|
log.info(
|
|
f"CivitAI resolved urn={model_id}@{version_id} file={result['expected_filename']}"
|
|
)
|
|
return result
|
|
|
|
except Exception as e:
|
|
log.error(f"CivitAI URN resolve error for {model_id}@{version_id}: {e}")
|
|
_urn_cache[cache_key] = None
|
|
return None
|
|
|
|
|
|
|
|
|
|
def get_model_info_by_hash(
|
|
file_hash: str, api_key: Optional[str] = None, use_cache: bool = True
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Look up a model on CivitAI using its sha256 hash.
|
|
Uses the CivitAI API endpoint: /api/v1/model-versions/by-hash/{hash}
|
|
|
|
Args:
|
|
file_hash: SHA256 hash of the model file
|
|
api_key: Optional CivitAI API key
|
|
use_cache: Whether to use cached results
|
|
|
|
Returns:
|
|
Dict with model info from CivitAI, or None if not found
|
|
"""
|
|
global _hash_cache
|
|
|
|
if not file_hash:
|
|
return None
|
|
|
|
cache_key = f"hash_{file_hash}"
|
|
|
|
if use_cache and cache_key in _hash_cache:
|
|
return _hash_cache[cache_key]
|
|
|
|
api_url = f"{CIVITAI_API_URL}/model-versions/by-hash/{file_hash}"
|
|
|
|
try:
|
|
data = execute_provider_json_request(
|
|
"CivitAI by-hash",
|
|
api_url,
|
|
api_key=api_key,
|
|
timeout=15,
|
|
raise_on_error=True,
|
|
)
|
|
|
|
if isinstance(data, dict) and data:
|
|
# Extract useful info from the response
|
|
model_info = data.get("model", {})
|
|
version_info = data
|
|
|
|
# Get model_id from top level or nested model object
|
|
model_id = data.get("modelId") or model_info.get("id")
|
|
|
|
# Extract images with metadata
|
|
images = _extract_model_images(
|
|
version_info,
|
|
model_info=model_info,
|
|
model_id=model_id,
|
|
)
|
|
|
|
result = {
|
|
"source": "civitai",
|
|
"model_id": model_id,
|
|
"model_name": model_info.get("name"),
|
|
"model_type": model_info.get("type"),
|
|
"version_id": version_info.get("id"),
|
|
"version_name": version_info.get("name"),
|
|
"sha256": file_hash,
|
|
"url": f"https://civitai.com/models/{model_id}" if model_id else None,
|
|
"version_url": f"https://civitai.com/models/{model_id}?modelVersionId={version_info.get('id')}"
|
|
if model_id
|
|
else None,
|
|
"download_url": version_info.get("downloadUrl"),
|
|
"base_model": version_info.get("baseModel"),
|
|
"tags": model_info.get("tags", []),
|
|
"trained_words": extract_trained_words(version_info),
|
|
"images": images,
|
|
"clip_skip": version_info.get("clipSkip"),
|
|
"description": version_info.get("description", ""),
|
|
"model_description": model_info.get("description", ""),
|
|
}
|
|
|
|
result = _enrich_model_info_with_details(
|
|
result,
|
|
model_id,
|
|
version_info.get("id"),
|
|
api_key,
|
|
)
|
|
_hash_cache[cache_key] = result
|
|
log.info(f"Found model by hash {file_hash}: {result.get('model_name')}")
|
|
return result
|
|
|
|
log.warning(
|
|
f"CivitAI hash lookup returned an unexpected JSON payload for {file_hash}"
|
|
)
|
|
return None
|
|
|
|
except requests.HTTPError as e:
|
|
status_code = getattr(e.response, "status_code", None)
|
|
if status_code == 404:
|
|
log.info(f"Model not found on CivitAI for hash {file_hash}")
|
|
_hash_cache[cache_key] = None
|
|
return None
|
|
log.warning(
|
|
f"CivitAI hash lookup returned HTTP {status_code or 'error'} for {file_hash}"
|
|
)
|
|
return None
|
|
|
|
except Exception as e:
|
|
log.error(f"Error looking up model by hash {file_hash}: {e}")
|
|
return None
|
|
|
|
def _extract_model_images(
|
|
version_info: Dict[str, Any],
|
|
*,
|
|
model_info: Optional[Dict[str, Any]] = None,
|
|
model_id: Any = None,
|
|
) -> List[Dict[str, Any]]:
|
|
"""
|
|
Extract images with metadata from model version info.
|
|
Each image can have: url, civitaiUrl, seed, steps, cfg, sampler, model, positive, negative
|
|
"""
|
|
images = []
|
|
image_version_info = version_info
|
|
|
|
# Get images from modelVersions or direct images field
|
|
images_data = version_info.get("images", [])
|
|
if not images_data:
|
|
# Check nested modelVersions
|
|
model_versions = version_info.get("modelVersions", [])
|
|
if model_versions and len(model_versions) > 0:
|
|
image_version_info = model_versions[0]
|
|
images_data = image_version_info.get("images", [])
|
|
|
|
resolved_model_info = model_info
|
|
if not isinstance(resolved_model_info, dict):
|
|
nested_model_info = version_info.get("model")
|
|
resolved_model_info = (
|
|
nested_model_info
|
|
if isinstance(nested_model_info, dict)
|
|
else version_info
|
|
if version_info.get("modelVersions")
|
|
else {}
|
|
)
|
|
|
|
resolved_model_id = (
|
|
model_id
|
|
or image_version_info.get("modelId")
|
|
or version_info.get("modelId")
|
|
or resolved_model_info.get("id")
|
|
)
|
|
version_id = image_version_info.get("id")
|
|
model_name = resolved_model_info.get("name") or ""
|
|
version_name = image_version_info.get("name") or ""
|
|
model_type = resolved_model_info.get("type") or ""
|
|
primary_resource = None
|
|
if resolved_model_id and version_id and model_name:
|
|
primary_resource = {
|
|
"name": model_name,
|
|
"versionName": version_name,
|
|
"type": model_type,
|
|
"modelId": resolved_model_id,
|
|
"modelVersionId": version_id,
|
|
"url": (
|
|
f"https://civitai.com/models/{resolved_model_id}"
|
|
f"?modelVersionId={version_id}"
|
|
),
|
|
"primary": True,
|
|
}
|
|
|
|
for img in images_data:
|
|
if not isinstance(img, dict):
|
|
continue
|
|
|
|
img_info = normalize_model_image(img)
|
|
if img_info.get("url"):
|
|
resources = img_info.get("resources")
|
|
if not isinstance(resources, list):
|
|
resources = []
|
|
has_linked_primary = any(
|
|
isinstance(resource, dict)
|
|
and str(resource.get("modelId") or resource.get("model_id") or "")
|
|
== str(resolved_model_id)
|
|
and str(
|
|
resource.get("modelVersionId")
|
|
or resource.get("versionId")
|
|
or resource.get("model_version_id")
|
|
or ""
|
|
)
|
|
== str(version_id)
|
|
for resource in resources
|
|
)
|
|
if primary_resource and not has_linked_primary:
|
|
matching_resource = next(
|
|
(
|
|
resource
|
|
for resource in resources
|
|
if _resource_matches_parent_model(
|
|
resource,
|
|
image_version_info,
|
|
model_name,
|
|
)
|
|
),
|
|
None,
|
|
)
|
|
enriched_primary = dict(primary_resource)
|
|
if matching_resource:
|
|
for weight_key in ("strength", "weight", "value"):
|
|
if weight_key in matching_resource:
|
|
enriched_primary[weight_key] = matching_resource[weight_key]
|
|
resources = [
|
|
resource
|
|
for resource in resources
|
|
if resource is not matching_resource
|
|
]
|
|
img_info["resources"] = [enriched_primary, *resources]
|
|
images.append(img_info)
|
|
|
|
return images
|
|
|
|
|
|
def _resource_matches_parent_model(
|
|
resource: Any,
|
|
version_info: Dict[str, Any],
|
|
model_name: str,
|
|
) -> bool:
|
|
"""Match an unlinked image resource to the model version that owns the image."""
|
|
if not isinstance(resource, dict):
|
|
return False
|
|
|
|
resource_name = str(
|
|
resource.get("name")
|
|
or resource.get("modelName")
|
|
or resource.get("model")
|
|
or ""
|
|
).strip()
|
|
if not resource_name:
|
|
return False
|
|
|
|
for file_info in version_info.get("files") or []:
|
|
if not isinstance(file_info, dict):
|
|
continue
|
|
file_name = str(file_info.get("name") or "").strip()
|
|
if file_name and calculate_filename_confidence(resource_name, file_name) >= 95:
|
|
return True
|
|
|
|
title_prefix = re.split(r"[\(\[\|]", str(model_name or ""), maxsplit=1)[0]
|
|
compact_title = re.sub(r"[^a-z0-9]+", "", title_prefix.lower())
|
|
compact_resource = re.sub(r"[^a-z0-9]+", "", resource_name.lower())
|
|
return len(compact_title) >= 8 and compact_title in compact_resource
|
|
|
|
|
|
def _read_model_metadata(
|
|
metadata_path: str,
|
|
model_path: str = "",
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Read model metadata from a JSON file.
|
|
Returns the metadata if found and valid, None otherwise.
|
|
"""
|
|
data = (
|
|
read_merged_model_metadata(model_path, None)
|
|
if model_path
|
|
else read_json_safe(metadata_path, None)
|
|
)
|
|
if not isinstance(data, dict):
|
|
return None
|
|
|
|
try:
|
|
files = as_list(data.get("files"))
|
|
file_hashes = [
|
|
as_dict(file_info).get("hashes")
|
|
for file_info in files
|
|
if isinstance(file_info, dict)
|
|
]
|
|
has_file_sha = any(
|
|
isinstance(hashes, dict)
|
|
and (hashes.get("SHA256") or hashes.get("sha256"))
|
|
for hashes in file_hashes
|
|
)
|
|
has_civitai_shape = bool(
|
|
data.get("civitai")
|
|
or data.get("modelId")
|
|
or data.get("model")
|
|
or data.get("images")
|
|
or data.get("downloadUrl")
|
|
)
|
|
|
|
# Check if it has the needed info. Older sidecars use a top-level
|
|
# CivitAI model-version shape and are commonly named *.civitai.info.
|
|
if (
|
|
not data.get("sha256")
|
|
and not data.get("hash")
|
|
and not data.get("hashes")
|
|
and not has_file_sha
|
|
and not has_civitai_shape
|
|
):
|
|
return None
|
|
|
|
log.info(f"Successfully read metadata from: {metadata_path}")
|
|
return data
|
|
except Exception as e:
|
|
log.debug(f"Error parsing metadata file {metadata_path}: {e}")
|
|
return None
|
|
|
|
|
|
def _format_model_location(file_path: str) -> str:
|
|
if not file_path:
|
|
return ""
|
|
|
|
location = os.path.dirname(file_path).replace("\\", "/")
|
|
if location and not location.endswith("/"):
|
|
location += "/"
|
|
return location
|
|
|
|
|
|
def _has_metadata_value(value: Any) -> bool:
|
|
if value is None:
|
|
return False
|
|
if isinstance(value, str):
|
|
return bool(value.strip())
|
|
if isinstance(value, (list, tuple, dict, set)):
|
|
return bool(value)
|
|
return True
|
|
|
|
|
|
def _merge_unique_values(existing: Any, incoming: Any) -> List[Any]:
|
|
existing_list = existing if isinstance(existing, list) else ([existing] if existing else [])
|
|
incoming_list = incoming if isinstance(incoming, list) else ([incoming] if incoming else [])
|
|
result: List[Any] = []
|
|
seen = set()
|
|
for item in [*existing_list, *incoming_list]:
|
|
if item in (None, ""):
|
|
continue
|
|
if isinstance(item, dict):
|
|
key = json.dumps(item, sort_keys=True, default=str)
|
|
else:
|
|
key = str(item).strip().lower()
|
|
if not key or key in seen:
|
|
continue
|
|
seen.add(key)
|
|
result.append(item)
|
|
return result
|
|
|
|
|
|
def _apply_safetensors_header_metadata(
|
|
result: Dict[str, Any],
|
|
file_path: str,
|
|
) -> Dict[str, Any]:
|
|
if not isinstance(result, dict):
|
|
return result
|
|
|
|
header_metadata = extract_safetensors_header_metadata(file_path)
|
|
if not header_metadata:
|
|
return result
|
|
|
|
filename = get_filename_from_path(file_path)
|
|
stem = os.path.splitext(filename)[0]
|
|
header_name = header_metadata.get("model_name")
|
|
current_name = str(result.get("model_name") or "").strip()
|
|
if header_name and (
|
|
not current_name
|
|
or current_name in (stem, filename)
|
|
or result.get("source") == "local"
|
|
):
|
|
result["model_name"] = header_name
|
|
|
|
for key in (
|
|
"description",
|
|
"model_description",
|
|
"author",
|
|
"creator",
|
|
"license",
|
|
"usage_hint",
|
|
"usage_tips",
|
|
"preview_url",
|
|
"clip_skip",
|
|
"sha256",
|
|
"hash",
|
|
"hashes",
|
|
"sha256_source",
|
|
"hash_status",
|
|
"model_type",
|
|
"base_model_raw",
|
|
):
|
|
if _has_metadata_value(header_metadata.get(key)) and not _has_metadata_value(result.get(key)):
|
|
result[key] = header_metadata[key]
|
|
|
|
for key in ("tags", "trained_words", "images"):
|
|
if _has_metadata_value(header_metadata.get(key)):
|
|
if _has_metadata_value(result.get(key)):
|
|
result[key] = _merge_unique_values(result.get(key), header_metadata.get(key))
|
|
else:
|
|
result[key] = header_metadata[key]
|
|
|
|
if header_metadata.get("base_model") and not result.get("base_model"):
|
|
result["base_model"] = header_metadata["base_model"]
|
|
result["base_model_source"] = header_metadata.get("base_model_source")
|
|
result["base_model_inferred"] = bool(header_metadata.get("base_model_inferred"))
|
|
|
|
for key in (
|
|
"metadata_summary",
|
|
"header_metadata_keys",
|
|
"metadata_source",
|
|
"from_safetensors_header",
|
|
"local_metadata_available",
|
|
):
|
|
if _has_metadata_value(header_metadata.get(key)) and not _has_metadata_value(result.get(key)):
|
|
result[key] = header_metadata[key]
|
|
|
|
if (
|
|
result.get("source") == "local"
|
|
and header_metadata.get("details_source")
|
|
):
|
|
result["details_source"] = header_metadata["details_source"]
|
|
|
|
return result
|
|
|
|
|
|
def _apply_inferred_base_model(result: Dict[str, Any], file_path: str) -> Dict[str, Any]:
|
|
if not isinstance(result, dict) or result.get("base_model"):
|
|
return result
|
|
|
|
inferred_base_model = infer_safetensors_base_model(file_path)
|
|
if inferred_base_model:
|
|
result["base_model"] = inferred_base_model
|
|
result["base_model_inferred"] = True
|
|
result["base_model_source"] = "safetensors_header"
|
|
log.info(
|
|
f"Inferred base model from safetensors tensor fingerprint: "
|
|
f"{inferred_base_model} ({file_path})"
|
|
)
|
|
return result
|
|
|
|
|
|
def _has_local_header_details(result: Optional[Dict[str, Any]]) -> bool:
|
|
if not isinstance(result, dict):
|
|
return False
|
|
return bool(
|
|
result.get("local_metadata_available")
|
|
or result.get("from_safetensors_header")
|
|
or result.get("base_model_inferred")
|
|
or result.get("metadata_summary")
|
|
)
|
|
|
|
|
|
def _build_local_model_info(
|
|
file_path: str,
|
|
*,
|
|
local_only: bool = False,
|
|
require_inferred: bool = False,
|
|
) -> Optional[Dict[str, Any]]:
|
|
filename = get_filename_from_path(file_path)
|
|
stem = os.path.splitext(filename)[0]
|
|
result = {
|
|
"source": "local",
|
|
"details_source": "local",
|
|
"filename": filename,
|
|
"file_path": file_path,
|
|
"resolved_path": file_path,
|
|
"model_name": stem,
|
|
"model_type": "",
|
|
"location": _format_model_location(file_path),
|
|
"from_metadata": False,
|
|
"local_only": local_only,
|
|
}
|
|
try:
|
|
result["size"] = os.path.getsize(file_path)
|
|
except Exception:
|
|
pass
|
|
|
|
_apply_safetensors_header_metadata(result, file_path)
|
|
_apply_inferred_base_model(result, file_path)
|
|
if require_inferred and not _has_local_header_details(result):
|
|
return None
|
|
return result
|
|
|
|
|
|
def _metadata_to_model_info(metadata: Dict[str, Any]) -> Dict[str, Any]:
|
|
"""
|
|
Convert metadata file format to model info format used by our extension.
|
|
"""
|
|
embedded_civitai_data = as_dict(metadata.get("civitai"))
|
|
looks_like_civitai_version = bool(
|
|
metadata.get("modelId")
|
|
or metadata.get("files")
|
|
or metadata.get("images")
|
|
or metadata.get("downloadUrl")
|
|
)
|
|
civitai_data = embedded_civitai_data or (
|
|
metadata if looks_like_civitai_version else {}
|
|
)
|
|
selected_version = as_dict(metadata.get("selected_version"))
|
|
if not selected_version and looks_like_civitai_version:
|
|
selected_version = civitai_data
|
|
path_metadata = as_dict(metadata.get("path_metadata"))
|
|
|
|
# Extract images with metadata
|
|
images = []
|
|
for img in (
|
|
civitai_data.get("images")
|
|
or selected_version.get("images")
|
|
or metadata.get("images")
|
|
or []
|
|
):
|
|
if isinstance(img, dict) and img.get("url"):
|
|
img_info = normalize_model_image(img)
|
|
if img_info.get("url"):
|
|
images.append(img_info)
|
|
|
|
# Get trained words from common sidecar shapes, including LoRA Manager metadata.
|
|
trained_words = extract_trained_words(
|
|
metadata.get("trained_words"),
|
|
metadata.get("trainedWords"),
|
|
civitai_data.get("trainedWords"),
|
|
selected_version.get("trained_words"),
|
|
selected_version.get("trainedWords"),
|
|
path_metadata.get("trained_words"),
|
|
path_metadata.get("trainedWords"),
|
|
)
|
|
|
|
# Get model info
|
|
model_info = as_dict(civitai_data.get("model"))
|
|
file_infos = as_list(civitai_data.get("files"))
|
|
file_info = as_dict(file_infos[0]) if file_infos else {}
|
|
metadata_hashes = as_dict(metadata.get("hashes"))
|
|
file_hashes = as_dict(file_info.get("hashes"))
|
|
|
|
# Build model_id from CivitAI data
|
|
model_id = (
|
|
civitai_data.get("modelId")
|
|
or metadata.get("model_id")
|
|
or metadata.get("modelId")
|
|
)
|
|
version_id = (
|
|
civitai_data.get("id")
|
|
or metadata.get("version_id")
|
|
or metadata.get("versionId")
|
|
)
|
|
details_source = (
|
|
metadata.get("details_source")
|
|
or metadata.get("source")
|
|
or metadata.get("metadata_source")
|
|
or "metadata"
|
|
)
|
|
stored_page_url = first_non_empty(
|
|
metadata.get("version_url"),
|
|
metadata.get("model_url"),
|
|
metadata.get("page_url"),
|
|
metadata.get("source_url"),
|
|
metadata.get("url"),
|
|
metadata.get("platform_url"),
|
|
path_metadata.get("version_url"),
|
|
path_metadata.get("model_url"),
|
|
path_metadata.get("source_url"),
|
|
path_metadata.get("url"),
|
|
path_metadata.get("platform_url"),
|
|
)
|
|
civitai_model_url = f"https://civitai.com/models/{model_id}" if model_id else None
|
|
civitai_version_url = (
|
|
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
|
|
if model_id and version_id
|
|
else civitai_model_url
|
|
)
|
|
page_url = stored_page_url or civitai_version_url
|
|
model_description = first_non_empty(
|
|
metadata.get("modelDescription"),
|
|
metadata.get("model_description"),
|
|
metadata.get("description"),
|
|
model_info.get("description"),
|
|
) or ""
|
|
version_description = first_non_empty(
|
|
metadata.get("versionDescription"),
|
|
metadata.get("version_description"),
|
|
selected_version.get("description"),
|
|
civitai_data.get("description"),
|
|
) or ""
|
|
if version_description == model_description:
|
|
version_description = ""
|
|
|
|
return {
|
|
"source": details_source,
|
|
"details_source": details_source,
|
|
"model_id": model_id,
|
|
"model_name": first_non_empty(
|
|
metadata.get("model_name"),
|
|
metadata.get("modelName"),
|
|
model_info.get("name"),
|
|
metadata.get("name") if not model_id else None,
|
|
metadata.get("file_name"),
|
|
)
|
|
or "",
|
|
"model_type": model_info.get("type", "") or civitai_data.get("type", ""),
|
|
"version_id": version_id,
|
|
"version_name": first_non_empty(
|
|
metadata.get("version_name"),
|
|
metadata.get("versionName"),
|
|
selected_version.get("name"),
|
|
civitai_data.get("name"),
|
|
)
|
|
or "",
|
|
"sha256": first_non_empty(
|
|
metadata.get("sha256"),
|
|
metadata.get("hash"),
|
|
metadata_hashes.get("SHA256"),
|
|
metadata_hashes.get("sha256"),
|
|
file_hashes.get("SHA256"),
|
|
file_hashes.get("sha256"),
|
|
)
|
|
or "",
|
|
"size": parse_size_to_bytes(
|
|
first_non_empty(
|
|
metadata.get("size"),
|
|
metadata.get("file_size"),
|
|
metadata.get("fileSize"),
|
|
metadata.get("sizeBytes"),
|
|
path_metadata.get("size"),
|
|
path_metadata.get("file_size"),
|
|
path_metadata.get("fileSize"),
|
|
path_metadata.get("sizeBytes"),
|
|
(file_info.get("sizeKB") * 1024)
|
|
if file_info.get("sizeKB") is not None
|
|
else None,
|
|
)
|
|
),
|
|
"url": page_url,
|
|
"version_url": page_url,
|
|
"model_url": stored_page_url or civitai_model_url,
|
|
"source_url": stored_page_url,
|
|
"platform_url": metadata.get("platform_url") or path_metadata.get("platform_url"),
|
|
"download_url": first_non_empty(
|
|
metadata.get("download_url"),
|
|
civitai_data.get("downloadUrl"),
|
|
file_info.get("downloadUrl"),
|
|
),
|
|
"base_model": first_non_empty(
|
|
metadata.get("base_model"),
|
|
metadata.get("baseModel"),
|
|
selected_version.get("base_model"),
|
|
selected_version.get("baseModel"),
|
|
civitai_data.get("baseModel"),
|
|
)
|
|
or "",
|
|
"tags": as_list(metadata.get("tags") or model_info.get("tags")),
|
|
"trained_words": trained_words,
|
|
"images": images,
|
|
"preview_url": first_non_empty(
|
|
metadata.get("preview_url"),
|
|
metadata.get("previewUrl"),
|
|
metadata.get("thumbnail_url"),
|
|
metadata.get("thumbnailUrl"),
|
|
)
|
|
or "",
|
|
"clip_skip": civitai_data.get("clipSkip"),
|
|
"description": model_description,
|
|
"model_description": model_description,
|
|
"version_description": version_description,
|
|
"from_metadata": True,
|
|
}
|
|
|
|
|
|
def get_model_info_for_file(
|
|
file_path: str, api_key: Optional[str] = None, local_only: bool = False
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""
|
|
Get model info from CivitAI by computing file hash and looking up by hash.
|
|
First checks if there's a metadata file in the same folder.
|
|
|
|
Args:
|
|
file_path: Full path to the model file
|
|
api_key: Optional CivitAI API key
|
|
local_only: If true, never calculate a hash or query CivitAI when
|
|
metadata is missing. Return only local file information.
|
|
|
|
Returns:
|
|
Dict with model info from CivitAI, or None if not found
|
|
"""
|
|
log.info(f"get_model_info_for_file called with: {file_path}")
|
|
|
|
# First check for metadata file
|
|
metadata_path = find_metadata_sidecar_path(file_path)
|
|
log.info(f"Looking for metadata file, checked: {metadata_path}")
|
|
|
|
if metadata_path:
|
|
metadata = _read_model_metadata(metadata_path, file_path)
|
|
if metadata:
|
|
log.info(f"Using metadata file for {file_path}")
|
|
result = _metadata_to_model_info(metadata)
|
|
result["file_path"] = file_path
|
|
result["metadata_path"] = metadata_path
|
|
result["location"] = _format_model_location(file_path)
|
|
_apply_safetensors_header_metadata(result, file_path)
|
|
_apply_inferred_base_model(result, file_path)
|
|
if not result.get("size"):
|
|
try:
|
|
result["size"] = os.path.getsize(file_path)
|
|
except Exception:
|
|
pass
|
|
return result
|
|
else:
|
|
# Debug: list all files in the same directory
|
|
directory = os.path.dirname(file_path)
|
|
if directory and os.path.exists(directory):
|
|
files = os.listdir(directory)
|
|
metadata_files = [
|
|
f for f in files if "metadata" in f.lower() or f.endswith(".json")
|
|
]
|
|
log.info(f"Files in {directory}: {metadata_files}")
|
|
|
|
if local_only:
|
|
return _build_local_model_info(file_path, local_only=True)
|
|
|
|
# If no metadata file, compute hash and look up on CivitAI
|
|
file_hash = calculate_file_sha256(file_path)
|
|
if not file_hash:
|
|
return _build_local_model_info(file_path, require_inferred=True)
|
|
|
|
try:
|
|
result = get_model_info_by_hash(file_hash, api_key)
|
|
except Exception as e:
|
|
log.warning(
|
|
f"CivitAI hash lookup failed; trying local tensor fingerprint fallback: {e}"
|
|
)
|
|
return _build_local_model_info(file_path, require_inferred=True)
|
|
|
|
if result:
|
|
result["file_path"] = file_path
|
|
result["location"] = _format_model_location(file_path)
|
|
_apply_safetensors_header_metadata(result, file_path)
|
|
_apply_inferred_base_model(result, file_path)
|
|
if not result.get("size"):
|
|
try:
|
|
result["size"] = os.path.getsize(file_path)
|
|
except Exception:
|
|
pass
|
|
return result
|
|
|
|
return _build_local_model_info(file_path, require_inferred=True)
|
|
|
|
|
|
def build_civitai_custom_result(
|
|
details: Dict[str, Any],
|
|
expected_filename: str = "",
|
|
api_key: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
if not isinstance(details, dict):
|
|
return None
|
|
|
|
selected_version = (
|
|
details.get("selected_version")
|
|
if isinstance(details.get("selected_version"), dict)
|
|
else {}
|
|
)
|
|
version_id = details.get("version_id") or selected_version.get("id")
|
|
model_id = details.get("model_id")
|
|
file_info = select_primary_model_file(
|
|
selected_version.get("files") or [],
|
|
expected_filename, require_download=True
|
|
) or {}
|
|
filename = (
|
|
file_info.get("name")
|
|
or file_info.get("filename")
|
|
or get_filename_from_path(expected_filename)
|
|
or details.get("name")
|
|
or f"civitai-{version_id or model_id}"
|
|
)
|
|
download_url = (
|
|
file_info.get("download_url")
|
|
or file_info.get("downloadUrl")
|
|
or (get_civitai_download_url(version_id, api_key) if version_id else "")
|
|
)
|
|
if not download_url:
|
|
return None
|
|
|
|
hashes = (
|
|
file_info.get("hashes")
|
|
if isinstance(file_info.get("hashes"), dict)
|
|
else {}
|
|
)
|
|
sha256 = (
|
|
file_info.get("sha256")
|
|
or file_info.get("hash")
|
|
or hashes.get("SHA256")
|
|
or hashes.get("sha256")
|
|
)
|
|
version_url = (
|
|
details.get("version_url")
|
|
or selected_version.get("url")
|
|
or (
|
|
f"https://civitai.com/models/{model_id}?modelVersionId={version_id}"
|
|
if model_id and version_id
|
|
else details.get("url")
|
|
)
|
|
)
|
|
return build_custom_url_result(
|
|
source="civitai",
|
|
model_id=model_id,
|
|
version_id=version_id,
|
|
name=details.get("name") or filename,
|
|
version_name=selected_version.get("name") or "",
|
|
type=details.get("type") or file_info.get("type") or "",
|
|
filename=filename,
|
|
url=version_url or details.get("url"),
|
|
download_url=download_url,
|
|
size=file_info.get("size"),
|
|
base_model=selected_version.get("base_model"),
|
|
tags=details.get("tags") or [],
|
|
trained_words=selected_version.get("trained_words") or [],
|
|
images=details.get("images") or selected_version.get("images") or [],
|
|
description=selected_version.get("description")
|
|
or details.get("description")
|
|
or "",
|
|
sha256=sha256,
|
|
hashes=hashes,
|
|
model_description=details.get("description") or "",
|
|
)
|
|
|
|
|
|
def resolve_civitai_version_custom_result(
|
|
version_id: int,
|
|
expected_filename: str = "",
|
|
api_key: Optional[str] = None,
|
|
) -> Optional[Dict[str, Any]]:
|
|
if not version_id:
|
|
return None
|
|
|
|
headers = {}
|
|
if api_key:
|
|
headers["Authorization"] = f"Bearer {api_key}"
|
|
|
|
try:
|
|
response = requests.get(
|
|
f"https://civitai.com/api/v1/model-versions/{version_id}",
|
|
headers=headers,
|
|
timeout=20,
|
|
)
|
|
try:
|
|
if response.status_code != 200:
|
|
return None
|
|
data = response.json()
|
|
finally:
|
|
response.close()
|
|
except Exception as e:
|
|
log.warning(
|
|
f"CivitAI custom URL version lookup failed: version_id={version_id}, error={e}"
|
|
)
|
|
return None
|
|
|
|
model_block = data.get("model") if isinstance(data.get("model"), dict) else {}
|
|
model_id = data.get("modelId") or data.get("model_id") or model_block.get("id")
|
|
if model_id:
|
|
try:
|
|
details = get_civitai_model_details(model_id, version_id, api_key)
|
|
result = build_civitai_custom_result(
|
|
details,
|
|
expected_filename=expected_filename,
|
|
api_key=api_key,
|
|
)
|
|
if result:
|
|
return result
|
|
except Exception as e:
|
|
log.warning(
|
|
f"CivitAI custom URL details lookup failed: model_id={model_id}, version_id={version_id}, error={e}"
|
|
)
|
|
return None
|
|
|