bridge.diffusion.data.flux.flux_energon_datamodule#

Module Contents#

Classes#

FluxDatasetConfig

Unified FLUX dataset config: mock vs real is decided at runtime in build_datasets() based on whether path is set (same pattern as Gemma/LLM recipes with dataset.blend).

FluxDataModuleConfig

Data#

logger

API#

bridge.diffusion.data.flux.flux_energon_datamodule.logger#

‘getLogger(…)’

class bridge.diffusion.data.flux.flux_energon_datamodule.FluxDatasetConfig#

Bases: megatron.bridge.data.utils.DatasetProvider

Unified FLUX dataset config: mock vs real is decided at runtime in build_datasets() based on whether path is set (same pattern as Gemma/LLM recipes with dataset.blend).

Use this in the recipe with path=None by default. Override with dataset.path=/path/to/wds to load real data; no separate –data_paths or mock flag needed.

path: Optional[str]#

None

seq_length: int#

1024

packing_buffer_size: Optional[int]#

None

micro_batch_size: int#

1

global_batch_size: int#

4

num_workers: int#

16

dataloader_type: str#

‘external’

vae_scale_factor: int#

8

latent_channels: int#

16

image_H: int#

1024

image_W: int#

1024

prompt_seq_len: int#

512

context_dim: int#

4096

pooled_prompt_dim: int#

768

__post_init__()#
build_datasets(
context: megatron.bridge.data.utils.DatasetBuildContext,
)#
class bridge.diffusion.data.flux.flux_energon_datamodule.FluxDataModuleConfig#

Bases: megatron.bridge.diffusion.data.common.diffusion_energon_datamodule.DiffusionDataModuleConfig

path: str#

None

seq_length: int#

None

packing_buffer_size: int#

None

micro_batch_size: int#

None

global_batch_size: int#

None

num_workers: int#

None

dataloader_type: str#

‘external’

vae_scale_factor: int#

8

latent_channels: int#

16

__post_init__()#
build_datasets(
context: megatron.bridge.data.utils.DatasetBuildContext,
)#