tar.gz Working

API/app.py

@@ -1,5 +1,5 @@
-from fastapi import FastAPI, Query
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, Query, BackgroundTasks
+from fastapi.responses import FileResponse
 import aiohttp
 import asyncio
 import requests.utils
@@ -8,7 +8,7 @@ import os
 import io
 from pydantic import BaseModel
 from typing import List, Dict
-import zipfile
+import tarfile
 from dotenv import load_dotenv
 
 img_base = 'https://i.pximg.net/img-original/img/'
@@ -224,39 +224,57 @@ async def pixif(
     return image_exifs
 
 
-async def generate_zip(posts, session):
-    zip_buffer = io.BytesIO()
-    zipf = zipfile.ZipFile(zip_buffer, mode="w", compression=zipfile.ZIP_DEFLATED)
+import tempfile
+from fastapi import BackgroundTasks
 
+async def generate_tar_gz(posts, session):
     semaphore = asyncio.Semaphore(100)
+    images = {}
 
-    async def add_to_zip(post_id, image_url):
+    async def fetch_image(post_id, image_url):
         async with semaphore:
             url = f"{img_base}{image_url}"
             async with session.get(url) as response:
                 image_data = await response.read()
-                image_name = f"{post_id}.png"
-                zipf.writestr(image_name, image_data)
+                return post_id, image_data
 
-    tasks = [add_to_zip(post_id, image_url) for post_id, image_url in posts.items()]
-    await asyncio.gather(*tasks)
+    tasks = [fetch_image(post_id, image_url) for post_id, image_url in posts.items()]
+    results = await asyncio.gather(*tasks)
 
-    zipf.close()
-    zip_buffer.seek(0)
-    return zip_buffer
+    images = {post_id: image_data for post_id, image_data in results}
+
+    temp_file = tempfile.NamedTemporaryFile(delete=False, suffix=".tar.gz")
+    with tarfile.open(fileobj=temp_file, mode="w:gz") as tar:
+        for post_id, image_data in images.items():
+            image_name = f"{post_id}.png"
+            file_info = tarfile.TarInfo(name=image_name)
+            file_info.size = len(image_data)
+            tar.addfile(tarinfo=file_info, fileobj=io.BytesIO(image_data))
+
+    temp_file.seek(0)
+    return temp_file
 
 @app.post("/download")
 async def download(
-    items: PixifDownloadModel    
+    items: PixifDownloadModel,
+    background_tasks: BackgroundTasks
 ):
     posts = items.posts
     async with aiohttp.ClientSession(cookies=cookies, headers=headers) as session:
-        zip_buffer = await generate_zip(posts, session)
+        temp_file = await generate_tar_gz(posts, session)
+
+    filename = f"{base26_time()}.tar.gz"
+
+    def cleanup_temp_file(file_path):
+        os.unlink(file_path)
+
+    background_tasks.add_task(cleanup_temp_file, temp_file.name)
 
-    return StreamingResponse(
-        zip_buffer,
-        media_type="application/zip",
-        headers={"Content-Disposition": f"attachment; filename={base26_time()}.zip"}
+    return FileResponse(
+        path=temp_file.name,
+        media_type="application/gzip",
+        filename=filename,
+        background=background_tasks
     )
 
 @app.get("/")

Client/Extract Pixiv/ai_search.py

@@ -1,7 +1,7 @@
 import requests
 import os
 
-local = False
+local = 0
 if local:
     endpoint = "http://127.0.0.1:7860"
 else:

Client/Extract Pixiv/user.py

@@ -2,7 +2,7 @@ import requests
 import re
 import os
 
-local = False
+local = 0
 if local:
     endpoint = "http://127.0.0.1:7860"
 else:

Client/hunt.py

@@ -2,9 +2,9 @@ import os
 import lmdb
 import requests
 import io
-import zipfile
+import tarfile
 
-local = False
+local = 0
 if local:
     endpoint = "http://127.0.0.1:7860"
 else:
@@ -62,12 +62,12 @@ for index in indexs:
     to_download = {post_id: url for post_id, url in post_ids.items() if url not in blacklist and f"{post_id}.png" not in images_cache}
     print(f"Downloading {len(to_download)} images...")
     if to_download:
-        images_zip = requests.post(f'{endpoint}/download', json={'posts': to_download}, stream=True)
+        images_tar_gz = requests.post(f'{endpoint}/download', json={'posts': to_download}, stream=True)
 
         print("Extracting images...")
-        with io.BytesIO(images_zip.content) as f:
-            with zipfile.ZipFile(f, mode='r') as zipf:
-                zipf.extractall("images/Stash")
+        with io.BytesIO(images_tar_gz.content) as f:
+            with tarfile.open(fileobj=f, mode='r:gz') as tarf:
+                tarf.extractall("images/Stash")
         
 
     images_cache = os.listdir("images/Stash")