diff --git a/modules/civitai/client_civitai.py b/modules/civitai/client_civitai.py index e62a8eb08..f98db6c2e 100644 --- a/modules/civitai/client_civitai.py +++ b/modules/civitai/client_civitai.py @@ -1,5 +1,7 @@ import os import time +import threading +from types import SimpleNamespace from modules.logger import log from modules.civitai.models_civitai import CivitModel, CivitVersion, CivitVersionMini, CivitImage, CivitSearchResponse, CivitTagResponse, CivitCreatorResponse, CivitUserProfile @@ -10,6 +12,43 @@ OPTIONS_TTL = 3600 # 1 hour # Civitai nsfwLevel bitmask: 1=PG/None 2=PG-13/Soft 4=R/Mature 8=X 16=XXX 32=Blocked NSFW_LEVEL_SFW = 3 # None + Soft: Civitai's SFW browsing boundary NSFW_LEVEL_ALL = 63 # every level set: disables filtering +BY_HASH_IDS_LIMIT = 10000 # POST /model-versions/by-hash/ids request cap +BY_HASH_LIMIT = 100 # POST /model-versions/by-hash request cap +MODEL_IDS_LIMIT = 100 # GET /models page cap; longer ids lists paginate +RETRY_LIMIT = 4 # retries after HTTP 429 +RETRY_DELAY_MAX = 60 # seconds +request_slots: threading.BoundedSemaphore | None = None +request_slots_lock = threading.Lock() + + +def get_request_slots() -> threading.BoundedSemaphore: + """Process-wide cap on concurrent CivitAI API requests, sized to shared.max_workers.""" + global request_slots # pylint: disable=global-statement + with request_slots_lock: + if request_slots is None: + from modules.shared import max_workers + request_slots = threading.BoundedSemaphore(max_workers) + return request_slots + + +def retry_delay(response, attempt: int) -> float: + """Seconds before retrying a 429: Retry-After when given in seconds, otherwise exponential.""" + headers = getattr(response, 'headers', None) or {} + try: + delay = float(headers.get('Retry-After')) + except (TypeError, ValueError): + delay = 2 ** attempt + return min(max(delay, 0.0), RETRY_DELAY_MAX) + + +def post_json(url: str, body, headers: dict): + """POST with the timeout, TLS and failure shape of shared.req.""" + import requests + try: + return requests.post(url, json=body, timeout=30, headers=headers, verify=False, allow_redirects=True) + except Exception as e: + log.error(f'HTTP request error: url={url} {e}') + return SimpleNamespace(status_code=500, text=f'HTTP request error: url={url} {e}') class CivitaiClient: @@ -24,7 +63,7 @@ class CivitaiClient: return tok return os.environ.get('CIVITAI_TOKEN', None) - def _get(self, path: str, params: dict | None = None, token: str | None = None, stream: bool = False): + def send(self, method: str, path: str, params: dict | None = None, body=None, token: str | None = None, stream: bool = False): from modules import shared url = f"{self.BASE_URL}{path}" headers = {} @@ -36,7 +75,22 @@ class CivitaiClient: query = urlencode({k: v for k, v in params.items() if v is not None and v != ''}, doseq=True) if query: url = f"{url}?{query}" - return shared.req(url, headers=headers if headers else None, stream=stream) + attempt = 0 + while True: + with get_request_slots(): + if method == 'POST': + r = post_json(url, body, headers) + else: + r = shared.req(url, headers=headers if headers else None, stream=stream) + if r.status_code != 429 or attempt >= RETRY_LIMIT: + return r + delay = retry_delay(r, attempt) + log.warning(f'CivitAI rate limited: path={path} attempt={attempt + 1} delay={delay:.0f}s') + time.sleep(delay) + attempt += 1 + + def _get(self, path: str, params: dict | None = None, token: str | None = None, stream: bool = False): + return self.send('GET', path, params=params, token=token, stream=stream) def search_models(self, *, query: str = "", tag: str = "", types: str = "", sort: str = "", period: str = "", base_models: list[str] | None = None, nsfw: bool | None = None, limit: int = 20, @@ -133,6 +187,59 @@ class CivitaiClient: log.error(f'CivitAI get version mini parse error: id={version_id} {e}') return None + def get_version_ids_by_hash(self, hashes: list[str], *, token: str | None = None) -> tuple[list[dict], dict[str, int]]: + """{modelVersionId, modelId, hash} rows for SHA256 hashes, plus the status code for each hash whose request failed.""" + rows, failed = [], {} + for i in range(0, len(hashes), BY_HASH_IDS_LIMIT): + chunk = hashes[i:i + BY_HASH_IDS_LIMIT] + r = self.send('POST', '/model-versions/by-hash/ids', body=chunk, token=token) + if r.status_code != 200: + log.error(f'CivitAI version ids by hash: count={len(chunk)} code={r.status_code}') + failed.update(dict.fromkeys(chunk, r.status_code)) + continue + try: + rows.extend(r.json()) + except Exception as e: + log.error(f'CivitAI version ids by hash parse error: count={len(chunk)} {e}') + failed.update(dict.fromkeys(chunk, 500)) + return rows, failed + + def get_versions_by_hash(self, hashes: list[str], *, token: str | None = None) -> tuple[list[CivitVersion], dict[str, int]]: + """Full versions for SHA256 hashes, plus the status code for each hash whose request failed.""" + versions, failed = [], {} + for i in range(0, len(hashes), BY_HASH_LIMIT): + chunk = hashes[i:i + BY_HASH_LIMIT] + r = self.send('POST', '/model-versions/by-hash', body=chunk, token=token) + if r.status_code != 200: + log.error(f'CivitAI versions by hash: count={len(chunk)} code={r.status_code}') + failed.update(dict.fromkeys(chunk, r.status_code)) + continue + try: + versions.extend([CivitVersion.parse_obj(v) for v in r.json()]) + except Exception as e: + log.error(f'CivitAI versions by hash parse error: count={len(chunk)} {e}') + failed.update(dict.fromkeys(chunk, 500)) + return versions, failed + + def get_models_raw(self, model_ids: list[int], *, token: str | None = None) -> tuple[dict[int, dict], dict[int, int]]: + """Unparsed /models items keyed by id, plus the status code for each id whose request failed.""" + models, failed = {}, {} + for i in range(0, len(model_ids), MODEL_IDS_LIMIT): + chunk = model_ids[i:i + MODEL_IDS_LIMIT] + params = {'ids': ','.join(str(m) for m in chunk), 'limit': MODEL_IDS_LIMIT, 'nsfw': 'true'} # ids query drops NSFW models unless nsfw=true + r = self.send('GET', '/models', params=params, token=token) + if r.status_code != 200: + log.error(f'CivitAI models by id: count={len(chunk)} code={r.status_code}') + failed.update(dict.fromkeys(chunk, r.status_code)) + continue + try: + for item in r.json().get('items', []): + models[item['id']] = item + except Exception as e: + log.error(f'CivitAI models by id parse error: count={len(chunk)} {e}') + failed.update(dict.fromkeys(chunk, 500)) + return models, failed + def get_images(self, *, model_version_id: int | None = None, limit: int | None = None, token: str | None = None) -> list[CivitImage]: params: dict = {} if model_version_id is not None: