Skip to content

DATA.md

Latest commit

 

History

History
92 lines (70 loc) · 5.08 KB

DATA.md

File metadata and controls

92 lines (70 loc) · 5.08 KB
Raw

Dataset Preparation for Latent Diffusion Model Training

This document provides details on the datasets selected and the preparation methods used for training a Latent Diffusion Model (LDM). The dataset selection and preparation strategies outlined below are tailored for a limited storage and compute environment, making them suitable for local desktop experiments.

Datasets Overview

LAION Datasets

To effectively train and reproduce latent diffusion models such as Stable Diffusion (SD), the LAION family of datasets is extensively utilized:

  • LAION-400M: The initial latent diffusion model research utilized this extensive dataset (Paper Reference).
  • LAION-2B-en: The Stable Diffusion V1 model was initially trained for around 250k steps with images resized to 256x256 resolution, subsequently fine-tuned at 512x512.
  • LAION-5B: The Stable Diffusion 2 iteration scaled further using the larger LAION-5B dataset. However, given the substantial size, this dataset exceeded the constraints of local training resources.

Dataset Subsets and Rationale

Given the computational and storage constraints, the following curated subsets were selected for practicality and quality:

  • LAION Aesthetics: A highly curated 50M+ subset of LAION-5B chosen based on visual quality criteria, aiming to achieve optimal aesthetic results while maintaining manageability for local training scenarios.

  • opendiffusionai/laion2b-en-aesthetic-square-cleaned: Contains over 250k high-quality images sampled from LAION-2B-en Aesthetic, filtered for images with square to nearly-square aspect ratios (1:1 to 5:4) at a resolution of at least 1024x1024 pixels. These images were re-captioned using the Moondream2B open-source vision-language model, significantly enhancing caption quality and expressiveness compared to the original dataset.

  • LAION-POP: Features approximately 600k images, specially curated to reflect popular concepts from platforms like Midjourney. These images have superior descriptive captions generated by advanced vision-language models (CogVLM and LLaVa-v1.5-13b). According to insights from the KOALA paper, training models from scratch with detailed captions and high-resolution images significantly improves the final performance compared to using only standard LAION subsets.

Dataset Downloading Procedure

To efficiently download and prepare datasets, the img2dataset Python package was utilized. This tool supports real-time resizing and various storage formats, optimizing the dataset preparation workflow.

Example download commands:

# Downloading LAION aesthetic square-cleaned subset
NO_ALBUMENTATIONS_UPDATE=1 img2dataset \
    --url_list laion-square-cleaned.jsonl.gz \
    --output_folder data \
    --input_format "jsonl.gz" \
    --encode_format png \
    --encode_quality 9 \
    --url_col "url" \
    --caption_col "moondream" \
    --output_format files \
    --skip_reencode true \
    --processes_count 8 \
    --thread_count 8 \
    --enable_wandb False \
    --resize_mode no

# Downloading LAION-POP dataset
NO_ALBUMENTATIONS_UPDATE=1 img2dataset \
    --url_list laion_pop.parquet \
    --output_folder data \
    --processes_count 16 \
    --thread_count 16 \
    --resize_mode no \
    --encode_quality 9 \
    --encode_format png \
    --skip_reencode true \
    --output_format webdataset \
    --input_format "parquet" \
    --url_col "url" \
    --extract_exif True \
    --timeout 10 \
    --retries 5 \
    --enable_wandb False

WebDataset Format Conversion

Initial experiments highlighted inefficiencies and image loss during simultaneous downloading and resizing operations. To mitigate this issue, images were downloaded in their original resolution first. A subsequent resizing step was performed using a custom Python script, prepare_webdataset.py.

Why WebDataset?

WebDataset offers an efficient streaming-based dataset format compatible with PyTorch. WebDataset documentation

Key advantages include:

  • Sequential Data Access: Optimizes local storage utilization, significantly improving I/O throughput (typically 3x–10x faster than random access).
  • Scalability: Streamlined for both local and cloud storage environments, simplifying dataset management for large-scale deep learning tasks.
  • PyTorch Integration: Implements a PyTorch-compatible IterableDataset, seamlessly integrating into existing training workflows.

Conversion Procedure

The prepare_webdataset.py script performs:

  • Image resizing to 256 pixels (smallest dimension).
  • Packaging resized images into POSIX-compliant tar archives.
  • Maintaining structured file naming conventions for streamlined WebDataset usage.