nemo_automodel.components.models.flux3.preprocessing

FLUX.3 image-edit (img2img) cache encoder.

Encodes generic image-edit manifests (one target image, one context/source image, one instruction prompt per row) into the cached contract read by ImageEditDataset:

target_latent: patchified + BN-normalized FLUX.2-VAE latent [128, H/16, W/16]
context_latents: single-entry list with the conditioning image latent in the same space (fed to the transformer’s image_cond stream with l = -1 position ids and timestep 0)
prompt_embeddings / prompt_attention_mask: stacked Qwen3-VL hidden states [sequence, joint_attention_dim] with an all-ones mask — FLUX.3 conditions on every padded position, matching the reference pipeline.

Module Contents

Classes

Name	Description
`Flux3ImageEditCacheEncoder`	Encode image-edit manifests into the FLUX.3 img2img cached contract.
`_ManifestSample`	-
`_MediaReference`	-

Functions

Name	Description
`_encode_rows_multiprocess`	Split samples across CUDA workers and merge their index records.
`_encode_rows_worker`	Encode one worker’s shard in a spawned subprocess.
`_read_manifest`	Read and validate every image-edit JSONL row.
`_resize_image`	Resize an image for VAE encoding with 16-aligned dimensions.
`_write_json_atomic`	Write one JSON value through a same-directory temporary file.
`_write_metadata_shards`	Write deterministic cache-index shards and return relative names.
`patchify_latents`	Convert latents to patchified format (2x2 spatial patch -> channel).

Data

_METADATA_SHARD_SIZE

_RESOLUTION_PRESETS

_SPATIAL_ALIGNMENT

logger

API

class nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder(
    model_name: str | None = None,
    max_sequence_length: int = MAX_SEQUENCE_LENGTH,
    device: str | None = None
)

Encode image-edit manifests into the FLUX.3 img2img cached contract.

Loads only the frozen FLUX.2 image VAE and Qwen3-VL text stack from the modular checkpoint; the trainable transformer is never loaded.

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder._encode_latent(
    models: nemo_automodel.components.models.flux3.types.Flux3Models,
    image: PIL.Image.Image,
    device: torch.device
) -> torch.Tensor

Encode one RGB image into the patchified + BN-normalized latent space.

Parameters:

models

Flux3Models

Dict containing ‘vae’, ‘bn_mean’, ‘bn_std’.

image

Image

PIL RGB image, already resized to 16-aligned dimensions.

device

torch.device

Encoding device.

Returns: torch.Tensor

Latent tensor of shape [128, H/16, W/16], FP32, on CPU.

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder._encode_rows(
    samples: list[nemo_automodel.components.models.flux3.preprocessing._ManifestSample],
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    worker_index: int,
    verify: bool
) -> list[nemo_automodel.components.models.flux3.types.Flux3EditRecord]

Encode one worker’s manifest shard on its assigned device.

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder._encode_sample(
    models: nemo_automodel.components.models.flux3.types.Flux3Models,
    sample: nemo_automodel.components.models.flux3.preprocessing._ManifestSample,
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    device: torch.device,
    verify: bool
) -> nemo_automodel.components.models.flux3.types.Flux3EditRecord

Encode and persist one image-edit example.

Parameters:

models

Flux3Models

Loaded VAE + text stack from _load_models.

sample

_ManifestSample

Validated target/context media references and prompt.

output_dir

Path

Cache root for the emitted tensor payload.

max_pixels

int

Maximum pixel area without a preset.

resolution_preset

str | None

Optional fixed square preset.

device

torch.device

Runtime device for VAE and text encoding.

verify

bool

Whether to run finite-value checks on the cached latents.

Returns: Flux3EditRecord

Index record with cache path, compound bucket signature, and token

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder._load_models(
    device: torch.device
) -> nemo_automodel.components.models.flux3.types.Flux3Models

Load the frozen FLUX.2 image VAE and Qwen3-VL text stack.

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder._worker_device(
    worker_index: int
) -> torch.device

Resolve the runtime device owned by an encoder worker.

nemo_automodel.components.models.flux3.preprocessing.Flux3ImageEditCacheEncoder.encode_manifest(
    manifest_path: pathlib.Path,
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    num_gpus: int,
    verify: bool
) -> pathlib.Path

Encode a generic image-edit manifest into cached FLUX.3 tensors.

Parameters:

manifest_path

Path

JSONL manifest whose rows contain an ID, prompt, ordered role-to-file media mappings, and provenance metadata.

output_dir

Path

Destination for cache tensors and indexes.

max_pixels

int

Maximum pixel area when no fixed preset is selected.

resolution_preset

str | None

Optional fixed square preset such as "512p"; target and context images are center-cropped to that shape.

num_gpus

int

Number of encoder workers (one CUDA device each).

verify

bool

Whether to run finite-value checks on each cached latent.

Returns: Path

Path to metadata.json whose shards index target latents of

class nemo_automodel.components.models.flux3.preprocessing._ManifestSample(
    identifier: str,
    prompt: str,
    target: nemo_automodel.components.models.flux3.preprocessing._MediaReference,
    context: nemo_automodel.components.models.flux3.preprocessing._MediaReference,
    metadata: dict[str, object],
    row_index: int
)

Dataclass

context

_MediaReference

identifier

str

metadata

dict[str, object]

prompt

str

row_index

int

target

_MediaReference

class nemo_automodel.components.models.flux3.preprocessing._MediaReference(
    role: str,
    file_name: str,
    path: pathlib.Path
)

Dataclass

file_name

str

path

Path

role

str

nemo_automodel.components.models.flux3.preprocessing._encode_rows_multiprocess(
    encoder: 'Flux3ImageEditCacheEncoder',
    samples: list[nemo_automodel.components.models.flux3.preprocessing._ManifestSample],
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    num_workers: int,
    verify: bool
) -> list[nemo_automodel.components.models.flux3.types.Flux3EditRecord]

Split samples across CUDA workers and merge their index records.

nemo_automodel.components.models.flux3.preprocessing._encode_rows_worker(
    encoder: 'Flux3ImageEditCacheEncoder',
    worker_index: int,
    samples: list[nemo_automodel.components.models.flux3.preprocessing._ManifestSample],
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    verify: bool,
    results: multiprocessing.queues.Queue[tuple[int, list[nemo_automodel.components.models.flux3.types.Flux3EditRecord]]]
) -> None

Encode one worker’s shard in a spawned subprocess.

nemo_automodel.components.models.flux3.preprocessing._read_manifest(
    manifest_path: pathlib.Path
) -> list[nemo_automodel.components.models.flux3.preprocessing._ManifestSample]

Read and validate every image-edit JSONL row.

Each row must reference exactly one target and exactly one context image; condition entries (used by VL-conditioned models like Qwen-Image-Edit) are not part of the FLUX.3 contract and are rejected.

nemo_automodel.components.models.flux3.preprocessing._resize_image(
    image: PIL.Image.Image,
    max_pixels: int,
    resolution_preset: str | None
) -> PIL.Image.Image

Resize an image for VAE encoding with 16-aligned dimensions.

Parameters:

image

Image

PIL image (any mode; converted to RGB).

max_pixels

int

Maximum pixel area when no preset is selected.

resolution_preset

str | None

Optional fixed square preset key.

Returns: Image

RGB PIL image whose dimensions are multiples of 16. Presets produce an

nemo_automodel.components.models.flux3.preprocessing._write_json_atomic(
    path: pathlib.Path,
    value: object
) -> None

Write one JSON value through a same-directory temporary file.

nemo_automodel.components.models.flux3.preprocessing._write_metadata_shards(
    records: list[nemo_automodel.components.models.flux3.types.Flux3EditRecord],
    output_dir: pathlib.Path
) -> list[str]

Write deterministic cache-index shards and return relative names.

nemo_automodel.components.models.flux3.preprocessing.patchify_latents(
    latents: torch.Tensor
) -> torch.Tensor

Convert latents to patchified format (2x2 spatial patch -> channel).

Parameters:

latents

torch.Tensor

Tensor of shape [batch, channels, height, width].

Returns: torch.Tensor

Tensor of shape [batch, channels*4, height//2, width//2].

nemo_automodel.components.models.flux3.preprocessing._METADATA_SHARD_SIZE = 1000

nemo_automodel.components.models.flux3.preprocessing._RESOLUTION_PRESETS = {'256p': (256, 256), '512p': (512, 512), '768p': (768, 768), '1024p': (1024, 102...

nemo_automodel.components.models.flux3.preprocessing._SPATIAL_ALIGNMENT = 16

nemo_automodel.components.models.flux3.preprocessing.logger = logging.getLogger(__name__)