Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-19 09:14:35 +02:00
Code Issues Packages Projects Releases Wiki Activity

feat(civitai): enrich base-model discovery with github constants

Browse Source 
CivitAI removed enum validation from /api/v1/models?baseModels=X (now
returns 200 with empty items on unknown values), which left the
base-model filter dropdown empty. Fix in two parts.

1. Point the base_models discovery probe at /images instead of /models.
   /images still returns a ZodError with the full enum list (82 names
   at time of writing). The existing parser handles the shape unchanged.

2. Also fetch civitai/civitai's base-model.constants.ts from GitHub and
   merge per-name metadata (group, ecosystem, engine, family, hidden)
   into a new base_models_info field on the response. The live probe
   remains authoritative for which names exist; GitHub provides the
   metadata and doubles as a fallback name list if the probe comes back
   empty.

Backward compatible: the existing base_models list keeps the same
shape. Clients that want grouped dropdowns, hidden-flag filtering, or
image/video separation can consume base_models_info instead.

The GitHub fetch has a separate 6h cache since these constants change
much less often than the probes it sits alongside.
This commit is contained in:
CalamitousFelicitousness
2026-04-13 23:50:41 +01:00
parent 101bd64b8c
commit 5d9b8d2e34
1 changed files with 145 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/civitai/client_civitai.py
+145 -2
View File
@@ -1,4 +1,5 @@
import os
import re
import time
from modules.logger import log
from modules.civitai.models_civitai import CivitModel, CivitVersion, CivitImage, CivitSearchResponse, CivitTagResponse, CivitCreatorResponse, CivitUserProfile
@@ -8,6 +9,15 @@ options_cache: dict = {}
options_cache_time: float = 0
OPTIONS_TTL = 3600 # 1 hour
# 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.
# discover_options fetches both and merges them.
base_models_github_cache: list[dict] = []
base_models_github_cache_time: float = 0
BASE_MODELS_GITHUB_TTL = 6 * 3600 # 6 hours
BASE_MODELS_GITHUB_URL = 'https://raw.githubusercontent.com/civitai/civitai/main/src/shared/constants/base-model.constants.ts'
class CivitaiClient:
BASE_URL = "https://civitai.com/api/v1"
@@ -196,13 +206,133 @@ class CivitaiClient:
return None
return {"username": profile.username, "id": profile.id}
def fetch_github_base_models(self) -> 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 separately
from discover_options with a longer TTL.
"""
global base_models_github_cache, base_models_github_cache_time # pylint: disable=global-statement
now = time.time()
if base_models_github_cache and (now - base_models_github_cache_time) < BASE_MODELS_GITHUB_TTL:
return base_models_github_cache
try:
from modules import shared
r = shared.req(BASE_MODELS_GITHUB_URL)
if r.status_code != 200:
log.debug(f'CivitAI github constants: code={r.status_code}')
return []
parsed = self.parse_base_model_config(r.text)
if parsed:
base_models_github_cache = parsed
base_models_github_cache_time = now
return parsed
except Exception as e:
log.debug(f'CivitAI github constants fetch failed: {e}')
return []
@staticmethod
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 discover_options(self) -> dict:
global options_cache, options_cache_time # pylint: disable=global-statement
now = time.time()
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 +385,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 = self.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