diff --git a/modules/civitai/basemodels_civitai.py b/modules/civitai/basemodels_civitai.py new file mode 100644 index 000000000..4b9a5c072 --- /dev/null +++ b/modules/civitai/basemodels_civitai.py @@ -0,0 +1,134 @@ +import re +import time +from modules.logger import log + + +# Canonical base-model metadata lives in the civitai/civitai repo. The live +# /images validator is generated from it, so the file has metadata (group, +# ecosystem, engine, hidden) while the validator has the current name list. +# Callers merge both. +github_cache: list[dict] = [] +github_cache_time: float = 0 +GITHUB_TTL = 6 * 3600 # 6 hours +GITHUB_URL = 'https://raw.githubusercontent.com/civitai/civitai/main/src/shared/constants/base-model.constants.ts' + + +def parse_base_model_config(ts_source: str) -> list[dict]: + """Parse the baseModelConfig array from base-model.constants.ts. + + Uses character-by-character bracket walking rather than regex because + entries can span multiple lines. Returns list of dicts with keys + name/type/group/hidden plus optional ecosystem/engine/family. + """ + start_match = re.search(r'const\s+baseModelConfig\s*=\s*\[', ts_source) + if not start_match: + return [] + # Walk to the matching ] respecting string literals + pos = start_match.end() + depth = 1 + in_string: str | None = None + end_pos = -1 + while pos < len(ts_source): + ch = ts_source[pos] + if in_string is not None: + if ch == '\\': + pos += 2 + continue + if ch == in_string: + in_string = None + else: + if ch in ("'", '"', '`'): + in_string = ch + elif ch == '[': + depth += 1 + elif ch == ']': + depth -= 1 + if depth == 0: + end_pos = pos + break + pos += 1 + if end_pos < 0: + return [] + array_body = ts_source[start_match.end():end_pos] + # Extract top-level {...} entries, respecting strings and nested braces + entries: list[str] = [] + brace_start = -1 + brace_depth = 0 + in_string = None + i = 0 + while i < len(array_body): + ch = array_body[i] + if in_string is not None: + if ch == '\\': + i += 2 + continue + if ch == in_string: + in_string = None + else: + if ch in ("'", '"', '`'): + in_string = ch + elif ch == '{': + if brace_depth == 0: + brace_start = i + brace_depth += 1 + elif ch == '}': + brace_depth -= 1 + if brace_depth == 0 and brace_start >= 0: + entries.append(array_body[brace_start:i + 1]) + brace_start = -1 + i += 1 + # Per-entry field extraction (string + bool values only) + field_re = re.compile( + r"(\w+)\s*:\s*(?:'([^'\\]*(?:\\.[^'\\]*)*)'|\"([^\"\\]*(?:\\.[^\"\\]*)*)\"|(true|false))" + ) + parsed: list[dict] = [] + for entry in entries: + fields: dict = {} + for m in field_re.finditer(entry): + key = m.group(1) + if m.group(2) is not None: + fields[key] = m.group(2) + elif m.group(3) is not None: + fields[key] = m.group(3) + elif m.group(4) is not None: + fields[key] = m.group(4) == 'true' + if 'name' in fields and 'type' in fields and 'group' in fields: + item: dict = { + 'name': fields['name'], + 'type': fields['type'], + 'group': fields['group'], + 'hidden': bool(fields.get('hidden', False)), + } + for opt in ('ecosystem', 'engine', 'family'): + if opt in fields: + item[opt] = fields[opt] + parsed.append(item) + return parsed + + +def fetch_github_base_models() -> list[dict]: + """Fetch and parse civitai's base-model constants from GitHub. + + Returns list of metadata dicts (name, type, group, hidden, plus + optional ecosystem/engine/family). Returns [] on any failure; + callers fall back to the live /images probe. Cached with a longer + TTL than discover_options since these constants change rarely. + """ + global github_cache, github_cache_time # pylint: disable=global-statement + now = time.time() + if github_cache and (now - github_cache_time) < GITHUB_TTL: + return github_cache + try: + from modules import shared + r = shared.req(GITHUB_URL) + if r.status_code != 200: + log.debug(f'CivitAI github constants: code={r.status_code}') + return [] + parsed = parse_base_model_config(r.text) + if parsed: + github_cache = parsed + github_cache_time = now + return parsed + except Exception as e: + log.debug(f'CivitAI github constants fetch failed: {e}') + return [] diff --git a/modules/civitai/client_civitai.py b/modules/civitai/client_civitai.py index fc99e15be..c24306d90 100644 --- a/modules/civitai/client_civitai.py +++ b/modules/civitai/client_civitai.py @@ -1,6 +1,7 @@ import os import time from modules.logger import log +from modules.civitai.basemodels_civitai import fetch_github_base_models from modules.civitai.models_civitai import CivitModel, CivitVersion, CivitImage, CivitSearchResponse, CivitTagResponse, CivitCreatorResponse, CivitUserProfile @@ -202,7 +203,7 @@ class CivitaiClient: if options_cache and (now - options_cache_time) < OPTIONS_TTL: return options_cache from modules import shared - result: dict = {'types': [], 'sort': [], 'period': [], 'base_models': []} + result: dict = {'types': [], 'sort': [], 'period': [], 'base_models': [], 'base_models_info': []} # Send invalid params to trigger 400 with valid enum values in error response probes = [ ('types', '/models', {'types': '__invalid__'}), @@ -255,9 +256,22 @@ class CivitaiClient: break except Exception as e: log.debug(f'CivitAI discover options: key={key} {e}') + # Merge live probe names with github metadata. The probe is the source + # of truth for which names exist; github provides per-entry metadata and + # doubles as a fallback name list if the probe returned empty. + github_entries = fetch_github_base_models() + github_index: dict = {entry['name']: entry for entry in github_entries} + probe_names: list = result['base_models'] + if not probe_names and github_entries: + probe_names = [entry['name'] for entry in github_entries] + result['base_models'] = probe_names + result['base_models_info'] = [ + github_index.get(name, {'name': name, 'type': 'image', 'group': '', 'hidden': False}) + for name in probe_names + ] options_cache = result options_cache_time = now - log.debug(f'CivitAI options: types={len(result["types"])} sort={len(result["sort"])} period={len(result["period"])} base_models={len(result["base_models"])}') + log.debug(f'CivitAI options: types={len(result["types"])} sort={len(result["sort"])} period={len(result["period"])} base_models={len(result["base_models"])} (enriched={len(github_index)})') return result