A

Client/hunt.py

@@ -1,14 +1,20 @@
 import json
 import os
 import sqlite3
+from concurrent.futures import ThreadPoolExecutor, as_completed
+
 import requests
-import tempfile
-import zipfile
 from tqdm import tqdm
 
 local = 0
 dry_run = 0
 endpoint = "http://127.0.0.1:7860" if local else "https://q6-p.hf.space"
+IMG_BASE = "https://i.pximg.net/img-original/img/"
+REQUEST_HEADERS = {
+    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:126.0) Gecko/20100101 Firefox/126.0",
+    "Referer": "https://www.pixiv.net/",
+}
+MAX_WORKERS = min(16, os.cpu_count() or 8)
 
 def read_dotenv_value(path, key):
     try:
@@ -129,28 +135,57 @@ for inp in inputs:
     elif inp.isdigit():
         indexs.append(int(inp) - 1)
 
-def download_zip(to_download, dest_dir):
-    response = requests.post(f"{endpoint}/pixif_zip", json={"d": to_download}, stream=True)
-    response.raise_for_status()
-    total = int(response.headers.get("Content-Length", 0))
+def build_image_url(url):
+    if url.startswith("http"):
+        return url
+    return IMG_BASE + url
+
+def download_one_image(post_id, url, dest_dir, phpsessid):
+    dest_path = os.path.join(dest_dir, f"{post_id}.png")
+    if os.path.exists(dest_path):
+        return post_id, "exists", None
 
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".zip") as tmp_file:
-        tmp_path = tmp_file.name
-        if total > 0:
-            with tqdm(total=total, unit="B", unit_scale=True, desc="Downloading zip") as pbar:
+    tmp_path = dest_path + ".part"
+    full_url = build_image_url(url)
+    try:
+        with requests.get(
+            full_url,
+            headers=REQUEST_HEADERS,
+            cookies={"PHPSESSID": phpsessid},
+            stream=True,
+            timeout=60,
+        ) as response:
+            response.raise_for_status()
+            with open(tmp_path, "wb") as handle:
                 for chunk in response.iter_content(chunk_size=1024 * 1024):
                     if chunk:
-                        tmp_file.write(chunk)
-                        pbar.update(len(chunk))
-        else:
-            for chunk in response.iter_content(chunk_size=1024 * 1024):
-                if chunk:
-                    tmp_file.write(chunk)
-
-    with zipfile.ZipFile(tmp_path, "r") as zf:
-        zf.extractall(dest_dir)
+                        handle.write(chunk)
+        os.replace(tmp_path, dest_path)
+        return post_id, "ok", None
+    except Exception as exc:
+        if os.path.exists(tmp_path):
+            try:
+                os.remove(tmp_path)
+            except OSError:
+                pass
+        return post_id, "error", str(exc)
+
+def download_images(to_download, dest_dir, phpsessid, max_workers=MAX_WORKERS):
+    if not to_download:
+        return
 
-    os.remove(tmp_path)
+    os.makedirs(dest_dir, exist_ok=True)
+    futures = []
+    with ThreadPoolExecutor(max_workers=max_workers) as executor:
+        for post_id, url in to_download.items():
+            futures.append(executor.submit(download_one_image, post_id, url, dest_dir, phpsessid))
+
+        with tqdm(total=len(futures), unit="image", desc="Downloading images") as pbar:
+            for future in as_completed(futures):
+                post_id, status, detail = future.result()
+                if status == "error":
+                    tqdm.write(f"Failed {post_id}: {detail}")
+                pbar.update(1)
 
 def decode_if_binary(val):
     if type(val) is bytes:
@@ -185,7 +220,7 @@ for index in indexs:
     if to_download:
         print(f"Total images to download: {len(to_download)}")
         if not dry_run:
-            download_zip(to_download, "images/Stash")
+            download_images(to_download, "images/Stash", phpsessid)
 
     images_cache.update(os.listdir("images/Stash"))