Skip to content

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

Open
karljang wants to merge 21 commits into
NVIDIA:mainfrom
karljang:trtllm-14177-flux2-kontext
Open

[TRTLLM-14177][feat] support reference images in FLUX.2#16644
karljang wants to merge 21 commits into
NVIDIA:mainfrom
karljang:trtllm-14177-flux2-kontext

Conversation

@karljang

@karljang karljang commented Jul 20, 2026

Copy link
Copy Markdown
Collaborator

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.
Red vintage bicycle reference Watercolor FLUX.2 output preserving the bicycle
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.
Bicycle reference Orange tabby cat reference Blue-green botanical watercolor style reference FLUX.2 output combining the cat, bicycle, and watercolor style

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

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@karljang

Copy link
Copy Markdown
Collaborator Author

/bot kill

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Link to invocation

@karljang

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@karljang

Copy link
Copy Markdown
Collaborator Author

/bot kill

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Link to invocation

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

@karljang
karljang marked this pull request as ready for review July 22, 2026 16:57
@karljang
karljang requested a review from a team as a code owner July 22, 2026 16:57

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py (1)

773-790: 🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Reject references smaller than one latent-patch multiple.

A 1–15px-wide or -high image is rounded down to zero at Lines 781-782, then passed to preprocessing. Validate dimensions before rounding and raise a clear ValueError.

Proposed fix
             image_width, image_height = image.size
+            if image_width < multiple_of or image_height < multiple_of:
+                raise ValueError(
+                    f"Reference image must be at least {multiple_of}x{multiple_of} pixels; "
+                    f"got {image_width}x{image_height}."
+                )
             image_width = (image_width // multiple_of) * multiple_of
             image_height = (image_height // multiple_of) * multiple_of
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py` around lines
773 - 790, Update _preprocess_reference_images to validate each image’s width
and height before rounding them to multiples of multiple_of, raising a clear
ValueError when either dimension is smaller than one latent-patch multiple.
Preserve the existing rounding and preprocessing behavior for valid images.
🧹 Nitpick comments (1)
tensorrt_llm/_torch/visual_gen/pipeline.py (1)

239-245: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick win

Type the prepared-request contract.

The new extension points pass mutable request state through Any, preventing static validation of the reference-conditioning handoff. Define a concrete request protocol and typed prepared-input container.

  • tensorrt_llm/_torch/visual_gen/pipeline.py#L239-L245: type the request parameter and warmup-key return contract.
  • tensorrt_llm/_torch/visual_gen/pipeline.py#L364-L370: type the request-preparation hook parameter.
  • tensorrt_llm/_torch/visual_gen/executor.py#L244-L244: replace Dict[str, Any] with the typed prepared-input container.
  • tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py#L203-L213: consume the typed request/prepared-image contract.
  • tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py#L372-L385: populate the typed conditioning input field.

As per coding guidelines, “Annotate every function” and “avoid Any.”

🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/visual_gen/pipeline.py` around lines 239 - 245, Define a
concrete prepared-request protocol and typed prepared-input container for the
reference-conditioning handoff, eliminating Any. In
tensorrt_llm/_torch/visual_gen/pipeline.py lines 239-245, type
request_warmup_cache_key’s request parameter and tuple return contract; at lines
364-370, type the request-preparation hook parameter. In
tensorrt_llm/_torch/visual_gen/executor.py line 244, replace Dict[str, Any] with
the prepared-input container. In
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py lines 203-213,
consume the typed request and prepared-image contract, and at lines 372-385
populate the typed conditioning input field.

Source: Coding guidelines

🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Outside diff comments:
In `@tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py`:
- Around line 773-790: Update _preprocess_reference_images to validate each
image’s width and height before rounding them to multiples of multiple_of,
raising a clear ValueError when either dimension is smaller than one
latent-patch multiple. Preserve the existing rounding and preprocessing behavior
for valid images.

---

Nitpick comments:
In `@tensorrt_llm/_torch/visual_gen/pipeline.py`:
- Around line 239-245: Define a concrete prepared-request protocol and typed
prepared-input container for the reference-conditioning handoff, eliminating
Any. In tensorrt_llm/_torch/visual_gen/pipeline.py lines 239-245, type
request_warmup_cache_key’s request parameter and tuple return contract; at lines
364-370, type the request-preparation hook parameter. In
tensorrt_llm/_torch/visual_gen/executor.py line 244, replace Dict[str, Any] with
the prepared-input container. In
tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py lines 203-213,
consume the typed request and prepared-image contract, and at lines 372-385
populate the typed conditioning input field.

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 57207902-905e-4466-844d-3b1997334bec

📥 Commits

Reviewing files that changed from the base of the PR and between 5f1f79c and 3bd1147.

📒 Files selected for processing (6)
  • examples/visual_gen/models/flux2.py
  • tensorrt_llm/_torch/visual_gen/executor.py
  • tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py
  • tensorrt_llm/_torch/visual_gen/pipeline.py
  • tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py
  • tests/unittest/_torch/visual_gen/test_visual_gen_params.py
🚧 Files skipped from review as they are similar to previous changes (1)
  • examples/visual_gen/models/flux2.py

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62821 [ run ] triggered by Bot. Commit: 3bd1147 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #62821 [ run ] completed with state FAILURE. Commit: 3bd1147
/LLM/main/L0_MergeRequest_PR pipeline #50947 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

Copy link
Copy Markdown
Collaborator Author

@zhenhuaw-me , @yingguo-trt , could you please review again, although CI is pending and I'm waiting for #17071 to fix the ci.

karljang added 21 commits July 30, 2026 18:29
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
Signed-off-by: Kanghwan Jang <861393+karljang@users.noreply.github.com>
@karljang
karljang force-pushed the trtllm-14177-flux2-kontext branch from af87f6b to 0be9ed1 Compare July 31, 2026 01:39
@karljang

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

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

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants