bridge.diffusion.models.flux.flux_provider#

Module Contents#

Classes#

FluxProvider

FLUX model provider configuration.

Data#

logger

API#

bridge.diffusion.models.flux.flux_provider.logger#

‘getLogger(…)’

class bridge.diffusion.models.flux.flux_provider.FluxProvider#

Bases: megatron.bridge.models.transformer_config.TransformerConfig, megatron.bridge.models.model_provider.ModelProviderMixin[megatron.core.models.common.vision_module.vision_module.VisionModule]

FLUX model provider configuration.

Extends TransformerConfig with FLUX-specific parameters and provides model instantiation through the ModelProviderMixin interface.

.. attribute:: num_layers

Dummy setting (required by base class).

.. attribute:: num_joint_layers

Number of double (joint) transformer blocks.

.. attribute:: num_single_layers

Number of single transformer blocks.

.. attribute:: hidden_size

Hidden dimension size.

.. attribute:: num_attention_heads

Number of attention heads.

.. attribute:: activation_func

Activation function to use.

.. attribute:: add_qkv_bias

Whether to add bias to QKV projections.

.. attribute:: in_channels

Number of input channels (latent channels).

.. attribute:: context_dim

Text encoder context dimension.

.. attribute:: model_channels

Model channel dimension for timestep embedding.

.. attribute:: patch_size

Patch size for image embedding.

.. attribute:: guidance_embed

Whether to use guidance embedding (for FLUX-dev).

.. attribute:: vec_in_dim

Vector input dimension (CLIP pooled output dim).

.. attribute:: rotary_interleaved

Whether to use interleaved rotary embeddings.

.. attribute:: apply_rope_fusion

Whether to apply RoPE fusion.

.. attribute:: guidance_scale

Classifier-free guidance scale.

.. attribute:: ckpt_path

Path to checkpoint for loading weights.

.. attribute:: load_dist_ckpt

Whether to load distributed checkpoint.

.. attribute:: do_convert_from_hf

Whether to convert from HuggingFace format.

.. attribute:: save_converted_model_to

Path to save converted model.

num_layers: int#

1

hidden_size: int#

3072

ffn_hidden_size: int#

12288

num_attention_heads: int#

24

layernorm_epsilon: float#

1e-06

hidden_dropout: float#

0

attention_dropout: float#

0

num_joint_layers: int#

19

num_single_layers: int#

38

activation_func: Callable#

None

add_qkv_bias: bool#

True

in_channels: int#

64

context_dim: int#

4096

model_channels: int#

256

axes_dims_rope: List[int]#

‘field(…)’

patch_size: int#

1

guidance_embed: bool#

False

vec_in_dim: int#

768

rotary_interleaved: bool#

True

apply_rope_fusion: bool#

False

use_cpu_initialization: bool#

True

gradient_accumulation_fusion: bool#

False

enable_cuda_graph: bool#

False

cuda_graph_scope: Optional[str]#

None

use_te_rng_tracker: bool#

False

cuda_graph_warmup_steps: int#

2

guidance_scale: float#

3.5

ckpt_path: Optional[str]#

None

load_dist_ckpt: bool#

False

do_convert_from_hf: bool#

False

save_converted_model_to: Optional[str]#

None

seq_length: int#

1024

share_embeddings_and_output_weights: bool#

False

vocab_size: int#

None

make_vocab_size_divisible_by: int#

128

provide(
pre_process=None,
post_process=None,
vp_stage=None,
) → megatron.bridge.diffusion.models.flux.flux_model.Flux#

Create and return a Flux model with this configuration.

Parameters:

  • pre_process – Whether this is the first pipeline stage (unused for Flux).

  • post_process – Whether this is the last pipeline stage (unused for Flux).

  • vp_stage – Virtual pipeline stage (unused for Flux).

Returns:

Configured Flux model instance.