This dataset is the downloaded variant of Spawning/PD12M. More specifically, this dataset
is compatible with webdataset. It was made public after obtaining permission
from the original authors of the dataset.
You can use the following to explore the dataset with webdataset:
import webdataset as wds
dataset_path = "pipe:curl -s -f -L https://huggingface.co/datasets/sayakpaul/pd12m-full/resolve/main/{00155..02480}.tar"
dataset = (
wds.WebDataset(dataset_path, handler=wds.warn_and_continue)
.shuffle(690, handler=wds.warn_and_continue)
.decode("pil", handler=wds.warn_and_continue)
)
for sample in dataset:
print(sample.keys())
print(sample["jpg"].size)
print(sample["json"])
print(sample["txt"])
break
Additionally, this script provides a reference dataloader implementation.
The dataset was downloaded by using the img2dataset tool. The following command was used to perform the
download on a CPU cluster:
Code
img2dataset --url_list pd12m_full.parquet --input_format "parquet" \
--url_col "url" --caption_col "caption" --output_format webdataset \
--number_sample_per_shard=5000 --skip_reencode=True \
--output_folder s3://diffusion-datasets/pd12m \
--processes_count 16 --thread_count 64 \
--resize_mode no \
--enable_wandb True
The command above serializes the webdataset shards to an S3 bucket. Additionally, here is the wandb log of the run.
pd12m_full.parquet was obtained by collating all the parquet files from here
into a single pandas dataframe. It's available here.
To copy the files from the S3 bucket to this repository, the following script was used:
Code
from huggingface_hub import create_repo, upload_file, dataset_info
import ray
import os
# Change `_temp_dir` path accordingly.
ray.init(num_cpus=16, _temp_dir="/scratch")
def main():
s3_fs = s3fs.S3FileSystem()
bucket_path = "s3://diffusion-datasets/pd12m"
files = s3_fs.ls(bucket_path, detail=True)
files = sorted([f["name"] for f in files if f["name"].endswith(".tar") and f["size"] > 0.0])
@ray.remote
def fn(tar_file):
# Change the paths accordingly.
full_s3_tar_file = f"s3://{tar_file}"
local_path = f"/scratch/{tar_file}"
s3_fs.download(full_s3_tar_file, local_path)
# Adjust according to what your local storage allows for.
batch_size = 20
for i in range(0, len(files), batch_size):
batch = files[i : i + batch_size]
futures = [fn.remote(tar_file) for tar_file in batch]
ray.get(futures)
os.system(
"huggingface-cli upload-large-folder sayakpaul/pd12m-full --repo-type=dataset /scratch/diffusion-datasets/pd12m --num-workers=16"
)
os.system(f"rm -rf /scratch/diffusion-datasets/pd12m/*.tar")
print("All shards have been downloaded successfully.")
if __name__ == "__main__":
create_repo(repo_id="sayakpaul/pd12m-full", repo_type="dataset", private=True, exist_ok=True)
main()