Revert "Optimised storing"

This reverts commit e733149ccf89c86b7bc96ba9ec9d24c3f6a0d85c.

Client/Extract Pixiv/ai_search.py

@@ -11,7 +11,7 @@ os.chdir(os.path.dirname(os.path.abspath(__file__)))
 
 input_url = input("Enter the URL: ")
 
-pages = 300 // 60
+pages = 1200 // 60
 
 params = {
     'raw': input_url,

Client/hunt.py

@@ -5,13 +5,17 @@ from concurrent.futures import ThreadPoolExecutor
 from tqdm import tqdm
 
 local = 0
-endpoint = "http://127.0.0.1:7860" if local else "https://q6-p.hf.space"
+if local:
+    endpoint = "http://127.0.0.1:7860"
+else:
+    endpoint = "https://q6-p.hf.space"
+
 img_base = 'https://i.pximg.net/img-original/img/'
 
 os.chdir(os.path.dirname(os.path.abspath(__file__)))
 os.makedirs("images/Stash", exist_ok=True)
 
-images_cache = set(os.listdir("images/Stash"))
+images_cache = os.listdir("images/Stash")
 
 db = lmdb.open("db", subdir=True, map_size=1048576)
 valid = [f for f in os.listdir() if f.endswith(".txt")]
@@ -21,14 +25,14 @@ for idx, file in enumerate(valid):
 
 inputs = input("Enter the index of the file: ")
 inputs = inputs.split()
-indexes = []
+indexs = []
 
 for inp in inputs:
     if inp.isdigit():
-        indexes.append(int(inp) - 1)
+        indexs.append(int(inp) - 1)
     elif "-" in inp:
         start, end = map(int, inp.split("-"))
-        indexes.extend(range(start - 1, end))
+        indexs.extend(range(start - 1, end))
 
 def download_image(args):
     post_id, url = args
@@ -41,52 +45,47 @@ def download_image(args):
         img_file.write(response.content)
     return post_id, True
 
-for index in indexes:
+for index in indexs:
     group_name = valid[index].rsplit(".", 1)[0]
     os.makedirs(f"images/{group_name}", exist_ok=True)
-
     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:
-        urls = {}
-        for post_id in post_ids:
-            url = txn.get(post_id.encode())
-            urls[post_id] = url.decode() if url is not None and url != b'' else None
-
-        filtered = [post_id for post_id, url in urls.items() 
-                    if url is None and f"{post_id}.png" not in images_cache]
-
-        print(f"Group: {group_name}\nFiltered: {len(filtered)}/{len(urls)}")
+        post_ids_dict = {post_id: txn.get(post_id.encode()) for post_id in post_ids}
+        post_ids_dict = {post_id: url.decode() if url != None else None for post_id, url in post_ids_dict.items()}
 
+        filtered = [post_id for post_id, url in post_ids_dict.items() if url == None and f"{post_id}.png" not in images_cache]
+        print(f"Group: {group_name}\nFiltered: {len(filtered)}/{len(post_ids_dict)}")
         if filtered:
             data = requests.post(f'{endpoint}/pixif', json={"post_ids": filtered}).json()
 
             for post_id, url in data.items():
                 txn.put(post_id.encode(), url.encode())
-                urls[post_id] = url
+                post_ids_dict[post_id] = url
 
             no_exif = set(filtered) - set(data.keys())
             for post_id in no_exif:
                 txn.put(post_id.encode(), b'')
-                urls[post_id] = ''
 
-    to_download = {post_id: url for post_id, url in urls.items() 
+    to_download = {post_id: url for post_id, url in post_ids_dict.items() 
                    if url and f"{post_id}.png" not in images_cache}
 
     print(f"Total images to download: {len(to_download)}")
 
-    if to_download:
+    max_workers = 30
+    to_download_items = list(to_download.items())
+
+    if to_download_items:
         print("Starting download...")
-        with ThreadPoolExecutor(max_workers=30) as executor:
-            list(tqdm(executor.map(download_image, to_download.items()),
-                      total=len(to_download),
-                      desc="Downloading"))
+        with ThreadPoolExecutor(max_workers=max_workers) as executor:
+            for result in tqdm(executor.map(download_image, to_download_items), total=len(to_download_items), desc="Downloading"):
+                pass
 
-        images_cache = set(os.listdir("images/Stash"))
+    images_cache = os.listdir("images/Stash")
 
     print("Linking images to the group directory...")
-    for i, post_id in enumerate(urls.keys()):
+    for i, post_id in enumerate(post_ids_dict.keys()):
         if f"{post_id}.png" in images_cache:
             stash_path = f"images/Stash/{post_id}.png"
             dest_path = f"images/{group_name}/{i}_{post_id}.png"