nemo_automodel.components.models.bagel

BAGEL mixed-modal LLM (Qwen2 + SigLIP-NaViT + MoT shell).

Submodules

Package Contents

Classes

Name	Description
`AutoEncoder`	BAGEL Stage 2 autoencoder wrapper.
`AutoEncoderParams`	Architecture parameters for the BAGEL/FLUX autoencoder.
`BagelConfig`	Top-level BAGEL config.
`BagelForUnifiedMultimodal`	BAGEL mixed-modal LLM wrapper for understanding and optional generation.

Functions

Name	Description
`load_bagel_autoencoder`	Load the BAGEL autoencoder from `ae.safetensors`.

API

class nemo_automodel.components.models.bagel.autoencoder.AutoEncoder(
    params: nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams
)

Bases: Module

BAGEL Stage 2 autoencoder wrapper.

decoder

encoder

reg

= DiagonalGaussian()

scale_factor

= params.scale_factor

shift_factor

= params.shift_factor

nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.decode(
    z: torch.Tensor
) -> torch.Tensor

Decode scaled latents to image tensors.

nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.encode(
    x: torch.Tensor
) -> torch.Tensor

Encode image tensors to scaled latents.

nemo_automodel.components.models.bagel.autoencoder.AutoEncoder.forward(
    x: torch.Tensor
) -> torch.Tensor

Encode and decode image tensors.

class nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams(
    resolution: int,
    in_channels: int,
    downsample: int,
    ch: int,
    out_ch: int,
    ch_mult: list[int],
    num_res_blocks: int,
    z_channels: int,
    scale_factor: float,
    shift_factor: float
)

Dataclass

Architecture parameters for the BAGEL/FLUX autoencoder.

int

ch_mult

list[int]

downsample

int

in_channels

int

num_res_blocks

int

out_ch

int

resolution

int

scale_factor

float

shift_factor

float

z_channels

int

class nemo_automodel.components.models.bagel.configuration.BagelConfig(
    vision_config: typing.Union[typing.Dict[str, typing.Any], nemo_automodel.components.models.bagel.modeling_siglip_navit.SiglipVisionConfig, None] = None,
    text_config: typing.Union[typing.Dict[str, typing.Any], transformers.Qwen2Config, None] = None,
    visual_und: bool = True,
    visual_gen: bool = False,
    stage: typing.Union[int, str, None] = None,
    llm_path: str = '',
    vit_path: str = '',
    vae_path: str = '',
    max_latent_size: int = 64,
    latent_patch_size: int = 2,
    vit_patch_size: int = 14,
    vit_max_num_patch_per_side: int = 70,
    connector_act: str = 'gelu_pytorch_tanh',
    interpolate_pos: bool = False,
    vit_select_layer: int = -2,
    vit_rope: bool = False,
    text_cond_dropout_prob: float = 0.1,
    vae_cond_dropout_prob: float = 0.3,
    vit_cond_dropout_prob: float = 0.3,
    timestep_shift: float = 1.0,
    pad_token_id: int = 151643,
    llm_config: typing.Union[typing.Dict[str, typing.Any], transformers.Qwen2Config, None] = None,
    vit_config: typing.Union[typing.Dict[str, typing.Any], nemo_automodel.components.models.bagel.modeling_siglip_navit.SiglipVisionConfig, None] = None,
    vae_config: typing.Union[typing.Dict[str, typing.Any], None] = None,
    kwargs: typing.Any = {}
)

Bases: PretrainedConfig

Top-level BAGEL config.

The text and vision sub-configs are nested PretrainedConfig instances (not bare dicts) so callers can introspect them the same way they would with any other HF config.

Attribute aliases:

text_config <-> llm_config (both point at the same object)
vision_config <-> vit_config (ditto)

architectures

= ['BagelForUnifiedMultimodal']

llm_config

Qwen2Config

model_type

= 'bagel'

text_config

= _coerce_text_config(text_config)

vae_config

= vae_config or {}

vision_config

= _coerce_vision_config(vision_config)

vit_config

SiglipVisionConfig | None

nemo_automodel.components.models.bagel.configuration.BagelConfig.to_dict() -> typing.Dict[str, typing.Any]

class nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal(
    config: nemo_automodel.components.models.bagel.configuration.BagelConfig,
    backend: nemo_automodel.components.models.bagel.configuration.BagelBackendConfig | None = None
)

Bases: HFCheckpointingMixin, Module

BAGEL mixed-modal LLM wrapper for understanding and optional generation.

visual_gen=False gives the Stage 1 understanding-only path. Stage 2 sets visual_gen=True and uses the MoT *_moe_gen parameter siblings, VAE latents prepared by the recipe, and the flow-matching MSE head.

backend

= resolve_bagel_backend(backend)

backend_config_resolver

= staticmethod(resolve_bagel_backend)

hidden_size

= config.text_config.hidden_size

model

= BagelModel(config, backend=(self.backend))

num_heads

= config.text_config.num_attention_heads

state_dict_adapter

tie_word_embeddings_support

TieSupport = TieSupport.UNTIED_ONLY

use_moe

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.forward(
    sequence_length: int,
    packed_text_ids: torch.LongTensor,
    packed_text_indexes: torch.LongTensor,
    sample_lens: typing.List[int],
    packed_position_ids: torch.LongTensor,
    nested_attention_masks: typing.List[torch.Tensor] | None = None,
    split_lens: typing.List[int] | None = None,
    attn_modes: typing.List[str] | None = None,
    packed_vit_tokens: torch.Tensor | None = None,
    packed_vit_token_indexes: torch.LongTensor | None = None,
    packed_vit_position_ids: torch.LongTensor | None = None,
    vit_token_seqlens: torch.Tensor | None = None,
    padded_latent: torch.Tensor | None = None,
    patchified_vae_latent_shapes: typing.List[typing.Tuple[int, int]] | None = None,
    packed_latent_position_ids: torch.LongTensor | None = None,
    packed_vae_token_indexes: torch.LongTensor | None = None,
    packed_timesteps: torch.Tensor | None = None,
    mse_loss_indexes: torch.Tensor | None = None,
    ce_loss_indexes: torch.Tensor | None = None,
    packed_label_ids: torch.Tensor | None = None,
    ce_loss_weights: torch.Tensor | None = None
) -> typing.Dict[str, torch.Tensor | None]

Run the BAGEL mixed-modal forward.

Stage 1 (visual_gen=False) skips the flow-matching branch and MSE computation; Stage 2 activates both. ce_loss_weights is accepted for data-pipeline compatibility but not consumed here - CE is returned per-token (reduction="none") and the trainer may apply weights downstream.

Stage 2 inputs (padded_latent, patchified_vae_latent_shapes, packed_latent_position_ids, packed_vae_token_indexes, packed_timesteps, mse_loss_indexes) are produced by the BAGEL collator when a pack contains t2i/edit samples. padded_latent is the VAE-encoded latent tensor; recipe must call vae_model.encode on the raw padded_images before forward; this module does not own the VAE.

Returns: Dict[str, torch.Tensor | None]

dict(ce=Tensor|None, mse=Tensor|None) - both can be None

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.from_pretrained(
    pretrained_model_name_or_path: typing.Union[str, os.PathLike],
    stage: typing.Union[int, str] = 1,
    strict: bool = False,
    kwargs: typing.Any = {}
) -> 'BagelForUnifiedMultimodal'

classmethod

Load a BAGEL-7B-MoT checkpoint directory into this class.

Reads config.json via BagelConfig.from_pretrained, constructs an empty model, and then loads ema.safetensors filtered by BagelStateDictAdapter. Stage 2 VAE weights are loaded by the recipe because the VAE is not owned by this module tree.

Parameters:

pretrained_model_name_or_path

Union[str, os.PathLike]

Directory containing the HF-layout BAGEL checkpoint.

stage

Union[int, str]Defaults to 1

1 (UND only) or 2 (UND + GEN). Strings "stage1" / "stage2" are also accepted.

strict

boolDefaults to False

If True, raise on state-dict keys that don’t match the adapter patterns. Defaults to False for compatibility with checkpoint sidecar files.

**kwargs

AnyDefaults to {}

Model overrides. backend is passed to the BAGEL model constructor; the remaining values are forwarded to BagelConfig.from_pretrained.

Returns: 'BagelForUnifiedMultimodal'

A fully-initialized BagelForUnifiedMultimodal with weights

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.get_input_embeddings() -> torch.nn.Module

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.get_output_embeddings() -> torch.nn.Module

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.initialize_weights() -> None

Initialize BAGEL weights after AM materializes a from_config model.

BagelForUnifiedMultimodal is an nn.Module root, not a HF PreTrainedModel root. AM’s meta-device from_config path materializes parameters after sharding and then calls this method. Delegate Qwen/SigLIP subtrees to their HF-style initializers, then initialize BAGEL-only connector and generation modules.

nemo_automodel.components.models.bagel.model.BagelForUnifiedMultimodal.supports_config(
    config: typing.Any
) -> bool

classmethod

Return True if this custom class supports config.

nemo_automodel.components.models.bagel.autoencoder.load_bagel_autoencoder(
    local_path: str | None
) -> tuple[nemo_automodel.components.models.bagel.autoencoder.AutoEncoder, nemo_automodel.components.models.bagel.autoencoder.AutoEncoderParams]

Load the BAGEL autoencoder from ae.safetensors.

Parameters:

local_path

str | None

Local path to ae.safetensors. If None, the module is returned with randomly initialized weights.

Returns: tuple[AutoEncoder, AutoEncoderParams]

The autoencoder module and its architecture parameters.