[civitai] ensure image files have an 'id' (#6251)

mikf · Sep 29, 2024 · 987f6c3 · 987f6c3
1 parent 53ae5e3
commit 987f6c3
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 14 deletions.
diff --git a/gallery_dl/extractor/civitai.py b/gallery_dl/extractor/civitai.py
@@ -93,7 +93,7 @@ class CivitaiModelExtractor(CivitaiExtractor):
     directory_fmt = ("{category}", "{user[username]}",
                      "{model[id]}{model[name]:? //}",
                      "{version[id]}{version[name]:? //}")
-    filename_fmt = "{filename}.{extension}"
+    filename_fmt = "{file[id]}.{extension}"
     archive_fmt = "{file[hash]}"
     pattern = BASE_PATTERN + r"/models/(\d+)(?:/?\?modelVersionId=(\d+))?"
     example = "https://civitai.com/models/12345/TITLE"
@@ -183,23 +183,22 @@ def _extract_files_image(self, model, version, user):
             }
             images = self.api.images(params, defaults=False)
 
-        return [
-            text.nameext_from_url(file["url"], {
-                "num" : num,
-                "file": file,
-                "url" : self._url(file),
-            })
-            for num, file in enumerate(images, 1)
-        ]
+        return self._image_results(images)
 
     def _extract_files_gallery(self, model, version, user):
         images = self.api.images_gallery(model, version, user)
+        return self._image_results(images)
+
+    def _image_results(self, images):
         for num, file in enumerate(images, 1):
-            yield text.nameext_from_url(file["url"], {
+            data = text.nameext_from_url(file["url"], {
                 "num" : num,
                 "file": file,
                 "url" : self._url(file),
             })
+            if "id" not in file and data["filename"].isdecimal():
+                file["id"] = text.parse_int(data["filename"])
+            yield data
 
     def _validate_file_model(self, response):
         if response.headers.get("Content-Type", "").startswith("text/html"):

diff --git a/test/results/civitai.py b/test/results/civitai.py
@@ -81,22 +81,26 @@
         "image"   : None,
         "username": "bolero537"
     },
-    "file"   : dict,
-    "model"  : dict,
+    "file"   : {
+        "id": {26887862, 26887856, 26887852},
+    },
+    "model"  : {
+        "id": 703211,
+    },
     "num"    : range(1, 3),
 },
 
 {
     "#url"  : "https://civitai.com/images/26962948",
     "#class": civitai.CivitaiImageExtractor,
-    "#options"     : {"image-flags": "w"},
+    "#options"     : {"quality": "w"},
     "#urls"        : "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/69bf3279-df2c-4ec8-b795-479e9cd3db1b/w/69bf3279-df2c-4ec8-b795-479e9cd3db1b.jpeg",
     "#sha1_content": "a9a9d08f5fcdbc1e1eec7f203717f9df97b7a671",
 
     "browsingLevel": 1,
     "createdAt": "2024-08-31T01:11:47.021Z",
     "date"     : "dt:2024-08-31 01:11:47",
-    "extension": "jpeg",
+    "extension": "jpg",
     "filename" : "69bf3279-df2c-4ec8-b795-479e9cd3db1b",
     "hash"     : "ULN0-w?b4nRjxGM{-;t7M_t7NGae~qRjMyt7",
     "height"   : 1536,