license:openrail region:us
---
license: openrail
---

This is a dataset of 4.4M images generated with Stable Diffusion 2 for Kaggle's stable diffusion image to prompt competition.

Prompts were extracted from public databases:

- mp: [Magic Prompt - 1M](https://www.kaggle.com/datasets/leonidkulyk/magic-prompt-1m)
- db: [DiffusionDB](https://poloclub.github.io/diffusiondb/)
- op: [Open Prompts](https://github.com/krea-ai/open-prompts/tree/main)
- co: [COCO](https://cocodataset.org/)
- cc: [Conceptual Captions](https://ai.google.com/research/ConceptualCaptions/)
- l0: [LAION-2B-en-aesthetic](https://huggingface.co/datasets/laion/laion2B-en-aesthetic)

The following prompts were filtered out:

- those with token length >77 CLIP tokens
- those whose [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) embedding have a cosine similarity >0.9 to any other prompt

Samples were clustered by their embedding using KMeans with `n_clusters=100`.

Data is avaiable in tar files in [webdataset](https://github.com/webdataset/webdataset) format. Example to read the dataset and autodecode all the keys:

```python
>>> ds = wds.DataPipeline(
...     wds.SimpleShardList([str(p) for p in PATH_TO_TARS.glob('*/*.tar')]),
...     wds.tarfile_to_samples(),
...     wds.decode('pil'),
...     wds.batched(2, collation_fn=None, partial=False),
... )
>>> batch = next(iter(ds))
>>> batch[0].keys()
dict_keys(['__key__', '__url__', 'embs.npy', 'json', 'txt', 'webp'])
>>> batch[0]['txt']
'a very typical bus station'
>>> batch[0]['embs.npy'].shape
(384,)
>>> batch[0]['webp'].size
(512, 512)
>>> batch[0]['json']
{'ntokens': 7, 'cluster': 96}
```

The keys available are:

- embs.npy: the prompt's [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) embedding as np.float16 array.
- json: extra metadata:
  - ntokens: the prompt's length in CLIP tokens
  - cluster: cluster id, useful for grouped train/validation splitting
- txt: prompt used to generate the image
- webp: the image in webp format

Additionally, metadata is available in a single metadata.parquet file:

- original_dataset: source dataset (one of: 'mp', 'db', 'op', etc.)
- original_index: image index
- prompt: prompt used to generate the image
- embs: the prompt's [all-MiniLM-L6-v2](https://huggingface.co/sentence-transformers/all-MiniLM-L6-v2) embedding 
- ntokens: the prompt's length in CLIP tokens
- cluster: cluster id, useful for grouped train/validation splitting

If you prefer to decompress the tar files for random access or further filtering, you can map metadata to image path with the following code:

```python
    img_path = Path(original_dataset / f'{original_index:08d}.webp')
```

Images were generated at 768px resolution with the following code:

```python
    pipe = StableDiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2", torch_dtype=torch.float16, revision="fp16")
    pipe.to(torch.device("cuda"))
    prompts = [...]
    images = pipe(prompts, guidance_scale=9.0).images
```

Then, resized to 512px with Lanczos interpolation and saved as lossless webp:

```python
    for img, idx in (images, indices):
        img = img.resize((512, 512), resample=Image.Resampling.LANCZOS)
        img.save("{idx:08}.webp", lossless=True)
```

Use entirely at your own risk. Prompts usage may be subjected to their original dataset terms and conditions.
 
