[civitai] fix & update 'search' extractor (#7609)

- rename to 'search-models'
- use SearchAPI
This commit is contained in:
Mike Fährmann
2025-06-01 10:28:03 +02:00
parent dc467c7790
commit 9ed31cf992
2 changed files with 43 additions and 21 deletions

View File

@@ -379,14 +379,15 @@ class CivitaiTagExtractor(CivitaiExtractor):
return self.api.models_tag(tag)
class CivitaiSearchExtractor(CivitaiExtractor):
subcategory = "search"
class CivitaiSearchModelsExtractor(CivitaiExtractor):
subcategory = "search-models"
pattern = BASE_PATTERN + r"/search/models\?([^#]+)"
example = "https://civitai.com/search/models?query=QUERY"
def models(self):
params = self._parse_query(self.groups[0])
return self.api.models(params)
return CivitaiSearchAPI(self).search_models(
params.get("query"), params.get("sortBy"), self.api.nsfw)
class CivitaiSearchImagesExtractor(CivitaiExtractor):
@@ -396,10 +397,8 @@ class CivitaiSearchImagesExtractor(CivitaiExtractor):
def images(self):
params = self._parse_query(self.groups[0])
return CivitaiSearchAPI(self).search(
params.get("query"),
params.get("sortBy") or "images_v6",
self.api.nsfw)
return CivitaiSearchAPI(self).search_images(
params.get("query"), params.get("sortBy"), self.api.nsfw)
class CivitaiModelsExtractor(CivitaiExtractor):
@@ -812,22 +811,13 @@ class CivitaiSearchAPI():
"Priority": "u=4",
}
def search(self, query, type, nsfw=31):
def search(self, query, type, facets, nsfw=31):
endpoint = "/multi-search"
query = {
"q": query,
"q" : query,
"indexUid": type,
"facets" : (
"aspectRatio",
"baseModel",
"createdAtUnix",
"tagNames",
"techniqueNames",
"toolNames",
"type",
"user.username",
),
"facets" : facets,
"attributesToHighlight": (),
"highlightPreTag" : "__ais-highlight__",
"highlightPostTag": "__/ais-highlight__",
@@ -838,6 +828,32 @@ class CivitaiSearchAPI():
return self._pagination(endpoint, query)
def search_models(self, query, type=None, nsfw=31):
facets = (
"category.name",
"checkpointType",
"fileFormats",
"lastVersionAtUnix",
"tags.name",
"type",
"user.username",
"version.baseModel",
)
return self.search(query, type or "models_v9", facets, nsfw)
def search_images(self, query, type=None, nsfw=31):
facets = (
"aspectRatio",
"baseModel",
"createdAtUnix",
"tagNames",
"techniqueNames",
"toolNames",
"type",
"user.username",
)
return self.search(query, type or "images_v6", facets, nsfw)
def _call(self, endpoint, query):
url = self.root + endpoint
params = util.json_dumps({"queries": (query,)})