From 03030149077761dc0d920d00af8a541d508ea2fe Mon Sep 17 00:00:00 2001 From: Zhenhua Wang Date: Wed, 15 Apr 2026 06:51:50 +0000 Subject: [PATCH 1/2] [TRTLLM-11373][refactor] Embed VisualGenParams in DiffusionRequest and simplify generate() inputs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Embed VisualGenParams directly in DiffusionRequest (req.params) instead of flattening fields, eliminating field-by-field copy in VisualGen.generate(). DiffusionRequest becomes a thin envelope (request_id + prompt + params), while pipelines unpack individual values in infer() before calling forward(). DiffusionRequest.params defaults to None. When None, the executor's _merge_defaults() constructs VisualGenParams directly from pipeline defaults in one shot, avoiding ambiguity between VisualGenParams general defaults and pipeline-specific defaults. When the user provides params, only None fields are filled from pipeline defaults. Also simplify generate()/generate_async() inputs from VisualGenInputs (Union of str, dict, token IDs, sequences) to plain Union[str, List[str]]. Drop VisualGenTextPrompt, VisualGenTokensPrompt, VisualGenPromptInputs, VisualGenInputs, and visual_gen_inputs() — token ID input is not a meaningful concept for diffusion text encoders. Params default to None when omitted, and negative_prompt is passed exclusively via VisualGenParams. Updates openai_server, examples, and tests accordingly. Signed-off-by: Zhenhua Wang --- .gitignore | 1 + examples/visual_gen/visual_gen_ltx2.py | 4 +- examples/visual_gen/visual_gen_wan_i2v.py | 2 +- examples/visual_gen/visual_gen_wan_t2v.py | 2 +- tensorrt_llm/_torch/visual_gen/executor.py | 78 +++++++--------- .../visual_gen/models/flux/pipeline_flux.py | 12 +-- .../visual_gen/models/flux/pipeline_flux2.py | 12 +-- .../visual_gen/models/ltx2/pipeline_ltx2.py | 24 ++--- .../models/ltx2/pipeline_ltx2_two_stages.py | 24 ++--- .../visual_gen/models/wan/pipeline_wan.py | 18 ++-- .../visual_gen/models/wan/pipeline_wan_i2v.py | 22 ++--- tensorrt_llm/_torch/visual_gen/pipeline.py | 4 +- tensorrt_llm/inputs/data.py | 79 +--------------- tensorrt_llm/serve/openai_server.py | 56 +++-------- tensorrt_llm/serve/visual_gen_utils.py | 53 ++++++----- tensorrt_llm/visual_gen/visual_gen.py | 62 +++---------- .../visual_gen/test_trtllm_serve_endpoints.py | 8 ++ .../_torch/visual_gen/test_visual_gen_args.py | 93 ++++--------------- .../visual_gen/test_visual_gen_params.py | 77 +++++++++++---- 19 files changed, 241 insertions(+), 390 deletions(-) diff --git a/.gitignore b/.gitignore index 7cdbea45dfb6..bf32438a0f41 100644 --- a/.gitignore +++ b/.gitignore @@ -1,4 +1,5 @@ __pycache__/ +.mypy_cache/ .vscode .cursor *.engine diff --git a/examples/visual_gen/visual_gen_ltx2.py b/examples/visual_gen/visual_gen_ltx2.py index bf3815391ef7..1b753217c520 100755 --- a/examples/visual_gen/visual_gen_ltx2.py +++ b/examples/visual_gen/visual_gen_ltx2.py @@ -383,8 +383,6 @@ def main(): start_time = time.time() - inputs = {"prompt": args.prompt} - extra_params = { "guidance_rescale": args.guidance_rescale, "stg_scale": args.stg_scale, @@ -411,7 +409,7 @@ def main(): extra_params=extra_params, ) - output = visual_gen.generate(inputs=inputs, params=params) + output = visual_gen.generate(inputs=args.prompt, params=params) end_time = time.time() logger.info(f"Generation completed in {end_time - start_time:.2f}s") diff --git a/examples/visual_gen/visual_gen_wan_i2v.py b/examples/visual_gen/visual_gen_wan_i2v.py index 0433bd46a0b1..f2ac5d76c4ce 100644 --- a/examples/visual_gen/visual_gen_wan_i2v.py +++ b/examples/visual_gen/visual_gen_wan_i2v.py @@ -336,7 +336,7 @@ def main(): extra_params["boundary_ratio"] = args.boundary_ratio output = visual_gen.generate( - inputs={"prompt": args.prompt}, + inputs=args.prompt, params=VisualGenParams( height=args.height, width=args.width, diff --git a/examples/visual_gen/visual_gen_wan_t2v.py b/examples/visual_gen/visual_gen_wan_t2v.py index 3c89277d5f7a..22c62f608838 100755 --- a/examples/visual_gen/visual_gen_wan_t2v.py +++ b/examples/visual_gen/visual_gen_wan_t2v.py @@ -340,7 +340,7 @@ def main(): extra_params["boundary_ratio"] = args.boundary_ratio output = visual_gen.generate( - inputs={"prompt": args.prompt}, + inputs=args.prompt, params=VisualGenParams( height=args.height, width=args.width, diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 7c8bf053276c..753e20312204 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -3,7 +3,7 @@ import threading import traceback from dataclasses import dataclass -from typing import List, Optional, Union +from typing import TYPE_CHECKING, List, Optional import torch import torch.distributed as dist @@ -15,42 +15,24 @@ from tensorrt_llm.executor.ipc import ZeroMqQueue from tensorrt_llm.logger import logger +if TYPE_CHECKING: + from tensorrt_llm.visual_gen.params import VisualGenParams + @dataclass class DiffusionRequest: """Request for diffusion inference. - Universal parameters are top-level fields with ``None`` meaning - "use model default" (resolved by the executor before calling - ``pipeline.infer()``). Model-specific parameters live in - ``extra_params`` and are passed through to the pipeline. + Generation parameters live in the optional ``params`` object + (a :class:`~tensorrt_llm.visual_gen.params.VisualGenParams` instance). + When ``params`` is ``None`` (the default), the executor creates a + ``VisualGenParams()`` and fills it with pipeline-specific defaults + before calling ``pipeline.infer()``. """ request_id: int prompt: List[str] - negative_prompt: Optional[str] = None - - # Core — None means "use model default" (resolved by executor) - height: Optional[int] = None - width: Optional[int] = None - num_inference_steps: Optional[int] = None - guidance_scale: Optional[float] = None - max_sequence_length: Optional[int] = None - seed: int = 42 - - # Video - num_frames: Optional[int] = None - frame_rate: Optional[float] = None - - # Image - num_images_per_prompt: int = 1 - - # Conditioning inputs - image: Optional[Union[str, bytes, List[Union[str, bytes]]]] = None - image_cond_strength: Optional[float] = None - - # Model-specific overflow (from VisualGenParams.extra_params) - extra_params: Optional[dict] = None + params: Optional["VisualGenParams"] = None @dataclass @@ -216,29 +198,40 @@ def serve_forever(self): self.process_request(req) def _merge_defaults(self, req: DiffusionRequest): - """Fill ``None`` fields in *req* with pipeline-specific defaults. + """Fill ``None`` fields in *req.params* with pipeline-specific defaults. Merges both universal defaults (from ``default_generation_params``) and extra_param defaults (from ``extra_param_specs``). """ + if req.params is None: + from tensorrt_llm.visual_gen.params import VisualGenParams + + kwargs = dict(self.pipeline.default_generation_params) + specs = self.pipeline.extra_param_specs + if specs: + kwargs["extra_params"] = {key: spec.default for key, spec in specs.items()} + req.params = VisualGenParams(**kwargs) + return + + params = req.params # Universal field defaults for field_name, default_value in self.pipeline.default_generation_params.items(): - if hasattr(req, field_name) and getattr(req, field_name) is None: - setattr(req, field_name, default_value) + if hasattr(params, field_name) and getattr(params, field_name) is None: + setattr(params, field_name, default_value) # Extra param defaults — fill all declared keys so infer() can use direct access specs = self.pipeline.extra_param_specs if specs: - if req.extra_params is None: - req.extra_params = {} + if params.extra_params is None: + params.extra_params = {} for key, spec in specs.items(): - if key not in req.extra_params: - req.extra_params[key] = spec.default + if key not in params.extra_params: + params.extra_params[key] = spec.default self._validate_request(req) def _validate_request(self, req: DiffusionRequest): - """Validate *req* against the loaded pipeline's declared parameters. + """Validate *req.params* against the loaded pipeline's declared parameters. Raises ``VisualGenParamsError`` on: - Unknown ``extra_params`` keys @@ -251,14 +244,15 @@ def _validate_request(self, req: DiffusionRequest): # (executor → visual_gen.visual_gen → _torch.visual_gen → executor) from tensorrt_llm.visual_gen.visual_gen import VisualGenParamsError + params = req.params errors: list[str] = [] pipeline_name = self.pipeline.__class__.__name__ declared_defaults = self.pipeline.default_generation_params specs = self.pipeline.extra_param_specs # --- unknown extra_params keys --- - if req.extra_params: - unknown = set(req.extra_params.keys()) - set(specs.keys()) + if params.extra_params: + unknown = set(params.extra_params.keys()) - set(specs.keys()) if unknown: errors.append( f"Unknown extra_params {sorted(unknown)} for {pipeline_name}. " @@ -271,7 +265,7 @@ def _validate_request(self, req: DiffusionRequest): # Conditioning inputs (image, negative_prompt, mask) are excluded — # they are validated at runtime by the pipeline's infer(). for field_name in _GENERATION_CONFIG_FIELDS: - value = getattr(req, field_name, None) + value = getattr(params, field_name, None) if value is not None and field_name not in declared_defaults: errors.append( f"Parameter '{field_name}' is set but {pipeline_name} does " @@ -280,8 +274,8 @@ def _validate_request(self, req: DiffusionRequest): ) # --- extra_params type and range checks --- - if req.extra_params: - for key, value in req.extra_params.items(): + if params.extra_params: + for key, value in params.extra_params.items(): if key not in specs: continue # already reported as unknown above spec = specs[key] @@ -315,7 +309,7 @@ def process_request(self, req: DiffusionRequest): try: self._merge_defaults(req) cache_key = self.pipeline.warmup_cache_key( - req.height, req.width, num_frames=req.num_frames + 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: logger.warning( 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 8e496cb4e40e..23dbe4620886 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux.py @@ -246,12 +246,12 @@ def infer(self, req): """Run inference from DiffusionRequest.""" return self.forward( prompt=req.prompt, - height=req.height, - width=req.width, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, - seed=req.seed, - max_sequence_length=req.max_sequence_length, + height=req.params.height, + width=req.params.width, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, + max_sequence_length=req.params.max_sequence_length, ) @torch.inference_mode() 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 a1fbb19f18fc..ef9445ba7fe6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -339,12 +339,12 @@ def infer(self, req): """Run inference from DiffusionRequest.""" return self.forward( prompt=req.prompt, - height=req.height, - width=req.width, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, - seed=req.seed, - max_sequence_length=req.max_sequence_length, + height=req.params.height, + width=req.params.width, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, + max_sequence_length=req.params.max_sequence_length, ) @torch.inference_mode() diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py index 5fa495481e93..aaf37afed602 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1170,22 +1170,22 @@ def extra_param_specs(self): def infer(self, req): """Run inference with request parameters.""" - extra = req.extra_params or {} + extra = req.params.extra_params or {} return self.forward( prompt=req.prompt, - negative_prompt=req.negative_prompt, - height=req.height, - width=req.width, - num_frames=req.num_frames, - frame_rate=req.frame_rate, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, - seed=req.seed, + negative_prompt=req.params.negative_prompt, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + frame_rate=req.params.frame_rate, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], - max_sequence_length=req.max_sequence_length, - image=req.image, - image_cond_strength=req.image_cond_strength, + max_sequence_length=req.params.max_sequence_length, + image=req.params.image, + image_cond_strength=req.params.image_cond_strength, stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], modality_scale=extra["modality_scale"], diff --git a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py index 40887e77515a..c92080d2a70c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2_two_stages.py @@ -565,22 +565,22 @@ def load_standard_components( # ------------------------------------------------------------------ def infer(self, req): - extra = req.extra_params or {} + extra = req.params.extra_params or {} return self.forward( prompt=req.prompt, - negative_prompt=req.negative_prompt, - height=req.height, - width=req.width, - num_frames=req.num_frames, - frame_rate=req.frame_rate, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, - seed=req.seed, + negative_prompt=req.params.negative_prompt, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + frame_rate=req.params.frame_rate, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], - max_sequence_length=req.max_sequence_length, - image=req.image, - image_cond_strength=req.image_cond_strength, + max_sequence_length=req.params.max_sequence_length, + image=req.params.image, + image_cond_strength=req.params.image_cond_strength, stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], modality_scale=extra["modality_scale"], diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py index afa1510d3413..b851f142e1de 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -325,19 +325,19 @@ def extra_param_specs(self): def infer(self, req): """Run inference with request parameters.""" - extra = req.extra_params or {} + extra = req.params.extra_params or {} return self.forward( prompt=req.prompt, - negative_prompt=req.negative_prompt, - height=req.height, - width=req.width, - num_frames=req.num_frames, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, + negative_prompt=req.params.negative_prompt, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, guidance_scale_2=extra.get("guidance_scale_2"), boundary_ratio=extra.get("boundary_ratio"), - seed=req.seed, - max_sequence_length=req.max_sequence_length, + seed=req.params.seed, + max_sequence_length=req.params.max_sequence_length, ) @nvtx_range("WanPipeline.forward") diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py index edc74783721b..920c94f0faba 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py @@ -393,11 +393,11 @@ def extra_param_specs(self): def infer(self, req): """Run inference with request parameters.""" # Extract image from request (can be path, PIL Image, or torch.Tensor) - if req.image is None: + if req.params.image is None: raise ValueError("I2V pipeline requires 'image' parameter") - image = req.image[0] if isinstance(req.image, list) else req.image - extra = req.extra_params or {} + image = req.params.image[0] if isinstance(req.params.image, list) else req.params.image + extra = req.params.extra_params or {} last_image = extra.get("last_image") if last_image is not None and isinstance(last_image, list): @@ -406,16 +406,16 @@ def infer(self, req): return self.forward( image=image, prompt=req.prompt, - negative_prompt=req.negative_prompt, - height=req.height, - width=req.width, - num_frames=req.num_frames, - num_inference_steps=req.num_inference_steps, - guidance_scale=req.guidance_scale, + negative_prompt=req.params.negative_prompt, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, guidance_scale_2=extra.get("guidance_scale_2"), boundary_ratio=extra.get("boundary_ratio"), - seed=req.seed, - max_sequence_length=req.max_sequence_length, + seed=req.params.seed, + max_sequence_length=req.params.max_sequence_length, last_image=last_image, ) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 947b29efb547..cd9d59b547f0 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -256,8 +256,8 @@ def extra_param_specs(self) -> Dict[str, ExtraParamSchema]: def default_generation_params(self) -> dict: """Model-specific defaults for ``None`` fields in ``VisualGenParams``. - Keys should match ``DiffusionRequest`` field names. The executor - merges these into the request before calling ``infer()``. + Keys should match ``VisualGenParams`` field names. The executor + merges these into ``request.params`` before calling ``infer()``. """ return {} diff --git a/tensorrt_llm/inputs/data.py b/tensorrt_llm/inputs/data.py index 48e3441df665..615043fe4878 100644 --- a/tensorrt_llm/inputs/data.py +++ b/tensorrt_llm/inputs/data.py @@ -1,6 +1,6 @@ # Adapt from # https://github.com/vllm-project/vllm/blob/2e33fe419186c65a18da6668972d61d7bbc31564/vllm/inputs/data.py -from typing import Any, Dict, List, Sequence, Union +from typing import Any, Dict, List, Union from typing_extensions import NotRequired, TypedDict @@ -85,80 +85,3 @@ def prompt_inputs(inputs: PromptInputs, ) -> Union[TextPrompt, TokensPrompt]: f"Invalid type of inputs for llm.generate: {type(inputs)}") return prompt_inputs - - -class VisualGenTextPrompt(TypedDict): - prompt: str - negative_prompt: NotRequired[str] - - -class VisualGenTokensPrompt(TypedDict): - prompt_token_ids: List[int] - negative_prompt_token_ids: NotRequired[List[int]] - - -VisualGenPromptInputs = Union[ - str, - List[int], - VisualGenTextPrompt, - VisualGenTokensPrompt, -] - -VisualGenInputs = Union[ - VisualGenPromptInputs, - Sequence[VisualGenPromptInputs], -] - - -def visual_gen_inputs( - inputs: "VisualGenPromptInputs", -) -> Union["VisualGenTextPrompt", "VisualGenTokensPrompt"]: - # str -> text prompt - if isinstance(inputs, str): - return VisualGenTextPrompt(prompt=inputs) - - # list[int] -> token prompt - if isinstance(inputs, list): - if len(inputs) == 0: - raise ValueError("`inputs` token list cannot be empty.") - if not all(isinstance(t, int) for t in inputs): - raise TypeError( - "`inputs` list must contain only ints when used as token IDs.") - return VisualGenTokensPrompt(prompt_token_ids=inputs) - - # dict form - if isinstance(inputs, dict): - has_prompt = "prompt" in inputs - has_prompt_token_ids = "prompt_token_ids" in inputs - - if has_prompt == has_prompt_token_ids: - raise ValueError( - "VisualGen prompt dict must contain exactly one of " - "`prompt` or `prompt_token_ids`.") - - if has_prompt: - prompt = inputs.get("prompt") - if not isinstance(prompt, str) or prompt == "": - raise TypeError("`prompt` must be a non-empty string.") - if "negative_prompt" in inputs and not isinstance( - inputs["negative_prompt"], str): - raise TypeError("`negative_prompt` must be a string.") - return inputs # VisualGenTextPrompt - - token_ids = inputs.get("prompt_token_ids") - if not isinstance(token_ids, list) or len(token_ids) == 0: - raise TypeError("`prompt_token_ids` must be a non-empty list[int].") - if not all(isinstance(t, int) for t in token_ids): - raise TypeError("`prompt_token_ids` must contain only ints.") - if "negative_prompt_token_ids" in inputs: - neg_ids = inputs["negative_prompt_token_ids"] - if not isinstance(neg_ids, list) or not all( - isinstance(t, int) for t in neg_ids): - raise TypeError( - "`negative_prompt_token_ids` must be a list[int].") - return inputs # VisualGenTokensPrompt - - raise TypeError( - "Invalid `inputs` for VisualGen.generate. " - "Expected one of: str, list[int], VisualGenTextPrompt, VisualGenTokensPrompt." - ) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index e248ac96b1fa..2b0077ad57e9 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -33,7 +33,7 @@ from tensorrt_llm.executor import CppExecutorError from tensorrt_llm.executor.postproc_worker import PostprocParams from tensorrt_llm.inputs import prompt_inputs -from tensorrt_llm.inputs.data import TokensPrompt, visual_gen_inputs +from tensorrt_llm.inputs.data import TokensPrompt from tensorrt_llm.inputs.multimodal import MultimodalServerConfig from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams @@ -1771,21 +1771,13 @@ async def openai_image_generation(self, request: ImageGenerationRequest, """ try: image_id = f"image_{uuid.uuid4().hex}" - params = parse_visual_gen_params(request, image_id) + params = parse_visual_gen_params(request, image_id, self.generator) logger.info( f"Generating image: {image_id} with params: {params} and prompt: {request.prompt}" ) - if request.negative_prompt is not None: - inputs = visual_gen_inputs({ - "prompt": - request.prompt, - "negative_prompt": - request.negative_prompt - }) - else: - inputs = visual_gen_inputs(request.prompt) - output = self.generator.generate(inputs=inputs, params=params) + output = self.generator.generate(inputs=request.prompt, + params=params) if output.image is None: return self.create_error_response( message="Image generation failed", @@ -1836,21 +1828,13 @@ async def openai_image_edit(self, request: ImageEditRequest, """ try: image_id = f"image_{uuid.uuid4().hex}" - params = parse_visual_gen_params(request, image_id) + params = parse_visual_gen_params(request, image_id, self.generator) logger.info( f"Editing image: {image_id} with params: {params} and prompt: {request.prompt}" ) - if request.negative_prompt is not None: - inputs = visual_gen_inputs({ - "prompt": - request.prompt, - "negative_prompt": - request.negative_prompt - }) - else: - inputs = visual_gen_inputs(request.prompt) - output = self.generator.generate(inputs=inputs, params=params) + output = self.generator.generate(inputs=request.prompt, + params=params) if output.image is None: return self.create_error_response( message="Image editing failed", @@ -1905,22 +1889,15 @@ async def openai_video_generation_sync(self, video_id = f"video_{uuid.uuid4().hex}" params = parse_visual_gen_params(request, video_id, + self.generator, media_storage_path=str( self.media_storage_path)) logger.info( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) - if request.negative_prompt is not None: - inputs = visual_gen_inputs({ - "prompt": - request.prompt, - "negative_prompt": - request.negative_prompt - }) - else: - inputs = visual_gen_inputs(request.prompt) - output = self.generator.generate(inputs=inputs, params=params) + output = self.generator.generate(inputs=request.prompt, + params=params) if output.video is None: return self.create_error_response( message="Video generation failed", @@ -2039,6 +2016,7 @@ async def openai_video_generation_async( video_id = f"video_{uuid.uuid4().hex}" params = parse_visual_gen_params(request, video_id, + self.generator, media_storage_path=str( self.media_storage_path)) logger.info( @@ -2087,16 +2065,8 @@ async def _generate_video_background( resolved_fmt, resolved_ext = resolve_video_format( request.output_format) - if request.negative_prompt is not None: - inputs = visual_gen_inputs({ - "prompt": - request.prompt, - "negative_prompt": - request.negative_prompt - }) - else: - inputs = visual_gen_inputs(request.prompt) - future = self.generator.generate_async(inputs=inputs, params=params) + future = self.generator.generate_async(inputs=request.prompt, + params=params) output = await future.result() if output.video is None: diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 27aacf33766a..f231156170fe 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -9,47 +9,52 @@ ImageGenerationRequest, VideoGenerationRequest, ) -from tensorrt_llm.visual_gen import VisualGenParams +from tensorrt_llm.visual_gen import VisualGen, VisualGenParams def parse_visual_gen_params( request: ImageGenerationRequest | VideoGenerationRequest | ImageEditRequest, id: str, + generator: VisualGen, media_storage_path: Optional[str] = None, ) -> VisualGenParams: - kwargs: Dict[str, Any] = {} - extra: Dict[str, Any] = {} - - kwargs["negative_prompt"] = request.negative_prompt + # Start from the pipeline's resolved defaults so unspecified request + # fields keep the model's defaults instead of being overwritten with None. + params = generator.default_params + if params.extra_params is None: + params.extra_params = {} + + if request.negative_prompt is not None: + params.negative_prompt = request.negative_prompt if request.size is not None and request.size != "auto": - kwargs["width"], kwargs["height"] = map(int, request.size.split("x")) + params.width, params.height = map(int, request.size.split("x")) if request.guidance_scale is not None: - kwargs["guidance_scale"] = request.guidance_scale + params.guidance_scale = request.guidance_scale if request.guidance_rescale is not None: - extra["guidance_rescale"] = request.guidance_rescale + params.extra_params["guidance_rescale"] = request.guidance_rescale if isinstance(request, (ImageGenerationRequest, ImageEditRequest)): if request.num_inference_steps is not None: - kwargs["num_inference_steps"] = request.num_inference_steps + params.num_inference_steps = request.num_inference_steps elif isinstance(request, ImageGenerationRequest) and request.quality == "hd": - kwargs["num_inference_steps"] = 30 + params.num_inference_steps = 30 if request.n is not None: - kwargs["num_images_per_prompt"] = request.n + params.num_images_per_prompt = request.n if isinstance(request, ImageEditRequest): if request.image is not None: if isinstance(request.image, list): - kwargs["image"] = [base64.b64decode(image) for image in request.image] + params.image = [base64.b64decode(image) for image in request.image] else: - kwargs["image"] = [base64.b64decode(request.image)] + params.image = [base64.b64decode(request.image)] if request.mask is not None: if isinstance(request.mask, list): - kwargs["mask"] = [base64.b64decode(mask) for mask in request.mask] + params.mask = [base64.b64decode(mask) for mask in request.mask] else: - kwargs["mask"] = base64.b64decode(request.mask) + params.mask = base64.b64decode(request.mask) elif isinstance(request, VideoGenerationRequest): if request.num_inference_steps is not None: - kwargs["num_inference_steps"] = request.num_inference_steps + params.num_inference_steps = request.num_inference_steps if request.input_reference is not None: if media_storage_path is None: raise ValueError("media_storage_path is required when input_reference is provided") @@ -60,18 +65,20 @@ def parse_visual_gen_params( else: with open(ref_path, "wb") as f: shutil.copyfileobj(request.input_reference.file, f) - kwargs["image"] = ref_path + params.image = ref_path - kwargs["frame_rate"] = request.fps - kwargs["num_frames"] = int(request.seconds * request.fps) + params.frame_rate = request.fps + params.num_frames = int(request.seconds * request.fps) if request.seed is not None: - kwargs["seed"] = int(request.seed) + params.seed = int(request.seed) - if extra: - kwargs["extra_params"] = extra + # Drop extra_params if we didn't end up with any — matches VisualGenParams + # convention where None means "no extras" for pipelines that declare none. + if not params.extra_params: + params.extra_params = None - return VisualGenParams(**kwargs) + return params class AsyncDictStore: diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 80d5e5974d0d..9d1d3cfa762c 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -45,7 +45,6 @@ "VisualGenResult", ] from tensorrt_llm.executor.ipc import ZeroMqQueue -from tensorrt_llm.inputs.data import VisualGenInputs from tensorrt_llm.llmapi.utils import set_api_status from tensorrt_llm.logger import logger @@ -538,13 +537,14 @@ def default_params(self) -> "VisualGenParams": @set_api_status("prototype") def generate( self, - inputs: VisualGenInputs, - params: VisualGenParams, + inputs: Union[str, List[str]], + params: Optional[VisualGenParams] = None, ) -> MediaOutput: """Synchronous generation. Blocks until complete. Args: - params: Generation parameters. + inputs: Text prompt string or list of prompt strings. + params: Generation parameters (optional; uses model defaults when None). Returns: MediaOutput: Generated media with model-specific fields populated: @@ -566,70 +566,36 @@ def generate( @set_api_status("prototype") def generate_async( self, - inputs: VisualGenInputs, - params: VisualGenParams, + inputs: Union[str, List[str]], + params: Optional[VisualGenParams] = None, ) -> VisualGenResult: """Async generation. Returns immediately with future-like object. Args: - params: Generation parameters. + inputs: Text prompt string or list of prompt strings. + params: Generation parameters (optional; uses model defaults when None). Returns: VisualGenResult: Call result() to get output dict. """ req_id = next(self._req_counter) - # Normalize inputs to (prompt: List[str], negative_prompt: Optional[str]) - # so DiffusionRequest.prompt is always a list. - if isinstance(inputs, dict): - prompt = [inputs.get("prompt")] - negative_prompt = inputs.get("negative_prompt", None) - elif isinstance(inputs, str): + # Normalize to List[str] for DiffusionRequest.prompt + if isinstance(inputs, str): prompt = [inputs] - negative_prompt = None elif isinstance(inputs, (list, tuple)): - # Batch generation: list of prompts if not inputs: raise ValueError("Batch inputs must contain at least one item") - - prompt = [] - negative_prompts = [] - for idx, inp in enumerate(inputs): - if isinstance(inp, str): - prompt.append(inp) - negative_prompts.append(None) - elif isinstance(inp, dict): - item_prompt = inp.get("prompt") - if item_prompt is None: - raise ValueError(f"Batch input at index {idx} is missing 'prompt'") - prompt.append(item_prompt) - negative_prompts.append(inp.get("negative_prompt")) - else: - raise ValueError(f"Invalid batch item type at index {idx}: {type(inp)}") - - unique_negatives = {p for p in negative_prompts if p is not None} - if len(unique_negatives) > 1: - raise ValueError("Per-item negative_prompt is not supported for batch inputs") - negative_prompt = next(iter(unique_negatives), None) + if not all(isinstance(item, str) for item in inputs): + raise ValueError("Batch inputs must contain only strings (prompt text)") + prompt = list(inputs) else: raise ValueError(f"Invalid inputs type: {type(inputs)}") request = DiffusionRequest( request_id=req_id, prompt=prompt, - negative_prompt=negative_prompt, - height=params.height, - width=params.width, - num_inference_steps=params.num_inference_steps, - guidance_scale=params.guidance_scale, - max_sequence_length=params.max_sequence_length, - seed=params.seed, - num_frames=params.num_frames, - frame_rate=params.frame_rate, - num_images_per_prompt=params.num_images_per_prompt, - image=params.image, - image_cond_strength=params.image_cond_strength, - extra_params=params.extra_params, + params=params, ) self.executor.enqueue_requests([request]) 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 c161d05ebdba..0ce92f9b7243 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -109,6 +109,14 @@ def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": should_fail=self._should_fail, ) + @property + def default_params(self): + """Stand-in for VisualGen.default_params — parse_visual_gen_params + seeds request params from this, so it must return a fresh instance.""" + from tensorrt_llm.visual_gen import VisualGenParams + + return VisualGenParams() + def _check_health(self) -> bool: return self._healthy diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py index 5a3f83a3b7c0..551fe39a2a93 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_args.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_args.py @@ -239,8 +239,8 @@ def test_list_prompt(self): assert len(req.prompt) == 2 -class TestVisualGenBatchInputParsing: - """Test that VisualGen.generate_async() correctly parses batch inputs. +class TestVisualGenInputParsing: + """Test that VisualGen.generate_async() correctly parses inputs. Uses mocking to avoid spawning GPU worker processes. """ @@ -258,116 +258,59 @@ def _make_visual_gen_with_mock_executor(self): vg.executor = MagicMock() return vg - def _make_params(self): - from tensorrt_llm.visual_gen import VisualGenParams - - return VisualGenParams() - def test_string_input(self): """String input → single-element list in DiffusionRequest.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - vg.generate_async(inputs="a cat", params=params) + vg.generate_async(inputs="a cat") - # Check the DiffusionRequest passed to enqueue_requests call_args = vg.executor.enqueue_requests.call_args[0][0] assert len(call_args) == 1 assert call_args[0].prompt == ["a cat"] - def test_dict_input(self): - """Dict input → prompt wrapped in list + negative_prompt extracted.""" - vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - - vg.generate_async( - inputs={"prompt": "a cat", "negative_prompt": "blurry"}, - params=params, - ) - - call_args = vg.executor.enqueue_requests.call_args[0][0] - assert call_args[0].prompt == ["a cat"] - assert call_args[0].negative_prompt == "blurry" - def test_list_of_strings_input(self): """List of strings → batch prompt in single DiffusionRequest.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - vg.generate_async(inputs=["a sunset", "a city"], params=params) + vg.generate_async(inputs=["a sunset", "a city"]) call_args = vg.executor.enqueue_requests.call_args[0][0] assert len(call_args) == 1 req = call_args[0] assert req.prompt == ["a sunset", "a city"] - assert req.negative_prompt is None - def test_list_of_dicts_input(self): - """List of dicts → batch prompts with negative_prompt from first dict.""" + def test_params_default_none(self): + """Omitting params passes None; executor materializes defaults later.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - - vg.generate_async( - inputs=[ - {"prompt": "a sunset", "negative_prompt": "dark"}, - {"prompt": "a city"}, - ], - params=params, - ) + + vg.generate_async(inputs="a cat") call_args = vg.executor.enqueue_requests.call_args[0][0] req = call_args[0] - assert req.prompt == ["a sunset", "a city"] - assert req.negative_prompt == "dark" + assert req.params is None + + def test_negative_prompt_via_params(self): + """negative_prompt is passed through params, not inputs.""" + from tensorrt_llm.visual_gen import VisualGenParams - def test_mixed_list_input(self): - """Mixed list of strings and dicts → batch prompts extracted.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() + params = VisualGenParams(negative_prompt="blurry") - vg.generate_async(inputs=["a sunset", {"prompt": "a city"}], params=params) + vg.generate_async(inputs="a cat", params=params) call_args = vg.executor.enqueue_requests.call_args[0][0] - req = call_args[0] - assert req.prompt == ["a sunset", "a city"] - - def test_conflicting_negative_prompt_raises(self): - """Conflicting per-item negative_prompt raises ValueError.""" - vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - - with pytest.raises(ValueError, match="Per-item negative_prompt is not supported"): - vg.generate_async( - inputs=[ - {"prompt": "a sunset", "negative_prompt": "dark"}, - {"prompt": "a city", "negative_prompt": "light"}, - ], - params=params, - ) + assert call_args[0].params.negative_prompt == "blurry" def test_empty_batch_raises(self): """Empty batch input raises ValueError.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() with pytest.raises(ValueError, match="at least one item"): - vg.generate_async(inputs=[], params=params) - - def test_missing_prompt_in_dict_raises(self): - """Dict without 'prompt' key raises ValueError.""" - vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() - - with pytest.raises(ValueError, match="missing 'prompt'"): - vg.generate_async( - inputs=[{"negative_prompt": "dark"}], - params=params, - ) + vg.generate_async(inputs=[]) def test_invalid_input_raises(self): """Invalid input type raises ValueError.""" vg = self._make_visual_gen_with_mock_executor() - params = self._make_params() with pytest.raises(ValueError, match="Invalid inputs type"): - vg.generate_async(inputs=12345, params=params) + vg.generate_async(inputs=12345) 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 8521592a26af..4eb70522226e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -300,8 +300,9 @@ def _make_mock_executor(self, pipeline_cls, mock_self=None): def _make_request(self, **kwargs): from tensorrt_llm._torch.visual_gen.executor import DiffusionRequest + from tensorrt_llm.visual_gen.params import VisualGenParams - return DiffusionRequest(request_id=0, prompt=["test"], **kwargs) + return DiffusionRequest(request_id=0, prompt=["test"], params=VisualGenParams(**kwargs)) def _merge(self, executor, req): from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor @@ -313,12 +314,12 @@ def test_universal_defaults_merged(self): executor = self._make_mock_executor(WanPipeline, _wan_mock(is_wan22=False, num_heads=12)) req = self._make_request() - assert req.height is None + assert req.params.height is None self._merge(executor, req) - assert req.height == 480 - assert req.width == 832 - assert req.num_inference_steps == 50 + assert req.params.height == 480 + assert req.params.width == 832 + assert req.params.num_inference_steps == 50 def test_user_values_not_overwritten(self): from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan import WanPipeline @@ -327,9 +328,9 @@ def test_user_values_not_overwritten(self): req = self._make_request(height=1080, width=1920) self._merge(executor, req) - assert req.height == 1080 # User value preserved - assert req.width == 1920 - assert req.num_inference_steps == 50 # Default filled + assert req.params.height == 1080 # User value preserved + assert req.params.width == 1920 + assert req.params.num_inference_steps == 50 # Default filled def test_extra_params_defaults_merged(self): from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -338,12 +339,12 @@ def test_extra_params_defaults_merged(self): req = self._make_request() self._merge(executor, req) - assert req.extra_params is not None - assert req.extra_params["stg_scale"] == 0.0 - assert req.extra_params["output_type"] == "pt" - assert req.extra_params["enhance_prompt"] is False + assert req.params.extra_params is not None + assert req.params.extra_params["stg_scale"] == 0.0 + assert req.params.extra_params["output_type"] == "pt" + assert req.params.extra_params["enhance_prompt"] is False # None defaults are also filled - assert req.extra_params["stg_blocks"] is None + assert req.params.extra_params["stg_blocks"] is None def test_user_extra_params_not_overwritten(self): from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline @@ -352,8 +353,8 @@ def test_user_extra_params_not_overwritten(self): req = self._make_request(extra_params={"stg_scale": 0.5}) self._merge(executor, req) - assert req.extra_params["stg_scale"] == 0.5 # User value preserved - assert req.extra_params["output_type"] == "pt" # Default filled + assert req.params.extra_params["stg_scale"] == 0.5 # User value preserved + assert req.params.extra_params["output_type"] == "pt" # Default filled def test_no_extra_params_for_flux(self): from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux import FluxPipeline @@ -362,7 +363,7 @@ def test_no_extra_params_for_flux(self): req = self._make_request() self._merge(executor, req) - assert req.extra_params is None # Flux has no extra specs + assert req.params.extra_params is None # Flux has no extra specs def test_all_declared_keys_present_after_merge(self): """After merge, all extra_param_specs keys are in extra_params.""" @@ -374,7 +375,30 @@ def test_all_declared_keys_present_after_merge(self): self._merge(executor, req) ltx2_specs = LTX2Pipeline.extra_param_specs.fget(None) for key in ltx2_specs: - assert key in req.extra_params, f"Missing key: {key}" + assert key in req.params.extra_params, f"Missing key: {key}" + + def test_params_none_materializes_defaults(self): + """req.params=None is the default path from generate_async(params=None); + _merge_defaults should materialize a VisualGenParams from pipeline defaults.""" + from tensorrt_llm._torch.visual_gen.executor import DiffusionRequest + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from tensorrt_llm.visual_gen.params import VisualGenParams + + executor = self._make_mock_executor(LTX2Pipeline) + req = DiffusionRequest(request_id=0, prompt=["test"], params=None) + + self._merge(executor, req) + + assert isinstance(req.params, VisualGenParams) + # Universal defaults are filled from the pipeline + assert req.params.height == 512 + assert req.params.width == 768 + assert req.params.num_inference_steps == 40 + # Extra-param defaults are filled for all declared keys + assert req.params.extra_params is not None + assert req.params.extra_params["stg_scale"] == 0.0 + assert req.params.extra_params["output_type"] == "pt" + assert "stg_blocks" in req.params.extra_params # ============================================================================= @@ -658,8 +682,9 @@ def _make_mock_executor(self, pipeline_cls, mock_self=None): def _make_request(self, **kwargs): from tensorrt_llm._torch.visual_gen.executor import DiffusionRequest + from tensorrt_llm.visual_gen.params import VisualGenParams - return DiffusionRequest(request_id=0, prompt=["test"], **kwargs) + return DiffusionRequest(request_id=0, prompt=["test"], params=VisualGenParams(**kwargs)) def _validate(self, executor, req): from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor @@ -757,6 +782,22 @@ def test_none_fields_not_flagged(self): req = self._make_request() # all None self._merge_and_validate(executor, req) + def test_params_none_merge_and_validate_ok(self): + """req.params=None must merge + validate cleanly (VisualGen.generate_async + defaults to params=None, so this is the canonical call path).""" + from tensorrt_llm._torch.visual_gen.executor import DiffusionRequest + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline + from tensorrt_llm.visual_gen.params import VisualGenParams + + executor = self._make_mock_executor(LTX2Pipeline) + req = DiffusionRequest(request_id=0, prompt=["test"], params=None) + + self._merge_and_validate(executor, req) # should not raise + + assert isinstance(req.params, VisualGenParams) + assert req.params.height == 512 + assert req.params.extra_params["stg_scale"] == 0.0 + # --- type validation on extra_params --- def test_wrong_type_extra_param_raises(self): From 8b6ee1c82029a970fd34d59d25405e0c0cb91d0d Mon Sep 17 00:00:00 2001 From: Zhenhua Wang Date: Fri, 24 Apr 2026 02:33:50 +0000 Subject: [PATCH 2/2] [TRTLLM-11373][fix] Snapshot VisualGenParams on enqueue and tighten endpoint tests Address review feedback on PR #13313: - generate_async now deep-copies caller-provided VisualGenParams before enqueueing, so later mutations by the caller cannot alter the in-flight request (the dispatcher thread serializes it lazily via zmq). - MockVisualGen in the serve-endpoint tests now records the last inputs and params; tests for image generation, image edit (with/without mask), sync video (with params / with reference file), and async video assert that openai_server and parse_visual_gen_params forward every field (width/height, guidance_scale, num_inference_steps, negative_prompt, seed, frame_rate/num_frames, image/mask, input_reference). Signed-off-by: Zhenhua Wang --- tensorrt_llm/visual_gen/visual_gen.py | 4 +- .../visual_gen/test_trtllm_serve_endpoints.py | 96 ++++++++++++++++++- 2 files changed, 98 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 9d1d3cfa762c..aac89e0e77fb 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -592,10 +592,12 @@ def generate_async( else: raise ValueError(f"Invalid inputs type: {type(inputs)}") + # Snapshot caller-provided params so later mutations don't affect + # the queued request (the dispatcher thread serializes it lazily). request = DiffusionRequest( request_id=req_id, prompt=prompt, - params=params, + params=params.model_copy(deep=True) if params is not None else None, ) self.executor.enqueue_requests([request]) 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 0ce92f9b7243..b047271f59eb 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -89,10 +89,16 @@ def __init__( self._should_fail = should_fail self._healthy = True self._req_counter = 0 + # Captured arguments of the most recent generate / generate_async call, + # used by tests to assert forwarded VisualGenParams fields. + self.last_inputs = None + self.last_params = None # --- VisualGen interface --- def generate(self, inputs=None, params=None) -> MediaOutput: + self.last_inputs = inputs + self.last_params = params if self._should_fail: raise RuntimeError("Generation intentionally failed") return MediaOutput( @@ -102,6 +108,8 @@ def generate(self, inputs=None, params=None) -> MediaOutput: ) def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": + self.last_inputs = inputs + self.last_params = params return MockVisualGenResult( image=self._image, video=self._video, @@ -174,7 +182,10 @@ def _create_server(generator: MockVisualGen, model_name: str = "test-model") -> server_role=ServerRole.VISUAL_GEN, metadata_server_cfg=None, ) - return TestClient(server.app) + client = TestClient(server.app) + # Expose the mock so tests can assert captured generate() arguments. + client.mock_gen = generator + return client # --------------------------------------------------------------------------- @@ -298,6 +309,15 @@ def test_image_generation_with_optional_params(self, image_client): data = resp.json() assert data["size"] == "128x64" + # Verify openai_server/parse_visual_gen_params forwarded every field. + params = image_client.mock_gen.last_params + assert image_client.mock_gen.last_inputs == "Sunset over ocean" + assert params.width == 128 + assert params.height == 64 + assert params.num_inference_steps == 20 + assert params.guidance_scale == 7.5 + assert params.negative_prompt == "blurry" + def test_image_generation_url_format_not_supported(self, image_client): resp = image_client.post( "/v1/images/generations", @@ -437,6 +457,13 @@ def test_image_edit_with_mask(self, image_client): ) assert resp.status_code == 200 + # Verify image + mask were base64-decoded and forwarded. + params = image_client.mock_gen.last_params + expected_bytes = base64.b64decode(b64_img) + assert params.image == [expected_bytes] + assert params.mask == base64.b64decode(b64_mask) + assert params.num_inference_steps == 10 + def test_image_edit_with_optional_params(self, image_client): b64_img = _b64_white_png_1x1() resp = image_client.post( @@ -455,6 +482,14 @@ def test_image_edit_with_optional_params(self, image_client): data = resp.json() assert data["size"] == "128x128" + params = image_client.mock_gen.last_params + assert params.width == 128 + assert params.height == 128 + assert params.guidance_scale == 8.0 + assert params.num_inference_steps == 15 + assert params.negative_prompt == "dark" + assert params.image == [base64.b64decode(b64_img)] + def test_image_edit_failure(self, failing_client): b64_img = _b64_white_png_1x1() resp = failing_client.post( @@ -518,6 +553,17 @@ def test_sync_video_generation_with_params(self, video_client): assert resp.status_code == 200 assert len(resp.content) > 0 + params = video_client.mock_gen.last_params + assert video_client.mock_gen.last_inputs == "Ocean waves" + assert params.width == 64 + assert params.height == 64 + assert params.num_inference_steps == 10 + assert params.guidance_scale == 5.0 + assert params.seed == 42 + assert params.negative_prompt == "blurry" + assert params.frame_rate == 8 + assert params.num_frames == int(2.0 * 8) + def test_sync_video_generation_multipart(self, video_client): # Use files={} with a dummy file to ensure multipart/form-data dummy_file = BytesIO(b"") @@ -554,6 +600,13 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert resp.status_code == 200 assert len(resp.content) > 0 + # input_reference should have been written to media storage and passed + # through as params.image (a filesystem path). + params = video_client.mock_gen.last_params + assert isinstance(params.image, str) + assert params.image.endswith("_reference.png") + assert os.path.exists(params.image) + def test_sync_video_failure(self, failing_client): resp = failing_client.post( "/v1/videos/generations", @@ -692,6 +745,47 @@ def test_async_video_invalid_fps(self, video_client): ) assert resp.status_code == 400 + def test_async_video_forwards_params(self, video_client): + """Ensure async video endpoint forwards VisualGenParams to generate_async.""" + resp = video_client.post( + "/v1/videos", + json={ + "prompt": "Rainy street", + "size": "128x64", + "seconds": 2.0, + "fps": 10, + "num_inference_steps": 12, + "guidance_scale": 6.0, + "seed": 7, + "negative_prompt": "noise", + }, + headers={"content-type": "application/json"}, + ) + assert resp.status_code == 202 + video_id = resp.json()["id"] + + # The background task calls generate_async lazily — drive the event + # loop via status polling until the job completes. + import time as _time + + deadline = _time.time() + 5 + while _time.time() < deadline: + meta = video_client.get(f"/v1/videos/{video_id}").json() + if meta.get("status") in ("completed", "failed"): + break + _time.sleep(0.05) + + params = video_client.mock_gen.last_params + assert video_client.mock_gen.last_inputs == "Rainy street" + assert params.width == 128 + assert params.height == 64 + assert params.num_inference_steps == 12 + assert params.guidance_scale == 6.0 + assert params.seed == 7 + assert params.negative_prompt == "noise" + assert params.frame_rate == 10 + assert params.num_frames == int(2.0 * 10) + # ========================================================================= # GET /v1/videos (list)