[TRTLLM-14730][feat] Add image edit serving endpoint for visual generation models

docs/source/models/supported-models.md

@@ -216,15 +216,15 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page
 
 | Model | FP8 blockwise | NVFP4 | TeaCache | CFG Parallelism | Ulysses Parallelism | Parallel VAE | CUDA Graph | torch.compile | trtllm-serve | Attention2D | Ring Attention | Tensor Parallelism |
 |---|---|---|---|---|---|---|---|---|---|--|--|--|
-| **FLUX.1** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
-| **FLUX.2** | Yes | Yes | Yes | No [^1] | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
+| **FLUX.1** | Yes | Yes | Yes | No [^vg1] | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
+| **FLUX.2** | Yes | Yes | Yes | No [^vg1] | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes |
 | **Wan 2.1** | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | **Wan 2.2** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes | Yes |
 | **LTX-2** | Yes | Yes | No | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No |
-| **Qwen-Image-Layered** [^3] | No | No | No | No | No | No | Yes | Yes | No | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | No | No | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Layered** [^vg2] | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | Yes | No | No | Yes | Yes | Yes | No | No | No |
 | **Cosmos3** | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes |
 
 [^vg1]: FLUX models use embedded guidance and do not have a separate negative prompt path, so CFG parallelism is not applicable.
-[^3]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition. FP8 blockwise, NVFP4, `trtllm-serve` image-edit routing, and attention-parallel backends are not enabled yet.
+[^vg2]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition through `trtllm-serve` image-edit routing. By default it returns one RGBA image per generated layer; set `extra_params.save_layers_to_grid` to `true` to pack layers into one saveable image grid. FP8 blockwise, NVFP4, and attention-parallel backends are not enabled yet.

docs/source/models/visual-generation.md

@@ -63,8 +63,8 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 | **FastWan 2.2** | Yes | Yes | No | No | No [^7] | No | No | Yes | Yes | Yes | No | No | No | No |
 | **LTX-2** | Yes | Yes | Yes [^4] | Yes | Yes | Yes | No | No | Yes | Yes | Yes | Yes | No | No |
 | **Qwen-Image** | Yes | Yes | Yes | Yes | Yes | Yes | No | Yes | Yes | Yes | Yes | Yes | No | No |
-| **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | Yes | Yes | No | No | No | No | No |
-| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | Yes | No | No | Yes | Yes | No | No | No | No | No |
+| **Qwen-Image-Layered** [^6] | No | No | No | No | No | No | No | Yes | Yes | Yes | No | No | No | No |
+| **Qwen-Image-Edit-2511** | Yes | Yes | No | No | Yes | No | No | Yes | Yes | Yes | No | No | No | No |
 | **Cosmos3** | Yes | Yes | No | No | Yes | Yes | Yes | Yes | Yes | Yes | No | No | Yes | No |
 | **HunyuanVideo 1.5** | Yes | Yes | No | No | No | No | No | No | No | Yes | No | No | No | No |
 
@@ -76,7 +76,7 @@ Models are auto-detected from the checkpoint directory. Diffusers-format models
 
 [^4]: LTX-2 has no built-in TeaCache coefficient table in TRT-LLM; set `teacache.coefficients` explicitly when enabling TeaCache.
 
-[^6]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition and returns the generated RGBA layer stack as a saveable image grid. FP8 blockwise, NVFP4, cache acceleration, attention-parallel/Sage/VSA backends, Tensor Parallelism, and `trtllm-serve` image-edit routing are not enabled for this pipeline yet.
+[^6]: Qwen-Image-Layered supports baseline BF16 image-conditioned layer decomposition through `trtllm-serve` image-edit routing and returns one RGBA image per generated layer by default. Set `extra_params.save_layers_to_grid` to `true` to pack layers into one saveable image grid. FP8 blockwise, NVFP4, cache acceleration, attention-parallel/Sage/VSA backends, and Tensor Parallelism are not enabled for this pipeline yet.
 
 [^7]: `FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers` — a distilled version of Wan2.2-TI2V-5B with 3 denoising steps. CFG parallelism, TeaCache, and Cache-DiT are not applicable.
 

examples/visual_gen/models/qwen_image_layered.py

@@ -21,6 +21,7 @@
 """
 
 import argparse
+from pathlib import Path
 
 from tensorrt_llm import VisualGen, VisualGenArgs
 
@@ -51,7 +52,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--output_path",
         default="qwen_image_layered_output.png",
-        help="Path to save the layer grid image.",
+        help="Path to save the output image.",
     )
     return parser.parse_args()
 
@@ -65,8 +66,17 @@ def main() -> None:
     params.image = args.image
 
     output = visual_gen.generate(inputs=args.prompt, params=params)
-    saved = output.save(args.output_path)
-    print(f"Saved image to {saved}")
+    if output.image is not None and output.image.shape[0] > 1:
+        output_path = Path(args.output_path)
+        paths = [
+            output_path.with_name(f"{output_path.stem}_layer_{i}{output_path.suffix}")
+            for i in range(output.image.shape[0])
+        ]
+        saved = output.save(paths)
+        print(f"Saved images to {saved}")
+    else:
+        saved = output.save(args.output_path)
+        print(f"Saved image to {saved}")
 
 
 if __name__ == "__main__":

tensorrt_llm/_torch/visual_gen/executor.py

@@ -369,6 +369,7 @@ def _load_pipeline(self):
                         "status": "READY",
                         "default_generation_params": self.pipeline.default_generation_params,
                         "extra_param_specs": self.pipeline.extra_param_specs,
+                        "supports_image_edit": self.pipeline.supports_image_edit,
                     },
                 )
             )
@@ -679,6 +680,7 @@ def __init__(
         # Pipeline metadata — populated by _wait_ready from the READY signal.
         self.default_generation_params: Dict = {}
         self.extra_param_specs: Dict = {}
+        self.supports_image_edit: bool = False
 
         # --- Launch workers ---
         self.worker_processes = []
@@ -1034,6 +1036,7 @@ async def _wait_ready_async(self):
                             "default_generation_params", {}
                         )
                         self.extra_param_specs = payload.get("extra_param_specs", {})
+                        self.supports_image_edit = bool(payload.get("supports_image_edit", False))
                     elapsed = time.time() - start_time
                     logger.info(f"DiffusionClient: Workers ready ({elapsed:.1f}s)")
                     return

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

@@ -109,6 +109,7 @@ class Flux2Pipeline(BasePipeline):
     Follows WAN pipeline pattern for DiffusionModelLoader integration.
     """
 
+    supports_image_edit = True
     derive_output_size_from_reference = True
 
     # Hidden state layers per text encoder type (auto-detected at load time)

tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py

@@ -86,6 +86,7 @@ class QwenImageEditPlusPipeline(QwenImagePipeline):
     concatenated sequence, and the scheduler only steps the generated prefix.
     """
 
+    supports_image_edit = True
     DEFAULT_GENERATION_PARAMS = _EDIT_DEFAULT_GENERATION_PARAMS
 
     def load_standard_components(

tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py

@@ -133,6 +133,7 @@ def _retrieve_latents(
 class QwenImageLayeredPipeline(BasePipeline):
     """Qwen-Image-Layered image decomposition pipeline."""
 
+    supports_image_edit = True
     DEFAULT_GENERATION_PARAMS = _LAYERED_DEFAULT_GENERATION_PARAMS
 
     def __init__(self, pipeline_config):
@@ -237,6 +238,14 @@ def extra_param_specs(self) -> dict:
                 default=False,
                 description="Use English auto-caption prompt when prompt is empty.",
             ),
+            "save_layers_to_grid": ExtraParamSchema(
+                type="bool",
+                default=False,
+                description=(
+                    "Pack generated layers into one image grid. By default the pipeline "
+                    "returns one image per layer."
+                ),
+            ),
         }
 
     def load_standard_components(
@@ -458,6 +467,31 @@ def _layer_stack_to_image_grid(layer_stack: torch.Tensor) -> torch.Tensor:
         grid = grid.permute(0, 1, 3, 2, 4, 5)
         return grid.reshape(batch_size, grid_rows * height, grid_cols * width, channels)
 
+    @staticmethod
+    def _validate_save_layers_to_grid(save_layers_to_grid: bool) -> bool:
+        if not isinstance(save_layers_to_grid, bool):
+            raise ValueError(
+                "save_layers_to_grid must be a bool, "
+                f"got {type(save_layers_to_grid).__name__}: {save_layers_to_grid!r}."
+            )
+        return save_layers_to_grid
+
+    @staticmethod
+    def _format_layer_output(
+        layer_stack: torch.Tensor,
+        save_layers_to_grid: bool,
+    ) -> torch.Tensor:
+        if layer_stack.ndim != 5:
+            raise ValueError(
+                "Qwen-Image-Layered output must have shape (B, layers, H, W, C), "
+                f"got {tuple(layer_stack.shape)}."
+            )
+
+        if QwenImageLayeredPipeline._validate_save_layers_to_grid(save_layers_to_grid):
+            return QwenImageLayeredPipeline._layer_stack_to_image_grid(layer_stack)
+        batch_size, layers, height, width, channels = layer_stack.shape
+        return layer_stack.reshape(batch_size * layers, height, width, channels)
+
     @staticmethod
     def _extract_masked_hidden(
         hidden_states: torch.Tensor, mask: torch.Tensor
@@ -738,6 +772,7 @@ def infer(self, req):
             resolution=extra.get("resolution", 640),
             cfg_normalize=extra.get("cfg_normalize", False),
             use_en_prompt=extra.get("use_en_prompt", False),
+            save_layers_to_grid=extra.get("save_layers_to_grid", False),
         )
 
     @torch.inference_mode()
@@ -756,6 +791,7 @@ def forward(
         resolution: int = 640,
         cfg_normalize: bool = False,
         use_en_prompt: bool = False,
+        save_layers_to_grid: bool = False,
         sigmas: Optional[list] = None,
         latents: Optional[torch.Tensor] = None,
     ) -> PipelineOutput:
@@ -765,6 +801,7 @@ def forward(
             raise ValueError(f"resolution must be 640 or 1024, got {resolution}")
         if layers < 1:
             raise ValueError(f"layers must be >= 1, got {layers}")
+        save_layers_to_grid = self._validate_save_layers_to_grid(save_layers_to_grid)
         if (height is None) != (width is None):
             raise ValueError("height and width must be set together for QwenImageLayeredPipeline.")
 
@@ -952,5 +989,5 @@ def forward(
             logger.info("Layered pipeline total: %.2fs", time.time() - pipeline_start)
 
         timer.mark_end()
-        image_grid = self._layer_stack_to_image_grid(layer_stack)
-        return timer.fill(PipelineOutput(image=image_grid))
+        image = self._format_layer_output(layer_stack, save_layers_to_grid)
+        return timer.fill(PipelineOutput(image=image))

tensorrt_llm/_torch/visual_gen/pipeline.py

@@ -70,6 +70,8 @@ class BasePipeline(nn.Module):
     Base class for diffusion pipelines.
     """
 
+    supports_image_edit: bool = False
+
     @classmethod
     def resolve_variant(cls, config: "DiffusionPipelineConfig") -> Type["BasePipeline"]:
         """Return *cls* or a more specialized subclass based on *config*.

tensorrt_llm/serve/openai_protocol.py

@@ -47,7 +47,7 @@
 from openai.types.responses.tool import Tool
 from openai.types.shared import Metadata, Reasoning
 from openai_harmony import ReasoningEffort
-from pydantic import (BaseModel, ConfigDict, Field, PositiveInt,
+from pydantic import (AliasChoices, BaseModel, ConfigDict, Field, PositiveInt,
                       field_validator, model_validator)
 from typing_extensions import Annotated, Required, TypeAlias, TypedDict
 
@@ -1693,6 +1693,71 @@ def _check_paired_dimensions(self):
         return self
 
 
+class ImageEditRequest(OpenAIBaseModel):
+    """OpenAI-compatible image editing request.
+
+    The server accepts the OpenAI multipart shape and a JSON/base64
+    shape for tests and non-SDK clients. Model-specific knobs travel
+    through ``extra_params`` and are validated by the loaded visual
+    generation pipeline.
+    """
+
+    prompt: str
+    image: Union[str, UploadFile, List[Union[str, UploadFile]]] = Field(
+        description="Input image or images to edit.")
+    mask: Optional[Union[str, UploadFile]] = Field(
+        default=None,
+        description=
+        "Optional edit mask. Currently accepted for compatibility but unsupported.",
+    )
+    response_format: Literal["url", "b64_json"] = "url"
+    output_format: Literal["png", "webp", "jpeg"] = Field(
+        default="png",
+        validation_alias=AliasChoices("output_format", "format"),
+        description="Edited image content encoding format.",
+    )
+    seed: Optional[int] = Field(default=None,
+                                ge=0,
+                                description="Random seed for reproducibility.")
+
+    size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$")
+    width: Optional[int] = Field(default=None, gt=0)
+    height: Optional[int] = Field(default=None, gt=0)
+
+    num_inference_steps: Optional[int] = Field(default=None, gt=0)
+    guidance_scale: Optional[float] = Field(default=None, gt=0)
+    max_sequence_length: Optional[int] = Field(default=None, gt=0)
+    negative_prompt: Optional[str] = None
+    n: Optional[int] = Field(
+        default=None,
+        gt=0,
+        le=10,
+        description=("Number of edited images to generate. Capped at 10 to "
+                     "match the OpenAI images API."),
+    )
+
+    extra_params: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description=(
+            "Model-specific parameters forwarded to the underlying pipeline. "
+            "See per-model docs for accepted keys."),
+    )
+
+    model: Optional[str] = None
+    quality: Optional[Literal["standard", "hd"]] = None
+    user: Optional[str] = None
+
+    @model_validator(mode="after")
+    def _check_paired_dimensions(self):
+        if isinstance(self.image, list) and not self.image:
+            raise ValueError("image must contain at least one input image")
+        if (self.width is None) != (self.height is None):
+            raise ValueError(
+                "width and height must be sent together; got width="
+                f"{self.width!r}, height={self.height!r}")
+        return self
+
+
 class ImageObject(OpenAIBaseModel):
     """Generated image object in the response."""
     b64_json: Optional[str] = None