Skip to content

Commit

Permalink
[civitai] improve model 'gallery' extraction (#3706)
Browse files Browse the repository at this point in the history
  • Loading branch information
mikf committed Sep 27, 2024
1 parent bc11dc0 commit 1ab90dd
Showing 1 changed file with 27 additions and 7 deletions.
34 changes: 27 additions & 7 deletions gallery_dl/extractor/civitai.py
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -193,12 +193,8 @@ def _extract_files_image(self, model, version, user):
]

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,
Expand Down Expand Up @@ -321,6 +317,15 @@ def images(self, params):
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)
Expand Down Expand Up @@ -398,6 +403,21 @@ def images(self, params, defaults=True):

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)}
Expand Down

0 comments on commit 1ab90dd

Please sign in to comment.