Data preprocessing guide

Data preprocessing guide#

This guide explains how to prepare datasets for Flux and other diffusion models in Primus, including pre-encoding of VAE latents and text embeddings into Energon WebDataset format.

Table of contents#

  1. Overview

  2. Quick start

  3. Two pipelines

  4. Running preprocessing

  5. Configuration

  6. Authentication

  7. Finalization

  8. Output format

  9. Validation

  10. Troubleshooting

Overview#

Diffusion models require three types of encodings:

  1. VAE latents: Images encoded to latent space

  2. Text embeddings: Captions encoded with T5-XXL (sequence)

  3. Pooled embeddings: Captions encoded with CLIP-L (pooled)

Why pre-encode?#

Benefits:

  • 5-10x faster training (no online encoding)

  • Lower GPU memory usage (encoders not loaded during training)

  • Deterministic inputs (same preprocessing for all runs)

  • Eliminates encoder differences as a variable

When to Use:

  • Training (highly recommended)

  • Fine-tuning on fixed datasets

  • Benchmarking and reproducibility

When NOT to Use:

  • Interactive data augmentation needed

  • Dataset too large to store pre-encoded

  • Rapid prototyping with changing data

Quick start#

Preprocess the Pokemon dataset with a single command:

primus-cli direct -- data diffusion-encoded \
  --config primus/configs/data/megatron/diffusion/preprocessing/quickstart_pokemon.yaml \
  --hf-token-file /path/to/.hf_token

This will:

  1. Download the diffusers/pokemon-gpt4-captions dataset from HuggingFace

  2. Encode all images with VAE and captions with T5/CLIP

  3. Write Energon WebDataset tar shards to /workspace/Primus/data/quickstart_pokemon

  4. Automatically finalize the dataset (create dataset.yaml, run energon prepare, validate)

The default encoder model (black-forest-labs/FLUX.1-dev) is gated and requires a HuggingFace token. Get one at https://huggingface.co/settings/tokens and save it to a file.

Two pipelines#

Primus provides two preprocessing pipelines via the primus data CLI:

diffusion-raw#

Stores raw images and captions without encoding. Smaller datasets but encoding happens on-the-fly during training (requires GPU + encoders loaded in memory).

primus-cli direct -- data diffusion-raw \
  --source-type huggingface \
  --hf-dataset diffusers/pokemon-gpt4-captions \
  --output-dir /workspace/Primus/data/raw_pokemon \
  --hf-token-file /path/to/.hf_token

Output: WebDataset shards containing jpg (or png/webp) and txt files.

Use when: Storage is limited, experimenting with different encoders, rapid prototyping.

Note: diffusion-raw does not support --config files. All options must be passed as CLI arguments.

Running preprocessing#

Using CLI arguments directly#

All config values can be provided as CLI arguments:

primus-cli direct -- data diffusion-encoded \
  --source-type huggingface \
  --hf-dataset diffusers/pokemon-gpt4-captions \
  --output-dir /workspace/Primus/data/encoded_pokemon \
  --model-path black-forest-labs/FLUX.1-dev \
  --batch-size 8 \
  --precision bf16 \
  --hf-token-file /path/to/.hf_token

CLI overrides config values#

When using both --config and CLI arguments, CLI arguments take priority:

primus-cli direct -- data diffusion-encoded \
  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
  --hf-token-file /path/to/.hf_token \
  --output-dir /my/custom/path \
  --batch-size 16 \
  --max-samples 1000

Priority order (highest to lowest):

  1. Explicitly provided CLI arguments

  2. YAML config file values

  3. CLI default values

Multi-GPU processing#

Use --nproc-per-node for data-parallel preprocessing across multiple GPUs:

primus-cli direct -- --nproc-per-node=8 data diffusion-encoded \
  --config primus/configs/data/megatron/diffusion/preprocessing/example_huggingface.yaml \
  --hf-token-file /path/to/.hf_token

Each GPU processes a subset of the data. Shards are named to avoid conflicts across ranks.

Configuration#

YAML config structure#

Preprocessing configs have four sections:

source:
  type: huggingface          # huggingface | directory | webdataset
  hf_dataset: diffusers/pokemon-gpt4-captions
  hf_split: train

output:
  output_dir: /workspace/Primus/data/encoded_pokemon
  shard_size: 1000           # samples per tar shard
  max_samples: null          # null = process all
  compress: false

model:
  model_path: black-forest-labs/FLUX.1-dev   # HF repo or local path
  precision: bf16            # bf16 | fp16 | fp32
  batch_size: 8
  # Optional per-encoder overrides:
  vae_path: null
  t5_path: null
  clip_path: null

image:
  image_size: 1024
  center_crop: false

Source types#

HuggingFace (type: huggingface):

source:
  type: huggingface
  hf_dataset: diffusers/pokemon-gpt4-captions
  hf_split: train
  hf_data_files: null        # optional: specific files within dataset

Local Directory (type: directory):

source:
  type: directory
  input_dir: /data/my_images

Expected directory structure for directory source:

my_images/
├── images/
│   ├── 00001.jpg
│   ├── 00002.png
│   └── ...
└── captions/
    ├── 00001.txt
    ├── 00002.txt
    └── ...

WebDataset (type: webdataset):

source:
  type: webdataset
  input_path: /data/existing_shards/*.tar

Model configuration#

The model_path defaults to black-forest-labs/FLUX.1-dev, which downloads VAE, T5-XXL, and CLIP-L encoders from HuggingFace. This model is gated and requires authentication (see Authentication).

Individual encoder paths can be overridden:

model:
  model_path: black-forest-labs/FLUX.1-dev
  vae_path: /local/models/vae        # use local VAE instead
  t5_path: null                       # falls back to model_path
  clip_path: null                     # falls back to model_path

Authentication#

The default encoder model (FLUX.1-dev) is gated on HuggingFace and requires authentication. Primus supports three authentication methods, checked in priority order:

2. Environment variable#

export HF_TOKEN=hf_your_token_here
primus-cli direct -- data diffusion-encoded --config your_config.yaml

3. HuggingFace CLI login#

huggingface-cli login
primus-cli direct -- data diffusion-encoded --config your_config.yaml

If authentication fails, Primus provides a clear error message indicating which encoder failed and how to fix it.

Finalization#

Finalization is automatic by default. After preprocessing completes, Primus automatically:

  1. Creates .nv-meta/dataset.yaml with CrudeWebdataset sample type and encoding subflavor

  2. Runs energon prepare to index the tar shards and create split assignments

  3. Validates the dataset using Primus’s custom validation (metadata checks, sample count verification, energon API spot-check)

Skipping finalization#

To skip automatic finalization (e.g., for manual post-processing):

primus-cli direct -- data diffusion-encoded \
  --config your_config.yaml \
  --hf-token-file /path/to/.hf_token \
  --no-finalize

Custom train/val/test splits#

By default, 100% of data goes to the training split. To create validation and test splits:

primus-cli direct -- data diffusion-encoded \
  --config your_config.yaml \
  --hf-token-file /path/to/.hf_token \
  --train-split 0.8

This creates an 80% train / 10% val / 10% test split.

Output format#

Directory structure#

After preprocessing and finalization, the output directory contains:

encoded_pokemon/
├── 000000.tar              # WebDataset shard
├── 000001.tar
├── 000000.tar.idx          # Energon index files
├── 000001.tar.idx
└── .nv-meta/               # Energon metadata
    ├── dataset.yaml         # Dataset type configuration
    ├── split.yaml           # Train/val/test split assignments
    └── .info.json           # Shard counts and sample counts

Pre-encoded shard contents (diffusion-encoded)#

Each tar shard contains samples with these keys:

000000.tar:
├── 0000000000.latents.pth              # VAE latents tensor
├── 0000000000.prompt_embeds.pth        # T5-XXL embeddings tensor
├── 0000000000.pooled_prompt_embeds.pth # CLIP-L pooled embeddings tensor
├── 0000000000.caption.txt              # Original caption text
├── 0000000001.latents.pth
├── 0000000001.prompt_embeds.pth
├── ...

Tensor shapes:

  • latents.pth: [64, H/8, W/8] (e.g., [64, 128, 128] for 1024x1024 images)

  • prompt_embeds.pth: [seq_len, 4096] (T5-XXL hidden dim)

  • pooled_prompt_embeds.pth: [768] (CLIP-L pooled dim)

Raw shard contents (diffusion-raw)#

000000.tar:
├── 0000000000.jpg       # Preprocessed image
├── 0000000000.txt       # Caption text
├── 0000000001.jpg
├── 0000000001.txt
├── ...

Dataset YAML#

The auto-generated .nv-meta/dataset.yaml uses CrudeWebdataset format:

__module__: megatron.energon
__class__: CrudeWebdataset
subflavors:
  encoding: preencoded    # or 'raw' for diffusion-raw

Validation#

Automatic validation#

Validation runs automatically as part of finalization. It performs four checks:

  1. Metadata check: Verifies .nv-meta/ files exist and are valid (.info.json, split.yaml, dataset.yaml)

  2. Sample count check: Spot-checks that tar shard entry counts match .info.json

  3. Sample load check: Loads one sample through Energon’s Python API (same code path as training)

  4. Summary report: Prints dataset statistics (encoding, total samples, splits, data shapes, size)

Standalone validation#

To validate a dataset independently:

python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset

# For raw datasets:
python -m primus.backends.megatron.data.diffusion.preprocessing.validate /path/to/dataset --encoding raw

Programmatic validation#

from primus.backends.megatron.data.diffusion.preprocessing.validate import validate_energon_dataset

ok = validate_energon_dataset('/path/to/dataset', encoding='preencoded')

Known Energon CLI limitations#

The standard Energon CLI tools (energon info, energon preview, energon lint) do not work correctly with CrudeWebdataset format. Primus uses custom validation instead:

  • energon info raises KeyError: 'sample_type'

  • energon preview raises TypeError (expects dataclass, CrudeSample is a dict)

  • energon lint raises AssertionError (expects registered cookers)

Use Primus’s built-in validation or the standalone script above.

Troubleshooting#

HuggingFace authentication failure#

Symptoms: Error mentioning “token”, “gated”, “401”, or “403” when downloading encoders.

Solutions:

  1. Provide a token: --hf-token-file /path/to/.hf_token

  2. Set environment variable: export HF_TOKEN=hf_xxx

  3. Run huggingface-cli login

  4. Accept the model’s license on https://huggingface.co/black-forest-labs/FLUX.1-dev

Out of memory during preprocessing#

Symptoms: CUDA out of memory error during encoding.

Solutions:

  1. Reduce batch size: --batch-size 4 or --batch-size 1

  2. Use smaller image size: --image-size 512

  3. Use fp16 precision: --precision fp16

  4. Use multi-GPU to distribute work: --nproc-per-node=8

Missing dependencies#

Symptoms: ModuleNotFoundError for webdataset, megatron-energon, tqdm, etc.

Solution:

pip install -r requirements.txt

Or set PRIMUS_AUTO_INSTALL=1 in the container to auto-install missing packages.

Finalization fails#

Symptoms: Error during energon prepare or validation after preprocessing.

Solutions:

  1. Check that tar shards exist in the output directory

  2. Ensure megatron-energon is installed (pip install megatron-energon)

  3. Re-run with --no-finalize, then manually inspect the output before finalizing

Slow preprocessing#

Symptoms: Low throughput (< 10 samples/sec).

Solutions:

  1. Increase batch size (limited by VRAM): --batch-size 16

  2. Use multiple GPUs: --nproc-per-node=8

  3. Use bf16 precision (faster on supported hardware): --precision bf16

  4. For raw pipeline, reduce image quality: --image-quality 85

Best practices#

  1. Always pre-encode for production training: 5-10x speedup is worth the storage

  2. Test on small dataset first: Use --max-samples 100 to verify the pipeline works

  3. Use bf16 precision: Good balance of speed, storage, and quality

  4. Start with quickstart_pokemon.yaml: Verify your setup before processing large datasets

  5. Keep raw data: Pre-encoding is a one-way transformation

  6. Version your configs: Track which preprocessing config produced each dataset

Flux-specific data preparation (torchrun)#

This section covers preparing datasets for Flux training using torchrun directly inside a Docker container, as an alternative to the primus-cli workflow above.

Prerequisites#

Start the container and optionally install requirements:

bash tools/docker/start_container.sh
docker exec dev_primus bash -c 'pip install -r /workspace/Primus/requirements.txt'

This mounts the repository to /workspace/Primus inside the container. Override the image with DOCKER_IMAGE:

DOCKER_IMAGE=docker.io/rocm/primus:v26.1 bash tools/docker/start_container.sh

Input directory structure#

When using --source-type directory, organize your data as follows:

dataset/
├── images/
│   ├── 0000000.png
│   ├── 0000001.png
│   └── ...
└── captions/
    ├── 0000000.txt
    ├── 0000001.txt
    └── ...

The file stem links each image to its caption (e.g., images/0000000.png pairs with captions/0000000.txt). Images can be .jpg, .jpeg, .png, or .webp. Captions must be UTF-8 .txt files. Samples without a matching caption are skipped.

Other supported source types:

  • huggingface – Load directly from HuggingFace Hub. Requires --hf-dataset.

  • webdataset – Read from existing WebDataset tar archives. Requires --input-path.

Image sizing#

Fixed size (default): Every image is resized to --image-size pixels square (default 1024). Use --center-crop to control center-cropping.

Variable size (--variable-size): Preserves aspect ratio by scaling the longest side to --max-size (default 1024) and rounding dimensions to multiples of 16. Only use this when all images share the same dimensions – mixed tensor sizes cause load imbalance across GPUs.

Key parameters#

Parameter

Default

Description

--source-type

—

directory, huggingface, or webdataset

--input-dir

—

Dataset root (for directory source)

--image-size

1024

Square target size

--variable-size

off

Preserve aspect ratio

--max-size

1024

Maximum dimension in variable-size mode

--model-path

black-forest-labs/FLUX.1-dev

Base model for encoder weights

--t5-max-length

512

T5 token limit (use 256 for FLUX.1-schnell)

--batch-size

8

Encoding batch size per GPU

--output-dir

—

Destination for encoded WebDataset shards

--shard-size

1000

Samples per tar shard

Examples#

From host (via docker exec):

docker exec dev_primus bash -c '\
  export HF_HOME=/workspace/Primus/checkpoints/flux; \
PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \
  torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \
    data diffusion-encoded \
    --source-type directory \
    --input-dir /workspace/Primus/data/dataset \
    --output-dir /workspace/Primus/data/dataset_encoded_256 \
    --image-size 256  \
    --t5-max-length 256'

Inside the container:

export HF_HOME=/workspace/Primus/checkpoints/flux
PYTHONPATH=/workspace/Primus:/workspace/Primus/third_party/Megatron-LM:$PYTHONPATH \
  torchrun --nproc_per_node=8 /workspace/Primus/primus/cli/main.py \
    data diffusion-encoded \
    --source-type directory \
    --input-dir /workspace/Primus/data/dataset \
    --output-dir /workspace/Primus/data/dataset_encoded_256 \
    --image-size 256 \
    --t5-max-length 256

Both commands encode a local directory dataset at 256px on 8 GPUs, then create an Energon dataset by default.

Limitations#

Uniform output size only. When using --variable-size, images may produce tensors of different shapes. The current data pipeline does not support mixed tensor sizes because samples are distributed evenly across GPUs and unequal shapes cause load imbalance and eventual timeout. Use fixed --image-size if your images have various sizes.

Next steps#

After preprocessing:

  1. Training: Use the preprocessed dataset path in your training config

  2. Validation: The dataset is ready for training immediately after finalization

See:

  • Energon Integration for TaskEncoder and dataloader details

  • Config Directory Guide for config file reference

  • Example configs in primus/configs/data/megatron/diffusion/preprocessing/

Last Updated: June 2026

Contents