From daa220370c7099b237e6c1f769953143ef09a7e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mike=20F=C3=A4hrmann?= <mike_faehrmann@web.de>
Date: Wed, 25 Sep 2024 07:39:52 +0200
Subject: [PATCH] [civitai] improve image URL handling (#3706)

- support building URLs from URL-hash values
- rename '_orig' to '_url'
---
 gallery_dl/extractor/civitai.py | 28 +++++++++++++++++++++-------
 1 file changed, 21 insertions(+), 7 deletions(-)

diff --git a/gallery_dl/extractor/civitai.py b/gallery_dl/extractor/civitai.py
index dd95b4bd..a7c29e47 100644
--- a/gallery_dl/extractor/civitai.py
+++ b/gallery_dl/extractor/civitai.py
@@ -33,8 +33,10 @@ class CivitaiExtractor(Extractor):
             if not isinstance(quality, str):
                 quality = ",".join(quality)
             self._image_quality = quality
+            self._image_ext = ("png" if quality == "original=true" else "jpg")
         else:
             self._image_quality = "original=true"
+            self._image_ext = "png"
 
     def items(self):
         models = self.models()
@@ -48,10 +50,11 @@ class CivitaiExtractor(Extractor):
         images = self.images()
         if images:
             for image in images:
-                url = self._orig(image["url"])
+                url = self._url(image)
                 image["date"] = text.parse_datetime(
                     image["createdAt"], "%Y-%m-%dT%H:%M:%S.%fZ")
                 text.nameext_from_url(url, image)
+                image["extension"] = self._image_ext
                 yield Message.Directory, image
                 yield Message.Url, url, image
             return
@@ -62,10 +65,21 @@ class CivitaiExtractor(Extractor):
     def images(self):
         return ()
 
-    def _orig(self, url):
-        parts = url.rsplit("/", 2)
-        parts[1] = self._image_quality
-        return "/".join(parts)
+    def _url(self, image):
+        url = image["url"]
+        if "/" in url:
+            parts = url.rsplit("/", 2)
+            parts[1] = self._image_quality
+            return "/".join(parts)
+
+        name = image.get("name")
+        if not name:
+            mime = image.get("mimeType") or self._image_ext
+            name = "{}.{}".format(image.get("id"), mime.rpartition("/")[2])
+        return (
+            "https://image.civitai.com/xG1nkqKTMzGDvpLrqFT7WA/{}/{}/{}".format(
+                url, self._image_quality, name)
+        )
 
 
 class CivitaiModelExtractor(CivitaiExtractor):
@@ -150,7 +164,7 @@ class CivitaiModelExtractor(CivitaiExtractor):
             text.nameext_from_url(file["url"], {
                 "num" : num,
                 "file": file,
-                "url" : self._orig(file["url"]),
+                "url" : self._url(file),
             })
             for num, file in enumerate(version["images"], 1)
         ]
@@ -165,7 +179,7 @@ class CivitaiModelExtractor(CivitaiExtractor):
             yield text.nameext_from_url(file["url"], {
                 "num" : num,
                 "file": file,
-                "url" : self._orig(file["url"]),
+                "url" : self._url(file),
             })
 
     def _validate_file_model(self, response):