Skip to content

feat(flux): diffusion data pipeline (energon/synthetic providers, encoders)#812

Open
jasainio wants to merge 7 commits into
mainfrom
feat/flux/data
Open

feat(flux): diffusion data pipeline (energon/synthetic providers, encoders)#812
jasainio wants to merge 7 commits into
mainfrom
feat/flux/data

Conversation

@jasainio

Copy link
Copy Markdown
Contributor

Part of an 18-PR series splitting the Flux diffusion-training feature (training Flux, a DiT text-to-image diffusion model, on Primus/Megatron) out of one large branch for reviewability. Targets feat/flux/flux — review after it.

What this changes

The diffusion data layer: dataloader + CUDA prefetch, energon and synthetic dataset providers, the image (VAE) and text (CLIP-L, T5-XXL) encoders, and the image task encoder. Also moves the energon/webdataset pins into requirements.txt here, since this layer's source and tests need them at runtime.

Why it's stacked here

One synthetic-dataset module top-level-imports flux.utils, so it bases on feat/flux/flux (this also transitively re-parents the prep/trainers/mlperf layers onto the Flux branch).

Dependencies

Sequenced after the CI-pins PR (feat/flux/ci-env); builds on feat/flux/flux.

Test plan

pytest tests/unit_tests/backends/megatron/diffusion/data (needs megatron-energon/webdataset from the moved requirements.txt). Validated locally on an AMD GPU container: 50 passed.

Files

32 (dataloader/prefetch, dataset providers, VAE/CLIP/T5 encoders, task encoder, requirements.txt + tests).

Flux Split Trial and others added 2 commits July 10, 2026 02:08
@luiza-amd luiza-amd changed the base branch from feat/flux/flux to main July 10, 2026 07:12
@luiza-amd luiza-amd marked this pull request as ready for review July 10, 2026 07:22
Comment thread primus/backends/megatron/data/diffusion/encoders/__init__.py Dismissed
Comment thread primus/backends/megatron/data/energon_dataset_provider.py Fixed
@luiza-amd luiza-amd marked this pull request as draft July 10, 2026 08:13
luiza-amd and others added 3 commits July 10, 2026 11:37
Unused import logging

Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
Co-authored-by: Copilot Autofix powered by AI <223894421+github-code-quality[bot]@users.noreply.github.com>
@luiza-amd luiza-amd marked this pull request as ready for review July 10, 2026 08:46
@luiza-amd luiza-amd marked this pull request as draft July 10, 2026 12:22
@luiza-amd luiza-amd marked this pull request as ready for review July 10, 2026 13:00
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants