Server zip

API/app.py

@@ -1,7 +1,11 @@
-from fastapi import FastAPI, Query
+from fastapi import FastAPI, Query, BackgroundTasks
+from fastapi.responses import FileResponse
 import aiohttp
 import asyncio
 import time
+import tempfile
+import zipfile
+import os
 from pydantic import BaseModel
 from typing import List, Dict
 
@@ -15,6 +19,9 @@ class pixifModel(BaseModel):
 class PixifDownloadModel(BaseModel):
     posts: Dict[str, str]
 
+class PixifZipModel(BaseModel):
+    d: Dict[str, str]
+
 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/',
@@ -123,6 +130,11 @@ async def process_post(post_id, session, semaphore):
         except Exception:
             return post_id, None
 
+async def fetch_image_bytes(session, url, semaphore):
+    async with semaphore:
+        async with session.get(url) as response:
+            return await response.read()
+
 @app.get("/allimages")
 async def all_images(
     only_first = Query("0", description="Only fetch the first image of each post."),
@@ -179,6 +191,39 @@ async def pixif(
 
     return image_exifs
 
+@app.post("/pixif_zip")
+async def pixif_zip(items: PixifZipModel, background_tasks: BackgroundTasks):
+    downloads = items.d
+    if not downloads:
+        return {"detail": "No downloads requested."}
+
+    tmp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".zip")
+    tmp_path = tmp_file.name
+    tmp_file.close()
+
+    connector = aiohttp.TCPConnector(limit=20)
+    semaphore = asyncio.Semaphore(20)
+    async with aiohttp.ClientSession(headers=headers, connector=connector) as session:
+        tasks = {}
+        for post_id, url in downloads.items():
+            full_url = url if url.startswith("http") else img_base + url
+            task = asyncio.create_task(fetch_image_bytes(session, full_url, semaphore))
+            tasks[task] = post_id
+
+        with zipfile.ZipFile(tmp_path, "w", compression=zipfile.ZIP_DEFLATED) as zf:
+            for task in asyncio.as_completed(tasks):
+                post_id = tasks[task]
+                try:
+                    data = await task
+                except Exception:
+                    continue
+                if data:
+                    zf.writestr(f"{post_id}.png", data)
+
+    background_tasks.add_task(os.remove, tmp_path)
+    filename = f"pixif_{base26_time()}.zip"
+    return FileResponse(tmp_path, media_type="application/zip", filename=filename)
+
 @app.get("/")
 async def read_root():
     return {"message": "Hello, World!"}

Client/hunt.py

@@ -1,12 +1,12 @@
 import os
 import lmdb
 import requests
-from concurrent.futures import ThreadPoolExecutor
+import tempfile
+import zipfile
 from tqdm import tqdm
 
 local = 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/'
 
 def read_dotenv_value(path, key):
     try:
@@ -54,15 +54,28 @@ for inp in inputs:
     elif inp.isdigit():
         indexs.append(int(inp) - 1)
 
-def download_image(args):
-    post_id, url = args
-    response = requests.get(img_base + url, headers={
-        "Referer": "https://www.pixiv.net/",
-        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64; rv:126.0) Gecko/20100101 Firefox/126.0"
-    })
-    with open(f"images/Stash/{post_id}.png", "wb") as img_file:
-        img_file.write(response.content)
-    return post_id, True
+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))
+
+    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:
+                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)
+
+    os.remove(tmp_path)
 
 def decode_if_binary(val):
     if type(val) is bytes:
@@ -94,8 +107,7 @@ for index in indexs:
     
     if to_download:
         print(f"Total images to download: {len(to_download)}")
-        with ThreadPoolExecutor(max_workers=8) as executor:
-            list(tqdm(executor.map(download_image, to_download.items()), total=len(to_download), desc="Downloading"))
+        download_zip(to_download, "images/Stash")
 
     images_cache.update(os.listdir("images/Stash"))