mirror of
https://github.com/willmiao/ComfyUI-Lora-Manager.git
synced 2026-08-08 23:10:15 -03:00
8e724538bd
Fix ~950 basedpyright errors across the backend: - Convert ineffective # type: ignore comments to # pyright: ignore[rule] - Add missing generic type arguments (Dict[str, Any], list[Any], ...) - Annotate dynamic dict literals and runtime-initialized attributes - Widen CivitAI provider tuple signatures in recipe parsers - Remove dead LoraRoutes handlers calling nonexistent LoraService methods - Suppress unavoidable ServiceRegistry import cycles (basedpyright counts function-local imports as cycle edges)
871 lines
32 KiB
Python
871 lines
32 KiB
Python
# pyright: reportImportCycles=false
|
|
# Lazy (function-local) imports still count as static edges in basedpyright's
|
|
# reportImportCycles, so the ServiceRegistry singleton pattern necessarily forms
|
|
# import cycles. Breaking them would require an architectural refactor.
|
|
import asyncio
|
|
import copy
|
|
import logging
|
|
import os
|
|
import time
|
|
from collections import OrderedDict
|
|
from typing import Any, Optional, Dict, Tuple, List, Sequence, cast
|
|
from .connectivity_guard import (
|
|
OFFLINE_FRIENDLY_MESSAGE,
|
|
is_expected_offline_error,
|
|
is_offline_cooldown_error,
|
|
)
|
|
from .model_metadata_provider import (
|
|
CivitaiModelMetadataProvider,
|
|
ModelMetadataProviderManager,
|
|
)
|
|
from .downloader import get_downloader
|
|
from .errors import RateLimitError, ResourceNotFoundError
|
|
from ..utils.civitai_utils import resolve_license_payload
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Best-effort cache for creator model counts, keyed by lowercase username.
|
|
# Values are (monotonic timestamp, count or None); None results are cached
|
|
# too so repeated failures don't hammer the API.
|
|
_CREATOR_COUNT_CACHE_TTL_SECONDS = 600
|
|
_creator_model_count_cache: Dict[str, Tuple[float, Optional[int]]] = {}
|
|
|
|
|
|
class CivitaiClient:
|
|
_instance = None
|
|
_lock = asyncio.Lock()
|
|
|
|
@classmethod
|
|
async def get_instance(cls):
|
|
"""Get singleton instance of CivitaiClient"""
|
|
async with cls._lock:
|
|
if cls._instance is None:
|
|
cls._instance = cls()
|
|
|
|
# Register this client as a metadata provider
|
|
provider_manager = await ModelMetadataProviderManager.get_instance()
|
|
provider_manager.register_provider(
|
|
"civitai", CivitaiModelMetadataProvider(cls._instance), True
|
|
)
|
|
|
|
return cls._instance
|
|
|
|
def __init__(self):
|
|
# Check if already initialized for singleton pattern
|
|
if hasattr(self, "_initialized"):
|
|
return
|
|
self._initialized = True
|
|
|
|
self.base_url = "https://civitai.red/api/v1"
|
|
# In-memory cache to avoid redundant get_model_version_info calls
|
|
# within the same import/scan flow. Only successful results are cached.
|
|
# Uses OrderedDict with LRU eviction at MAX_CACHE_ENTRIES to prevent
|
|
# unbounded growth in long-running server processes.
|
|
self._version_info_cache: OrderedDict[
|
|
str, Tuple[Optional[Dict[str, Any]], Optional[str]]
|
|
] = OrderedDict()
|
|
self._MAX_CACHE_ENTRIES = 500
|
|
|
|
def _build_image_info_url(self, image_id: str) -> str:
|
|
return f"{self.base_url}/images?imageId={image_id}&nsfw=X&withMeta=true"
|
|
|
|
async def _make_request(
|
|
self,
|
|
method: str,
|
|
url: str,
|
|
*,
|
|
use_auth: bool = False,
|
|
**kwargs,
|
|
) -> Tuple[bool, Dict[str, Any] | str]:
|
|
"""Wrapper around downloader.make_request that surfaces rate limits,
|
|
with retry for transient server errors (5xx, Cloudflare 524, network flakiness)."""
|
|
|
|
max_retries = 3
|
|
for attempt in range(max_retries):
|
|
downloader = await get_downloader()
|
|
success, result = await downloader.make_request(
|
|
method,
|
|
url,
|
|
use_auth=use_auth,
|
|
**kwargs,
|
|
)
|
|
if success:
|
|
# RateLimitError is raised below; a successful result is dict or str.
|
|
return True, cast(Dict[str, Any] | str, result)
|
|
|
|
if isinstance(result, RateLimitError):
|
|
if result.provider is None:
|
|
result.provider = "civitai_api"
|
|
raise result
|
|
|
|
if is_offline_cooldown_error(result):
|
|
return False, OFFLINE_FRIENDLY_MESSAGE
|
|
|
|
# Transient server error — retry with exponential backoff
|
|
if self._is_transient_server_error(str(result)):
|
|
if attempt < max_retries - 1:
|
|
wait = 2**attempt # 1s, 2s, 4s
|
|
logger.info(
|
|
"Transient error on %s %s, retrying in %ds "
|
|
"(attempt %d/%d): %s",
|
|
method,
|
|
url,
|
|
wait,
|
|
attempt + 1,
|
|
max_retries,
|
|
result,
|
|
)
|
|
await asyncio.sleep(wait)
|
|
continue
|
|
logger.warning(
|
|
"All %d retries exhausted for %s %s: %s",
|
|
max_retries,
|
|
method,
|
|
url,
|
|
result,
|
|
)
|
|
return False, result
|
|
|
|
return False, result
|
|
|
|
return False, "Unexpected error in _make_request"
|
|
|
|
@staticmethod
|
|
def _remove_comfy_metadata(model_version: Optional[Dict[str, Any]]) -> None:
|
|
"""Remove Comfy-specific metadata from model version images."""
|
|
if not isinstance(model_version, dict):
|
|
return
|
|
|
|
images = model_version.get("images")
|
|
if not isinstance(images, list):
|
|
return
|
|
|
|
for image in images:
|
|
if not isinstance(image, dict):
|
|
continue
|
|
|
|
meta = image.get("meta")
|
|
if isinstance(meta, dict) and "comfy" in meta:
|
|
meta.pop("comfy", None)
|
|
|
|
async def download_file(
|
|
self, url: str, save_dir: str, default_filename: str, progress_callback=None
|
|
) -> Tuple[bool, str]:
|
|
"""Download file with resumable downloads and retry mechanism
|
|
|
|
Args:
|
|
url: Download URL
|
|
save_dir: Directory to save the file
|
|
default_filename: Fallback filename if none provided in headers
|
|
progress_callback: Optional async callback function for progress updates (0-100)
|
|
|
|
Returns:
|
|
Tuple[bool, str]: (success, save_path or error message)
|
|
"""
|
|
downloader = await get_downloader()
|
|
save_path = os.path.join(save_dir, default_filename)
|
|
|
|
# Use unified downloader with CivitAI authentication
|
|
success, result = await downloader.download_file(
|
|
url=url,
|
|
save_path=save_path,
|
|
progress_callback=progress_callback,
|
|
use_auth=True, # Enable CivitAI authentication
|
|
allow_resume=True,
|
|
)
|
|
|
|
return success, result
|
|
|
|
async def get_model_by_hash(
|
|
self, model_hash: str
|
|
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
try:
|
|
success, version = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/model-versions/by-hash/{model_hash}",
|
|
use_auth=True,
|
|
)
|
|
if not success:
|
|
message = str(version)
|
|
if is_expected_offline_error(message):
|
|
return None, OFFLINE_FRIENDLY_MESSAGE
|
|
if "not found" in message.lower():
|
|
return None, "Model not found"
|
|
|
|
logger.error(
|
|
"Failed to fetch model info for %s: %s", model_hash[:10], message
|
|
)
|
|
return None, message
|
|
|
|
if isinstance(version, dict):
|
|
model_id = version.get("modelId")
|
|
if model_id:
|
|
model_data = await self._fetch_model_data(model_id)
|
|
if model_data:
|
|
self._enrich_version_with_model_data(version, model_data)
|
|
|
|
self._remove_comfy_metadata(version)
|
|
return version, None
|
|
else:
|
|
return None, "Invalid response format"
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.error("API Error: %s", exc)
|
|
return None, str(exc)
|
|
|
|
async def download_preview_image(self, image_url: str, save_path: str):
|
|
try:
|
|
downloader = await get_downloader()
|
|
success, content, headers = await downloader.download_to_memory(
|
|
image_url,
|
|
use_auth=False, # Preview images don't need auth
|
|
)
|
|
if success:
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(save_path), exist_ok=True)
|
|
with open(save_path, "wb") as f:
|
|
f.write(content if isinstance(content, bytes) else content.encode("utf-8"))
|
|
return True
|
|
return False
|
|
except Exception as e:
|
|
if is_expected_offline_error(str(e)):
|
|
logger.debug("Preview download skipped due to offline state.")
|
|
return False
|
|
logger.error(f"Download Error: {str(e)}")
|
|
return False
|
|
|
|
@staticmethod
|
|
def _extract_error_message(payload: Any) -> str:
|
|
"""Return a human-readable error message from an API payload."""
|
|
|
|
def _from_value(value: Any) -> str:
|
|
if isinstance(value, str):
|
|
return value
|
|
if isinstance(value, dict):
|
|
for key in ("message", "error", "detail", "details"):
|
|
if key in value:
|
|
candidate = _from_value(value[key])
|
|
if candidate:
|
|
return candidate
|
|
if isinstance(value, list):
|
|
for item in value:
|
|
candidate = _from_value(item)
|
|
if candidate:
|
|
return candidate
|
|
return ""
|
|
|
|
return _from_value(payload)
|
|
|
|
@staticmethod
|
|
def _is_transient_server_error(message: str) -> bool:
|
|
"""Return True when the message indicates a transient upstream failure.
|
|
|
|
Recognises Cloudflare 524, generic 5xx, and connectivity-level flakiness
|
|
that should not be treated as a permanent failure.
|
|
"""
|
|
normalized = message.lower()
|
|
if "status 5" in normalized or "status 524" in normalized:
|
|
return True
|
|
if any(
|
|
keyword in normalized
|
|
for keyword in (
|
|
"connection refused",
|
|
"connection reset",
|
|
"temporary failure",
|
|
"name resolution",
|
|
"connection closed",
|
|
)
|
|
):
|
|
return True
|
|
return False
|
|
|
|
async def get_model_versions(self, model_id: str) -> Optional[Dict[str, Any]]:
|
|
"""Get all versions of a model with local availability info"""
|
|
try:
|
|
success, result = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/models/{model_id}",
|
|
use_auth=True,
|
|
)
|
|
if success and isinstance(result, dict):
|
|
# Also return model type along with versions
|
|
return {
|
|
"modelVersions": result.get("modelVersions", []),
|
|
"type": result.get("type", ""),
|
|
"name": result.get("name", ""),
|
|
}
|
|
message = self._extract_error_message(result)
|
|
if message and "not found" in message.lower():
|
|
raise ResourceNotFoundError(f"Resource not found for model {model_id}")
|
|
if is_expected_offline_error(message):
|
|
logger.info("Civitai request skipped: %s", OFFLINE_FRIENDLY_MESSAGE)
|
|
return None
|
|
if message:
|
|
if self._is_transient_server_error(message):
|
|
logger.info(
|
|
"Transient server error for model %s: %s",
|
|
model_id,
|
|
message,
|
|
)
|
|
return None
|
|
raise RuntimeError(message)
|
|
return None
|
|
except RateLimitError:
|
|
raise
|
|
except ResourceNotFoundError as exc:
|
|
logger.info("Model %s is no longer available on Civitai: %s", model_id, exc)
|
|
raise
|
|
except Exception as e:
|
|
logger.error("Error fetching model versions: %s", e, exc_info=True)
|
|
raise
|
|
|
|
async def get_model_versions_bulk(
|
|
self, model_ids: Sequence[int]
|
|
) -> Optional[Dict[int, Dict[str, Any]]]:
|
|
"""Fetch model metadata for multiple ids using the batch API."""
|
|
|
|
deduped: Dict[int, None] = {}
|
|
for raw_id in model_ids:
|
|
try:
|
|
normalized = int(raw_id)
|
|
except (TypeError, ValueError):
|
|
continue
|
|
deduped.setdefault(normalized, None)
|
|
|
|
normalized_ids = [str(model_id) for model_id in deduped.keys()]
|
|
if not normalized_ids:
|
|
return {}
|
|
|
|
try:
|
|
query = ",".join(normalized_ids)
|
|
success, result = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/models",
|
|
use_auth=True,
|
|
params={"ids": query, "nsfw": "true"},
|
|
)
|
|
if not success:
|
|
return None
|
|
|
|
items = result.get("items") if isinstance(result, dict) else None
|
|
if not isinstance(items, list):
|
|
return {}
|
|
|
|
payload: Dict[int, Dict[str, Any]] = {}
|
|
for item in items:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
model_id = item.get("id")
|
|
try:
|
|
normalized_id = int(cast(Any, model_id))
|
|
except (TypeError, ValueError):
|
|
continue
|
|
payload[normalized_id] = {
|
|
"modelVersions": item.get("modelVersions", []),
|
|
"type": item.get("type", ""),
|
|
"name": item.get("name", ""),
|
|
"allowNoCredit": item.get("allowNoCredit"),
|
|
"allowCommercialUse": item.get("allowCommercialUse"),
|
|
"allowDerivatives": item.get("allowDerivatives"),
|
|
"allowDifferentLicense": item.get("allowDifferentLicense"),
|
|
}
|
|
return payload
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as exc:
|
|
logger.error(f"Error fetching model versions in bulk: {exc}")
|
|
return None
|
|
|
|
async def get_model_version(
|
|
self, model_id: int | None = None, version_id: int | None = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Get specific model version with additional metadata."""
|
|
try:
|
|
if model_id is None and version_id is not None:
|
|
return await self._get_version_by_id_only(version_id)
|
|
|
|
if model_id is not None:
|
|
return await self._get_version_with_model_id(model_id, version_id)
|
|
|
|
logger.error("Either model_id or version_id must be provided")
|
|
return None
|
|
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as e:
|
|
logger.error(f"Error fetching model version: {e}")
|
|
return None
|
|
|
|
async def _get_version_by_id_only(self, version_id: int) -> Optional[Dict[str, Any]]:
|
|
version = await self._fetch_version_by_id(version_id)
|
|
if version is None:
|
|
return None
|
|
|
|
model_id = version.get("modelId")
|
|
if not model_id:
|
|
logger.error(f"No modelId found in version {version_id}")
|
|
return None
|
|
|
|
model_data = await self._fetch_model_data(model_id)
|
|
if model_data:
|
|
self._enrich_version_with_model_data(version, model_data)
|
|
|
|
self._remove_comfy_metadata(version)
|
|
return version
|
|
|
|
async def _get_version_with_model_id(
|
|
self, model_id: int, version_id: Optional[int]
|
|
) -> Optional[Dict[str, Any]]:
|
|
model_data = await self._fetch_model_data(model_id)
|
|
if not model_data:
|
|
return None
|
|
|
|
target_version = self._select_target_version(model_data, model_id, version_id)
|
|
|
|
# If modelVersions is empty (e.g. CivitAI cache lag for newly published
|
|
# models) but a specific version_id is known, fall back to fetching the
|
|
# version directly via the individual model-versions endpoint, then
|
|
# enrich it with the model-level data we already have.
|
|
if target_version is None and version_id is not None:
|
|
logger.info(
|
|
"modelVersions empty for model %s; falling back to direct "
|
|
"version lookup for %s",
|
|
model_id,
|
|
version_id,
|
|
)
|
|
version = await self._fetch_version_by_id(version_id)
|
|
if version:
|
|
self._enrich_version_with_model_data(version, model_data)
|
|
self._remove_comfy_metadata(version)
|
|
return version
|
|
return None
|
|
|
|
if target_version is None:
|
|
return None
|
|
|
|
target_version_id = target_version.get("id")
|
|
version = (
|
|
await self._fetch_version_by_id(target_version_id)
|
|
if target_version_id
|
|
else None
|
|
)
|
|
|
|
if version is None:
|
|
model_hash = self._extract_primary_model_hash(target_version)
|
|
if model_hash:
|
|
version = await self._fetch_version_by_hash(model_hash)
|
|
else:
|
|
logger.warning(
|
|
f"No primary model hash found for model {model_id} version {target_version_id}"
|
|
)
|
|
|
|
if version is None:
|
|
version = self._build_version_from_model_data(
|
|
target_version, model_id, model_data
|
|
)
|
|
|
|
self._enrich_version_with_model_data(version, model_data)
|
|
self._remove_comfy_metadata(version)
|
|
return version
|
|
|
|
async def _fetch_model_data(self, model_id: int) -> Optional[Dict[str, Any]]:
|
|
success, data = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/models/{model_id}",
|
|
use_auth=True,
|
|
)
|
|
if success and isinstance(data, dict):
|
|
return data
|
|
if is_expected_offline_error(data):
|
|
return None
|
|
logger.warning(f"Failed to fetch model data for model {model_id}")
|
|
return None
|
|
|
|
async def _fetch_version_by_id(self, version_id: Optional[int]) -> Optional[Dict[str, Any]]:
|
|
if version_id is None:
|
|
return None
|
|
|
|
success, version = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/model-versions/{version_id}",
|
|
use_auth=True,
|
|
)
|
|
if success and isinstance(version, dict):
|
|
return version
|
|
if is_expected_offline_error(version):
|
|
return None
|
|
|
|
logger.warning(f"Failed to fetch version by id {version_id}")
|
|
return None
|
|
|
|
async def _fetch_version_by_hash(self, model_hash: Optional[str]) -> Optional[Dict[str, Any]]:
|
|
if not model_hash:
|
|
return None
|
|
|
|
success, version = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/model-versions/by-hash/{model_hash}",
|
|
use_auth=True,
|
|
)
|
|
if success and isinstance(version, dict):
|
|
return version
|
|
if is_expected_offline_error(version):
|
|
return None
|
|
|
|
logger.warning(f"Failed to fetch version by hash {model_hash}")
|
|
return None
|
|
|
|
def _select_target_version(
|
|
self, model_data: Dict[str, Any], model_id: int, version_id: Optional[int]
|
|
) -> Optional[Dict[str, Any]]:
|
|
model_versions = model_data.get("modelVersions", [])
|
|
if not model_versions:
|
|
logger.warning(f"No model versions found for model {model_id}")
|
|
return None
|
|
|
|
if version_id is not None:
|
|
target_version = next(
|
|
(item for item in model_versions if item.get("id") == version_id), None
|
|
)
|
|
if target_version is None:
|
|
logger.warning(
|
|
f"Version {version_id} not found for model {model_id}, defaulting to first version"
|
|
)
|
|
return model_versions[0]
|
|
return target_version
|
|
|
|
return model_versions[0]
|
|
|
|
def _extract_primary_model_hash(self, version_entry: Dict[str, Any]) -> Optional[str]:
|
|
for file_info in version_entry.get("files", []):
|
|
if file_info.get("type") == "Model" and file_info.get("primary"):
|
|
hashes = file_info.get("hashes", {})
|
|
model_hash = hashes.get("SHA256")
|
|
if model_hash:
|
|
return model_hash
|
|
return None
|
|
|
|
def _build_version_from_model_data(
|
|
self, version_entry: Dict[str, Any], model_id: int, model_data: Dict[str, Any]
|
|
) -> Dict[str, Any]:
|
|
version = copy.deepcopy(version_entry)
|
|
version.pop("index", None)
|
|
version["modelId"] = model_id
|
|
version["model"] = {
|
|
"name": model_data.get("name"),
|
|
"type": model_data.get("type"),
|
|
"nsfw": model_data.get("nsfw"),
|
|
"poi": model_data.get("poi"),
|
|
}
|
|
return version
|
|
|
|
def _enrich_version_with_model_data(self, version: Dict[str, Any], model_data: Dict[str, Any]) -> None:
|
|
model_info = version.get("model")
|
|
if not isinstance(model_info, dict):
|
|
model_info = {}
|
|
version["model"] = model_info
|
|
|
|
model_info["description"] = model_data.get("description")
|
|
model_info["tags"] = model_data.get("tags", [])
|
|
version["creator"] = model_data.get("creator")
|
|
|
|
license_payload = resolve_license_payload(model_data)
|
|
for field, value in license_payload.items():
|
|
model_info[field] = value
|
|
|
|
async def get_model_version_info(
|
|
self, version_id: str
|
|
) -> Tuple[Optional[Dict[str, Any]], Optional[str]]:
|
|
"""Fetch model version metadata from Civitai
|
|
|
|
Args:
|
|
version_id: The Civitai model version ID
|
|
|
|
Returns:
|
|
Tuple[Optional[Dict], Optional[str]]: A tuple containing:
|
|
- The model version data or None if not found
|
|
- An error message if there was an error, or None on success
|
|
"""
|
|
# In-memory cache avoids redundant API calls within the same
|
|
# import/scan flow (e.g. _resolve_base_model_from_checkpoint
|
|
# followed by _resolve_and_populate_checkpoint with the same id).
|
|
if version_id in self._version_info_cache:
|
|
logger.debug("Cache hit for model version info: %s", version_id)
|
|
self._version_info_cache.move_to_end(version_id) # LRU bump
|
|
return self._version_info_cache[version_id]
|
|
|
|
try:
|
|
url = f"{self.base_url}/model-versions/{version_id}"
|
|
|
|
logger.debug("Resolving Civitai model version info: %s", url)
|
|
success, result = await self._make_request("GET", url, use_auth=True)
|
|
|
|
if success and isinstance(result, dict):
|
|
logger.debug("Successfully fetched model version info for: %s", version_id)
|
|
self._remove_comfy_metadata(result)
|
|
self._version_info_cache[version_id] = (result, None)
|
|
self._version_info_cache.move_to_end(version_id)
|
|
# Evict oldest entry when over capacity
|
|
if len(self._version_info_cache) > self._MAX_CACHE_ENTRIES:
|
|
self._version_info_cache.popitem(last=False)
|
|
return result, None
|
|
|
|
# Handle specific error cases
|
|
if is_expected_offline_error(result):
|
|
return None, OFFLINE_FRIENDLY_MESSAGE
|
|
if "not found" in str(result):
|
|
error_msg = f"Model not found"
|
|
logger.warning(f"Model version not found: {version_id} - {error_msg}")
|
|
return None, error_msg
|
|
|
|
# Other error cases
|
|
logger.error(f"Failed to fetch model info for {version_id}: {result}")
|
|
return None, str(result)
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as e:
|
|
error_msg = f"Error fetching model version info: {e}"
|
|
logger.error(error_msg)
|
|
return None, error_msg
|
|
|
|
async def get_image_info(
|
|
self, image_id: str, source_url: str | None = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch image information from Civitai API
|
|
|
|
Args:
|
|
image_id: The Civitai image ID
|
|
source_url: Original image page URL. Accepted for caller compatibility;
|
|
API requests always target ``civitai.red``.
|
|
|
|
Returns:
|
|
Optional[Dict]: The image data or None if not found
|
|
"""
|
|
try:
|
|
requested_id = int(image_id)
|
|
url = self._build_image_info_url(image_id)
|
|
success, result = await self._make_request("GET", url, use_auth=True)
|
|
|
|
if not success:
|
|
if is_expected_offline_error(result):
|
|
return None
|
|
if self._is_transient_server_error(str(result)):
|
|
logger.info(
|
|
"Transient server error fetching image info for ID %s: %s",
|
|
image_id,
|
|
result,
|
|
)
|
|
return None
|
|
logger.error(
|
|
"Failed to fetch image info for ID %s from civitai.red: %s",
|
|
image_id,
|
|
result,
|
|
)
|
|
return None
|
|
|
|
if isinstance(result, dict) and "items" in result and isinstance(result["items"], list):
|
|
items = result["items"]
|
|
|
|
for item in items:
|
|
if isinstance(item, dict) and item.get("id") == requested_id:
|
|
logger.debug(
|
|
"Successfully fetched image info for ID %s from civitai.red",
|
|
image_id,
|
|
)
|
|
return item
|
|
|
|
returned_ids = [
|
|
item.get("id")
|
|
for item in items
|
|
if isinstance(item, dict) and "id" in item
|
|
]
|
|
|
|
logger.warning(
|
|
"CivitAI API returned no matching image for requested ID %s from civitai.red. Returned %d item(s) with IDs: %s. This may indicate the image was deleted, hidden, or there is a database lag.",
|
|
image_id,
|
|
len(items),
|
|
returned_ids,
|
|
)
|
|
return None
|
|
|
|
logger.warning("No image found with ID: %s", image_id)
|
|
return None
|
|
except RateLimitError:
|
|
raise
|
|
except ValueError as e:
|
|
error_msg = f"Invalid image ID format: {image_id}"
|
|
logger.error(error_msg)
|
|
return None
|
|
except Exception as e:
|
|
error_msg = f"Error fetching image info: {e}"
|
|
logger.error(error_msg)
|
|
return None
|
|
|
|
async def get_model_versions_by_hashes(
|
|
self, hashes: List[str]
|
|
) -> Optional[List[Dict[str, Any]]]:
|
|
"""Fetch full version details for up to 100 SHA256 hashes via the batch endpoint.
|
|
|
|
Uses POST /api/v1/model-versions/by-hash which returns full version
|
|
details including ``usageControl`` and ``earlyAccessEndsAt`` that are
|
|
not available from the model-level API.
|
|
|
|
Args:
|
|
hashes: List of SHA256 hashes (max 100 per batch; auto-split).
|
|
|
|
Returns:
|
|
List of version dicts or None on failure.
|
|
"""
|
|
if not hashes:
|
|
return []
|
|
|
|
BATCH_SIZE = 100
|
|
all_versions: List[Dict[str, Any]] = []
|
|
|
|
for start in range(0, len(hashes), BATCH_SIZE):
|
|
batch = hashes[start : start + BATCH_SIZE]
|
|
try:
|
|
success, result = await self._make_request(
|
|
"POST",
|
|
f"{self.base_url}/model-versions/by-hash",
|
|
use_auth=True,
|
|
json=batch,
|
|
)
|
|
if not success:
|
|
logger.warning(
|
|
"Batch by-hash request failed for %d hashes: %s",
|
|
len(batch),
|
|
result,
|
|
)
|
|
continue
|
|
|
|
if isinstance(result, list):
|
|
all_versions.extend(cast(Any, result))
|
|
else:
|
|
logger.debug(
|
|
"Unexpected by-hash response type: %s", type(result)
|
|
)
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as exc: # pragma: no cover - defensive logging
|
|
logger.error(
|
|
"Error fetching model versions by hashes: %s", exc
|
|
)
|
|
|
|
return all_versions if all_versions else None
|
|
|
|
async def get_user_models(
|
|
self, username: str, cursor: Optional[str] = None
|
|
) -> Optional[Dict[str, Any]]:
|
|
"""Fetch one page (up to 100 models) for a specific Civitai user.
|
|
|
|
Returns ``{"items": [...], "nextCursor": <str|None>}`` on success,
|
|
or None on failure. Pass ``cursor`` (from a previous response's
|
|
``nextCursor``) to fetch subsequent pages.
|
|
"""
|
|
if not username:
|
|
return None
|
|
|
|
params: Dict[str, Any] = {
|
|
"username": username,
|
|
"nsfw": "true",
|
|
"limit": 100,
|
|
"sort": "Newest",
|
|
"period": "AllTime",
|
|
}
|
|
if cursor:
|
|
params["cursor"] = cursor
|
|
|
|
try:
|
|
success, result = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/models",
|
|
use_auth=True,
|
|
params=params,
|
|
)
|
|
|
|
if not success:
|
|
if is_expected_offline_error(result):
|
|
logger.info("User model fetch skipped: %s", OFFLINE_FRIENDLY_MESSAGE)
|
|
return None
|
|
logger.error("Failed to fetch models for %s: %s", username, result)
|
|
return None
|
|
|
|
items = result.get("items") if isinstance(result, dict) else None
|
|
if not isinstance(items, list):
|
|
items = []
|
|
|
|
for model in items:
|
|
versions = model.get("modelVersions")
|
|
if not isinstance(versions, list):
|
|
continue
|
|
for version in versions:
|
|
self._remove_comfy_metadata(version)
|
|
|
|
next_cursor: Optional[str] = None
|
|
metadata = result.get("metadata") if isinstance(result, dict) else None
|
|
if isinstance(metadata, dict):
|
|
raw_cursor = metadata.get("nextCursor")
|
|
if raw_cursor is not None:
|
|
next_cursor = str(raw_cursor)
|
|
|
|
return {"items": items, "nextCursor": next_cursor}
|
|
except RateLimitError:
|
|
raise
|
|
except Exception as exc: # pragma: no cover - defensive logging
|
|
logger.error("Error fetching models for %s: %s", username, exc)
|
|
return None
|
|
|
|
async def get_creator_model_count(self, username: str) -> Optional[int]:
|
|
"""Best-effort lookup of a creator's published model count.
|
|
|
|
Uses the ``/creators`` endpoint (a contains-match query), picking the
|
|
entry whose username matches exactly (case-insensitive). Returns None
|
|
on any failure; never raises. Results (including None) are cached
|
|
for ``_CREATOR_COUNT_CACHE_TTL_SECONDS``.
|
|
"""
|
|
if not username:
|
|
return None
|
|
|
|
cache_key = username.lower()
|
|
cached = _creator_model_count_cache.get(cache_key)
|
|
if cached is not None:
|
|
cached_at, cached_count = cached
|
|
if time.monotonic() - cached_at < _CREATOR_COUNT_CACHE_TTL_SECONDS:
|
|
return cached_count
|
|
|
|
count: Optional[int] = None
|
|
try:
|
|
success, result = await self._make_request(
|
|
"GET",
|
|
f"{self.base_url}/creators",
|
|
use_auth=True,
|
|
params={"query": username, "limit": 10},
|
|
)
|
|
|
|
if success and isinstance(result, dict):
|
|
creators = result.get("items")
|
|
if isinstance(creators, list):
|
|
for creator in creators:
|
|
if not isinstance(creator, dict):
|
|
continue
|
|
creator_name = creator.get("username")
|
|
if not isinstance(creator_name, str):
|
|
continue
|
|
if creator_name.lower() != cache_key:
|
|
continue
|
|
model_count = creator.get("modelCount")
|
|
if isinstance(model_count, (int, float)) and not isinstance(
|
|
model_count, bool
|
|
):
|
|
count = int(model_count)
|
|
break
|
|
except Exception as exc: # best-effort only, never propagate
|
|
logger.debug(
|
|
"Failed to fetch creator model count for %s: %s", username, exc
|
|
)
|
|
|
|
_creator_model_count_cache[cache_key] = (time.monotonic(), count)
|
|
return count
|