poloclub/diffusiondb

1# Copyright 2022  Jay Wang, Evan Montoya, David Munechika, Alex Yang, Ben Hoover, Polo Chau2# MIT License3"""Loading script for DiffusionDB."""4 5import re6import numpy as np7import pandas as pd8 9from json import load, dump10from os.path import join, basename11from huggingface_hub import hf_hub_url12 13import datasets14 15# Find for instance the citation on arxiv or on the dataset repo/website16_CITATION = """\17@article{wangDiffusionDBLargescalePrompt2022,18  title = {{{DiffusionDB}}: {{A}} Large-Scale Prompt Gallery Dataset for Text-to-Image Generative Models},19  author = {Wang, Zijie J. and Montoya, Evan and Munechika, David and Yang, Haoyang and Hoover, Benjamin and Chau, Duen Horng},20  year = {2022},21  journal = {arXiv:2210.14896 [cs]},22  url = {https://arxiv.org/abs/2210.14896}23}24"""25 26# You can copy an official description27_DESCRIPTION = """28DiffusionDB is the first large-scale text-to-image prompt dataset. It contains 229million images generated by Stable Diffusion using prompts and hyperparameters30specified by real users. The unprecedented scale and diversity of this31human-actuated dataset provide exciting research opportunities in understanding32the interplay between prompts and generative models, detecting deepfakes, and33designing human-AI interaction tools to help users more easily use these models.34"""35 36_HOMEPAGE = "https://poloclub.github.io/diffusiondb"37_LICENSE = "CC0 1.0"38_VERSION = datasets.Version("0.9.1")39 40# Programmatically generate the URLs for different parts41# hf_hub_url() provides a more flexible way to resolve the file URLs42# https://huggingface.co/datasets/poloclub/diffusiondb/resolve/main/images/part-000001.zip43_URLS = {}44_URLS_LARGE = {}45_PART_IDS = range(1, 2001)46_PART_IDS_LARGE = range(1, 14001)47 48for i in _PART_IDS:49    _URLS[i] = hf_hub_url(50        "poloclub/diffusiondb",51        filename=f"images/part-{i:06}.zip",52        repo_type="dataset",53    )54 55for i in _PART_IDS_LARGE:56    if i < 10001:57        _URLS_LARGE[i] = hf_hub_url(58            "poloclub/diffusiondb",59            filename=f"diffusiondb-large-part-1/part-{i:06}.zip",60            repo_type="dataset",61        )62    else:63        _URLS_LARGE[i] = hf_hub_url(64            "poloclub/diffusiondb",65            filename=f"diffusiondb-large-part-2/part-{i:06}.zip",66            repo_type="dataset",67        )68 69# Add the metadata parquet URL as well70_URLS["metadata"] = hf_hub_url(71    "poloclub/diffusiondb", filename="metadata.parquet", repo_type="dataset"72)73 74_URLS_LARGE["metadata"] = hf_hub_url(75    "poloclub/diffusiondb",76    filename="metadata-large.parquet",77    repo_type="dataset",78)79 80_SAMPLER_DICT = {81    1: "ddim",82    2: "plms",83    3: "k_euler",84    4: "k_euler_ancestral",85    5: "ddik_heunm",86    6: "k_dpm_2",87    7: "k_dpm_2_ancestral",88    8: "k_lms",89    9: "others",90}91 92 93class DiffusionDBConfig(datasets.BuilderConfig):94    """BuilderConfig for DiffusionDB."""95 96    def __init__(self, part_ids, is_large, **kwargs):97        """BuilderConfig for DiffusionDB.98        Args:99          part_ids([int]): A list of part_ids.100          is_large(bool): If downloading data from DiffusionDB Large (14 million)101          **kwargs: keyword arguments forwarded to super.102        """103        super(DiffusionDBConfig, self).__init__(version=_VERSION, **kwargs)104        self.part_ids = part_ids105        self.is_large = is_large106 107 108class DiffusionDB(datasets.GeneratorBasedBuilder):109    """A large-scale text-to-image prompt gallery dataset based on Stable Diffusion."""110 111    BUILDER_CONFIGS = []112 113    # Programmatically generate configuration options (HF requires to use a string114    # as the config key)115    for num_k in [1, 5, 10, 50, 100, 500, 1000]:116        for sampling in ["first", "random"]:117            for is_large in [False, True]:118                num_k_str = f"{num_k}k" if num_k < 1000 else f"{num_k // 1000}m"119                subset_str = "large_" if is_large else "2m_"120 121                if sampling == "random":122                    # Name the config123                    cur_name = subset_str + "random_" + num_k_str124 125                    # Add a short description for each config126                    cur_description = (127                        f"Random {num_k_str} images with their prompts and parameters"128                    )129 130                    # Sample part_ids131                    total_part_ids = _PART_IDS_LARGE if is_large else _PART_IDS132                    part_ids = np.random.choice(133                        total_part_ids, num_k, replace=False134                    ).tolist()135                else:136                    # Name the config137                    cur_name = subset_str + "first_" + num_k_str138 139                    # Add a short description for each config140                    cur_description = f"The first {num_k_str} images in this dataset with their prompts and parameters"141 142                    # Sample part_ids143                    total_part_ids = _PART_IDS_LARGE if is_large else _PART_IDS144                    part_ids = total_part_ids[1 : num_k + 1]145 146                # Create configs147                BUILDER_CONFIGS.append(148                    DiffusionDBConfig(149                        name=cur_name,150                        part_ids=part_ids,151                        is_large=is_large,152                        description=cur_description,153                    ),154                )155 156    # Add few more options for Large only157    for num_k in [5000, 10000]:158        for sampling in ["first", "random"]:159            num_k_str = f"{num_k // 1000}m"160            subset_str = "large_"161 162            if sampling == "random":163                # Name the config164                cur_name = subset_str + "random_" + num_k_str165 166                # Add a short description for each config167                cur_description = (168                    f"Random {num_k_str} images with their prompts and parameters"169                )170 171                # Sample part_ids172                total_part_ids = _PART_IDS_LARGE173                part_ids = np.random.choice(174                    total_part_ids, num_k, replace=False175                ).tolist()176            else:177                # Name the config178                cur_name = subset_str + "first_" + num_k_str179 180                # Add a short description for each config181                cur_description = f"The first {num_k_str} images in this dataset with their prompts and parameters"182 183                # Sample part_ids184                total_part_ids = _PART_IDS_LARGE185                part_ids = total_part_ids[1 : num_k + 1]186 187            # Create configs188            BUILDER_CONFIGS.append(189                DiffusionDBConfig(190                    name=cur_name,191                    part_ids=part_ids,192                    is_large=True,193                    description=cur_description,194                ),195            )196 197    # Need to manually add all (2m) and all (large)198    BUILDER_CONFIGS.append(199        DiffusionDBConfig(200            name="2m_all",201            part_ids=_PART_IDS,202            is_large=False,203            description="All images with their prompts and parameters",204        ),205    )206 207    BUILDER_CONFIGS.append(208        DiffusionDBConfig(209            name="large_all",210            part_ids=_PART_IDS_LARGE,211            is_large=True,212            description="All images with their prompts and parameters",213        ),214    )215 216    # We also prove a text-only option, which loads the meatadata parquet file217    BUILDER_CONFIGS.append(218        DiffusionDBConfig(219            name="2m_text_only",220            part_ids=[],221            is_large=False,222            description="Only include all prompts and parameters (no image)",223        ),224    )225 226    BUILDER_CONFIGS.append(227        DiffusionDBConfig(228            name="large_text_only",229            part_ids=[],230            is_large=True,231            description="Only include all prompts and parameters (no image)",232        ),233    )234 235    # Add a random 1k from 2M as the first entry point to show on HF data viewer236    # Sample part_ids237    part_ids = np.random.choice(_PART_IDS, 1000, replace=False).tolist()238    BUILDER_CONFIGS.append(239        DiffusionDBConfig(240            name="1k_random_2m",241            part_ids=part_ids,242            is_large=False,243            description="Another random 1k images with meta data from DiffusionDB 2M",244        ),245    )246 247    # Default to only load 1k random images248    DEFAULT_CONFIG_NAME = "2m_random_1k"249 250    def _info(self):251        """Specify the information of DiffusionDB."""252 253        if "text_only" in self.config.name:254            features = datasets.Features(255                {256                    "image_name": datasets.Value("string"),257                    "prompt": datasets.Value("string"),258                    "part_id": datasets.Value("uint16"),259                    "seed": datasets.Value("uint32"),260                    "step": datasets.Value("uint16"),261                    "cfg": datasets.Value("float32"),262                    "sampler": datasets.Value("string"),263                    "width": datasets.Value("uint16"),264                    "height": datasets.Value("uint16"),265                    "user_name": datasets.Value("string"),266                    "timestamp": datasets.Value("timestamp[us, tz=UTC]"),267                    "image_nsfw": datasets.Value("float32"),268                    "prompt_nsfw": datasets.Value("float32"),269                },270            )271 272        else:273            features = datasets.Features(274                {275                    "image": datasets.Image(),276                    "prompt": datasets.Value("string"),277                    "seed": datasets.Value("uint32"),278                    "step": datasets.Value("uint16"),279                    "cfg": datasets.Value("float32"),280                    "sampler": datasets.Value("string"),281                    "width": datasets.Value("uint16"),282                    "height": datasets.Value("uint16"),283                    "user_name": datasets.Value("string"),284                    "timestamp": datasets.Value("timestamp[us, tz=UTC]"),285                    "image_nsfw": datasets.Value("float32"),286                    "prompt_nsfw": datasets.Value("float32"),287                },288            )289 290        return datasets.DatasetInfo(291            description=_DESCRIPTION,292            features=features,293            supervised_keys=None,294            homepage=_HOMEPAGE,295            license=_LICENSE,296            citation=_CITATION,297        )298 299    def _split_generators(self, dl_manager):300        # If several configurations are possible (listed in BUILDER_CONFIGS),301        # the configuration selected by the user is in self.config.name302 303        # dl_manager is a datasets.download.DownloadManager that can be used to304        # download and extract URLS It can accept any type or nested list/dict305        # and will give back the same structure with the url replaced with path306        # to local files. By default the archives will be extracted and a path307        # to a cached folder where they are extracted is returned instead of the308        # archive309 310        # Download and extract zip files of all sampled part_ids311        data_dirs = []312        json_paths = []313 314        # Resolve the urls315        if self.config.is_large:316            urls = _URLS_LARGE317        else:318            urls = _URLS319 320        for cur_part_id in self.config.part_ids:321            cur_url = urls[cur_part_id]322            data_dir = dl_manager.download_and_extract(cur_url)323 324            data_dirs.append(data_dir)325            json_paths.append(join(data_dir, f"part-{cur_part_id:06}.json"))326 327        # Also download the metadata table328        metadata_path = dl_manager.download(urls["metadata"])329 330        return [331            datasets.SplitGenerator(332                name=datasets.Split.TRAIN,333                # These kwargs will be passed to _generate_examples334                gen_kwargs={335                    "data_dirs": data_dirs,336                    "json_paths": json_paths,337                    "metadata_path": metadata_path,338                },339            ),340        ]341 342    def _generate_examples(self, data_dirs, json_paths, metadata_path):343        # This method handles input defined in _split_generators to yield344        # (key, example) tuples from the dataset.345        # The `key` is for legacy reasons (tfds) and is not important in itself,346        # but must be unique for each example.347 348        # Load the metadata parquet file if the config is text_only349        if "text_only" in self.config.name:350            metadata_df = pd.read_parquet(metadata_path)351            for _, row in metadata_df.iterrows():352                yield row["image_name"], {353                    "image_name": row["image_name"],354                    "prompt": row["prompt"],355                    "part_id": row["part_id"],356                    "seed": row["seed"],357                    "step": row["step"],358                    "cfg": row["cfg"],359                    "sampler": _SAMPLER_DICT[int(row["sampler"])],360                    "width": row["width"],361                    "height": row["height"],362                    "user_name": row["user_name"],363                    "timestamp": None364                    if pd.isnull(row["timestamp"])365                    else row["timestamp"],366                    "image_nsfw": row["image_nsfw"],367                    "prompt_nsfw": row["prompt_nsfw"],368                }369 370        else:371            num_data_dirs = len(data_dirs)372            assert num_data_dirs == len(json_paths)373 374            # Read the metadata table (only rows with the needed part_ids)375            part_ids = []376            for path in json_paths:377                cur_id = int(re.sub(r"part-(\d+)\.json", r"\1", basename(path)))378                part_ids.append(cur_id)379 380            # We have to use pandas here to make the dataset preview work (it381            # uses streaming mode)382            metadata_table = pd.read_parquet(383                metadata_path,384                filters=[("part_id", "in", part_ids)],385            )386 387            # Iterate through all extracted zip folders for images388            for k in range(num_data_dirs):389                cur_data_dir = data_dirs[k]390                cur_json_path = json_paths[k]391 392                json_data = load(open(cur_json_path, "r", encoding="utf8"))393 394                for img_name in json_data:395                    img_params = json_data[img_name]396                    img_path = join(cur_data_dir, img_name)397 398                    # Query the metadata399                    query_result = metadata_table.query(f'`image_name` == "{img_name}"')400 401                    # Yields examples as (key, example) tuples402                    yield img_name, {403                        "image": {404                            "path": img_path,405                            "bytes": open(img_path, "rb").read(),406                        },407                        "prompt": img_params["p"],408                        "seed": int(img_params["se"]),409                        "step": int(img_params["st"]),410                        "cfg": float(img_params["c"]),411                        "sampler": img_params["sa"],412                        "width": query_result["width"].to_list()[0],413                        "height": query_result["height"].to_list()[0],414                        "user_name": query_result["user_name"].to_list()[0],415                        "timestamp": None416                        if pd.isnull(query_result["timestamp"].to_list()[0])417                        else query_result["timestamp"].to_list()[0],418                        "image_nsfw": query_result["image_nsfw"].to_list()[0],419                        "prompt_nsfw": query_result["prompt_nsfw"].to_list()[0],420                    }421