p / Client /download_all.py
q6's picture
F
728c581
import os
import requests
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"
img_base = 'https://i.pximg.net/img-original/img/'
os.chdir(os.path.dirname(os.path.abspath(__file__)))
valid = [f for f in os.listdir() if f.endswith(".txt")]
for idx, file in enumerate(valid):
print(f"{idx + 1}: {file}")
inputs = input("Enter the index of the file: ").split()
indexes = []
for inp in inputs:
if "-" in inp:
start, end = map(int, inp.split("-"))
indexes.extend(range(start - 1, end))
elif inp.isdigit():
indexes.append(int(inp) - 1)
def download_image(args):
idx, url, group_name = args
r = requests.get(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/{group_name}/{idx}.png", "wb") as f:
f.write(r.content)
for i in indexes:
group_name = valid[i].rsplit(".", 1)[0]
os.makedirs(f"images/{group_name}", exist_ok=True)
with open(valid[i], "r") as f:
post_ids = [line.strip() for line in f if line.strip()]
if not post_ids:
continue
params = [("post_ids", pid) for pid in post_ids] + [("only_first", "1")]
data = requests.get(f"{endpoint}/allimages", params=params).json()
image_urls = data.get("image_urls", [])
if not image_urls:
continue
to_download = [(idx, url, group_name) for idx, url in enumerate(image_urls)]
with ThreadPoolExecutor(max_workers=20) as executor:
list(tqdm(executor.map(download_image, to_download), total=len(to_download), desc=f"Downloading {group_name}"))
print(f"Saved {len(image_urls)} images to images/{group_name}")