nemo_automodel.components.models.qwen_image_edit.preprocessing

Offline cache encoder for Qwen/Qwen-Image-Edit-2511.

Module Contents

Classes

Name	Description
`QwenImageEditCacheEncoder`	Encode generic image-edit manifests into the cached contract.
`_CacheRecord`	-
`_CompoundBucketSignature`	-
`_ManifestSample`	-
`_MediaReference`	-
`_WorkerSettings`	-

Functions

Name	Description
`_encode_rows_multiprocess`	Encode round-robin manifest shards in spawned GPU workers.
`_encode_rows_worker`	Construct a private encoder and run one spawned GPU shard.
`_encode_vae_image`	Encode one deterministic Qwen VAE latent.
`_image_to_tensor`	Convert an RGB image to a normalized channels-first tensor.
`_load_unique_images`	Load each manifest-relative media file once as an RGB image.
`_read_manifest`	Read and validate every generic image-edit JSONL row.
`_resize_condition_image`	Match the upstream Qwen edit vision-conditioning resize.
`_resize_to_area`	Resize an RGB image to an aspect-preserving, 32-aligned pixel area.
`_resize_vae_image`	Resize an RGB image for a cached Qwen VAE latent.
`_shared_metadata_optional_string`	Return an optional provenance string shared by all manifest rows.
`_shared_metadata_string`	Return a required provenance string shared by all manifest rows.
`_validate_output_directory`	Require a fresh cache while allowing CLI-owned materialized input.
`_verify_vae_latent`	Decode a cached latent and require a finite RGB result.
`_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.

Data

_CONDITION_IMAGE_AREA

_DEFAULT_MODEL_NAME

_DTYPES

_METADATA_SHARD_SIZE

_RESOLUTION_PRESETS

logger

API

class nemo_automodel.components.models.qwen_image_edit.preprocessing.QwenImageEditCacheEncoder(
    model_name: str | None = None,
    max_sequence_length: int = 512,
    device: str | None = None,
    torch_dtype: str = 'bfloat16'
)

Encode generic image-edit manifests into the cached contract.

The encoder loads only the upstream Diffusers VAE and Qwen2.5-VL conditioning stack. The trainable transformer is explicitly omitted.

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

Encode one worker’s deterministic manifest shard.

Parameters:

samples

list[_ManifestSample]

Validated manifest rows referencing local RGB images.

output_dir

Path

Cache root shared by workers; sample filenames are derived from globally unique row indices.

max_pixels

int

Maximum target/context pixel area without a preset.

resolution_preset

str | None

Optional fixed square spatial preset.

worker_index

int

CUDA device index owned by this worker.

verify

bool

Whether to decode and validate each target latent tensor of shape [channels, target_height, target_width].

Returns: list[_CacheRecord]

Metadata records describing every emitted target/context tensor

nemo_automodel.components.models.qwen_image_edit.preprocessing.QwenImageEditCacheEncoder._encode_sample(
    pipeline,
    sample: nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample,
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    device: torch.device,
    verify: bool
) -> nemo_automodel.components.models.qwen_image_edit.preprocessing._CacheRecord

Encode and persist one instruction-based editing example.

Parameters:

pipeline

Upstream pipeline holding the frozen VAE and Qwen2.5-VL conditioning stack.

sample

_ManifestSample

Validated target/context/condition media references.

output_dir

Path

Cache root for the emitted tensor payload.

max_pixels

int

Maximum target/context pixel area without a preset.

resolution_preset

str | None

Optional fixed square target/context preset.

device

torch.device

Runtime device for VAE and text encoding.

verify

bool

Whether to decode the target latent tensor of shape [channels, target_height, target_width].

Returns: _CacheRecord

Index record containing the target/context tensor layouts and

nemo_automodel.components.models.qwen_image_edit.preprocessing.QwenImageEditCacheEncoder._load_pipeline(
    device: torch.device
)

Load only upstream VAE and Qwen2.5-VL conditioning components.

nemo_automodel.components.models.qwen_image_edit.preprocessing.QwenImageEditCacheEncoder._worker_device(
    worker_index: int
) -> torch.device

Resolve the runtime device owned by an encoder worker.

nemo_automodel.components.models.qwen_image_edit.preprocessing.QwenImageEditCacheEncoder.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 manifest into cached target/context/text 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. It may contain the materialized source tree that owns manifest_path but no prior cache artifacts.

max_pixels

int

Maximum target/context pixel area when no fixed preset is selected.

resolution_preset

str | None

Optional fixed square benchmark preset such as "1024p". When set, target and context images are center- cropped to that exact spatial shape.

num_gpus

int

Number of independent encoder workers. Each worker owns one CUDA device and a private VAE/text encoder.

verify

bool

Whether to decode each target latent and validate its finite RGB tensor output.

Returns: Path

Path to metadata.json, whose shards index target tensors of

class nemo_automodel.components.models.qwen_image_edit.preprocessing._CacheRecord

Bases: typing.TypedDict

cache_file

str

compound_bucket_signature

_CompoundBucketSignature

context_token_lengths

list[int]

original_id

str

row_index

int

target_token_length

int

text_token_length

int

class nemo_automodel.components.models.qwen_image_edit.preprocessing._CompoundBucketSignature

Bases: typing.TypedDict

contexts

list[list[int]]

target

list[int]

class nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample(
    identifier: str,
    prompt: str,
    media: tuple[nemo_automodel.components.models.qwen_image_edit.preprocessing._MediaReference, ...],
    metadata: dict[str, object],
    row_index: int
)

Dataclass

identifier

str

media

tuple[_MediaReference, ...]

metadata

dict[str, object]

prompt

str

row_index

int

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

Dataclass

file_name

str

path

Path

role

str

class nemo_automodel.components.models.qwen_image_edit.preprocessing._WorkerSettings(
    model_name: str,
    max_sequence_length: int,
    device: str | None,
    torch_dtype: str
)

Dataclass

device

str | None

max_sequence_length

int

model_name

str

torch_dtype

str

nemo_automodel.components.models.qwen_image_edit.preprocessing._encode_rows_multiprocess(
    settings: nemo_automodel.components.models.qwen_image_edit.preprocessing._WorkerSettings,
    samples: list[nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample],
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    num_workers: int,
    verify: bool
) -> list[nemo_automodel.components.models.qwen_image_edit.preprocessing._CacheRecord]

Encode round-robin manifest shards in spawned GPU workers.

nemo_automodel.components.models.qwen_image_edit.preprocessing._encode_rows_worker(
    settings: nemo_automodel.components.models.qwen_image_edit.preprocessing._WorkerSettings,
    samples: list[nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample],
    output_dir: pathlib.Path,
    max_pixels: int,
    resolution_preset: str | None,
    worker_index: int,
    verify: bool
) -> list[nemo_automodel.components.models.qwen_image_edit.preprocessing._CacheRecord]

Construct a private encoder and run one spawned GPU shard.

nemo_automodel.components.models.qwen_image_edit.preprocessing._encode_vae_image(
    vae: typing.Any,
    image: Image.Image,
    device: torch.device,
    storage_dtype: torch.dtype
) -> torch.Tensor

Encode one deterministic Qwen VAE latent.

Parameters:

vae

Any

Upstream Qwen VAE. Its input layout is [batch, 3, frames, height, width], and its latent layout is [batch, channels, frames, latent_height, latent_width].

image

Image.Image

RGB image in channels-last layout [height, width, channels].

device

torch.device

VAE execution device.

storage_dtype

torch.dtype

Floating-point dtype for the returned cache tensor.

Returns: torch.Tensor

Normalized latent mode tensor of shape [channels, latent_height,

nemo_automodel.components.models.qwen_image_edit.preprocessing._image_to_tensor(
    image: Image.Image,
    device: torch.device,
    dtype: torch.dtype
) -> torch.Tensor

Convert an RGB image to a normalized channels-first tensor.

Parameters:

image

Image.Image

RGB image in channels-last layout [height, width, channels].

device

torch.device

Output tensor device.

dtype

torch.dtype

Output floating-point dtype.

Returns: torch.Tensor

Tensor of shape [1, 3, height, width] in the range [-1, 1].

nemo_automodel.components.models.qwen_image_edit.preprocessing._load_unique_images(
    media: tuple[nemo_automodel.components.models.qwen_image_edit.preprocessing._MediaReference, ...]
) -> dict[str, Image.Image]

Load each manifest-relative media file once as an RGB image.

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

Read and validate every generic image-edit JSONL row.

nemo_automodel.components.models.qwen_image_edit.preprocessing._resize_condition_image(
    image: Image.Image
) -> Image.Image

Match the upstream Qwen edit vision-conditioning resize.

Diffusers derives an aspect-preserving shape at the fixed 384-squared area and rounds both axes to the nearest multiple of 32. Matching that policy is required because the Qwen2.5-VL image grid changes the prompt embeddings.

Parameters:

image

Image.Image

RGB image in channels-last pixel layout [height, width, channels].

Returns: Image.Image

RGB image with the upstream 32-aligned conditioning dimensions.

nemo_automodel.components.models.qwen_image_edit.preprocessing._resize_to_area(
    image: Image.Image,
    target_area: int
) -> Image.Image

Resize an RGB image to an aspect-preserving, 32-aligned pixel area.

Parameters:

image

Image.Image

RGB image in channels-last pixel layout [height, width, channels].

target_area

int

Approximate output pixel count.

Returns: Image.Image

RGB image in channels-last layout [target_height, target_width,

nemo_automodel.components.models.qwen_image_edit.preprocessing._resize_vae_image(
    image: Image.Image,
    max_pixels: int,
    resolution_preset: str | None
) -> Image.Image

Resize an RGB image for a cached Qwen VAE latent.

Parameters:

image

Image.Image

RGB image in channels-last pixel layout [height, width, channels].

max_pixels

int

Maximum output pixel area without a preset.

resolution_preset

str | None

Optional fixed square output preset.

Returns: Image.Image

RGB image in channels-last layout [target_height, target_width,

nemo_automodel.components.models.qwen_image_edit.preprocessing._shared_metadata_optional_string(
    samples: list[nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample],
    field_name: str
) -> str | None

Return an optional provenance string shared by all manifest rows.

nemo_automodel.components.models.qwen_image_edit.preprocessing._shared_metadata_string(
    samples: list[nemo_automodel.components.models.qwen_image_edit.preprocessing._ManifestSample],
    field_name: str
) -> str

Return a required provenance string shared by all manifest rows.

nemo_automodel.components.models.qwen_image_edit.preprocessing._validate_output_directory(
    output_dir: pathlib.Path,
    manifest_path: pathlib.Path
) -> None

Require a fresh cache while allowing CLI-owned materialized input.

The image-edit CLI may place its materialized source manifest inside the requested output directory before invoking the configured encoder. That source tree is intentionally retained for provenance. Cache-owned paths must not exist so encoding never overwrites prior outputs.

nemo_automodel.components.models.qwen_image_edit.preprocessing._verify_vae_latent(
    vae: typing.Any,
    latent: torch.Tensor,
    device: torch.device
) -> None

Decode a cached latent and require a finite RGB result.

Parameters:

vae

Any

Upstream Qwen VAE whose decoded tensor layout is [batch, 3, frames, height, width].

latent

torch.Tensor

Normalized cache tensor of shape [channels, latent_height, latent_width].

device

torch.device

VAE execution device.

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

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

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

Write deterministic cache-index shards and return relative names.

nemo_automodel.components.models.qwen_image_edit.preprocessing._CONDITION_IMAGE_AREA = 384 * 384

nemo_automodel.components.models.qwen_image_edit.preprocessing._DEFAULT_MODEL_NAME = 'Qwen/Qwen-Image-Edit-2511'

nemo_automodel.components.models.qwen_image_edit.preprocessing._DTYPES = {'bfloat16': torch.bfloat16, 'float16': torch.float16, 'float32': torch.float32}

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

nemo_automodel.components.models.qwen_image_edit.preprocessing._RESOLUTION_PRESETS = {'256p': 256, '512p': 512, '768p': 768, '1024p': 1024, '1536p': 1536}

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