File size: 1,465 Bytes
69aa3f7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 | import os, time
import requests
from zipfile import ZipFile
output_dir = r"X:\AI Data\AI Training\Training Data\diffusiondb"
os.makedirs(output_dir, exist_ok=True)
total_parts = 2000
start_part = 1691
end_part = total_parts
base_url = "https://huggingface.co/datasets/poloclub/diffusiondb/resolve/main/images"
def is_zip_valid(file_path):
"""Check if a zip file can be opened successfully."""
try:
with ZipFile(file_path, 'r') as zf:
bad_file = zf.testzip()
return bad_file is None
except Exception:
return False
for i in range(start_part, end_part + 1):
part_name = f"part-{i:06d}.zip"
out_path = os.path.join(output_dir, part_name)
if os.path.exists(out_path) and is_zip_valid(out_path):
print(f"Skipping {part_name}, already downloaded and valid.")
continue
url = f"{base_url}/{part_name}"
print(f"Downloading {part_name} ...")
try:
r = requests.get(url, stream=True)
r.raise_for_status()
with open(out_path, "wb") as f:
for chunk in r.iter_content(chunk_size=1024*1024):
f.write(chunk)
print(f"Downloaded {part_name}")
if not is_zip_valid(out_path):
print(f"Warning: {part_name} appears corrupted, will re-download next time.")
except Exception as e:
print(f"Failed to download {part_name}: {e}")
time.sleep(10)
|