X

API/app.py

@@ -6,12 +6,19 @@ import time
 import uvicorn
 from dotenv import load_dotenv
 import os
+from aiofiles import open as aio_open
 from pydantic import BaseModel
-from typing import List
+from typing import List, Dict
 
+img_base = 'https://i.pximg.net/img-original/'
 class pixifModel(BaseModel):
     post_ids: List[int]
 
+class PixifDownloadModel(BaseModel):
+    posts: Dict[str, str]
+
+os.makedirs('Stash', exist_ok=True)
+
 env_path = os.path.dirname(os.path.realpath(__file__)) + "/../.env"
 # config = dotenv_values(env_path)
 # PHPSESSID = config.get("PHPSESSID")
@@ -206,9 +213,28 @@ async def pixif(
         tasks = [process_post(post_id, session, semaphore) for post_id in post_ids]
         results = await asyncio.gather(*tasks)
 
-    image_exifs = {post_id: image_url for post_id, image_url in results if image_url}
+    image_exifs = {post_id: image_url.replace('https://i.pximg.net/img-original/', '', 1) for post_id, image_url in results if image_url}
     return image_exifs
 
+async def download_image(session, post_id, post_url):
+    url = f"{img_base}{post_url}"
+    async with session.get(url) as response:
+        content = await response.read()
+        async with aio_open(f"Stash/{post_id}.png", "wb") as f:
+            await f.write(content)
+
+
+@app.post("download")
+async def download(
+    items: PixifDownloadModel    
+):
+    posts = items.posts
+    
+    async with aiohttp.ClientSession(cookies=cookies, headers=headers) as session:
+        tasks = [download_image(session, post_id, image_url) for post_id, image_url in posts.items()]
+        await asyncio.gather(*tasks)
+    print(os.listdir("Stash"))
+
 @app.get("/")
 async def read_root():
     return {"message": "Hello, World!"}

Client/hunt.py

@@ -2,11 +2,12 @@ import os
 import lmdb
 import requests
 
+img_base = 'https://i.pximg.net/img-original/'
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 os.makedirs("images/Stash", exist_ok=True)
 
-db = lmdb.open("db", subdir=True, map_size=10485760)
+db = lmdb.open("db", subdir=True, map_size=1048576)
 images_cache = os.listdir("images/Stash")
 valid = [f for f in os.listdir() if f.endswith(".txt")]
 
@@ -24,21 +25,24 @@ for inp in inputs:
         start, end = map(int, inp.split("-"))
         indexs.extend(range(start - 1, end))
 
-for index in indexs:
-    with open(valid[index], "r") as f:
-        post_ids = [x for x in f.read().split("\n") if x]
+with db.begin(write=True) as txn:
+    for index in indexs:
+        group_name = valid[index].rsplit(".", 1)[0]
 
-    with db.begin(write=False) as txn:
-        post_ids = [post_id for post_id in post_ids if not txn.get(post_id.encode())]
+        with open(valid[index], "r") as f:
+            post_ids = [x for x in f.read().split("\n") if x]
 
-    if not post_ids:
-        continue
-    data = requests.post('https://q6-p.hf.space/pixif', json={"post_ids": post_ids}).json()
+        filtered_post_ids = [post_id for post_id in post_ids if not txn.get(post_id.encode())]
+        data = requests.post('https://q6-p.hf.space/pixif', json={"post_ids": filtered_post_ids}).json()
 
-    with db.begin(write=True) as txn:
         for post_id, url in data.items():
             txn.put(post_id.encode(), url.encode())
 
-        missing_post_ids = set(post_ids) - set(data.keys())
+        missing_post_ids = set(filtered_post_ids) - set(data.keys())
         for post_id in missing_post_ids:
-            txn.put(post_id.encode(), b'')
+            txn.put(post_id.encode(), b'\x00')
+        
+        data.update({post_id: txn.get(post_id.encode()).decode() for post_id in set(post_ids) - missing_post_ids if txn.get(post_id.encode()) != b'\x00'})
+        
+        print(data)
+        images_zip = requests.post('https://q6-p.hf.space/download', json={'posts': data}).json()

Dockerfile

@@ -8,7 +8,7 @@ ENV PYTHONDONTWRITEBYTECODE=1
 ENV PYTHONUNBUFFERED=1
 WORKDIR /app
 
-RUN pip install --no-cache-dir fastapi aiohttp uvicorn python-dotenv pydantic requests
+RUN pip install --no-cache-dir fastapi aiohttp uvicorn python-dotenv pydantic requests aiofiles
 
 COPY --chown=user ./API /app
 CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]