[civitai] add 'generate' extractor (#7796)

docs/supportedsites.md

@@ -184,7 +184,7 @@ Consider all listed sites to potentially be NSFW.
 <tr>
     <td>Civitai</td>
     <td>https://www.civitai.com/</td>
-    <td>individual Images, Image Listings, Models, Model Listings, Posts, Post Listings, Image Searches, Model Searches, Tag Searches, User Profiles, User Images, Image Reactions, User Models, User Posts, User Videos, Video Reactions</td>
+    <td>Generated Files, individual Images, Image Listings, Models, Model Listings, Posts, Post Listings, Image Searches, Model Searches, Tag Searches, User Profiles, User Images, Image Reactions, User Models, User Posts, User Videos, Video Reactions</td>
     <td></td>
 </tr>
 <tr>

gallery_dl/extractor/civitai.py

@@ -186,10 +186,7 @@ class CivitaiExtractor(Extractor):
             yield data
 
     def _image_reactions(self):
-        if "Authorization" not in self.api.headers and \
+        self._require_auth()
-                not self.cookies.get(
-                "__Secure-civitai-token", domain=".civitai.com"):
-            raise exception.AuthorizationError("api-key or cookies required")
 
         params = self.params
         params["authed"] = True
@@ -198,6 +195,12 @@ class CivitaiExtractor(Extractor):
             params["reactions"] = ("Like", "Dislike", "Heart", "Laugh", "Cry")
         return self.api.images(params)
 
+    def _require_auth(self):
+        if "Authorization" not in self.api.headers and \
+                not self.cookies.get(
+                "__Secure-civitai-token", domain=".civitai.com"):
+            raise exception.LoginRequired("'api-key' or cookies needed")
+
     def _parse_query(self, value):
         return text.parse_query_list(
             value, {"tags", "reactions", "baseModels", "tools", "techniques",
@@ -525,6 +528,28 @@ class CivitaiUserVideosExtractor(CivitaiExtractor):
     images = CivitaiUserImagesExtractor.images
 
 
+class CivitaiGenerateExtractor(CivitaiExtractor):
+    """Extractor for your generated files feed"""
+    subcategory = "generate"
+    filename_fmt = "{filename}.{extension}"
+    directory_fmt = ("{category}", "generated")
+    pattern = f"{BASE_PATTERN}/generate"
+    example = "https://civitai.com/generate"
+
+    def items(self):
+        self._require_auth()
+
+        for gen in self.api.orchestrator_queryGeneratedImages():
+            gen["date"] = text.parse_datetime(
+                gen["createdAt"], "%Y-%m-%dT%H:%M:%S.%fZ")
+            yield Message.Directory, gen
+            for step in gen.pop("steps", ()):
+                for image in step.pop("images", ()):
+                    data = {"file": image, **step, **gen}
+                    url = image["url"]
+                    yield Message.Url, url, text.nameext_from_url(url, data)
+
+
 class CivitaiRestAPI():
     """Interface for the Civitai Public REST API
 
@@ -739,6 +764,15 @@ class CivitaiTrpcAPI():
         params = {"username": username}
         return (self._call(endpoint, params),)
 
+    def orchestrator_queryGeneratedImages(self):
+        endpoint = "orchestrator.queryGeneratedImages"
+        params = {
+            "ascending": False,
+            "tags"     : ("gen",),
+            "authed"   : True,
+        }
+        return self._pagination(endpoint, params)
+
     def _call(self, endpoint, params, meta=None):
         url = self.root + endpoint
         headers = self.headers

scripts/supportedsites.py

@@ -236,6 +236,7 @@ SUBCATEGORY_MAP = {
         "user-images": ("User Images", "Image Reactions"),
         "user-posts" : "User Posts",
         "user-videos": ("User Videos", "Video Reactions"),
+        "generate": "Generated Files",
     },
     "coomer": {
         "discord"       : "",

test/results/civitai.py

@@ -512,5 +512,10 @@ __tests__ = (
     "#exception": exception.AuthorizationError,
 },
 
+{
+    "#url"     : "https://civitai.com/generate",
+    "#class"   : civitai.CivitaiGenerateExtractor,
+    "#auth"    : True,
+},
 
 )