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

examples/visual_gen/models/flux2.py

@@ -13,11 +13,13 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""FLUX.2 text-to-image generation.
+"""FLUX.2 text-to-image and reference-image generation.
 
 Usage:
     python flux2.py
     python flux2.py --visual_gen_args ../configs/flux2-dev-fp4-1gpu.yaml
+    python flux2.py --image subject.png
+    python flux2.py --image subject.png --image style.png
 """
 
 import argparse
@@ -35,7 +37,9 @@ def _output_paths(output_path: str, num_images: int) -> str | list[str]:
 
 
 def main():
-    parser = argparse.ArgumentParser(description="FLUX.2 Text-to-Image example")
+    parser = argparse.ArgumentParser(
+        description="FLUX.2 text-to-image and reference-image generation example"
+    )
     parser.add_argument(
         "--model",
         type=str,
@@ -61,6 +65,42 @@ def main():
         default=1,
         help="Number of images to generate for the prompt",
     )
+    parser.add_argument(
+        "--image",
+        action="append",
+        default=None,
+        help="Reference image path; repeat for a shared set of images",
+    )
+    parser.add_argument(
+        "--height",
+        type=int,
+        default=None,
+        help="Output height; with references, omitted uses the first processed image",
+    )
+    parser.add_argument(
+        "--width",
+        type=int,
+        default=None,
+        help="Output width; with references, omitted uses the first processed image",
+    )
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=None,
+        help="Number of denoising steps; omitted uses the model default",
+    )
+    parser.add_argument(
+        "--guidance_scale",
+        type=float,
+        default=None,
+        help="Embedded guidance scale; omitted uses the model default",
+    )
+    parser.add_argument(
+        "--seed",
+        type=int,
+        default=None,
+        help="Random seed; omitted selects a fresh random seed",
+    )
     parser.add_argument(
         "--output_path",
         type=str,
@@ -75,10 +115,25 @@ def main():
     extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None
     visual_gen = VisualGen(model=args.model, args=extra_args)
 
-    # --- Model-specific: T2I request construction ---
-    # Start from per-model defaults (resolution, steps, guidance, seed, etc.) and set image count.
+    # Start from per-model defaults and override only user-provided request fields.
     params = visual_gen.default_params
     params.num_images_per_prompt = args.num_images_per_prompt
+    params.image = args.image
+    if args.image:
+        # Let FLUX.2 derive omitted dimensions from the first processed reference.
+        params.height = args.height
+        params.width = args.width
+    else:
+        if args.height is not None:
+            params.height = args.height
+        if args.width is not None:
+            params.width = args.width
+    if args.num_inference_steps is not None:
+        params.num_inference_steps = args.num_inference_steps
+    if args.guidance_scale is not None:
+        params.guidance_scale = args.guidance_scale
+    if args.seed is not None:
+        params.seed = args.seed
 
     output = visual_gen.generate(inputs=args.prompt, params=params)
 

tensorrt_llm/_torch/visual_gen/cache/teacache.py

@@ -1,3 +1,6 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 import inspect
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, Optional
@@ -71,6 +74,7 @@ class ExtractorConfig:
         guidance_param_name: Parameter name for guidance if used (default: None)
         forward_params: List of parameter names (None = auto-introspect from forward signature)
         return_dict_default: Default value for return_dict parameter (default: True)
+        return_tuple_when_return_dict_false: Whether return_dict=False uses a one-element tuple
         output_model_class: Output class name for return type (default: "Transformer2DModelOutput")
     """
 
@@ -80,6 +84,7 @@ class ExtractorConfig:
     guidance_param_name: Optional[str] = None
     forward_params: Optional[List[str]] = None
     return_dict_default: bool = True
+    return_tuple_when_return_dict_false: bool = False
     output_model_class: str = "Transformer2DModelOutput"
 
 
@@ -154,10 +159,8 @@ def postprocess(output):
                 if isinstance(output, tuple):
                     return output
                 return Transformer2DModelOutput(sample=output)
-            # For return_dict=False, unwrap single-element tuple to raw tensor
-            if isinstance(output, tuple) and len(output) == 1:
-                return output[0]
-            # Return raw tensor as-is (TeaCacheHook always passes tensors to postprocess)
+            if self.config.return_tuple_when_return_dict_false:
+                return (output,)
             return output
 
         return CacheContext(

tensorrt_llm/_torch/visual_gen/executor.py

@@ -6,9 +6,9 @@
 import time
 import traceback
 from collections import deque
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from datetime import datetime, timezone
-from typing import TYPE_CHECKING, Dict, List, Optional, Set, Tuple, Union
+from typing import TYPE_CHECKING, Any, Dict, List, Optional, Set, Tuple, Union
 
 import torch
 import torch.distributed as dist
@@ -241,6 +241,7 @@ class DiffusionRequest:
     request_id: int
     prompt: List[str]
     params: Optional["VisualGenParams"] = None
+    prepared_inputs: Dict[str, Any] = field(default_factory=dict, repr=False)
 
 
 @dataclass
@@ -253,11 +254,11 @@ class DiffusionResponse:
             model-specific fields populated. Set to ``None`` on the error
             path; on the READY signal it carries a ``dict`` instead.
         error_msg: Error message if generation failed.
-        generation: Wall-clock time the executor measured around the
-            engine's inference call (host ``time.perf_counter()``), in
-            seconds. Default ``0.0`` so the dataclass round-trips through
-            pickling across worker/client; the error path leaves it at
-            ``0.0``.
+        generation: Wall-clock time the executor measured around request
+            preparation and the engine's inference call (host
+            ``time.perf_counter()``), in seconds. Default ``0.0`` so the
+            dataclass round-trips through pickling across worker/client; the
+            error path leaves it at ``0.0``.
     """
 
     request_id: int
@@ -404,6 +405,12 @@ def _merge_defaults(self, req: DiffusionRequest):
         # Universal field defaults
         for field_name, default_value in self.pipeline.default_generation_params.items():
             if hasattr(params, field_name) and getattr(params, field_name) is None:
+                if (
+                    params.image is not None
+                    and getattr(self.pipeline, "derive_output_size_from_reference", False) is True
+                    and field_name in ("height", "width")
+                ):
+                    continue
                 setattr(params, field_name, default_value)
 
         # Extra param defaults — fill all declared keys so infer() can use direct access
@@ -419,21 +426,24 @@ def process_request(self, req: DiffusionRequest):
         """Process a single request."""
         try:
             self._merge_defaults(req)
-            cache_key = self.pipeline.warmup_cache_key(
-                req.params.height, req.params.width, num_frames=req.params.num_frames
-            )
-            if self.pipeline._warmed_up_shapes and cache_key not in self.pipeline._warmed_up_shapes:
+            # Include request preparation in executor-side generation latency.
+            # Model-specific preparation runs before the warmup lookup so it
+            # can resolve shape-dependent request fields such as output size.
+            generation_start = time.perf_counter()
+            self.pipeline.prepare_request(req)
+            cache_key = self.pipeline.request_warmup_cache_key(req)
+            cache_key_is_resolved = all(value is not None for value in cache_key)
+            if (
+                cache_key_is_resolved
+                and self.pipeline._warmed_up_shapes
+                and cache_key not in self.pipeline._warmed_up_shapes
+            ):
                 logger.warning(
                     f"Requested shape {cache_key} was not warmed up. "
                     f"First request with this shape will be slower due to "
                     f"torch.compile recompilation or CUDA graph capture. "
                     f"Warmed-up shapes: {self.pipeline._warmed_up_shapes}"
                 )
-            # Host wall-clock around pipeline.infer(). The pipeline already
-            # syncs at the end (decode_latents path), so this captures the
-            # full executor-side envelope including any pre/post-pipeline work
-            # that the per-phase CUDA-event timings on PipelineOutput do not.
-            generation_start = time.perf_counter()
             output = self.pipeline.infer(req)
             generation = time.perf_counter() - generation_start  # seconds
             if self.rank == 0:

tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py

@@ -239,6 +239,7 @@ def post_load_weights(self) -> None:
                         "return_dict",
                     ],
                     return_dict_default=False,
+                    return_tuple_when_return_dict_false=True,
                 )
             )