[civitai] add 'quality' option (#3706)

download 'original=true' files by default
mikf · Sep 25, 2024 · a2db0d5 · a2db0d5
1 parent 92bbb9b
commit a2db0d5
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 23 deletions.
diff --git a/docs/configuration.rst b/docs/configuration.rst
@@ -383,6 +383,7 @@ Type
     |Duration|_
 Default
     * ``"0.5-1.5"``
+        ``ao3``, ``civitai``,
         ``[Danbooru]``, ``[E621]``, ``[foolfuuka]:search``, ``itaku``,
         ``koharu``,
         ``newgrounds``, ``[philomena]``, ``pixiv:novel``, ``plurk``,
@@ -1533,14 +1534,34 @@ Type
 Default
     ``["image"]``
 Description
-    Determines the type and order of files to download.
+    Determines the type and order of files to download when processing models.
 
     Available types are
     ``model``,
     ``image``,
     ``gallery``.
 
 
+extractor.civitai.quality
+-------------------------
+Type
+    * ``string``
+    * ``list`` of ``strings``
+Default
+    ``"original=true"``
+Example
+    * ``"width=1280,quality=90"``
+    * ``["width=1280", "quality=90"]``
+Description
+    A (comma-separated) list of image quality options
+    to pass with every image URL.
+
+    Known available options include ``original``, ``quality``, ``width``
+
+    Note: Set this option to an arbitrary letter, e.g., ``"w"``,
+    to download images in JPEG format at their original resolution.
+
+
 extractor.cohost.asks
 ---------------------
 Type

diff --git a/gallery_dl/extractor/civitai.py b/gallery_dl/extractor/civitai.py
@@ -10,9 +10,7 @@
 
 from .common import Extractor, Message
 from .. import text
-import functools
 import itertools
-import re
 
 BASE_PATTERN = r"(?:https?://)?civitai\.com"
 USER_PATTERN = BASE_PATTERN + r"/user/([^/?#]+)"
@@ -30,6 +28,14 @@ class CivitaiExtractor(Extractor):
     def _init(self):
         self.api = CivitaiAPI(self)
 
+        quality = self.config("quality")
+        if quality:
+            if not isinstance(quality, str):
+                quality = ",".join(quality)
+            self._image_quality = quality
+        else:
+            self._image_quality = "original=true"
+
     def items(self):
         models = self.models()
         if models:
@@ -57,16 +63,17 @@ def images(self):
         return ()
 
     def _orig(self, url):
-        sub_width = functools.partial(re.compile(r"/width=\d*/").sub, "/w/")
-        CivitaiExtractor._orig = sub_width
-        return sub_width(url)
+        parts = url.rsplit("/", 2)
+        parts[1] = self._image_quality
+        return "/".join(parts)
 
 
 class CivitaiModelExtractor(CivitaiExtractor):
     subcategory = "model"
     directory_fmt = ("{category}", "{user[username]}",
                      "{model[id]}{model[name]:? //}",
                      "{version[id]}{version[name]:? //}")
+    filename_fmt = "{filename}.{extension}"
     archive_fmt = "{file[hash]}"
     pattern = BASE_PATTERN + r"/models/(\d+)(?:/?\?modelVersionId=(\d+))?"
     example = "https://civitai.com/models/12345/TITLE"

diff --git a/test/results/civitai.py b/test/results/civitai.py
@@ -9,15 +9,15 @@
 
 __tests__ = (
 {
-    "#url"     : "https://civitai.com/models/703211/maid-classic",
-    "#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",
+    "#url"  : "https://civitai.com/models/703211/maid-classic",
+    "#class": civitai.CivitaiModelExtractor,
+    "#urls" : [
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/5c4efa68-bb58-47c5-a716-98cd0f51f047/original=true/26962950.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/69bf3279-df2c-4ec8-b795-479e9cd3db1b/original=true/26962948.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/2dd1dc69-45a6-4beb-b36b-2e2bc65e3cda/original=true/26962957.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/52b6efa7-801c-4901-90b4-fa3964d23480/original=true/26887862.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/c4d3bcd5-0e23-4f4e-9f34-d13b2f2bf14c/original=true/26887856.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/68568d22-c4f3-45cb-ac32-82f1cedf968f/original=true/26887852.jpeg",
     ],
 
     "model"  : {
@@ -47,12 +47,12 @@
 },
 
 {
-    "#url"     : "https://civitai.com/models/703211?modelVersionId=786644",
-    "#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",
+    "#url"  : "https://civitai.com/models/703211?modelVersionId=786644",
+    "#class": civitai.CivitaiModelExtractor,
+    "#urls" : [
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/52b6efa7-801c-4901-90b4-fa3964d23480/original=true/26887862.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/c4d3bcd5-0e23-4f4e-9f34-d13b2f2bf14c/original=true/26887856.jpeg",
+        "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/68568d22-c4f3-45cb-ac32-82f1cedf968f/original=true/26887852.jpeg",
     ],
 
     "version": {
@@ -87,8 +87,9 @@
 },
 
 {
-    "#url"     : "https://civitai.com/images/26962948",
-    "#class"   : civitai.CivitaiImageExtractor,
+    "#url"  : "https://civitai.com/images/26962948",
+    "#class": civitai.CivitaiImageExtractor,
+    "#options"     : {"image-flags": "w"},
     "#urls"        : "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/69bf3279-df2c-4ec8-b795-479e9cd3db1b/w/69bf3279-df2c-4ec8-b795-479e9cd3db1b.jpeg",
     "#sha1_content": "a9a9d08f5fcdbc1e1eec7f203717f9df97b7a671",