Mirrors/gallery-dl
1
0
Fork 0
mirror of https://github.com/mikf/gallery-dl.git synced 2024-11-25 04:02:32 +01:00
Code Issues Releases Wiki Activity

[civitai] improve model 'gallery' extraction (#3706)

Browse Source
This commit is contained in:
Mike Fährmann 2024-09-27 22:07:54 +02:00
parent bc11dc0de2
commit 1ab90dd9e2
No known key found for this signature in database
GPG Key ID: 5680CA389D365A88
1 changed files with 27 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

34
gallery_dl/extractor/civitai.py
View File

@ -28,7 +28,7 @@ class CivitaiExtractor(Extractor):
def _init(self):
if self.config("api") == "trpc":
self.log.debug("Using TRPC API")
self.log.debug("Using tRPC API")
self.api = CivitaiTrpcAPI(self)
else:
self.log.debug("Using REST API")
@ -193,12 +193,8 @@ class CivitaiModelExtractor(CivitaiExtractor):
]
def _extract_files_gallery(self, model, version, user):
params = {
"modelId" : model["id"],
"modelVersionId": version["id"],
}
for num, file in enumerate(self.api.images(params), 1):
images = self.api.images_gallery(model, version, user)
for num, file in enumerate(images, 1):
yield text.nameext_from_url(file["url"], {
"num" : num,
"file": file,
@ -321,6 +317,15 @@ class CivitaiRestAPI():
endpoint = "/v1/images"
return self._pagination(endpoint, params)
def images_gallery(self, model, version, user):
endpoint = "/v1/images"
params = {
"modelId" : model["id"],
"modelVersionId": version["id"],
"nsfw" : "X",
}
return self._pagination(endpoint, params)
def model(self, model_id):
endpoint = "/v1/models/{}".format(model_id)
return self._call(endpoint)
@ -398,6 +403,21 @@ class CivitaiTrpcAPI():
return self._pagination(endpoint, params_)
def images_gallery(self, model, version, user):
endpoint = "image.getImagesAsPostsInfinite"
params = {
"period" : "AllTime",
"sort" : "Newest",
"modelVersionId": version["id"],
"modelId" : model["id"],
"hidden" : False,
"limit" : 50,
"browsingLevel" : 31,
}
for post in self._pagination(endpoint, params):
yield from post["images"]
def model(self, model_id):
endpoint = "model.getById"
params = {"id": int(model_id)}