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-encoded (recommended for training)

Pre-encodes images with VAE and text with T5/CLIP. Produces larger datasets but enables faster training since encoders are not needed at training time.

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

Output: WebDataset shards containing latents.pth, prompt_embeds.pth, pooled_prompt_embeds.pth, caption.txt

Use when: Training on production datasets, maximum training speed is needed, storage is available.

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 a config file (recommended)

The --config flag is supported by diffusion-encoded only. The simplest approach uses a YAML config file:

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

Available example configs in primus/configs/data/megatron/diffusion/preprocessing/:

Config	Source	Description
quickstart_pokemon.yaml	HuggingFace	Minimal config, 256px, fast
example_huggingface.yaml	HuggingFace	Full example with all options
example_directory.yaml	Local directory	Images + captions from disk
example_webdataset.yaml	WebDataset	Existing tar archives
example_base.yaml	N/A	Comprehensive reference with all fields
text_to_image_2m_10k.yaml	HuggingFace	10K subset of text-to-image-2M (1024px)

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:

1. Token file (recommended)

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

The token file must have secure permissions (600 or 400). Create it with:

echo "hf_your_token_here" > /path/to/.hf_token
chmod 600 /path/to/.hf_token

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