diff --git a/examples/visual_gen/models/flux2.py b/examples/visual_gen/models/flux2.py index 3b48a17923bd..a39044448b2f 100644 --- a/examples/visual_gen/models/flux2.py +++ b/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) diff --git a/tensorrt_llm/_torch/visual_gen/cache/teacache.py b/tensorrt_llm/_torch/visual_gen/cache/teacache.py index 45098a9e8f97..fdf9f827edbe 100644 --- a/tensorrt_llm/_torch/visual_gen/cache/teacache.py +++ b/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( diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 95a170120795..5431653e0804 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/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: diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py index e08d07047878..5881614c872b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/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, ) ) diff --git a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py index 1b2116f7f27f..85a87f05a631 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -18,15 +18,19 @@ - 4-axis RoPE: (32, 32, 32, 32) instead of 3-axis """ +import io import json import os import time -from typing import List, Optional, Tuple, Union +from contextlib import contextmanager +from typing import Any, Iterator, List, Optional, Tuple, Union import numpy as np +import PIL.Image import torch from diffusers import FlowMatchEulerDiscreteScheduler from diffusers.models.autoencoders.autoencoder_kl_flux2 import AutoencoderKLFlux2 +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor from diffusers.utils.torch_utils import randn_tensor from transformers import ( AutoModelForCausalLM, @@ -92,10 +96,10 @@ def format_input(prompts: List[str], system_message: str) -> List[List[dict]]: @register_pipeline( "Flux2Pipeline", hf_ids=["black-forest-labs/FLUX.2-dev"], - doc="Black Forest Labs FLUX.2 family (text-to-image).", + doc="Black Forest Labs FLUX.2 family (text-to-image and reference-image generation).", ) class Flux2Pipeline(BasePipeline): - """FLUX.2 Text-to-Image Pipeline. + """FLUX.2 text-to-image and reference-image pipeline. Supports FLUX.2 model variants: - FLUX.2-dev (35B): guidance_embeds=True, embedded guidance @@ -105,6 +109,8 @@ class Flux2Pipeline(BasePipeline): Follows WAN pipeline pattern for DiffusionModelLoader integration. """ + derive_output_size_from_reference = True + # Hidden state layers per text encoder type (auto-detected at load time) _TEXT_ENCODER_CONFIG = { "Mistral3ForConditionalGeneration": { @@ -194,6 +200,17 @@ def default_warmup_num_frames(self): def warmup_cache_key(self, height: int, width: int, **kwargs) -> tuple: return (height, width) + def request_warmup_cache_key(self, req: Any) -> tuple: + cache_key = super().request_warmup_cache_key(req) + condition_images = req.prepared_inputs.get("condition_images") + if condition_images is None: + return cache_key + + reference_shapes = tuple( + (int(image.shape[-2]), int(image.shape[-1])) for image in condition_images + ) + return (*cache_key, len(condition_images), reference_shapes) + def _init_transformer(self) -> None: """Initialize FLUX.2 transformer with quantization support.""" logger.info("Creating FLUX.2 transformer with quantization support...") @@ -284,6 +301,7 @@ def load_standard_components( ) self.vae_scale_factor = 8 # FLUX.2 uses scale_factor=8 + self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) # Scheduler if PipelineComponent.SCHEDULER not in skip_components: @@ -333,6 +351,7 @@ def post_load_weights(self) -> None: guidance_param_name=guidance_param, forward_params=forward_params, return_dict_default=False, + return_tuple_when_return_dict_false=True, ) ) @@ -350,6 +369,20 @@ def default_generation_params(self): "max_sequence_length": 512, } + def prepare_request(self, req: Any) -> None: + """Load and preprocess reference images before warmup bookkeeping.""" + if req.params.image is None: + return + + reference_images = self._load_reference_images(req.params.image) + condition_images = self._preprocess_reference_images(reference_images) + req.params.height, req.params.width = self._resolve_target_dimensions( + req.params.height, + req.params.width, + condition_images, + ) + req.prepared_inputs["condition_images"] = condition_images + def infer(self, req): """Run inference from DiffusionRequest.""" return self.forward( @@ -361,6 +394,8 @@ def infer(self, req): seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, num_images_per_prompt=req.params.num_images_per_prompt, + image=req.params.image, + _condition_images=req.prepared_inputs.get("condition_images"), ) @torch.inference_mode() @@ -368,21 +403,32 @@ def forward( self, prompt: Union[str, List[str]], seed: int, - height: int = 1024, - width: int = 1024, + height: Optional[int] = None, + width: Optional[int] = None, num_inference_steps: int = 50, guidance_scale: float = 3.5, max_sequence_length: int = 512, num_images_per_prompt: int = 1, + image: Optional[ + Union[ + PIL.Image.Image, + str, + bytes, + List[Union[PIL.Image.Image, str, bytes]], + ] + ] = None, + _condition_images: Optional[List[torch.Tensor]] = None, ): - """Generate image(s) from text prompt(s). + """Generate image(s) from text and optional reference images. Args: prompt: Text prompt or list of prompts for image generation. When a list is provided, generates one image per prompt in a single batched forward pass. - height: Output image height (default: 1024) - width: Output image width (default: 1024) + height: Output image height. Defaults to the first processed + reference image's height, or 1024 without a reference. + width: Output image width. Defaults to the first processed + reference image's width, or 1024 without a reference. num_inference_steps: Number of denoising steps guidance_scale: Embedded guidance scale seed: Random seed for reproducibility. @@ -390,6 +436,9 @@ def forward( num_images_per_prompt: Number of images to generate per prompt. Each prompt's embeddings are repeated and independent noise is sampled, producing N different images per prompt. + image: Reference image or shared list of reference images for + image conditioning. Public ``VisualGenParams`` requests use + file paths or encoded bytes; direct calls may also use PIL images. Returns: PipelineOutput with image tensor (B, H, W, C) where @@ -418,9 +467,35 @@ def forward( if num_images_per_prompt > 1: prompt_embeds = prompt_embeds.repeat_interleave(num_images_per_prompt, dim=0) + condition_images = _condition_images + if condition_images is None and image is not None: + reference_images = self._load_reference_images(image) + condition_images = self._preprocess_reference_images(reference_images) + + height, width = self._resolve_target_dimensions(height, width, condition_images) latents, latent_ids = self._prepare_latents(batch_size, height, width, generator) logger.info(f"Latents shape: {latents.shape}") + image_latents = None + image_latent_ids = None + if condition_images is not None: + image_latents, image_latent_ids = self._prepare_image_latents( + condition_images, + batch_size=batch_size, + ) + image_latents = image_latents.to(device=latents.device, dtype=latents.dtype) + image_latent_ids = image_latent_ids.to(device=latent_ids.device, dtype=latent_ids.dtype) + self._validate_reference_sequence_length( + target_seq_len=latents.shape[1], + reference_seq_len=image_latents.shape[1], + sharder=getattr(self.transformer, "sharder", None), + ) + logger.info( + "Prepared %d FLUX.2 reference image(s), %d tokens total", + len(condition_images), + image_latents.shape[1], + ) + # Prepare timesteps with dynamic shifting # Use explicit linear sigmas (matches HF diffusers exactly) # This is critical for step-distilled models like FLUX.2-klein @@ -456,25 +531,39 @@ def forward_fn( extra_tensors, ): """Forward function for FLUX.2 transformer.""" - return self.transformer( - hidden_states=latents, + transformer_latents = latents + transformer_latent_ids = latent_ids + if image_latents is not None: + transformer_latents = torch.cat([latents, image_latents], dim=1) + transformer_latent_ids = torch.cat([latent_ids, image_latent_ids], dim=0) + + noise_pred = self.transformer( + hidden_states=transformer_latents, encoder_hidden_states=encoder_hidden_states, timestep=timestep / 1000, # FLUX.2 expects normalized timesteps - img_ids=latent_ids, + img_ids=transformer_latent_ids, txt_ids=text_ids, guidance=guidance, return_dict=False, )[0] + if image_latents is None: + return noise_pred + return noise_pred[:, : latents.shape[1]] timer.mark_denoise_start() - latents = self.denoise( - latents=latents, - scheduler=self.scheduler, - prompt_embeds=prompt_embeds, - guidance_scale=1.0, # No CFG: guidance is embedded - forward_fn=forward_fn, - timesteps=timesteps, - ) + # Reference count and dimensions change the transformer sequence length. + # Run those requests outside CUDA graphs so a long-lived process does not + # retain one graph per reference shape. torch.compile remains active and + # may compile a new sequence shape; text-only requests still use graphs. + with self._temporarily_disable_cuda_graphs(disable=image_latents is not None): + latents = self.denoise( + latents=latents, + scheduler=self.scheduler, + prompt_embeds=prompt_embeds, + guidance_scale=1.0, # No CFG: guidance is embedded + forward_fn=forward_fn, + timesteps=timesteps, + ) timer.mark_post_start() # Decode @@ -489,6 +578,23 @@ def forward_fn( timer.mark_end() return timer.fill(PipelineOutput(image=image)) + @contextmanager + def _temporarily_disable_cuda_graphs(self, disable: bool) -> Iterator[None]: + """Bypass CUDA graphs for a request without discarding captured graphs.""" + runners = list(getattr(self, "_cuda_graph_runners", {}).values()) + if not disable or not runners: + yield + return + + previous_states = [runner.enabled for runner in runners] + for runner in runners: + runner.enabled = False + try: + yield + finally: + for runner, enabled in zip(runners, previous_states): + runner.enabled = enabled + def _encode_prompt( self, prompt: List[str], @@ -630,6 +736,151 @@ def _prepare_latent_ids(self, height: int, width: int) -> torch.Tensor: return latent_ids # [seq_len, 4] + @staticmethod + def _load_reference_images( + image: Union[ + PIL.Image.Image, + str, + bytes, + List[Union[PIL.Image.Image, str, bytes]], + ], + ) -> List[PIL.Image.Image]: + """Normalize supported reference-image inputs to materialized RGB images.""" + inputs = image if isinstance(image, list) else [image] + if not inputs: + raise ValueError("`image` must contain at least one reference image.") + + images = [] + for index, item in enumerate(inputs): + try: + if isinstance(item, PIL.Image.Image): + images.append(item.convert("RGB")) + elif isinstance(item, str): + with PIL.Image.open(item) as loaded: + images.append(loaded.convert("RGB")) + elif isinstance(item, bytes): + with PIL.Image.open(io.BytesIO(item)) as loaded: + images.append(loaded.convert("RGB")) + else: + raise ValueError( + "Reference images must be PIL images, file paths, or encoded bytes; " + f"item {index} has type {type(item).__name__}." + ) + except OSError as exc: + raise ValueError(f"Unable to load reference image {index}: {exc}") from exc + return images + + def _preprocess_reference_images(self, images: List[PIL.Image.Image]) -> List[torch.Tensor]: + """Apply the upstream FLUX.2 area, crop, and normalization rules.""" + condition_images = [] + multiple_of = self.vae_scale_factor * 2 + for image in images: + self.image_processor.check_image_input(image) + image = self.image_processor._resize_if_exceeds_area(image) + image_width, image_height = image.size + image_width = (image_width // multiple_of) * multiple_of + image_height = (image_height // multiple_of) * multiple_of + condition_images.append( + self.image_processor.preprocess( + image, + height=image_height, + width=image_width, + resize_mode="crop", + ) + ) + return condition_images + + @staticmethod + def _prepare_image_ids( + image_latents: List[torch.Tensor], + scale: int = 10, + ) -> torch.Tensor: + """Create reference-image position IDs with distinct FLUX.2 T offsets.""" + image_ids = [] + for index, latent in enumerate(image_latents): + _batch_size, _channels, height, width = latent.shape + device = latent.device + t_dim = torch.tensor([scale * (index + 1)], device=device) + h_dim = torch.arange(height, device=device) + w_dim = torch.arange(width, device=device) + l_dim = torch.arange(1, device=device) + image_ids.append(torch.cartesian_prod(t_dim, h_dim, w_dim, l_dim)) + return torch.cat(image_ids, dim=0).float() + + @staticmethod + def _validate_reference_sequence_length( + target_seq_len: int, + reference_seq_len: int, + sharder: Any, + ) -> None: + """Fail before denoising when sequence parallelism cannot shard image tokens.""" + if sharder is None or not sharder.is_active: + return + + combined_seq_len = target_seq_len + reference_seq_len + if combined_seq_len % sharder.size != 0: + raise ValueError( + "FLUX.2 reference-image conditioning produced " + f"{combined_seq_len} image tokens ({target_seq_len} target + " + f"{reference_seq_len} reference), which is not divisible by the configured " + f"sequence-parallel size {sharder.size}. Adjust the target/reference image " + "dimensions or disable sequence parallelism." + ) + + @staticmethod + def _resolve_target_dimensions( + height: Optional[int], + width: Optional[int], + condition_images: Optional[List[torch.Tensor]], + ) -> Tuple[int, int]: + """Match Diffusers defaults: first reference dimensions, then 1024 fallback.""" + if condition_images: + height = height or condition_images[0].shape[-2] + width = width or condition_images[0].shape[-1] + return height or 1024, width or 1024 + + @staticmethod + def _patchify_latents(latents: torch.Tensor) -> torch.Tensor: + """Patchify VAE latents from 32 channels to FLUX.2's packed 128 channels.""" + batch_size, channels, height, width = latents.shape + latents = latents.reshape(batch_size, channels, height // 2, 2, width // 2, 2) + latents = latents.permute(0, 1, 3, 5, 2, 4) + return latents.reshape(batch_size, channels * 4, height // 2, width // 2) + + def _encode_vae_image(self, image: torch.Tensor) -> torch.Tensor: + """Encode one preprocessed reference image deterministically.""" + if image.ndim != 4: + raise ValueError(f"Expected reference image rank 4, got {image.ndim}.") + + encoded = self.vae.encode(image) + image_latents = encoded.latent_dist.mode() + image_latents = self._patchify_latents(image_latents) + + bn_eps = getattr(self.vae.config, "batch_norm_eps", 1e-5) + latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to( + image_latents.device, image_latents.dtype + ) + latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + bn_eps).to( + image_latents.device, image_latents.dtype + ) + return (image_latents - latents_bn_mean) / latents_bn_std + + def _prepare_image_latents( + self, + images: List[torch.Tensor], + batch_size: int, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """VAE-encode, pack, and concatenate a shared reference-image set.""" + image_latents_4d = [] + for image in images: + image = image.to(device=self.device, dtype=self.vae.dtype) + image_latents_4d.append(self._encode_vae_image(image)) + + image_ids = self._prepare_image_ids(image_latents_4d) + packed_latents = [self._pack_latents(latent) for latent in image_latents_4d] + image_latents = torch.cat(packed_latents, dim=1).repeat(batch_size, 1, 1) + return image_latents, image_ids + def _prepare_latents( self, batch_size: int, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index eba5e4320772..c930b66d7c85 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -251,6 +251,14 @@ def warmup_cache_key(self, height: int, width: int, num_frames: int) -> tuple: """ return (height, width, num_frames) + def request_warmup_cache_key(self, req: Any) -> tuple: + """Return the warmup cache key for a prepared inference request.""" + return self.warmup_cache_key( + req.params.height, + req.params.width, + num_frames=req.params.num_frames, + ) + @property def default_warmup_resolutions(self) -> List[Tuple[int, int]]: """Model-specific default warmup resolutions (height, width). @@ -368,6 +376,14 @@ def default_generation_params(self) -> dict: """ return {} + def prepare_request(self, req: Any) -> None: + """Prepare model-specific inputs before warmup bookkeeping. + + Subclasses may mutate internal request state and resolve request + parameters needed by :meth:`request_warmup_cache_key`. The default + implementation is a no-op. + """ + def infer(self, req: Any): raise NotImplementedError diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 2a729d93c132..a5d29b537e27 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -1850,6 +1850,47 @@ def test_flux2_example(_visual_gen_deps, llm_root, llm_venv): assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" +def test_flux2_reference_image_example(_visual_gen_deps, llm_root, llm_venv, tmp_path): + """Run the FLUX.2 example with the existing reference-image request argument.""" + model_path = _lpips_model_path("FLUX.2-dev") + _skip_if_missing(model_path, "FLUX.2-dev checkpoint", is_dir=True) + reference_path = _golden_media_path( + tmp_path, "flux2_lpips_golden.png", "FLUX.2 reference image" + ) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "flux2_reference_image_example" + ) + os.makedirs(out_dir, exist_ok=True) + output_path = os.path.join(out_dir, "flux2_reference_image_output.png") + script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "flux2.py") + config_path = os.path.join( + llm_root, "examples", "visual_gen", "configs", "flux2-dev-fp4-1gpu.yaml" + ) + + _venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--visual_gen_args", + config_path, + "--image", + str(reference_path), + "--height", + "256", + "--width", + "256", + "--num_inference_steps", + "4", + "--output_path", + output_path, + ], + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + + def test_ltx2_example(_visual_gen_deps, llm_root, llm_venv): """Run examples/visual_gen/models/ltx2.py with NVFP4 config end-to-end. diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index a0f5de033f09..f1dc3996a74c 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -211,6 +211,7 @@ l0_b200: - unittest/_torch/visual_gen/test_flux_transformer.py - unittest/_torch/visual_gen/test_flux_attention.py - unittest/_torch/visual_gen/test_flux_pipeline.py + - unittest/_torch/visual_gen/test_flux2_image_conditioning.py - unittest/_torch/visual_gen/test_ltx2_transformer.py - unittest/_torch/visual_gen/test_ltx2_attention.py - unittest/_torch/visual_gen/test_ltx2_pipeline.py @@ -232,6 +233,7 @@ l0_b200: - examples/visual_gen/test_visual_gen.py::test_wan_t2v_example - examples/visual_gen/test_visual_gen.py::test_flux1_example - examples/visual_gen/test_visual_gen.py::test_flux2_example + - examples/visual_gen/test_visual_gen.py::test_flux2_reference_image_example - examples/visual_gen/test_visual_gen.py::test_ltx2_example - examples/visual_gen/test_visual_gen.py::test_wan_i2v_example - examples/visual_gen/test_visual_gen.py::test_cosmos3_example diff --git a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py new file mode 100644 index 000000000000..ac4a5ee84b14 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py @@ -0,0 +1,302 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for FLUX.2 reference-image conditioning.""" + +import io +from collections.abc import Iterator +from types import SimpleNamespace +from unittest.mock import MagicMock + +import PIL.Image +import pytest +import torch +from diffusers.pipelines.flux2.image_processor import Flux2ImageProcessor + +from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + + +def _png_bytes() -> bytes: + buffer = io.BytesIO() + PIL.Image.new("RGB", (64, 64), color=(10, 20, 30)).save(buffer, format="PNG") + return buffer.getvalue() + + +def test_load_reference_images_accepts_pil_path_and_bytes(tmp_path) -> None: + pil_image = PIL.Image.new("L", (64, 64), color=128) + image_path = tmp_path / "reference.png" + pil_image.save(image_path) + + images = Flux2Pipeline._load_reference_images([pil_image, str(image_path), _png_bytes()]) + + assert len(images) == 3 + assert all(image.mode == "RGB" for image in images) + assert [image.size for image in images] == [(64, 64), (64, 64), (64, 64)] + + +@pytest.mark.parametrize("image", [[], [object()]]) +def test_load_reference_images_rejects_invalid_inputs(image: list[object]) -> None: + with pytest.raises(ValueError): + Flux2Pipeline._load_reference_images(image) + + +def test_preprocess_reference_images_caps_area_and_aligns_to_16() -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + pipeline.vae_scale_factor = 8 + pipeline.image_processor = Flux2ImageProcessor(vae_scale_factor=16) + images = [PIL.Image.new("RGB", (81, 65)), PIL.Image.new("RGB", (2048, 1024))] + + processed = pipeline._preprocess_reference_images(images) + + assert processed[0].shape == (1, 3, 64, 80) + for tensor in processed: + height, width = tensor.shape[-2:] + assert height % 16 == 0 + assert width % 16 == 0 + assert height * width <= 1024 * 1024 + + +def test_target_dimensions_default_to_first_processed_reference() -> None: + condition_images = [torch.zeros(1, 3, 64, 80), torch.zeros(1, 3, 96, 112)] + + assert Flux2Pipeline._resolve_target_dimensions(None, None, condition_images) == (64, 80) + assert Flux2Pipeline._resolve_target_dimensions(128, None, condition_images) == (128, 80) + assert Flux2Pipeline._resolve_target_dimensions(None, None, None) == (1024, 1024) + + +def test_prepare_request_resolves_dimensions_and_infer_reuses_images() -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + pipeline.vae_scale_factor = 8 + pipeline.image_processor = Flux2ImageProcessor(vae_scale_factor=16) + req = SimpleNamespace( + prompt=["edit this image"], + params=SimpleNamespace( + image=_png_bytes(), + height=None, + width=None, + num_inference_steps=1, + guidance_scale=3.5, + seed=0, + max_sequence_length=512, + num_images_per_prompt=1, + num_frames=None, + ), + prepared_inputs={}, + ) + + pipeline.prepare_request(req) + + condition_images = req.prepared_inputs["condition_images"] + assert req.params.height == 64 + assert req.params.width == 64 + assert len(condition_images) == 1 + assert condition_images[0].shape == (1, 3, 64, 64) + assert pipeline.request_warmup_cache_key(req) == (64, 64, 1, ((64, 64),)) + + pipeline.forward = MagicMock(return_value=object()) + pipeline.infer(req) + + assert pipeline.forward.call_args.kwargs["_condition_images"] is condition_images + + +def test_text_only_request_preserves_existing_warmup_cache_key() -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + req = SimpleNamespace( + params=SimpleNamespace(height=1024, width=1024, num_frames=None), + prepared_inputs={}, + ) + + assert pipeline.request_warmup_cache_key(req) == (1024, 1024) + + +def test_reference_warmup_cache_key_preserves_count_and_ordered_shapes() -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + req = SimpleNamespace( + params=SimpleNamespace(height=128, width=256, num_frames=None), + prepared_inputs={ + "condition_images": [ + torch.zeros(1, 3, 32, 48), + torch.zeros(1, 3, 64, 80), + ] + }, + ) + + assert pipeline.request_warmup_cache_key(req) == ( + 128, + 256, + 2, + ((32, 48), (64, 80)), + ) + + +@pytest.mark.parametrize("cache_backend", ["teacache", "cache_dit"]) +@pytest.mark.parametrize("reference_count", [1, 3]) +def test_reference_images_run_with_cache_acceleration( + cache_backend: str, + reference_count: int, +) -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + pipeline.pipeline_config = SimpleNamespace(cache_backend=cache_backend) + pipeline._encode_prompt = lambda _prompt, _max_length: ( + torch.zeros(1, 2, 8), + torch.zeros(2, 4), + ) + pipeline._preprocess_reference_images = lambda images: [ + torch.zeros(1, 3, 16, 16) for _ in images + ] + pipeline._prepare_latents = lambda _batch_size, _height, _width, _generator: ( + torch.zeros(1, 4, 8), + torch.zeros(4, 4), + ) + pipeline._prepare_image_latents = lambda images, batch_size: ( + torch.zeros(batch_size, 4 * len(images), 8), + torch.zeros(4 * len(images), 4), + ) + + class Transformer: + guidance_embeds = False + sharder = SimpleNamespace(is_active=False) + + def __init__(self) -> None: + self.sequence_lengths: list[int] = [] + + def parameters(self) -> Iterator[torch.Tensor]: + return iter([torch.empty(0)]) + + def __call__(self, hidden_states: torch.Tensor, **_kwargs) -> tuple[torch.Tensor]: + self.sequence_lengths.append(hidden_states.shape[1]) + return (hidden_states + 1,) + + class Scheduler: + config = SimpleNamespace(use_flow_sigmas=False) + + def set_timesteps(self, *_args, **_kwargs) -> None: + self.timesteps = torch.tensor([1000.0]) + + def set_begin_index(self, _index: int) -> None: + return None + + transformer = Transformer() + pipeline.transformer = transformer + pipeline.scheduler = Scheduler() + graph_runner = SimpleNamespace(enabled=True) + pipeline._cuda_graph_runners = {"transformer": graph_runner} + + denoised_latents: list[torch.Tensor] = [] + graph_states_during_denoise: list[bool] = [] + + def denoise(**kwargs) -> torch.Tensor: + graph_states_during_denoise.append(graph_runner.enabled) + result = kwargs["forward_fn"]( + kwargs["latents"], + {}, + 0, + kwargs["timesteps"][0], + kwargs["prompt_embeds"], + {}, + ) + denoised_latents.append(result) + return result + + pipeline.denoise = denoise + pipeline.decode_latents = lambda _latents, _decode_fn: torch.zeros(1, 16, 16, 3) + + references = [_png_bytes() for _ in range(reference_count)] + result = pipeline.forward( + prompt="edit this image", + seed=0, + image=references[0] if reference_count == 1 else references, + num_inference_steps=1, + ) + + assert transformer.sequence_lengths == [4 + 4 * reference_count] + assert graph_states_during_denoise == [False] + assert graph_runner.enabled is True + assert denoised_latents[0].shape == (1, 4, 8) + assert result.image.shape == (1, 16, 16, 3) + + +def test_cuda_graph_bypass_preserves_text_only_state_and_restores_after_failure() -> None: + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + graph_runner = SimpleNamespace(enabled=True) + pipeline._cuda_graph_runners = {"transformer": graph_runner} + + with pipeline._temporarily_disable_cuda_graphs(disable=False): + assert graph_runner.enabled is True + + with pytest.raises(RuntimeError, match="denoise failed"): + with pipeline._temporarily_disable_cuda_graphs(disable=True): + assert graph_runner.enabled is False + raise RuntimeError("denoise failed") + + assert graph_runner.enabled is True + + +def test_prepare_image_ids_assigns_distinct_time_offsets() -> None: + first = torch.zeros(1, 128, 2, 3) + second = torch.zeros(1, 128, 1, 2) + + image_ids = Flux2Pipeline._prepare_image_ids([first, second]) + + assert image_ids.shape == (8, 4) + torch.testing.assert_close(image_ids[:6, 0], torch.full((6,), 10.0)) + torch.testing.assert_close(image_ids[6:, 0], torch.full((2,), 20.0)) + assert torch.count_nonzero(image_ids[:, 3]) == 0 + + +def test_reference_sequence_length_rejects_nondivisible_parallel_shape() -> None: + sharder = SimpleNamespace(is_active=True, size=4) + + with pytest.raises(ValueError, match="not divisible by the configured sequence-parallel"): + Flux2Pipeline._validate_reference_sequence_length( + target_seq_len=16, + reference_seq_len=6, + sharder=sharder, + ) + + +def test_reference_sequence_length_accepts_divisible_parallel_shape() -> None: + sharder = SimpleNamespace(is_active=True, size=4) + + Flux2Pipeline._validate_reference_sequence_length( + target_seq_len=16, + reference_seq_len=8, + sharder=sharder, + ) + + +def test_patchify_and_pack_reference_latents() -> None: + latents = torch.arange(1 * 2 * 4 * 6).reshape(1, 2, 4, 6) + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + + patchified = Flux2Pipeline._patchify_latents(latents) + packed = pipeline._pack_latents(patchified) + + assert patchified.shape == (1, 8, 2, 3) + assert packed.shape == (1, 6, 8) + torch.testing.assert_close(packed[0, 0], latents[0, :, :2, :2].reshape(-1)) + + +def test_encode_vae_image_uses_mode_patchify_and_batch_norm() -> None: + source_latents = torch.arange(1 * 2 * 4 * 6, dtype=torch.float32).reshape(1, 2, 4, 6) + + class LatentDistribution: + mode_called = False + + def mode(self) -> torch.Tensor: + self.mode_called = True + return source_latents + + latent_distribution = LatentDistribution() + vae = SimpleNamespace( + encode=lambda _image: SimpleNamespace(latent_dist=latent_distribution), + bn=SimpleNamespace(running_mean=torch.zeros(8), running_var=torch.ones(8)), + config=SimpleNamespace(batch_norm_eps=0.0), + ) + pipeline = Flux2Pipeline.__new__(Flux2Pipeline) + pipeline.vae = vae + + encoded = pipeline._encode_vae_image(torch.zeros(1, 3, 32, 48)) + + assert latent_distribution.mode_called + torch.testing.assert_close(encoded, Flux2Pipeline._patchify_latents(source_latents)) diff --git a/tests/unittest/_torch/visual_gen/test_flux_infer.py b/tests/unittest/_torch/visual_gen/test_flux_infer.py index 7b28ae2b42e0..87dd0f6a3d99 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_infer.py +++ b/tests/unittest/_torch/visual_gen/test_flux_infer.py @@ -19,6 +19,7 @@ def test_infer_forwards_num_images_per_prompt( pipeline.forward = Mock(return_value="image") request = SimpleNamespace( prompt="a cat", + prepared_inputs={}, params=SimpleNamespace( height=256, width=256, @@ -27,6 +28,7 @@ def test_infer_forwards_num_images_per_prompt( seed=42, max_sequence_length=512, num_images_per_prompt=2, + image=[b"reference"], ), ) @@ -35,3 +37,7 @@ def test_infer_forwards_num_images_per_prompt( assert result == "image" pipeline.forward.assert_called_once() assert pipeline.forward.call_args.kwargs["num_images_per_prompt"] == 2 + if pipeline_cls is Flux2Pipeline: + assert pipeline.forward.call_args.kwargs["image"] == [b"reference"] + else: + assert "image" not in pipeline.forward.call_args.kwargs diff --git a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py index 0455e25faff7..4e33fee0670f 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_flux_pipeline.py @@ -18,6 +18,7 @@ from pathlib import Path import numpy as np +import PIL.Image import pytest import torch import torch.distributed as dist @@ -772,6 +773,76 @@ def test_flux2_e2e_vs_hf(self, flux2_checkpoint_exists): gc.collect() torch.cuda.empty_cache() + @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") + def test_flux2_reference_image_e2e_vs_hf(self, flux2_checkpoint_exists): + """FLUX.2 reference-image generation matches the diffusers pipeline.""" + from diffusers import Flux2Pipeline as HFFlux2Pipeline + + reference_array = np.zeros((256, 256, 3), dtype=np.uint8) + reference_array[..., 0] = np.arange(256, dtype=np.uint8)[None, :] + reference_array[..., 1] = np.arange(256, dtype=np.uint8)[:, None] + reference_array[..., 2] = 127 + reference_image = PIL.Image.fromarray(reference_array) + cases = [ + { + "image": reference_image, + "prompt": "turn the reference into a detailed watercolor painting", + "height": 256, + "width": 256, + "num_images_per_prompt": 1, + }, + { + "image": [ + reference_image.resize((160, 128)), + PIL.Image.new("RGB", (112, 96), color=(30, 90, 180)), + PIL.Image.new("RGB", (80, 64), color=(180, 90, 30)), + ], + "prompt": [ + "combine the references into a watercolor scene", + "combine the references into a pencil illustration", + ], + "height": None, + "width": None, + "num_images_per_prompt": 2, + }, + ] + + hf_pipe = HFFlux2Pipeline.from_pretrained( + FLUX2_CHECKPOINT_PATH, torch_dtype=torch.bfloat16 + ).to("cuda") + hf_images = [] + for case in cases: + hf_result = hf_pipe( + **case, + num_inference_steps=4, + guidance_scale=4.0, + generator=torch.Generator("cuda").manual_seed(42), + ) + hf_images.append(np.stack([np.array(image) for image in hf_result.images])) + del hf_pipe + gc.collect() + torch.cuda.empty_cache() + + pipeline = PipelineLoader(VisualGenArgs(model=FLUX2_CHECKPOINT_PATH)).load() + for case, hf_image_batch in zip(cases, hf_images): + result = pipeline.forward( + **case, + num_inference_steps=4, + guidance_scale=4.0, + seed=42, + ) + native_image_batch = result.image.cpu().numpy() + + assert native_image_batch.shape == hf_image_batch.shape + for hf_image, native_image in zip(hf_image_batch, native_image_batch): + mse = ((hf_image.astype(float) - native_image.astype(float)) ** 2).mean() + psnr = 10 * np.log10(255**2 / mse) if mse > 0 else float("inf") + assert psnr > 20.0, f"PSNR too low: {psnr:.2f} dB (expected >20 dB)" + + del pipeline + gc.collect() + torch.cuda.empty_cache() + class TestFluxBatchGeneration: """Batch generation tests for FLUX pipelines. diff --git a/tests/unittest/_torch/visual_gen/test_teacache.py b/tests/unittest/_torch/visual_gen/test_teacache.py index 6deeb6b6bfab..aeb36cbb5377 100644 --- a/tests/unittest/_torch/visual_gen/test_teacache.py +++ b/tests/unittest/_torch/visual_gen/test_teacache.py @@ -19,8 +19,13 @@ from unittest.mock import MagicMock, patch import pytest +import torch -from tensorrt_llm._torch.visual_gen.cache.teacache import TeaCacheBackend +from tensorrt_llm._torch.visual_gen.cache.teacache import ( + ExtractorConfig, + TeaCacheBackend, + register_extractor_from_config, +) from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline from tensorrt_llm.visual_gen.args import TeaCacheConfig @@ -301,6 +306,107 @@ def test_no_backends_is_noop(self): acc.refresh(10) # should not raise +class _TupleTransformer(torch.nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + return_dict: bool = False, + ) -> tuple[torch.Tensor]: + assert not return_dict + return (hidden_states + 1,) + + +def _identity_timestep_embedding( + _module: torch.nn.Module, + timestep: torch.Tensor, + **_kwargs, +) -> torch.Tensor: + return timestep + + +def test_teacache_preserves_tuple_output_on_cache_miss_and_hit() -> None: + transformer = _TupleTransformer() + register_extractor_from_config( + ExtractorConfig( + model_class_name=transformer.__class__.__name__, + timestep_embed_fn=_identity_timestep_embedding, + forward_params=["hidden_states", "timestep", "return_dict"], + return_dict_default=False, + return_tuple_when_return_dict_false=True, + ) + ) + backend = TeaCacheBackend( + TeaCacheConfig( + coefficients=[0.0, 0.0], + teacache_thresh=0.2, + use_ret_steps=False, + ) + ) + backend.enable(transformer) + backend.refresh(num_inference_steps=4) + + hidden_states = torch.zeros(1, 8, 4) + timestep = torch.ones(1, 4) + try: + cache_miss = transformer(hidden_states, timestep, return_dict=False) + cache_hit = transformer(hidden_states, timestep, return_dict=False) + cache_stats = backend.get_stats() + finally: + backend.disable(transformer) + + assert isinstance(cache_miss, tuple) + assert isinstance(cache_hit, tuple) + torch.testing.assert_close(cache_miss[0], hidden_states + 1) + torch.testing.assert_close(cache_hit[0], hidden_states + 1) + assert cache_stats["cached"] == 1 + + +class _TensorTransformer(torch.nn.Module): + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + return hidden_states + timestep[:, :1].unsqueeze(-1) + + +def test_teacache_preserves_tensor_output_on_cache_miss_and_hit() -> None: + transformer = _TensorTransformer() + register_extractor_from_config( + ExtractorConfig( + model_class_name=transformer.__class__.__name__, + timestep_embed_fn=_identity_timestep_embedding, + forward_params=["hidden_states", "timestep"], + return_dict_default=False, + ) + ) + backend = TeaCacheBackend( + TeaCacheConfig( + coefficients=[0.0, 0.0], + teacache_thresh=0.2, + use_ret_steps=False, + ) + ) + backend.enable(transformer) + backend.refresh(num_inference_steps=4) + + hidden_states = torch.zeros(1, 8, 4) + timestep = torch.ones(1, 4) + try: + cache_miss = transformer(hidden_states, timestep) + cache_hit = transformer(hidden_states, timestep) + cache_stats = backend.get_stats() + finally: + backend.disable(transformer) + + assert isinstance(cache_miss, torch.Tensor) + assert isinstance(cache_hit, torch.Tensor) + torch.testing.assert_close(cache_miss, hidden_states + 1) + torch.testing.assert_close(cache_hit, hidden_states + 1) + assert cache_stats["cached"] == 1 + + class TestFlux2TeacacheTable: """FLUX.2 built-in coefficient table (dev variant).""" diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py index e69c81d21f1b..75930e827a7f 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -375,6 +375,31 @@ def test_user_values_not_overwritten(self): assert req.params.width == 1920 assert req.params.num_inference_steps == 50 # Default filled + def test_flux2_reference_dimensions_remain_unset_for_pipeline_resolution(self): + from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + + executor = self._make_mock_executor(Flux2Pipeline) + executor.pipeline.derive_output_size_from_reference = True + req = self._make_request(image=b"encoded image") + + self._merge(executor, req) + + assert req.params.height is None + assert req.params.width is None + assert req.params.num_inference_steps == 50 + + def test_flux2_reference_dimensions_preserve_explicit_values(self): + from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + + executor = self._make_mock_executor(Flux2Pipeline) + executor.pipeline.derive_output_size_from_reference = True + req = self._make_request(image=b"encoded image", height=768, width=512) + + self._merge(executor, req) + + assert req.params.height == 768 + assert req.params.width == 512 + def test_extra_params_defaults_merged(self): from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -1160,7 +1185,7 @@ def test_runtime_error_carried_on_response(self): executor = self._make_executor(FluxPipeline) executor._merge_defaults = lambda req: DiffusionExecutor._merge_defaults(executor, req) - executor.pipeline.warmup_cache_key = MagicMock(return_value=(1024, 1024, None)) + executor.pipeline.request_warmup_cache_key = MagicMock(return_value=(1024, 1024, None)) executor.pipeline._warmed_up_shapes = None executor.pipeline.infer = MagicMock(side_effect=RuntimeError("oops")) @@ -1176,3 +1201,41 @@ def test_runtime_error_carried_on_response(self): resp = executor.response_queue.put.call_args[0][0] assert isinstance(resp, DiffusionResponse) assert resp.error_msg == "oops" + + def test_reference_size_is_prepared_before_warmup_lookup(self): + from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor, DiffusionRequest + from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + from tensorrt_llm.visual_gen.params import VisualGenParams + + events = [] + executor = self._make_executor(Flux2Pipeline) + executor.rank = 1 + executor._merge_defaults = lambda req: DiffusionExecutor._merge_defaults(executor, req) + executor.pipeline.derive_output_size_from_reference = True + + def prepare_request(req): + events.append("prepare") + req.params.height = 64 + req.params.width = 80 + + def request_warmup_cache_key(req): + events.append("warmup_cache_key") + return (req.params.height, req.params.width) + + executor.pipeline.prepare_request = MagicMock(side_effect=prepare_request) + executor.pipeline.request_warmup_cache_key = MagicMock(side_effect=request_warmup_cache_key) + executor.pipeline._warmed_up_shapes = {(1024, 1024)} + executor.pipeline.infer = MagicMock( + side_effect=lambda _req: events.append("infer") or MagicMock() + ) + req = DiffusionRequest( + request_id=8, + prompt=["test"], + params=VisualGenParams(image=b"encoded image"), + ) + + DiffusionExecutor.process_request(executor, req) + + assert events == ["prepare", "warmup_cache_key", "infer"] + executor.pipeline.request_warmup_cache_key.assert_called_once_with(req) + executor.pipeline.infer.assert_called_once_with(req)