Update civitai_api.py

civitai_api.py

@@ -334,6 +334,41 @@ def search_civitai(query, types, base_model=[], sort=SORT[0], period=PERIOD[0],
            state
 
 
+def download_media(session, media_info, base_url, headers, is_video=False):
+    media_url = media_info["url"]
+    media_path = get_save_path(base_url, media_info["name"], is_video)
+
+    try:
+        response = session.get(media_url, headers=headers, timeout=(5.0, 15), stream=True)
+        response.raise_for_status()
+
+        with open(media_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+
+        if not is_video:
+            temp_image_path = media_path + ".tmp"
+            with Image.open(temp_image_path) as img:
+                img = img.convert('RGBA')
+                img.save(media_path)
+            os.remove(temp_image_path)
+
+        return media_path
+
+    except requests.RequestException as e:
+        print(f"Media download error: {e}")
+    return ""
+
+def get_save_path(base_url, file_name, is_video=False):
+    dir_path = Path(TEMP_DIR)
+    os.makedirs(dir_path, exist_ok=True)
+    name = Path(file_name).stem
+    ext = Path(file_name).suffix
+    if is_video:
+        return str(dir_path / f"{name}{ext}")
+    else:
+        return str(dir_path / f"{name}.png")
+
 def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", api_key=""):
     if not image_baseurl: image_baseurl = dl_url
     default = ("", "", "") if is_html else ""
@@ -360,15 +395,22 @@ def get_civitai_json(dl_url: str, is_html: bool=False, image_baseurl: str="", ap
             r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
             if not r.ok: return json, html, image
             html = r.text
-        if 'images' in json.keys() and len(json["images"]) != 0:
-            url = json["images"][0]["url"]
-            r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
-            if not r.ok: return json, html, image
-            image_temp = str(Path(TEMP_DIR, "image" + Path(url.split("/")[-1]).suffix))
-            image = str(Path(TEMP_DIR, Path(image_baseurl.split("/")[-1]).stem + ".png"))
-            with open(image_temp, 'wb') as f:
-                f.write(r.content)
-            Image.open(image_temp).convert('RGBA').save(image)
+        if "images" in json.keys() and len(json["images"]) != 0:
+            for image_info in json_data["images"]:
+                if image_info["type"] == "video":
+                    video_path = download_media(session, image_info, image_baseurl, headers, is_video=True)
+                elif image_info["type"] == "image":
+                    # image_path = download_media(session, image_info, image_baseurl, headers, is_video=False)
+                    url = json["images"][0]["url"]
+                    r = session.get(url, params=params, headers=headers, stream=True, timeout=(5.0, 15))
+                    if not r.ok: return json, html, image
+                    image_temp = str(Path(TEMP_DIR, "image" + Path(url.split("/")[-1]).suffix))
+                    image = str(Path(TEMP_DIR, Path(image_baseurl.split("/")[-1]).stem + ".png"))
+                    with open(image_temp, 'wb') as f:
+                        f.write(r.content)
+                    Image.open(image_temp).convert('RGBA').save(image)
+                if image_path and video_path:
+                    break
         return json, html, image
     except Exception as e:
         print(e)