diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index cf5339c4ef6f..7e8fb78e6e3a 100644 --- a/docs/source/models/supported-models.md +++ b/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. diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index dcfd46dc1375..1cf96b4e2a73 100644 --- a/docs/source/models/visual-generation.md +++ b/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. diff --git a/examples/visual_gen/models/qwen_image_layered.py b/examples/visual_gen/models/qwen_image_layered.py index 6fbcaf9e9f47..06af0d386914 100644 --- a/examples/visual_gen/models/qwen_image_layered.py +++ b/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__": diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 25edec04570a..60e8b16449b2 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/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 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 85a87f05a631..20420592caaa 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/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) diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py index 6293b28a2ddc..ddc0c453bf7c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image/pipeline_qwen_image_edit.py +++ b/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( diff --git a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py b/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py index 113a329a27f3..440979a5a287 100644 --- a/tensorrt_llm/_torch/visual_gen/models/qwen_image_layered/pipeline_qwen_image_layered.py +++ b/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)) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 089fdd5f8034..8512039760cb 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/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*. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index a2608d4ae199..824634970d96 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/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 diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 1c38c3c90634..ed819d60a1e0 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -75,7 +75,7 @@ ChatCompletionRequest, ChatCompletionResponse, ChatCompletionResponseChoice, ChatMessage, CompletionRequest, CompletionResponse, CompletionResponseChoice, EmbeddingRequest, EmbeddingResponse, - EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, + EmbeddingResponseData, EmbeddingUsageInfo, ErrorResponse, ImageEditRequest, ImageGenerationRequest, ImageGenerationResponse, ImageObject, MemoryUpdateRequest, ModelCard, ModelList, PromptTokensDetails, ResponseFormat, ResponsesRequest, ResponsesResponse, TokenizeRequest, @@ -104,7 +104,8 @@ from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import \ build_visual_gen_timing_headers -from tensorrt_llm.serve.visual_gen_utils import parse_visual_gen_params +from tensorrt_llm.serve.visual_gen_utils import ( + cleanup_materialized_conditioning_inputs, parse_visual_gen_params) from tensorrt_llm.version import __version__ as VERSION from .._utils import nvtx_mark, set_prometheus_multiproc_dir @@ -319,6 +320,30 @@ def _normalize_image_output(image) -> list: return [image] +def _image_output_size(image) -> Optional[str]: + pil_size = getattr(image, "size", None) + if isinstance(pil_size, tuple) and len(pil_size) >= 2: + width, height = pil_size[:2] + return f"{int(width)}x{int(height)}" + + shape = getattr(image, "shape", None) + if shape is None: + return None + + dims = tuple(int(dim) for dim in shape) + if len(dims) == 2: + height, width = dims + elif len(dims) == 3: + if dims[0] in (1, 3, 4) and dims[1] > 4 and dims[2] > 4: + height, width = dims[1], dims[2] + else: + height, width = dims[0], dims[1] + else: + return None + + return f"{width}x{height}" + + class OpenAIServer(_VideoRoutesMixin): @staticmethod @@ -583,6 +608,24 @@ def _init_visual_gen(self): self.media_storage_path.mkdir(exist_ok=True, parents=True) self.video_gen_tasks = {} + def _supports_image_edit(self) -> bool: + if not self._is_visual_gen: + return False + + executor = getattr(self.generator, "executor", None) + if executor is None: + return False + + pipeline = getattr(executor, "pipeline", None) + if pipeline is not None: + return bool(getattr(pipeline, "supports_image_edit", False)) + + supports_image_edit = getattr(executor, "supports_image_edit", None) + if supports_image_edit is not None: + return bool(supports_image_edit) + + return False + def _init_llm(self, chat_template: Optional[str] = None): self.tokenizer = self.generator.tokenizer hf_tokenizer_path = self.generator._hf_model_dir @@ -2631,18 +2674,160 @@ def _build_image_content_url(raw_request: Request, image_id: str, base = str(raw_request.base_url).rstrip("/") return f"{base}/v1/images/{image_id}/content?i={i}" + async def _parse_image_edit_request( + self, + raw_request: Request, + ) -> ImageEditRequest: + """Parse an image edit request from JSON or multipart form data.""" + content_type = raw_request.headers.get("content-type", "") + normalized_content_type = content_type.lower() + + if "application/json" in normalized_content_type: + body = await raw_request.json() + if not isinstance(body, dict): + raise ValueError( + "JSON image edit request body must be an object") + return ImageEditRequest(**body) + + if "multipart/form-data" in normalized_content_type: + form = await raw_request.form() + data = {} + for key in form: + values = form.getlist(key) + if key == "image": + values = [ + value for value in values + if not (isinstance(value, str) and value == "") + ] + if not values: + continue + data[key] = values if len(values) > 1 else values[-1] + continue + + value = values[-1] + if key == "extra_params": + if value == "": + continue + if not isinstance(value, str): + raise ValueError( + "'extra_params' must be a JSON object string") + try: + data[key] = json.loads(value) + except (json.JSONDecodeError, TypeError) as exc: + raise ValueError( + f"'extra_params' must be a JSON object string; {exc}" + ) from exc + continue + if value == "": + continue + data[key] = value + return ImageEditRequest(**data) + + raise ValueError( + f"Unsupported content-type: {content_type}. Use 'application/json' or 'multipart/form-data'" + ) + async def openai_image_edit(self, raw_request: Request) -> Response: - """OpenAI-compatible image editing endpoint — returns HTTP 501. - - No in-tree pipeline implements image editing today: Flux/Flux2 are - text-to-image only and ignore ``params.image``; Wan and LTX-2 produce - video, not edited images. The route is registered so callers get an - honest NotImplemented signal instead of a 404. The request body is - not parsed because no schema is committed for this endpoint yet — - bring a typed request model back when an edit-capable pipeline lands. - """ - return self._create_not_supported_error( - "Image editing is not supported by any in-tree pipeline yet.") + """OpenAI-compatible image editing endpoint.""" + if not self._supports_image_edit(): + return self._create_not_supported_error( + "Image editing is not supported by the loaded visual generation model." + ) + + try: + image_id = f"image_{uuid.uuid4().hex}" + input_paths = None + + try: + request = await self._parse_image_edit_request(raw_request) + params = parse_visual_gen_params( + request, + image_id, + self.generator, + media_storage_path=str(self.media_storage_path), + ) + input_paths = params.image + logger.info( + f"Editing image: {image_id} with params: {params} and prompt: {request.prompt}" + ) + image_edit_start = time.perf_counter() + try: + output = self.generator.generate(inputs=request.prompt, + params=params) + finally: + cleanup_materialized_conditioning_inputs(input_paths) + except ValidationError as exc: + return self._render_pydantic_validation_error(exc) + except ValueError as exc: + logger.error(f"Image edit request error: {exc}") + return self.create_error_response( + message=str(exc), + status_code=HTTPStatus.BAD_REQUEST, + ) + + if output.image is None: + return self.create_error_response( + message="Image editing failed", + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) + + # Model-specific ``extra_params`` stay pipeline-owned. Layered + # image-edit models such as Qwen-Image-Layered use + # ``save_layers_to_grid`` to pack all layers into one image here. + output_images = _normalize_image_output(output.image) + output_size = _image_output_size( + output_images[0]) if output_images else None + pil_format = request.output_format.upper() + ext = f".{request.output_format}" + if request.response_format == "b64_json": + data = [ + ImageObject( + b64_json=base64.b64encode( + image_to_bytes(image, + format=pil_format)).decode("utf-8"), + revised_prompt=request.prompt, + ) for image in output_images + ] + else: + data = [] + for i, image in enumerate(output_images): + path = self.media_storage_path / f"{image_id}_{i}{ext}" + path.write_bytes(image_to_bytes(image, format=pil_format)) + data.append( + ImageObject( + url=self._build_image_content_url( + raw_request, image_id, i), + revised_prompt=request.prompt, + )) + + response = ImageGenerationResponse( + created=int(time.time()), + data=data, + output_format=request.output_format, + size=output_size, + ) + + latency = time.perf_counter() - image_edit_start + metrics = output.metrics + generation = metrics.generation if metrics is not None else 0.0 + denoise = metrics.denoise if metrics is not None else 0.0 + logger.info(f"Image {image_id} edited and encoded: " + f"latency={latency:.3f}s generation={generation:.3f}s " + f"denoise={denoise:.3f}s") + headers = build_visual_gen_timing_headers(metrics) + + return JSONResponse(content=response.model_dump(), headers=headers) + + except ValidationError as exc: + return self._render_pydantic_validation_error(exc) + except Exception as e: + logger.error(traceback.format_exc()) + return self.create_error_response( + message=str(e), + err_type="InternalServerError", + status_code=HTTPStatus.INTERNAL_SERVER_ERROR, + ) async def __call__(self, host, diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index afd5d4d6e5a1..49f093a2297e 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -2,18 +2,34 @@ import asyncio import base64 +import binascii import os +from io import BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional +from urllib.parse import urlparse + +from PIL import Image, UnidentifiedImageError from tensorrt_llm.inputs.media_io import is_isobmff_image_bytes, sniff_media_kind from tensorrt_llm.logger import logger -from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest +from tensorrt_llm.serve.openai_protocol import ( + ImageEditRequest, + ImageGenerationRequest, + VideoGenerationRequest, +) if TYPE_CHECKING: # Type-only: importing tensorrt_llm.visual_gen at runtime would pull the # whole visual_gen tree into every LLM serving process. from tensorrt_llm.visual_gen import VisualGen, VisualGenParams +IMAGE_EDIT_MAX_IMAGES = 16 +IMAGE_EDIT_MAX_IMAGE_BYTES = 50 * 1024 * 1024 +IMAGE_EDIT_MAX_TOTAL_IMAGE_BYTES = 256 * 1024 * 1024 +IMAGE_EDIT_MAX_OUTPUT_IMAGES = 64 +_IMAGE_EDIT_INPUT_FORMATS = {"PNG", "JPEG"} +_INVALID_IMAGE_EDIT_INPUT_MESSAGE = "image edit input is not a PNG/JPEG image" + # Per-field warnings for OpenAI-shaped knobs that the engine has no # semantic for. Each entry maps the request attribute to the message # logged when the client sends a non-None value. @@ -107,8 +123,189 @@ def _read_reference_payload(reference) -> bytes: return reference.file.read() +def _decode_base64_media(value: str) -> Optional[bytes]: + payload = value + if value.startswith("data:"): + _, sep, payload = value.partition(",") + if not sep: + return None + if len(payload) > ((IMAGE_EDIT_MAX_IMAGE_BYTES + 2) // 3) * 4: + raise ValueError( + "Image edit input exceeds the per-image byte limit " + f"before decoding ({len(payload)} encoded bytes)." + ) + try: + return base64.b64decode(payload, validate=True) + except (binascii.Error, ValueError): + return None + + +def _write_bytes_with_limit(value: bytes, path: str) -> int: + size = len(value) + if size > IMAGE_EDIT_MAX_IMAGE_BYTES: + raise ValueError( + "Image edit input exceeds the per-image byte limit " + f"({size} > {IMAGE_EDIT_MAX_IMAGE_BYTES})." + ) + _validate_png_jpeg_image(value) + with open(path, "wb") as f: + f.write(value) + return size + + +def _validate_png_jpeg_image(value: bytes) -> None: + try: + with Image.open(BytesIO(value)) as image: + image_format = image.format + image.verify() + except (UnidentifiedImageError, OSError, SyntaxError, ValueError) as exc: + raise ValueError(_INVALID_IMAGE_EDIT_INPUT_MESSAGE) from exc + if image_format not in _IMAGE_EDIT_INPUT_FORMATS: + raise ValueError(_INVALID_IMAGE_EDIT_INPUT_MESSAGE) + + +def _copy_upload_with_limit(value: Any, path: str) -> int: + total = 0 + if hasattr(value.file, "seek"): + value.file.seek(0) + chunks = [] + while True: + chunk = value.file.read(1024 * 1024) + if not chunk: + break + total += len(chunk) + if total > IMAGE_EDIT_MAX_IMAGE_BYTES: + raise ValueError( + "Image edit input exceeds the per-image byte limit " + f"({total} > {IMAGE_EDIT_MAX_IMAGE_BYTES})." + ) + chunks.append(chunk) + return _write_bytes_with_limit(b"".join(chunks), path) + + +def _materialize_conditioning_input( + value: Any, + path: str, +) -> tuple[str, int]: + """Return a server-owned file path for upload or base64 inputs.""" + try: + if isinstance(value, str): + decoded = _decode_base64_media(value) + if decoded is None: + parsed = urlparse(value) + if parsed.scheme in ("file", "http", "https"): + raise ValueError( + "Image edit inputs must be uploaded files or base64-encoded images; " + "local paths and URLs are not supported." + ) + raise ValueError("String image edit inputs must be base64-encoded image data.") + return path, _write_bytes_with_limit(decoded, path) + + if isinstance(value, bytes): + return path, _write_bytes_with_limit(value, path) + + if hasattr(value, "file"): + return path, _copy_upload_with_limit(value, path) + except Exception: + try: + os.remove(path) + except FileNotFoundError: + pass + raise + + raise ValueError(f"Unsupported conditioning input type: {type(value)}") + + +def _resolve_image_edit_layer_multiplier( + request: ImageEditRequest, + generator: VisualGen, +) -> int: + extra = request.extra_params or {} + save_layers_to_grid = extra.get("save_layers_to_grid", False) + if save_layers_to_grid is True: + return 1 + if save_layers_to_grid not in (False, None): + raise ValueError( + "extra_params.save_layers_to_grid must be a bool when estimating image edit output count." + ) + + layer_spec = generator.extra_param_specs.get("layers") + if layer_spec is None: + return 1 + + layers = extra.get("layers", getattr(layer_spec, "default", 1)) + if layers is None: + return 1 + if isinstance(layers, bool) or not isinstance(layers, int): + raise ValueError( + "extra_params.layers must be an int when estimating image edit output count." + ) + return layers + + +def _validate_image_edit_request_limits( + request: ImageEditRequest, + generator: VisualGen, +) -> int: + image_count = len(request.image) if isinstance(request.image, list) else 1 + if image_count > IMAGE_EDIT_MAX_IMAGES: + raise ValueError( + f"Image edit accepts at most {IMAGE_EDIT_MAX_IMAGES} input images, got {image_count}." + ) + + output_count = (request.n or 1) * _resolve_image_edit_layer_multiplier(request, generator) + if output_count > IMAGE_EDIT_MAX_OUTPUT_IMAGES: + raise ValueError( + "Image edit request can produce at most " + f"{IMAGE_EDIT_MAX_OUTPUT_IMAGES} output images, got {output_count}." + ) + return image_count + + +def _materialize_conditioning_inputs( + value: Any, + *, + id: str, + field_name: str, + media_storage_path: str, +) -> str | List[str]: + values = value if isinstance(value, list) else [value] + paths = [] + total_bytes = 0 + try: + for i, item in enumerate(values): + path, size = _materialize_conditioning_input( + item, + os.path.join(media_storage_path, f"{id}_{field_name}_{i}.png"), + ) + paths.append(path) + total_bytes += size + if total_bytes > IMAGE_EDIT_MAX_TOTAL_IMAGE_BYTES: + raise ValueError( + "Image edit inputs exceed the total byte limit " + f"({total_bytes} > {IMAGE_EDIT_MAX_TOTAL_IMAGE_BYTES})." + ) + except Exception: + cleanup_materialized_conditioning_inputs(paths) + raise + return paths if isinstance(value, list) else paths[0] + + +def cleanup_materialized_conditioning_inputs(value: Any) -> None: + paths = value if isinstance(value, list) else [value] + for path in paths: + if not isinstance(path, str): + continue + try: + os.remove(path) + except FileNotFoundError: + pass + except OSError as exc: + logger.warning("Failed to remove temporary image edit input %r: %s", path, exc) + + def parse_visual_gen_params( - request: ImageGenerationRequest | VideoGenerationRequest, + request: ImageGenerationRequest | ImageEditRequest | VideoGenerationRequest, id: str, generator: VisualGen, media_storage_path: Optional[str] = None, @@ -133,6 +330,10 @@ def parse_visual_gen_params( params.width, params.height = request.width, request.height elif request.size is not None and request.size != "auto": params.width, params.height = map(int, request.size.split("x")) + elif isinstance(request, ImageEditRequest): + if request.width is None and request.height is None and request.size in (None, "auto"): + params.width = None + params.height = None # Universal per-request overlays — each guard is the "do not # override with None" rule in action. @@ -151,6 +352,21 @@ def parse_visual_gen_params( if request.n is not None: params.num_images_per_prompt = request.n + elif isinstance(request, ImageEditRequest): + if request.mask is not None: + raise ValueError("Image edit mask input is not supported yet.") + if request.n is not None: + params.num_images_per_prompt = request.n + if media_storage_path is None: + raise ValueError("media_storage_path is required when image edit inputs are provided") + _validate_image_edit_request_limits(request, generator) + params.image = _materialize_conditioning_inputs( + request.image, + id=id, + field_name="image", + media_storage_path=media_storage_path, + ) + elif isinstance(request, VideoGenerationRequest): if request.frame_rate is not None: params.frame_rate = request.frame_rate diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 0ca143964ea6..6d7ae0463c9e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -229,6 +229,7 @@ l0_b200: - unittest/_torch/visual_gen/test_attention_trtllm_sage.py - unittest/_torch/visual_gen/test_attention_integration.py - unittest/_torch/visual_gen/test_attention_perf.py + - unittest/_torch/visual_gen/test_qwen_image_layered_registry.py - unittest/_torch/visual_gen/test_trtllm_serve_e2e.py - unittest/_torch/visual_gen/test_model_loader.py - unittest/_torch/visual_gen/test_flux_transformer.py diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py b/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py index 79e61bb65067..671cf15b79b2 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_layered_registry.py @@ -120,6 +120,7 @@ def test_qwen_image_layered_default_params_match_runtime_inputs(): assert pipeline.default_generation_params["height"] is None assert pipeline.default_generation_params["width"] is None assert pipeline.extra_param_specs["resolution"].range is None + assert pipeline.extra_param_specs["save_layers_to_grid"].default is False assert pipeline.default_warmup_num_frames == [1] assert pipeline.warmup_cache_key(None, None, num_frames=1) == (640, 640) assert pipeline.warmup_cache_key(512, 768, num_frames=1) == (512, 768) @@ -167,6 +168,7 @@ def test_qwen_image_layered_layer_stack_to_image_grid(): grid = QwenImageLayeredPipeline._layer_stack_to_image_grid(layer_stack) assert grid.shape == (1, 4, 4, 1) + assert torch.equal(QwenImageLayeredPipeline._format_layer_output(layer_stack, True), grid) assert grid[0, :, :, 0].tolist() == [ [0, 1, 4, 5], [2, 3, 6, 7], @@ -175,6 +177,25 @@ def test_qwen_image_layered_layer_stack_to_image_grid(): ] +def test_qwen_image_layered_formats_multiple_layer_images_by_default(): + """Layer stacks can be returned as one image per generated layer.""" + layer_stack = torch.arange(24, dtype=torch.uint8).view(2, 3, 2, 2, 1) + + separate = QwenImageLayeredPipeline._format_layer_output(layer_stack, False) + + assert separate.shape == (6, 2, 2, 1) + assert torch.equal(separate[0], layer_stack[0, 0]) + assert torch.equal(separate[3], layer_stack[1, 0]) + + +def test_qwen_image_layered_rejects_non_bool_save_layers_to_grid(): + """Invalid output packing values fail before saving or serving the image.""" + layer_stack = torch.zeros(1, 1, 2, 2, 1, dtype=torch.uint8) + + with pytest.raises(ValueError, match="save_layers_to_grid"): + QwenImageLayeredPipeline._format_layer_output(layer_stack, "true") + + def test_transformer_constructs_with_layered_config(): """Layered config fields select layer-aware RoPE and additional time conditioning.""" model = QwenImageLayeredTransformer2DModel( diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py index a32aef4f3ca0..cf87cefa25c8 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -15,6 +15,7 @@ import asyncio import base64 +import json import os from io import BytesIO from pathlib import Path @@ -138,9 +139,14 @@ def __init__( batch_aware: bool = True, validation_error: Optional[ValueError] = None, generate_error: Optional[BaseException] = None, + extra_param_specs: Optional[dict] = None, + model: str = "test-model", + supports_image_edit: bool = False, ): from types import SimpleNamespace + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + self._image = image_output self._video = video_output self._audio = audio_output @@ -150,6 +156,8 @@ def __init__( # Raised out of generate(): models an engine-side failure class, # where validation_error models a coordinator preflight rejection. self._generate_error = generate_error + self._extra_param_specs = extra_param_specs or {} + self._model = model self._healthy = True self._req_counter = 0 # Captured arguments of the most recent generate / generate_async call, @@ -164,8 +172,6 @@ def __init__( # reject legitimate width/height/num_frames/... requests; # ``extra_param_specs`` lists a single known key so tests can # exercise both the accept-known and reject-unknown paths. - from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema - self.executor = SimpleNamespace( default_generation_params={ "height": 64, @@ -176,9 +182,9 @@ def __init__( "num_frames": 8, "frame_rate": 8.0, }, - extra_param_specs={ - "stg_scale": ExtraParamSchema(type="float", default=1.0), - }, + extra_param_specs=extra_param_specs + or {"stg_scale": ExtraParamSchema(type="float", default=1.0)}, + supports_image_edit=supports_image_edit, ) def _maybe_batch(self, tensor, n): @@ -232,7 +238,7 @@ def default_params(self): seeds request params from this, so it must return a fresh instance.""" from tensorrt_llm.visual_gen import VisualGenParams - return VisualGenParams() + return VisualGenParams(**self.executor.default_generation_params) @property def extra_param_specs(self): @@ -240,12 +246,12 @@ def extra_param_specs(self): every request ``extra_params`` key reaches the executor as ``unknown_extra_param`` (matches a pipeline with no model-specific knobs declared, like Flux or Wan 2.1).""" - return {} + return self._extra_param_specs @property def model(self): """Stand-in for VisualGen.model — used by warn-on-set logic.""" - return "test-model" + return self._model def _check_health(self) -> bool: return self._healthy @@ -700,38 +706,397 @@ def test_image_generation_b64_with_4d_batch_pipeline_output(self, tmp_path): class TestImageEdit: - """``/v1/images/edits`` returns 501 NotImplemented in the current release. + """``/v1/images/edits`` support is gated by the loaded visual model.""" - No in-tree pipeline implements image editing: Flux/Flux2 are - text-to-image only and ignore ``params.image``; Wan and LTX-2 produce - video, not edited images. Restore the full happy-path coverage when an - edit-capable pipeline lands. - """ + def _client( + self, + tmp_path, + monkeypatch, + *, + image_output: Optional[torch.Tensor] = None, + extra_param_specs: Optional[dict] = None, + model: str = "Qwen/Qwen-Image-Layered", + should_fail: bool = False, + supports_image_edit: bool = True, + ): + gen = MockVisualGen( + image_output=image_output if image_output is not None else _make_dummy_image_tensor(), + extra_param_specs=extra_param_specs, + model=model, + should_fail=should_fail, + supports_image_edit=supports_image_edit, + ) + monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) + return _create_server(gen, model_name=model), gen - def test_image_edit_returns_not_implemented(self, image_client): - """Valid request body short-circuits to 501 NotImplemented.""" - b64_img = _b64_white_png_1x1() - resp = image_client.post( + @pytest.mark.parametrize( + ("model", "supports_image_edit", "expected_status"), + [ + ("not-a-canonical-edit-model-id", True, 200), + ("Qwen/Qwen-Image-Layered", False, 501), + ], + ) + def test_image_edit_support_uses_loaded_pipeline_capability( + self, tmp_path, monkeypatch, model, supports_image_edit, expected_status + ): + client, gen = self._client( + tmp_path, + monkeypatch, + model=model, + supports_image_edit=supports_image_edit, + ) + + resp = client.post( "/v1/images/edits", json={ - "image": b64_img, - "prompt": "Make it blue", - "num_inference_steps": 10, + "prompt": "Make it red", + "image": _b64_white_png_1x1(), + "response_format": "b64_json", }, ) - assert resp.status_code == 501 + + assert resp.status_code == expected_status + if expected_status == 501: + assert gen.last_params is None + + def test_image_edit_accepts_json_base64_image(self, tmp_path, monkeypatch): + """JSON edit requests materialize inputs and map OpenAI-shaped fields.""" + client, gen = self._client( + tmp_path, + monkeypatch, + image_output=_make_dummy_image_tensor(4, 4), + ) + + resp = client.post( + "/v1/images/edits", + content=json.dumps( + { + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "n": 2, + "output_format": "webp", + "response_format": "b64_json", + } + ), + headers={"content-type": "Application/JSON"}, + ) + + assert resp.status_code == 200 + assert str(gen.last_params.image).startswith(str(tmp_path)) + assert not os.path.exists(gen.last_params.image) + assert gen.last_params.num_images_per_prompt == 2 body = resp.json() - assert body.get("type") == "NotImplementedError" - assert "not supported" in body.get("message", "").lower() - - def test_image_edit_no_body_returns_not_implemented(self, image_client): - """The route doesn't parse a typed body; any incoming request still - gets 501, including ones that would have failed schema validation - before. Restore typed-body coverage when an edit pipeline lands.""" - resp = image_client.post("/v1/images/edits", json={"prompt": "Edit without image"}) - assert resp.status_code == 501 + assert body["output_format"] == "webp" + assert body["size"] == "4x4" + assert len(body["data"]) == 2 + + @pytest.mark.parametrize( + ("size", "expected_dimensions"), + [ + ("auto", (None, None)), + ("32x48", (32, 48)), + ], + ) + def test_image_edit_auto_size_allows_reference_size_derivation( + self, tmp_path, monkeypatch, size, expected_dimensions + ): + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "use reference dimensions", + "image": _b64_white_png_1x1(), + "size": size, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 200 + assert (gen.last_params.width, gen.last_params.height) == expected_dimensions + + @pytest.mark.threadleak(enabled=False) # FileResponse spawns AnyIO worker threads + def test_image_edit_default_url_returns_fetchable_output(self, tmp_path, monkeypatch): + """The default edit response writes a fetchable image content URL.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + }, + ) + + assert resp.status_code == 200 body = resp.json() - assert body.get("type") == "NotImplementedError" + url = body["data"][0]["url"] + assert "/v1/images/" in url and "/content" in url + assert str(gen.last_params.image).startswith(str(tmp_path)) + assert not os.path.exists(gen.last_params.image) + + path = url.split("//", 1)[-1].split("/", 1)[1] + content = client.get("/" + path) + assert content.status_code == 200 + assert content.content.startswith(b"\x89PNG\r\n\x1a\n") + assert content.headers["content-type"] == "image/png" + + def test_image_edit_rejects_json_array_body(self, tmp_path, monkeypatch): + """Non-object JSON bodies are client errors, not server errors.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + content=json.dumps( + [ + { + "prompt": "split layers", + "image": _b64_white_png_1x1(), + } + ] + ), + headers={"content-type": "application/json"}, + ) + + assert resp.status_code == 400 + assert "must be an object" in resp.json()["message"] + assert gen.last_params is None + + def test_image_edit_rejects_file_extra_params(self, tmp_path, monkeypatch): + """Multipart extra_params must be a JSON string field.""" + client, gen = self._client(tmp_path, monkeypatch) + + image_bytes = BytesIO(base64.b64decode(_b64_white_png_1x1())) + resp = client.post( + "/v1/images/edits", + data={ + "prompt": "split layers", + "response_format": "b64_json", + }, + files={ + "image": ("input.png", image_bytes, "image/png"), + "extra_params": ("extra.json", b"{}", "application/json"), + }, + ) + + assert resp.status_code == 400 + assert "extra_params" in resp.json()["message"] + assert gen.last_params is None + + def test_image_edit_rejects_empty_image_list(self, tmp_path, monkeypatch): + """Empty image lists fail request validation before pipeline dispatch.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": [], + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 422 + _assert_llm_envelope(resp.json(), code=422, message_contains="image") + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_rejects_non_image_base64_input(self, tmp_path, monkeypatch): + """Decoded image-edit bytes must be a supported image before disk write.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": base64.b64encode(b"not an image").decode("utf-8"), + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + _assert_llm_envelope( + resp.json(), + code=400, + message_contains="image edit input is not a PNG/JPEG image", + ) + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_rejects_non_image_upload_input(self, tmp_path, monkeypatch): + """Multipart image-edit bytes are sniffed before materialization.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + data={ + "prompt": "split layers", + "response_format": "b64_json", + }, + files={"image": ("input.png", BytesIO(b"not an image"), "image/png")}, + ) + + assert resp.status_code == 400 + _assert_llm_envelope( + resp.json(), + code=400, + message_contains="image edit input is not a PNG/JPEG image", + ) + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_rejects_mask_with_clear_error(self, tmp_path, monkeypatch): + """Mask is OpenAI-shaped but not implemented by TRTLLM image edit yet.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "mask": _b64_white_png_1x1(), + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + assert "mask input is not supported" in resp.json()["message"] + assert gen.last_params is None + + def test_image_edit_rejects_too_many_input_images(self, tmp_path, monkeypatch): + """Input image count is capped before files are materialized.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": [_b64_white_png_1x1()] * 17, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + assert "at most 16 input images" in resp.json()["message"] + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_allows_max_input_images_without_output_fanout(self, tmp_path, monkeypatch): + """Multiple edit inputs are joint conditioning, not output fan-out.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": [_b64_white_png_1x1()] * 16, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 200 + assert len(gen.last_params.image) == 16 + assert len(resp.json()["data"]) == 1 + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_rejects_excessive_output_fanout(self, tmp_path, monkeypatch): + """Layered output fan-out is capped before files are materialized.""" + from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + + client, gen = self._client( + tmp_path, + monkeypatch, + extra_param_specs={ + "layers": ExtraParamSchema(type="int", default=4, range=(1, 16)), + }, + ) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "n": 5, + "extra_params": {"layers": 16}, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + assert "at most 64 output images" in resp.json()["message"] + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_rejects_oversized_base64_image_before_decode(self, tmp_path, monkeypatch): + """Base64 image size is capped before allocating decoded bytes.""" + from tensorrt_llm.serve import visual_gen_utils + + client, gen = self._client(tmp_path, monkeypatch) + monkeypatch.setattr(visual_gen_utils, "IMAGE_EDIT_MAX_IMAGE_BYTES", 8) + monkeypatch.setattr( + visual_gen_utils.base64, + "b64decode", + lambda *args, **kwargs: pytest.fail("oversized payload was decoded"), + ) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": "A" * 13, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + assert "per-image byte limit" in resp.json()["message"] + assert gen.last_params is None + assert list(tmp_path.iterdir()) == [] + + def test_image_edit_cleans_inputs_when_generation_fails(self, tmp_path, monkeypatch): + """Temporary edit inputs are removed even when generation raises.""" + client, gen = self._client( + tmp_path, + monkeypatch, + should_fail=True, + ) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": _b64_white_png_1x1(), + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 500 + assert gen.last_params is not None + assert list(tmp_path.iterdir()) == [] + + @pytest.mark.parametrize( + "image_value", + [ + "/tmp/server-local-image.png", + "file:///tmp/server-local-image.png", + "https://example.com/server-local-image.png", + ], + ) + def test_image_edit_rejects_json_path_or_url_image(self, tmp_path, monkeypatch, image_value): + """Serving image-edit input strings must be base64, not server paths or URLs.""" + client, gen = self._client(tmp_path, monkeypatch) + + resp = client.post( + "/v1/images/edits", + json={ + "prompt": "split layers", + "image": image_value, + "response_format": "b64_json", + }, + ) + + assert resp.status_code == 400 + assert gen.last_params is None # ========================================================================= @@ -1874,6 +2239,7 @@ def test_seconds_without_frame_rate_returns_400(self, video_client): declares a ``frame_rate``: the parser must reject the request with HTTP 400 instead of silently dropping the duration and returning the pipeline's default ``num_frames``.""" + video_client.mock_gen.executor.default_generation_params.pop("frame_rate", None) resp = video_client.post( "/v1/videos/generations", json={