From 1ab90dd9e252761515a12b5bd631c6fcea33f052 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20F=C3=A4hrmann?= <mike_faehrmann@web.de>
Date: Fri, 27 Sep 2024 22:07:54 +0200
Subject: [PATCH] [civitai] improve model 'gallery' extraction (#3706)

---
 gallery_dl/extractor/civitai.py | 34 ++++++++++++++++++++++++++-------
 1 file changed, 27 insertions(+), 7 deletions(-)

diff --git a/gallery_dl/extractor/civitai.py b/gallery_dl/extractor/civitai.py
index 00880e6246..d434cca54b 100644
--- a/gallery_dl/extractor/civitai.py
+++ b/gallery_dl/extractor/civitai.py
@@ -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 @@ 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,
@@ -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)
@@ -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)}