| 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) | |