Skip to navigation

nemo_automodel.components.models.flux3.text_encoding

View as Markdown

FLUX.3 Qwen3-VL text-encoding helpers shared by the preprocessing processors.

FLUX.3 conditions on stacked Qwen3-VL hidden states: the prompt goes through the chat template, is right-padded to a multiple of 80 tokens, and evenly spaced hidden layers (derived from the checkpoint geometry) are concatenated per token.

Module Contents

Classes

NameDescription
Flux3TextStackFrozen image-small text encoder, processor, and selected hidden layers.

Functions

NameDescription
encode_text_qwen3vlEncode a prompt with FLUX3’s training-time text-conditioning semantics.
load_qwen3vl_text_stackLoad the FLUX.3 Qwen3-VL text stack from a resolved modular snapshot dir.
resolve_text_encoder_out_layersDerive the text-encoder hidden layers to stack from checkpoint geometry.

Data

MAX_SEQUENCE_LENGTH

PAD_MULTIPLE

logger

API

class nemo_automodel.components.models.flux3.text_encoding.Flux3TextStack

Bases: typing.TypedDict

Frozen image-small text encoder, processor, and selected hidden layers.

text_encoder
Qwen3VLForConditionalGeneration
text_encoder_out_layers
tuple[int, ...]
tokenizer
Qwen3VLProcessor
nemo_automodel.components.models.flux3.text_encoding.encode_text_qwen3vl(
prompt: str,
device: str
) -> torch.Tensor

Encode a prompt with FLUX3’s training-time text-conditioning semantics.

Each prompt is right-padded to a multiple of 80 tokens. Explicit Qwen position IDs keep padding tokens at position 1, including with Transformers versions that would otherwise assign them position 0.

Parameters:

prompt
str

Caption used to condition image generation.

models
Flux3TextStack

Frozen text encoder, processor, and hidden-layer indices.

device
str

Device used for text encoding.

Returns: torch.Tensor

FP16 CPU tensor of shape [sequence, joint_attention_dim], where sequence

nemo_automodel.components.models.flux3.text_encoding.load_qwen3vl_text_stack(
model_name: str,
device: str

Load the FLUX.3 Qwen3-VL text stack from a resolved modular snapshot dir.

Parameters:

model_name
str

Resolved local snapshot dir of the FLUX.3 modular checkpoint.

device
str

Device to load the text encoder on.

Returns: Flux3TextStack

Dict with ‘text_encoder’ (Qwen3VLForConditionalGeneration), ‘tokenizer’

nemo_automodel.components.models.flux3.text_encoding.resolve_text_encoder_out_layers(
joint_attention_dim: int,
text_hidden_size: int,
text_num_hidden_layers: int
) -> tuple[int, ...]

Derive the text-encoder hidden layers to stack from checkpoint geometry.

FLUX.3 conditions on joint_attention_dim / hidden_size stacked hidden layers, evenly spaced through the encoder depth. For the released checkpoint (Qwen3-VL-4B: 36 layers, hidden 2560, joint 20480) this yields the pipeline default (4, 8, …, 32); for tiny test checkpoints it scales down consistently.

Parameters:

joint_attention_dim
int

Transformer’s text-embedding channel width.

text_hidden_size
int

Text encoder hidden size.

text_num_hidden_layers
int

Text encoder layer count.

Returns: tuple[int, ...]

Tuple of hidden-state indices to stack, in ascending order.

Raises:

  • ValueError: If the checkpoint geometry is inconsistent (dim not a multiple of hidden size, or more stacked layers than the encoder has).
nemo_automodel.components.models.flux3.text_encoding.MAX_SEQUENCE_LENGTH = 8192
nemo_automodel.components.models.flux3.text_encoding.PAD_MULTIPLE = 80
nemo_automodel.components.models.flux3.text_encoding.logger = logging.getLogger(__name__)