[TRTLLM-14177][feat] support reference images in FLUX.2

[karljang]

What changed

• add FLUX.2 reference-image conditioning through the existing VisualGenParams.image request argument
• prepare reference images before inference and distinguish reference count and processed shapes in warmup keys
• match the Diffusers FLUX.2 Kontext packing, positional IDs, and denoising behavior
• expose reference-image input in the FLUX.2 example without adding a new pipeline or public configuration surface
• support reference-image conditioning with both TeaCache and Cache-DiT
• preserve each transformer's output contract under TeaCache: a Diffusers-style tuple for FLUX and a tensor for Wan
• add focused unit, Diffusers parity, cache-combination, and example integration coverage

Why

The FLUX.2 implementation supported text-to-image generation but did not consume the existing reference-image request argument. This prevented FLUX.2 Kontext-style image-conditioned generation even though the shared VisualGen request API already represented the input.

The initial implementation conservatively rejected cache acceleration for reference-image requests. B200 validation showed that Cache-DiT already supports the combined target/reference sequence. TeaCache exposed a shared return-contract bug: FLUX transformers return a one-element tuple when return_dict=False, while Wan transformers return a tensor. Preserving the contract declared by each extractor fixes both paths and enables FLUX.2 reference conditioning with TeaCache.

OpenAI image-edit serving is intentionally out of scope. This PR keeps the existing offline VisualGen request contract; a shared, model-neutral image-edit endpoint can enable FLUX.2 and other image-conditioned pipelines separately.

Usage

The public API accepts one reference path or a shared list of reference paths:

from tensorrt_llm import VisualGen, VisualGenArgs
from tensorrt_llm.visual_gen.args import CompilationConfig, TorchCompileConfig

args = VisualGenArgs(
    quant_config={"quant_algo": "NVFP4", "dynamic": True},
    torch_compile_config=TorchCompileConfig(enable=False, enable_autotune=False),
    compilation_config=CompilationConfig(skip_warmup=True),
)
visual_gen = VisualGen(model="black-forest-labs/FLUX.2-dev", args=args)
params = visual_gen.default_params
params.image = [
    "bicycle.png",
    "cat.png",
    "watercolor_style.png",
]
params.height = 512
params.width = 512
params.num_inference_steps = 28
params.guidance_scale = 4.0
params.max_sequence_length = 256
params.seed = 43

output = visual_gen.generate(
    inputs=(
        "Create a detailed watercolor illustration of the orange tabby cat from "
        "reference 2 sitting in the tan wicker basket of the red vintage bicycle "
        "from reference 1. Preserve the bicycle side profile and render the scene "
        "in the loose blue-green botanical watercolor style and paper texture "
        "of reference 3."
    ),
    params=params,
)
output.save("flux2_kontext.png")

Use a single path in params.image for one-reference generation. If height or width is left unset, FLUX.2 derives it from the first processed reference. The examples/visual_gen/models/flux2.py CLI also supports repeatable --image arguments.

B200 sample outputs

These uncached NVFP4 outputs were generated end to end with the public VisualGen API on one B200 at 512×512, 28 denoising steps, guidance 4.0, maximum text length 256, and seeds 42 for one reference and 43 for three references.

Case	Reference input(s)	Output
One reference Transform the reference into a detailed watercolor illustration. Preserve the bright red vintage bicycle, tan wicker basket, full side profile, and composition.
Three references Create a detailed watercolor illustration of the orange tabby cat from reference 2 sitting in the tan wicker basket of the red vintage bicycle from reference 1. Preserve the bicycle side profile and render the scene in the loose blue-green botanical watercolor style and paper texture of reference 3.

Validation

• pre-commit hooks: passed
• B200 test_teacache.py: 28 passed
• B200 test_flux2_image_conditioning.py and test_visual_gen_params.py: 100 passed
• B200 Diffusers reference-image parity test: passed
• B200 FLUX.2 reference-image example integration test: passed
• B200 public VisualGen sample generation: one-reference and three-reference outputs completed
• 2×B200 Ulysses-size-2 qualification completed without detected failures for uncached, TeaCache, and Cache-DiT runs with one and three references at 256×256 and 16 steps

NVFP4 B200 smoke test on FLUX.2-dev, 256×256, 16 steps, seed 42. TeaCache used threshold 0.2; Cache-DiT used residual-difference threshold 0.24.

Cache backend	References	Cached steps	PSNR vs uncached
TeaCache	1	5/16	28.41 dB
TeaCache	3	5/16	32.32 dB
Cache-DiT	1	6/16	25.97 dB
Cache-DiT	3	6/16	29.80 dB

All four cache cases completed successfully. Cached outputs retained the uncached scene composition; the table reports their pixel-level difference from the uncached outputs.

Dev Engineer Review

• Implemented FLUX.2 reference-image conditioning end-to-end via VisualGenParams.image / request params.image, including:
  • Reference image loading from PIL / file path / PNG bytes, validation (now requires at least one reference), and preprocessing to RGB tensors aligned to spatial constraints.
  • VAE latent encoding + latent/conditioning packing, reference IDs/positional ID handling, and denoising integration by concatenating reference latents/IDs with the target stream and slicing transformer outputs back to the target portion.
  • Target height/width resolution derived from reference conditioning when Flux2Pipeline.derive_output_size_from_reference = True; executor default-merging avoids overwriting these derived dimensions.
  • CUDA graphs temporarily disabled when reference conditioning changes the transformer sequence shape.
• Removed the previous module-level MAX_REFERENCE_IMAGES cap; reference loading now enforces only minimum presence and supported input types.
• Added pipeline warmup/prepare request hooks:
  • BasePipeline.request_warmup_cache_key() and prepare_request() introduced; Flux2Pipeline implements both to preload/store condition_images in req.prepared_inputs before warmup-cache lookup.
  • DiffusionExecutor updated to start wall-clock timing before prepare_request, compute warmup using request_warmup_cache_key, and only warn on unresolved warmup keys after full resolution.
• TeaCache correctness for return_dict=False:
  • ExtractorConfig.return_tuple_when_return_dict_false added.
  • Flux TeaCache registration updated to request tuple returns for FluxTransformer2DModel.
  • Added tests ensuring tuple vs tensor output types are preserved across cache misses/hits.
• Preserved transformer/executor output contracts by explicitly tuning TeaCache postprocess behavior and maintaining denoise slicing semantics for target outputs.
• Updated FLUX.2 example CLI to accept repeatable reference-image input and override relevant generation parameters when provided.

QA Engineer Review

Test functions added/modified:

• tests/integration/defs/examples/visual_gen/test_visual_gen.py
  • Added test_flux2_reference_image_example: runs examples/visual_gen/models/flux2.py with a golden reference image and asserts an output image is produced.
• tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py
  • Updated/expanded unit coverage for reference-image input types (PIL/path/bytes), preprocessing/dimension resolution, request preparation + caching keys, denoise/packing behavior across TeaCache backends (teacache and cache_dit), CUDA graph disabling, and latent packing/encoding utilities.
• tests/unittest/_torch/visual_gen/test_flux_infer.py
  • Updated forwarding assertions to ensure infer() passes image only for Flux2Pipeline and not for FluxPipeline.
• tests/unittest/_torch/visual_gen/test_flux_pipeline.py
  • Added TestFluxE2E.test_flux2_reference_image_e2e_vs_hf: validates TRT-LLM output vs HuggingFace outputs (including PSNR threshold) for deterministic FLUX.2 reference-image generation.
• tests/unittest/_torch/visual_gen/test_teacache.py
  • Added CPU-only tests test_teacache_preserves_tuple_output_on_cache_miss_and_hit and test_teacache_preserves_tensor_output_on_cache_miss_and_hit.
• tests/unittest/_torch/visual_gen/test_visual_gen_params.py
  • Updated warmup/prepare ordering and _merge_defaults behavior tests for reference-derived dimensions.

Coverage vs CI/manual lists:

• These are test-code changes (all under tests/); corresponding entries in tests/integration/test_lists/ were not included in the provided context, so CI/CBTS registration coverage cannot be confirmed.
• Verdict: needs follow-up.

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60500 [ run ] triggered by Bot. Commit: 55ba2c4 Link to invocation

[tensorrt-cicd]

PR_Github #60500 [ run ] completed with state SUCCESS. Commit: 55ba2c4
/LLM/main/L0_MergeRequest_PR pipeline #48824 completed with status: 'SUCCESS'

CI Report

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60637 [ run ] triggered by Bot. Commit: 32dc658 Link to invocation

[tensorrt-cicd]

PR_Github #60637 [ run ] completed with state FAILURE. Commit: 32dc658
/LLM/main/L0_MergeRequest_PR pipeline #48939 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60736 [ run ] triggered by Bot. Commit: 465fe7e Link to invocation

[karljang]

/bot kill

[tensorrt-cicd]

PR_Github #60746 [ kill ] triggered by Bot. Commit: 742f2d3 Link to invocation

[tensorrt-cicd]

PR_Github #60736 [ run ] completed with state ABORTED. Commit: 465fe7e

Link to invocation

[tensorrt-cicd]

PR_Github #60746 [ kill ] completed with state SUCCESS. Commit: 742f2d3
Successfully killed previous jobs for commit 742f2d3

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60748 [ run ] triggered by Bot. Commit: 174cf9d Link to invocation

[karljang]

/bot kill

[tensorrt-cicd]

PR_Github #60758 [ kill ] triggered by Bot. Commit: 2a2dd5f Link to invocation

[tensorrt-cicd]

PR_Github #60748 [ run ] completed with state ABORTED. Commit: 174cf9d

Link to invocation

[tensorrt-cicd]

PR_Github #60758 [ kill ] completed with state SUCCESS. Commit: 2a2dd5f
Successfully killed previous jobs for commit 2a2dd5f

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60801 [ run ] triggered by Bot. Commit: 1328016 Link to invocation

[tensorrt-cicd]

PR_Github #60801 [ run ] completed with state FAILURE. Commit: 1328016
/LLM/main/L0_MergeRequest_PR pipeline #49078 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60881 [ run ] triggered by Bot. Commit: 1328016 Link to invocation

[tensorrt-cicd]

PR_Github #60881 [ run ] completed with state FAILURE. Commit: 1328016
/LLM/main/L0_MergeRequest_PR pipeline #49150 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61019 [ run ] triggered by Bot. Commit: 1328016 Link to invocation

[karljang]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62903 [ run ] triggered by Bot. Commit: 0be9ed1 Link to invocation

[tensorrt-cicd]

PR_Github #62903 [ run ] completed with state FAILURE. Commit: 0be9ed1
/LLM/main/L0_MergeRequest_PR pipeline #51025 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

• Please check the failed tests and fix your PR
• If you cannot view the failures, ask the CI triggerer to share details
• Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

[karljang]

/bot run

[tensorrt-cicd]

PR_Github #63000 [ run ] triggered by Bot. Commit: 0be9ed1 Link to invocation

[tensorrt-cicd]

PR_Github #63000 [ run ] completed with state SUCCESS. Commit: 0be9ed1
/LLM/main/L0_MergeRequest_PR pipeline #51108 completed with status: 'SUCCESS'

CI Report

Link to invocation