[TRTLLM-14177][feat] support reference images in FLUX.2 - #16644
[TRTLLM-14177][feat] support reference images in FLUX.2#16644karljang wants to merge 21 commits into
Conversation
|
/bot run --disable-fail-fast |
|
PR_Github #60500 [ run ] triggered by Bot. Commit: |
|
PR_Github #60500 [ run ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #60637 [ run ] triggered by Bot. Commit: |
|
PR_Github #60637 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60736 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #60746 [ kill ] triggered by Bot. Commit: |
|
PR_Github #60736 [ run ] completed with state |
|
PR_Github #60746 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #60748 [ run ] triggered by Bot. Commit: |
|
/bot kill |
|
PR_Github #60758 [ kill ] triggered by Bot. Commit: |
|
PR_Github #60748 [ run ] completed with state |
|
PR_Github #60758 [ kill ] completed with state |
|
/bot run --disable-fail-fast |
|
PR_Github #60801 [ run ] triggered by Bot. Commit: |
|
PR_Github #60801 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #60881 [ run ] triggered by Bot. Commit: |
|
PR_Github #60881 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #61019 [ run ] triggered by Bot. Commit: |
There was a problem hiding this comment.
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 winReject 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 winType 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: replaceDict[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
📒 Files selected for processing (6)
examples/visual_gen/models/flux2.pytensorrt_llm/_torch/visual_gen/executor.pytensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.pytensorrt_llm/_torch/visual_gen/pipeline.pytests/unittest/_torch/visual_gen/test_flux2_image_conditioning.pytests/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
|
/bot run --disable-fail-fast |
|
PR_Github #62821 [ run ] triggered by Bot. Commit: |
|
PR_Github #62821 [ run ] completed with state
|
|
@zhenhuaw-me , @yingguo-trt , could you please review again, although CI is pending and I'm waiting for #17071 to fix the ci. |
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>
af87f6b to
0be9ed1
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #62903 [ run ] triggered by Bot. Commit: |
What changed
VisualGenParams.imagerequest argumentWhy
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:
Use a single path in
params.imagefor one-reference generation. Ifheightorwidthis left unset, FLUX.2 derives it from the first processed reference. Theexamples/visual_gen/models/flux2.pyCLI also supports repeatable--imagearguments.B200 sample outputs
These uncached NVFP4 outputs were generated end to end with the public
VisualGenAPI 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.Transform the reference into a detailed watercolor illustration. Preserve the bright red vintage bicycle, tan wicker basket, full side profile, and composition.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
test_teacache.py: 28 passedtest_flux2_image_conditioning.pyandtest_visual_gen_params.py: 100 passedVisualGensample generation: one-reference and three-reference outputs completedNVFP4 B200 smoke test on FLUX.2-dev, 256×256, 16 steps, seed 42. TeaCache used threshold
0.2; Cache-DiT used residual-difference threshold0.24.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
VisualGenParams.image/ requestparams.image, including:height/widthresolution derived from reference conditioning whenFlux2Pipeline.derive_output_size_from_reference = True; executor default-merging avoids overwriting these derived dimensions.MAX_REFERENCE_IMAGEScap; reference loading now enforces only minimum presence and supported input types.BasePipeline.request_warmup_cache_key()andprepare_request()introduced;Flux2Pipelineimplements both to preload/storecondition_imagesinreq.prepared_inputsbefore warmup-cache lookup.DiffusionExecutorupdated to start wall-clock timing beforeprepare_request, compute warmup usingrequest_warmup_cache_key, and only warn on unresolved warmup keys after full resolution.return_dict=False:ExtractorConfig.return_tuple_when_return_dict_falseadded.FluxTransformer2DModel.QA Engineer Review
Test functions added/modified:
tests/integration/defs/examples/visual_gen/test_visual_gen.pytest_flux2_reference_image_example: runsexamples/visual_gen/models/flux2.pywith a golden reference image and asserts an output image is produced.tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.pyteacacheandcache_dit), CUDA graph disabling, and latent packing/encoding utilities.tests/unittest/_torch/visual_gen/test_flux_infer.pyinfer()passesimageonly forFlux2Pipelineand not forFluxPipeline.tests/unittest/_torch/visual_gen/test_flux_pipeline.pyTestFluxE2E.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.pytest_teacache_preserves_tuple_output_on_cache_miss_and_hitandtest_teacache_preserves_tensor_output_on_cache_miss_and_hit.tests/unittest/_torch/visual_gen/test_visual_gen_params.py_merge_defaultsbehavior tests for reference-derived dimensions.Coverage vs CI/manual lists:
tests/); corresponding entries intests/integration/test_lists/were not included in the provided context, so CI/CBTS registration coverage cannot be confirmed.