[civitai] add initial support (#3706, #3787, #4129, #5995)

docs/configuration.rst

@@ -1501,7 +1501,7 @@ Type
 Default
     ``["image", "video", "download", "gallery"]``
 Description
-    Determines the type and order of files to be downloaded.
+    Determines the type and order of files to download.
 
     Available types are
     ``image``,
@@ -1510,6 +1510,34 @@ Description
     ``gallery``.
 
 
+extractor.civitai.api-key
+-------------------------
+Type
+    ``string``
+Description
+    The API Key value generated in your
+    `User Account Settings <https://civitai.com/user/account>`__
+    to make authorized API requests.
+
+    See `API/Authorization <https://developer.civitai.com/docs/api/public-rest#authorization>`__
+    for details.
+
+
+extractor.civitai.files
+-----------------------
+Type
+    ``list`` of ``strings``
+Default
+    ``["image"]``
+Description
+    Determines the type and order of files to download.
+
+    Available types are
+    ``model``,
+    ``image``,
+    ``gallery``.
+
+
 extractor.cohost.asks
 ---------------------
 Type

docs/supportedsites.md

@@ -157,10 +157,16 @@ Consider all listed sites to potentially be NSFW.
     <td>Articles, Creators, Followed Users, Recent Images</td>
     <td></td>
 </tr>
+<tr>
+    <td>Civitai</td>
+    <td>https://www.civitai.com/</td>
+    <td>Models, Search Results, Tag Searches, User Profiles</td>
+    <td></td>
+</tr>
 <tr>
     <td>cohost!</td>
     <td>https://cohost.org/</td>
-    <td>Posts, Tag Searches, User Profiles</td>
+    <td>Likes, Posts, Tag Searches, User Profiles</td>
     <td></td>
 </tr>
 <tr>

gallery_dl/extractor/__init__.py

@@ -36,6 +36,7 @@
     "catbox",
     "chevereto",
     "cien",
+    "civitai",
     "cohost",
     "comicvine",
     "cyberdrop",

gallery_dl/extractor/civitai.py

@@ -0,0 +1,229 @@
+# -*- coding: utf-8 -*-
+
+# Copyright 2024 Mike Fährmann
+#
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+
+"""Extractors for https://www.civitai.com/"""
+
+from .common import Extractor, Message
+from .. import text
+import itertools
+import re
+
+BASE_PATTERN = r"(?:https?://)?civitai\.com"
+
+
+class CivitaiExtractor(Extractor):
+    """Base class for civitai extractors"""
+    category = "civitai"
+    root = "https://civitai.com"
+    request_interval = (0.5, 1.5)
+
+    def _init(self):
+        self.api = CivitaiAPI(self)
+
+    def items(self):
+        data = {"_extractor": CivitaiModelExtractor}
+        for model in self.models():
+            url = "{}/models/{}".format(self.root, model["id"])
+            yield Message.Queue, url, data
+
+
+class CivitaiModelExtractor(CivitaiExtractor):
+    subcategory = "model"
+    directory_fmt = ("{category}", "{user[username]}",
+                     "{model[id]}{model[name]:? //}",
+                     "{version[id]}{version[name]:? //}")
+    archive_fmt = "{file[hash]}"
+    pattern = BASE_PATTERN + r"/models/(\d+)(?:/?\?modelVersionId=(\d+))?"
+    example = "https://civitai.com/models/12345/TITLE"
+
+    def items(self):
+        model_id, version_id = self.groups
+        self._sub = re.compile(r"/width=\d*/").sub
+
+        model = self.api.model(model_id)
+        creator = model["creator"]
+        versions = model["modelVersions"]
+        del model["creator"]
+        del model["modelVersions"]
+
+        if version_id:
+            version_id = int(version_id)
+            for version in versions:
+                if version["id"] == version_id:
+                    break
+            else:
+                version = self.api.model_version(version_id)
+            versions = (version,)
+
+        for version in versions:
+            version["date"] = text.parse_datetime(
+                version["createdAt"], "%Y-%m-%dT%H:%M:%S.%fZ")
+
+            data = {
+                "model"  : model,
+                "version": version,
+                "user"   : creator,
+            }
+
+            yield Message.Directory, data
+            for file in self._extract_files(model, version):
+                file.update(data)
+                yield Message.Url, file["url"], file
+
+    def _extract_files(self, model, version):
+        filetypes = self.config("files")
+        if filetypes is None:
+            return self._extract_files_image(model, version)
+
+        generators = {
+            "model"   : self._extract_files_model,
+            "image"   : self._extract_files_image,
+            "gallery" : self._extract_files_gallery,
+            "gallerie": self._extract_files_gallery,
+        }
+        if isinstance(filetypes, str):
+            filetypes = filetypes.split(",")
+
+        return itertools.chain.from_iterable(
+            generators[ft.rstrip("s")](model, version)
+            for ft in filetypes
+        )
+
+    def _extract_files_model(self, model, version):
+        return [
+            {
+                "num"      : num,
+                "file"     : file,
+                "filename" : file["name"],
+                "extension": "bin",
+                "url"      : file["downloadUrl"],
+                "_http_headers" : {
+                    "Authorization": self.api.headers.get("Authorization")},
+                "_http_validate": self._validate_file_model,
+            }
+            for num, file in enumerate(version["files"], 1)
+        ]
+
+    def _extract_files_image(self, model, version):
+        return [
+            text.nameext_from_url(file["url"], {
+                "num" : num,
+                "file": file,
+                "url" : self._sub("/w/", file["url"]),
+            })
+            for num, file in enumerate(version["images"], 1)
+        ]
+
+    def _extract_files_gallery(self, model, version):
+        params = {
+            "modelId"       : model["id"],
+            "modelVersionId": version["id"],
+        }
+
+        for num, file in enumerate(self.api.images(params), 1):
+            yield text.nameext_from_url(file["url"], {
+                "num" : num,
+                "file": file,
+                "url" : self._sub("/w/", file["url"]),
+            })
+
+    def _validate_file_model(self, response):
+        if response.headers.get("Content-Type", "").startswith("text/html"):
+            alert = text.extr(
+                response.text, 'mantine-Alert-message">', "</div></div></div>")
+            if alert:
+                msg = "\"{}\" - 'api-key' required".format(
+                    text.remove_html(alert))
+            else:
+                msg = "'api-key' required to download this file"
+            self.log.warning(msg)
+            return False
+        return True
+
+
+class CivitaiTagExtractor(CivitaiExtractor):
+    subcategory = "tag"
+    pattern = BASE_PATTERN + r"/tag/([^?/#]+)"
+    example = "https://civitai.com/tag/TAG"
+
+    def models(self):
+        tag = text.unquote(self.groups[0])
+        return self.api.models({"tag": tag})
+
+
+class CivitaiSearchExtractor(CivitaiExtractor):
+    subcategory = "search"
+    pattern = BASE_PATTERN + r"/search/models\?([^#]+)"
+    example = "https://civitai.com/search/models?query=QUERY"
+
+    def models(self):
+        params = text.parse_query(self.groups[0])
+        return self.api.models(params)
+
+
+class CivitaiUserExtractor(CivitaiExtractor):
+    subcategory = "user"
+    pattern = BASE_PATTERN + r"/user/([^/?#]+)(?:/models)?/?(?:$|\?|#)"
+    example = "https://civitai.com/user/USER/models"
+
+    def models(self):
+        params = {"username": text.unquote(self.groups[0])}
+        return self.api.models(params)
+
+
+class CivitaiAPI():
+    """Interface for the Civitai Public REST API
+
+    https://developer.civitai.com/docs/api/public-rest
+    """
+
+    def __init__(self, extractor):
+        self.extractor = extractor
+        self.root = extractor.root + "/api"
+        self.headers = {"Content-Type": "application/json"}
+
+        api_key = extractor.config("api-key")
+        if api_key:
+            extractor.log.debug("Using api_key authentication")
+            self.headers["Authorization"] = "Bearer " + api_key
+
+    def images(self, params):
+        endpoint = "/v1/images"
+        return self._pagination(endpoint, params)
+
+    def model(self, model_id):
+        endpoint = "/v1/models/{}".format(model_id)
+        return self._call(endpoint)
+
+    def model_version(self, model_version_id):
+        endpoint = "/v1/model-versions/{}".format(model_version_id)
+        return self._call(endpoint)
+
+    def models(self, params):
+        return self._pagination("/v1/models", params)
+
+    def _call(self, endpoint, params=None):
+        if endpoint[0] == "/":
+            url = self.root + endpoint
+        else:
+            url = endpoint
+
+        response = self.extractor.request(
+            url, params=params, headers=self.headers)
+        return response.json()
+
+    def _pagination(self, endpoint, params):
+        while True:
+            data = self._call(endpoint, params)
+            yield from data["items"]
+
+            try:
+                endpoint = data["metadata"]["nextPage"]
+            except KeyError:
+                return
+            params = None

test/results/civitai.py

@@ -0,0 +1,91 @@
+# -*- coding: utf-8 -*-
+
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of the GNU General Public License version 2 as
+# published by the Free Software Foundation.
+
+from gallery_dl.extractor import civitai
+
+
+__tests__ = (
+{
+    "#url"     : "https://civitai.com/models/703211/maid-classic",
+    "#category": ("", "civitai", "model"),
+    "#class"   : civitai.CivitaiModelExtractor,
+    "#urls"    : [
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/5c4efa68-bb58-47c5-a716-98cd0f51f047/w/26962950.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/69bf3279-df2c-4ec8-b795-479e9cd3db1b/w/26962948.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/2dd1dc69-45a6-4beb-b36b-2e2bc65e3cda/w/26962957.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/52b6efa7-801c-4901-90b4-fa3964d23480/w/26887862.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/c4d3bcd5-0e23-4f4e-9f34-d13b2f2bf14c/w/26887856.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/68568d22-c4f3-45cb-ac32-82f1cedf968f/w/26887852.jpeg",
+    ],
+
+    "model"  : {
+        "allowCommercialUse": ["RentCivit"],
+        "allowDerivatives": True,
+        "allowDifferentLicense": True,
+        "allowNoCredit": True,
+        "cosmetic"   : None,
+        "description": "<p>The strength of Lora is recommended to be around 1.0.</p>",
+        "id"         : 703211,
+        "minor"      : False,
+        "name"       : "メイド　クラシック/maid classic",
+        "nsfw"       : False,
+        "nsfwLevel"  : 1,
+        "poi"        : False,
+        "stats"      : dict,
+        "tags"       : ["clothing"],
+        "type"       : "LORA"
+    },
+    "user"   : {
+        "image"   : None,
+        "username": "bolero537"
+    },
+    "file"   : dict,
+    "version": dict,
+    "num"    : range(1, 3),
+},
+
+{
+    "#url"     : "https://civitai.com/models/703211?modelVersionId=786644",
+    "#category": ("", "civitai", "model"),
+    "#class"   : civitai.CivitaiModelExtractor,
+    "#urls"    : [
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/52b6efa7-801c-4901-90b4-fa3964d23480/w/26887862.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/c4d3bcd5-0e23-4f4e-9f34-d13b2f2bf14c/w/26887856.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/68568d22-c4f3-45cb-ac32-82f1cedf968f/w/26887852.jpeg",
+    ],
+
+    "version": {
+        "availability": "Public",
+        "baseModel"   : "Pony",
+        "createdAt"   : "2024-08-30T15:28:47.661Z",
+        "date"        : "dt:2024-08-30 15:28:47",
+        "downloadUrl" : "https://civitai.com/api/download/models/786644",
+        "files"       : list,
+        "id"          : 786644,
+        "images"      : list,
+        "index"       : 1,
+        "name"        : "v1.0 pony",
+        "nsfwLevel"   : 1,
+        "publishedAt" : "2024-08-30T15:39:17.674Z",
+        "stats"       : dict,
+        "status"      : "Published",
+        "trainedWords": [
+            "maid",
+            "madi apron",
+            "maid headdress",
+            "long sleeves",
+        ],
+    },
+    "user"   : {
+        "image"   : None,
+        "username": "bolero537"
+    },
+    "file"   : dict,
+    "model"  : dict,
+    "num"    : range(1, 3),
+},
+
+)