From 7eb5340ccae958425848a47c02b18012cc7fca87 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:00:00 -0700 Subject: [PATCH 01/61] [TRTLLM-15277][feat] VisualGen media reference input API (Scheme C) Replace the untyped params.image field and the sniff-routed serve input_reference with typed per-modality reference inputs: image_reference / video_reference / audio_reference, each item optionally carrying a role. Add per-model ref_slot_specs (role membership + conditional-required role + per-role arity) declared by each pipeline and enforced by validate_visual_gen_params at both the engine entry and the serve preflight. Migrate all VisualGen pipelines (FLUX.2, Wan I2V/TI2V, LTX-2 x2, Qwen-edit, Qwen-Layered, Cosmos3) off params.image; serve materializes each transport (base64 / data-URI / upload) to a local path and hands the pipeline a typed *Ref. Remove params.image and input_reference (breaking). Export ImageRef / VideoRef / AudioRef publicly. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- examples/visual_gen/models/cosmos3/cosmos3.py | 4 +- examples/visual_gen/models/flux2.py | 2 +- examples/visual_gen/models/qwen_image_edit.py | 2 +- .../visual_gen/models/qwen_image_layered.py | 2 +- examples/visual_gen/models/wan_i2v.py | 8 +- examples/visual_gen/serve/README.md | 11 +- examples/visual_gen/serve/async_video_gen.py | 20 +- examples/visual_gen/serve/sync_video_gen.py | 24 +-- tensorrt_llm/__init__.py | 12 +- tensorrt_llm/_torch/visual_gen/executor.py | 5 +- .../visual_gen/models/cosmos3/defaults.py | 12 -- .../models/cosmos3/pipeline_cosmos3.py | 28 ++- .../visual_gen/models/flux/pipeline_flux2.py | 20 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 19 +- .../models/ltx2/pipeline_ltx2_two_stages.py | 3 +- .../qwen_image/pipeline_qwen_image_edit.py | 15 +- .../pipeline_qwen_image_layered.py | 19 +- .../visual_gen/models/wan/pipeline_wan.py | 21 ++- .../visual_gen/models/wan/pipeline_wan_i2v.py | 42 +++-- tensorrt_llm/_torch/visual_gen/pipeline.py | 35 ++++ tensorrt_llm/serve/openai_protocol.py | 74 +++++++- tensorrt_llm/serve/openai_video_routes.py | 9 +- tensorrt_llm/serve/visual_gen_utils.py | 177 +++++++++++++----- tensorrt_llm/visual_gen/__init__.py | 8 +- tensorrt_llm/visual_gen/params.py | 128 ++++++++++++- tensorrt_llm/visual_gen/visual_gen.py | 14 +- .../README_test_visual_gen_perf_sanity.md | 2 +- .../defs/perf/visual_gen_perf_utils.py | 4 +- .../visual_gen/ltx2_blackwell.yaml | 4 +- .../visual_gen/wan22_i2v_a14b_blackwell.yaml | 2 +- .../visual_gen/test_trtllm_serve_e2e.py | 4 +- .../visual_gen/test_trtllm_serve_endpoints.py | 46 +++-- .../visual_gen/test_visual_gen_params.py | 75 +++----- .../visual_gen/test_visual_gen_utils.py | 156 +++++++++------ .../references/trtllm_serve_api.yaml | 16 +- 35 files changed, 734 insertions(+), 289 deletions(-) diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 18b1cba051c8..4f35e67c1ca8 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -477,7 +477,7 @@ def main(): # Query per-model defaults (resolution, steps, guidance, seed, etc.). params = visual_gen.default_params if image_path is not None: - params.image = image_path + params.image_reference = image_path negative_prompt = resolve_negative_prompt( negative_prompt=args.negative_prompt, @@ -514,7 +514,7 @@ def main(): with open(args.action_json, encoding="utf-8") as f: params.extra_params["action"] = json.load(f) if args.video_path is not None: - params.extra_params["video"] = Path(args.video_path).read_bytes() + params.video_reference = args.video_path if args.extra_params: # Merged last: explicit JSON wins over flag-derived values. params.extra_params.update(args.extra_params) diff --git a/examples/visual_gen/models/flux2.py b/examples/visual_gen/models/flux2.py index a39044448b2f..ad983d72b321 100644 --- a/examples/visual_gen/models/flux2.py +++ b/examples/visual_gen/models/flux2.py @@ -118,7 +118,7 @@ def main(): # Start from per-model defaults and override only user-provided request fields. params = visual_gen.default_params params.num_images_per_prompt = args.num_images_per_prompt - params.image = args.image + params.image_reference = args.image if args.image: # Let FLUX.2 derive omitted dimensions from the first processed reference. params.height = args.height diff --git a/examples/visual_gen/models/qwen_image_edit.py b/examples/visual_gen/models/qwen_image_edit.py index 43f37a73d1c4..4d1740f340da 100644 --- a/examples/visual_gen/models/qwen_image_edit.py +++ b/examples/visual_gen/models/qwen_image_edit.py @@ -64,7 +64,7 @@ def main() -> None: extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None visual_gen = VisualGen(model=args.model, args=extra_args) params = visual_gen.default_params - params.image = args.image if len(args.image) > 1 else args.image[0] + params.image_reference = args.image output = visual_gen.generate(inputs=args.prompt, params=params) saved = output.save(args.output_path) print(f"Saved edited image to {saved}") diff --git a/examples/visual_gen/models/qwen_image_layered.py b/examples/visual_gen/models/qwen_image_layered.py index 06af0d386914..f5fc47d2be82 100644 --- a/examples/visual_gen/models/qwen_image_layered.py +++ b/examples/visual_gen/models/qwen_image_layered.py @@ -63,7 +63,7 @@ def main() -> None: visual_gen = VisualGen(model=args.model, args=extra_args) params = visual_gen.default_params - params.image = args.image + params.image_reference = args.image output = visual_gen.generate(inputs=args.prompt, params=params) if output.image is not None and output.image.shape[0] > 1: diff --git a/examples/visual_gen/models/wan_i2v.py b/examples/visual_gen/models/wan_i2v.py index 93b16d51310b..9876b6ebaf45 100644 --- a/examples/visual_gen/models/wan_i2v.py +++ b/examples/visual_gen/models/wan_i2v.py @@ -23,7 +23,7 @@ import argparse import os -from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm import ImageRef, VisualGen, VisualGenArgs _DEFAULT_IMAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cat_piano.png") @@ -62,9 +62,11 @@ def main(): visual_gen = VisualGen(model=args.model, args=extra_args) # --- Model-specific: I2V request construction --- - # Start from per-model defaults (steps, guidance, seed, etc.) and set the input image. + # Start from per-model defaults (steps, guidance, seed, etc.) and set the + # first-frame reference. Wan I2V also accepts a ``last_frame`` role, so the + # role must be given to disambiguate. params = visual_gen.default_params - params.image = args.image + params.image_reference = [ImageRef(image=args.image, role="first_frame")] output = visual_gen.generate( inputs="A cat presses the piano keys with its paws, soft notes filling the quiet room.", diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 1a8a3019bdd7..604942e4d86f 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,7 +286,7 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `input_reference`: Reference image (I2V/TI2V) or video (V2V), accepted as a base64-encoded string in JSON or as a file in multipart form-data +- `image_reference`: Reference image(s) for I2V/TI2V. `video_reference`: reference video(s) for V2V. `audio_reference`: reference audio(s). Each accepts a base64-encoded string (optionally a `data:` URI), a `{image|video|audio, role}` object, or a list of them in JSON; or a single uploaded file in multipart form-data. - **Supported formats**: PNG and JPEG images; MP4 and AVI video, with H.264 the tested codec and others best-effort. HEIF/AVIF are not supported. - `extra_params`: model-specific overflow (see below) - `response_format`: `"file"` (default; `FileResponse` byte download) or `"path"` (server-side output path JSON, for co-located clients) @@ -384,7 +384,7 @@ curl -X POST "http://localhost:8000/v1/videos" \ ```bash curl -X POST "http://localhost:8000/v1/videos" \ -F "prompt=She turns around and smiles" \ - -F "input_reference=@./media/woman_skyline_original_720p.jpeg" \ + -F "image_reference=@./media/woman_skyline_original_720p.jpeg" \ -F "seconds=4.0" \ -F "fps=24" \ -F "size=256x256" \ @@ -393,11 +393,12 @@ curl -X POST "http://localhost:8000/v1/videos" \ ### Video-to-Video (Multipart with File Upload, Cosmos3) ```bash -# The reference is classified by content: image -> I2V, video -> V2V. -# V2V conditioning knobs ride in extra_params (values below are the defaults). +# The modality is declared by the field name: image_reference -> I2V, +# video_reference -> V2V. V2V conditioning knobs ride in extra_params +# (values below are the defaults). curl -X POST "http://localhost:8000/v1/videos" \ -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \ - -F "input_reference=@./media/reference.mp4" \ + -F "video_reference=@./media/reference.mp4" \ -F "num_frames=189" \ -F "fps=24" \ -F 'extra_params={"condition_video_latent_indexes": [0, 1], "condition_video_keep": "first"}' diff --git a/examples/visual_gen/serve/async_video_gen.py b/examples/visual_gen/serve/async_video_gen.py index b884de0bbd22..2c90ed305349 100755 --- a/examples/visual_gen/serve/async_video_gen.py +++ b/examples/visual_gen/serve/async_video_gen.py @@ -28,7 +28,7 @@ def test_async_video_generation( base_url: str = "http://localhost:8000/v1", model: str = "wan", prompt: str = "A video of a cool cat on a motorcycle in the night", - input_reference: str = None, + image_reference: str = None, duration: float = 4.0, fps: int = 24, size: str = "256x256", @@ -41,7 +41,7 @@ def test_async_video_generation( base_url: Base URL of the API server model: Model name to use prompt: Text prompt for generation - input_reference: Path to reference image (optional, for TI2V mode) + image_reference: Path to reference image (optional, for TI2V mode) duration: Video duration in seconds fps: Frames per second size: Video resolution (WxH format) @@ -51,7 +51,7 @@ def test_async_video_generation( The server may return either MP4 (H.264) or AVI (MJPEG) format depending on the available encoder. The output filename extension will be adjusted to match. """ - mode = "TI2V" if input_reference else "T2V" + mode = "TI2V" if image_reference else "T2V" print("=" * 80) print(f"Testing Async Video Generation API - {mode} Mode") print("=" * 80) @@ -62,8 +62,8 @@ def test_async_video_generation( print("\n1. Creating video generation job...") print(f" Mode: {mode}") print(f" Prompt: {prompt}") - if input_reference: - print(f" Input Reference: {input_reference}") + if image_reference: + print(f" Input Reference: {image_reference}") print(f" Duration: {duration}s") print(f" FPS: {fps}") print(f" Size: {size}") @@ -82,11 +82,11 @@ def test_async_video_generation( } # Add input reference if provided (TI2V mode) - if input_reference: - if not Path(input_reference).exists(): - print(f"\n❌ Error: Input reference image not found: {input_reference}") + if image_reference: + if not Path(image_reference).exists(): + print(f"\n❌ Error: Input reference image not found: {image_reference}") return False - create_params["input_reference"] = open(input_reference, "rb") + create_params["image_reference"] = open(image_reference, "rb") # Create video generation job job = client.videos.create(**create_params) @@ -269,7 +269,7 @@ def test_async_video_generation( base_url=args.base_url, model=args.model, prompt=args.prompt, - input_reference=args.image, + image_reference=args.image, duration=args.duration, fps=args.fps, size=args.size, diff --git a/examples/visual_gen/serve/sync_video_gen.py b/examples/visual_gen/serve/sync_video_gen.py index 385369699f3c..52100a27ea2c 100755 --- a/examples/visual_gen/serve/sync_video_gen.py +++ b/examples/visual_gen/serve/sync_video_gen.py @@ -27,7 +27,7 @@ def test_sync_video_generation( base_url: str = "http://localhost:8000/v1", model: str = "wan", prompt: str = "A video of a cute cat playing with a ball in the park", - input_reference: str = None, + image_reference: str = None, duration: float = 4.0, fps: int = 24, size: str = "256x256", @@ -40,7 +40,7 @@ def test_sync_video_generation( base_url: Base URL of the API server model: Model name to use prompt: Text prompt for generation - input_reference: Path to reference image (optional, for TI2V mode) + image_reference: Path to reference image (optional, for TI2V mode) duration: Video duration in seconds fps: Frames per second size: Video resolution (WxH format) @@ -50,7 +50,7 @@ def test_sync_video_generation( The server may return either MP4 (H.264) or AVI (MJPEG) format depending on the available encoder. The output filename extension will be adjusted to match. """ - mode = "TI2V" if input_reference else "T2V" + mode = "TI2V" if image_reference else "T2V" print("=" * 80) print(f"Testing Sync Video Generation API - {mode} Mode") print("=" * 80) @@ -58,8 +58,8 @@ def test_sync_video_generation( print("\n1. Generating video (waiting for completion)...") print(f" Mode: {mode}") print(f" Prompt: {prompt}") - if input_reference: - print(f" Input Reference: {input_reference}") + if image_reference: + print(f" Input Reference: {image_reference}") print(f" Duration: {duration}s") print(f" FPS: {fps}") print(f" Size: {size}") @@ -67,10 +67,10 @@ def test_sync_video_generation( try: endpoint = f"{base_url}/videos/sync" - if input_reference: + if image_reference: # TI2V mode - Use multipart/form-data with file upload - if not Path(input_reference).exists(): - print(f"\n❌ Error: Input reference image not found: {input_reference}") + if not Path(image_reference).exists(): + print(f"\n❌ Error: Input reference image not found: {image_reference}") return False # Prepare form data (all values as strings for multipart) @@ -86,9 +86,9 @@ def test_sync_video_generation( # Add the file ## Note: The content-type must be multipart/form-data. files = { - "input_reference": ( - Path(input_reference).name, - open(input_reference, "rb"), + "image_reference": ( + Path(image_reference).name, + open(image_reference, "rb"), "multipart/form-data", ) } @@ -254,7 +254,7 @@ def test_sync_video_generation( base_url=args.base_url, model=args.model, prompt=args.prompt, - input_reference=args.image, + image_reference=args.image, duration=args.duration, fps=args.fps, size=args.size, diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index 516e31ad2c52..aefaeaf503c6 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -62,9 +62,9 @@ from .mapping import Mapping from .models.automodel import AutoConfig, AutoModelForCausalLM from .sampling_params import SamplingParams - from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, - VisualGenMetrics, VisualGenOutput, VisualGenParams, - VisualGenResult) + from .visual_gen import (AudioRef, ExtraParamSchema, ImageRef, VideoRef, + VisualGen, VisualGenArgs, VisualGenMetrics, + VisualGenOutput, VisualGenParams, VisualGenResult) # Public name -> (source module, attribute); attribute None = the module itself. _LAZY_ATTRS = { @@ -105,6 +105,9 @@ 'VisualGenMetrics': ('tensorrt_llm.visual_gen', 'VisualGenMetrics'), 'VisualGenOutput': ('tensorrt_llm.visual_gen', 'VisualGenOutput'), 'VisualGenParams': ('tensorrt_llm.visual_gen', 'VisualGenParams'), + 'ImageRef': ('tensorrt_llm.visual_gen', 'ImageRef'), + 'VideoRef': ('tensorrt_llm.visual_gen', 'VideoRef'), + 'AudioRef': ('tensorrt_llm.visual_gen', 'AudioRef'), 'VisualGenResult': ('tensorrt_llm.visual_gen', 'VisualGenResult'), } @@ -173,6 +176,9 @@ def __dir__(): 'math_utils', 'VisualGen', 'VisualGenParams', + 'ImageRef', + 'VideoRef', + 'AudioRef', '__version__', ] diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 050d09cdafbd..6dbdce267008 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -385,6 +385,7 @@ def _load_pipeline(self): "default_generation_params": self.pipeline.default_generation_params, "extra_param_specs": self.pipeline.extra_param_specs, "supports_image_edit": self.pipeline.supports_image_edit, + "ref_slot_specs": self.pipeline.ref_slot_specs, }, ) ) @@ -428,7 +429,7 @@ def _merge_defaults(self, req: DiffusionRequest): for field_name, default_value in self.pipeline.default_generation_params.items(): if hasattr(params, field_name) and getattr(params, field_name) is None: if ( - params.image is not None + params.image_reference and getattr(self.pipeline, "derive_output_size_from_reference", False) is True and field_name in ("height", "width") ): @@ -721,6 +722,7 @@ def __init__( self.default_generation_params: Dict = {} self.extra_param_specs: Dict = {} self.supports_image_edit: bool = False + self.ref_slot_specs: Dict = {} # --- Launch workers --- self.worker_processes = [] @@ -1077,6 +1079,7 @@ async def _wait_ready_async(self): ) self.extra_param_specs = payload.get("extra_param_specs", {}) self.supports_image_edit = bool(payload.get("supports_image_edit", False)) + self.ref_slot_specs = payload.get("ref_slot_specs", {}) elapsed = time.time() - start_time logger.info(f"DiffusionClient: Workers ready ({elapsed:.1f}s)") return diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index cf18b82bb498..0d04c85be07f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -646,18 +646,6 @@ def _resolve_field( default=None, description="Optional scheduler flow shift override. Uses the Cosmos3 mode default when omitted.", ), - "video": ExtraParamSchema( - type="bytes", - default=None, - description=( - "V2V reference: encoded MP4/AVI bytes (e.g. " - "Path(video).read_bytes()). Each worker rank demuxes them from " - "memory and NVDEC-decodes only the conditioning window per " - "condition_video_latent_indexes / condition_video_keep, resized " - "to the output resolution, then VAE-encodes it." - ), - validator=_validate_video_reference, - ), "action_mode": ExtraParamSchema( type="Literal['policy', 'forward_dynamics', 'inverse_dynamics']", default=None, diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index a78e1bfba4a5..00fa4c70049c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -37,7 +37,7 @@ def tqdm(iterable, **kwargs): from tensorrt_llm._torch.visual_gen.models.wan.vae_loader import load_wan_vae from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import ( classify_worker_error, @@ -73,6 +73,7 @@ def tqdm(iterable, **kwargs): _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, resolve_domain_action_config, + _validate_video_reference, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT @@ -602,6 +603,19 @@ def extra_param_specs(self): # ``default_use_system_prompt``. return dict(COSMOS3_EXTRA_SPECS) + @property + def ref_slot_specs(self): + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="first_frame", min=1, max=1)], + ), + "video_reference": RefSlotSpec( + modality="video", + roles=[RoleSpec(role="reference", min=1, max=1)], + ), + } + def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: # Checkpoint-aware guidance: distilled defaults carry a concrete 1.0; # base defaults leave it None ("by mode") — warmup runs the video mode. @@ -742,7 +756,14 @@ def as_given(field_name): value = getattr(req.params, field_name) return value if field_name in specified else None - video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) + refs_v = req.params.video_reference + video = refs_v[0].video if refs_v else None + if isinstance(video, str): + from pathlib import Path + + video = Path(video).read_bytes() # forward() NVDEC-demuxes from memory + if video is not None: + _validate_video_reference(video) is_action = extra_params.get("action_mode") is not None if is_action: # Action resolves its whole recipe in forward() -- the canvas from @@ -807,11 +828,12 @@ def as_given(field_name): width = resolved["width"] num_inference_steps = resolved["num_inference_steps"] guidance_scale = resolved["guidance_scale"] + refs_i = req.params.image_reference return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, - image=req.params.image, + image=refs_i[0].image if refs_i else None, height=height, width=width, num_frames=req.params.num_frames, 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 20420592caaa..0845beb6f070 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -44,7 +44,7 @@ register_extractor_from_config, ) from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm.logger import logger @@ -370,12 +370,23 @@ def default_generation_params(self): "max_sequence_length": 512, } + @property + def ref_slot_specs(self): + # Reference image(s): single "reference" role, count 1..N (multi-subject). + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="reference", min=1, max=None)], + ) + } + def prepare_request(self, req: Any) -> None: """Load and preprocess reference images before warmup bookkeeping.""" - if req.params.image is None: + refs = req.params.image_reference + if not refs: return - reference_images = self._load_reference_images(req.params.image) + reference_images = self._load_reference_images([r.image for r in refs]) condition_images = self._preprocess_reference_images(reference_images) req.params.height, req.params.width = self._resolve_target_dimensions( req.params.height, @@ -386,6 +397,7 @@ def prepare_request(self, req: Any) -> None: def infer(self, req): """Run inference from DiffusionRequest.""" + refs = req.params.image_reference return self.forward( prompt=req.prompt, height=req.params.height, @@ -395,7 +407,7 @@ def infer(self, req): seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, num_images_per_prompt=req.params.num_images_per_prompt, - image=req.params.image, + image=[r.image for r in refs] if refs else None, _condition_images=req.prepared_inputs.get("condition_images"), ) 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 2ebee01b7d35..27287bab5924 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -20,7 +20,12 @@ from tensorrt_llm._torch.visual_gen.checkpoints.prefetch import prefetch_files_to_host_cache from tensorrt_llm._torch.visual_gen.cuda_graph_runner import CUDAGraphRunner, CUDAGraphRunnerConfig from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline import ( + BasePipeline, + ExtraParamSchema, + RefSlotSpec, + RoleSpec, +) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm.logger import logger @@ -1374,9 +1379,19 @@ def extra_param_specs(self): ), } + @property + def ref_slot_specs(self): + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="first_frame", min=1, max=1)], + ), + } + def infer(self, req): """Run inference with request parameters.""" extra = req.params.extra_params or {} + refs = req.params.image_reference return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, @@ -1390,7 +1405,7 @@ def infer(self, req): output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], max_sequence_length=req.params.max_sequence_length, - image=req.params.image, + image=refs[0].image if refs else None, image_cond_strength=extra["image_cond_strength"], stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], 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 afb2bdd687e9..0bc6c374e0be 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 @@ -1226,6 +1226,7 @@ def _load_two_stage_components( def infer(self, req): extra = req.params.extra_params or {} + refs = req.params.image_reference return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, @@ -1239,7 +1240,7 @@ def infer(self, req): output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], max_sequence_length=req.params.max_sequence_length, - image=req.params.image, + image=refs[0].image if refs else None, image_cond_strength=extra["image_cond_strength"], stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], 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 ddc0c453bf7c..5f0d43cf2307 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 @@ -18,6 +18,7 @@ import torch.distributed as dist from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger @@ -131,6 +132,15 @@ def default_generation_params(self) -> dict: def default_warmup_resolutions(self) -> list[tuple[int, int]]: return [(1024, 1024)] + @property + def ref_slot_specs(self): + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="reference", min=1, max=None)], + ), + } + def warmup_cache_key(self, height: int | None, width: int | None, **kwargs) -> tuple: return (height, width) @@ -152,7 +162,7 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N @staticmethod def _load_edit_images(image: Any) -> list[Any]: if image is None: - raise ValueError("Qwen-Image-Edit requires params.image.") + raise ValueError("Qwen-Image-Edit requires image_reference.") images = image if isinstance(image, list) else [image] pil_images = [] for item in images: @@ -341,7 +351,8 @@ def infer(self, req: Any) -> PipelineOutput: "QwenImageEditPlusPipeline currently supports num_images_per_prompt=1 only." ) prompts = req.prompt if isinstance(req.prompt, list) else [req.prompt] - pil_images = self._load_edit_images(params.image) + refs = params.image_reference + pil_images = self._load_edit_images([r.image for r in refs] if refs else None) height = params.height width = params.width if height is None or width is None: 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 440979a5a287..4edf20cb4eee 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 @@ -23,7 +23,12 @@ import torch from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline import ( + BasePipeline, + ExtraParamSchema, + RefSlotSpec, + RoleSpec, +) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm.logger import logger @@ -248,6 +253,15 @@ def extra_param_specs(self) -> dict: ), } + @property + def ref_slot_specs(self): + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="reference", min=1, max=1)], + ), + } + def load_standard_components( self, checkpoint_dir: str, @@ -758,8 +772,9 @@ def infer(self, req): ) negative = [n for n in negatives for _ in range(num_per)] + refs = req.params.image_reference return self.forward( - image=req.params.image, + image=refs[0].image if refs else None, prompt=prompts, negative_prompt=negative, height=req.params.height, 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 148cf659da22..04c9cda12e66 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -39,7 +39,7 @@ ) from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import retrieve_latents from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm._utils import nvtx_range @@ -424,17 +424,22 @@ def default_generation_params(self): def extra_param_specs(self): return get_wan_extra_param_specs(self.is_wan22_14b) + @property + def ref_slot_specs(self): + # Optional single conditioning image (first frame); T2V when absent. + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[RoleSpec(role="first_frame", min=1, max=1)], + ) + } + def infer(self, req): """Run inference with request parameters.""" extra = req.params.extra_params or {} # Wan 2.2 TI2V-5B takes one conditioning image if provided - image = req.params.image - if isinstance(image, list): - if len(image) != 1: - raise ValueError( - f"WanPipeline I2V expects a single image, got list of {len(image)}." - ) - image = image[0] + refs = req.params.image_reference + image = refs[0].image if refs else None return self.forward( prompt=req.prompt, 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 d8318fe02fe8..7eda46df2842 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 @@ -36,7 +36,7 @@ ) from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_utils import retrieve_latents from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput -from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm.logger import logger @@ -402,26 +402,34 @@ def default_generation_params(self): @property def extra_param_specs(self): - specs = get_wan_extra_param_specs(self.is_wan22_14b) - specs["last_image"] = ExtraParamSchema( - type="str", - default=None, - description="Last frame path for video interpolation (Wan I2V).", - ) - return specs + return get_wan_extra_param_specs(self.is_wan22_14b) + + @property + def ref_slot_specs(self): + # I2V first frame (required) + optional last frame for interpolation. + return { + "image_reference": RefSlotSpec( + modality="image", + roles=[ + RoleSpec(role="first_frame", min=1, max=1), + RoleSpec(role="last_frame", min=0, max=1), + ], + ) + } def infer(self, req): """Run inference with request parameters.""" - # Extract image from request (can be path, PIL Image, or torch.Tensor) - if req.params.image is None: - raise ValueError("I2V pipeline requires 'image' parameter") - - image = req.params.image[0] if isinstance(req.params.image, list) else req.params.image + refs = req.params.image_reference + if not refs: + raise ValueError("I2V pipeline requires an image_reference (first_frame)") + by_role = {r.role: r for r in refs} + first = by_role.get("first_frame") or by_role.get(None) + if first is None: + raise ValueError("I2V pipeline requires a first_frame image_reference") + last = by_role.get("last_frame") + image = first.image + last_image = last.image if last is not None else None extra = req.params.extra_params or {} - last_image = extra.get("last_image") - - if last_image is not None and isinstance(last_image, list): - last_image = last_image[0] if last_image else None return self.forward( image=image, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index d4978872e471..7cb0fd815a81 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -70,6 +70,30 @@ class ExtraParamSchema(StrictBaseModel): ) +class RoleSpec(StrictBaseModel): + """One accepted role for a reference modality, with its count bounds.""" + + role: str = Field(description="'reference' | 'first_frame' | 'last_frame'.") + min: int = Field(default=1, description="Minimum count for this role.") + max: Optional[int] = Field( + default=1, description="Maximum count for this role (None = unbounded)." + ) + + +class RefSlotSpec(StrictBaseModel): + """Reference slot a pipeline accepts for one modality. + + A request item's ``role`` is required only when ``roles`` has more than one + entry (same modality carries multiple roles, e.g. first + last frame); + otherwise the single declared role is inferred. Exposed via + ``VisualGen.ref_slot_specs`` and enforced by ``validate_visual_gen_params``. + Pickled to the coordinator in the READY handshake, so keep it plain data. + """ + + modality: str = Field(description="'image' | 'video' | 'audio'.") + roles: List[RoleSpec] = Field(description="Accepted roles + counts for this modality.") + + if TYPE_CHECKING: from .cache import CacheAccelerator from .config import DiffusionPipelineConfig @@ -350,6 +374,17 @@ def extra_param_specs(self) -> Dict[str, ExtraParamSchema]: """ return {} + @property + def ref_slot_specs(self) -> Dict[str, RefSlotSpec]: + """Reference slots this pipeline accepts. + + Maps a ``VisualGenParams`` reference field name + (``image_reference`` / ``video_reference`` / ``audio_reference``) to a + :class:`RefSlotSpec` declaring the accepted roles and per-role counts. + Empty by default (pipeline takes no reference inputs). + """ + return {} + @property def default_generation_params(self) -> dict: """Model-specific defaults for ``None`` fields in ``VisualGenParams``. diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 84fed73f2b9f..b3f26c6df0ed 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2011,6 +2011,41 @@ class ImageGenerationResponse(OpenAIBaseModel): size: Optional[str] = None +class ImageReferenceItem(OpenAIBaseModel): + """One image reference for conditioning (mirrors ``ImageRef``). + + ``image`` carries base64-encoded bytes, optionally as a ``data:`` URI. + ``role`` is required only for models that accept more than one image + role (e.g. Wan first/last frame); single-role models infer it. + """ + + image: str = Field( + description="Base64-encoded image bytes, optionally as a ``data:`` URI." + ) + role: Optional[str] = Field( + default=None, + description=( + "Reference role (e.g. 'reference', 'first_frame', 'last_frame'). " + "Required only when the model accepts multiple image roles."), + ) + + +class VideoReferenceItem(OpenAIBaseModel): + """One video reference for conditioning (mirrors ``VideoRef``).""" + + video: str = Field( + description="Base64-encoded video bytes, optionally as a ``data:`` URI." + ) + + +class AudioReferenceItem(OpenAIBaseModel): + """One audio reference for conditioning (mirrors ``AudioRef``).""" + + audio: str = Field( + description="Base64-encoded audio bytes, optionally as a ``data:`` URI." + ) + + class VideoGenerationRequest(OpenAIBaseModel): """Video generation request (extended API). @@ -2034,15 +2069,36 @@ class VideoGenerationRequest(OpenAIBaseModel): seed: Optional[int] = Field(default=None, ge=0, description="Random seed for reproducibility.") - input_reference: Optional[Union[str, UploadFile]] = Field( - default=None, - description=( - "Optional image or video reference that guides generation. PNG or " - "JPEG images condition image-to-video; MP4 or AVI video conditions " - "video-to-video, with H.264 the tested codec and others " - "best-effort. HEIF/AVIF are not supported. JSON requests carry " - "base64 bytes; multipart requests upload the file."), - ) + image_reference: Optional[Union[ + str, UploadFile, ImageReferenceItem, + List[Union[str, ImageReferenceItem]]]] = Field( + default=None, + description= + ("Image reference(s) conditioning generation (e.g. image-to-video " + "first frame). JSON sends base64 bytes, an ``{image, role}`` " + "object, or a list of them; multipart uploads a single image file. " + "PNG or JPEG only — HEIF/AVIF are not supported."), + ) + video_reference: Optional[Union[ + str, UploadFile, VideoReferenceItem, + List[Union[str, VideoReferenceItem]]]] = Field( + default=None, + description= + ("Video reference(s) conditioning generation (video-to-video). JSON " + "sends base64 bytes, a ``{video}`` object, or a list of them; " + "multipart uploads a single video file. MP4 or AVI, with H.264 the " + "tested codec and others best-effort."), + ) + audio_reference: Optional[Union[ + str, UploadFile, AudioReferenceItem, + List[Union[str, AudioReferenceItem]]]] = Field( + default=None, + description= + ("Audio reference(s) conditioning generation. JSON sends base64 " + "bytes, an ``{audio}`` object, or a list of them; multipart uploads " + "a single audio file. Accepted only by models that declare an audio " + "reference slot."), + ) # Resolution size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 8fbaca541071..a0338e459f19 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -141,7 +141,7 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: Supports both JSON and multipart/form-data requests: - JSON: Send VideoGenerationRequest as application/json - - Multipart: Send form fields + optional input_reference file + - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time try: @@ -309,8 +309,8 @@ async def _parse_video_generation_request( for key in form: value = form[key] if hasattr(value, "file"): - # Uploaded file (``input_reference``) — pass through - # so the conversion layer reads ``.file``. + # Uploaded reference file (image_reference / video_reference) + # — pass through so the conversion layer reads ``.file``. data[key] = value continue if key == "extra_params": @@ -373,7 +373,7 @@ async def openai_video_generation_async( Supports both JSON and multipart/form-data requests: - JSON: Send VideoGenerationRequest as application/json - - Multipart: Send form fields + optional input_reference file + - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time try: @@ -398,6 +398,7 @@ async def openai_video_generation_async( params, declared_defaults=self.generator.executor.default_generation_params, extra_param_specs=self.generator.executor.extra_param_specs, + ref_slot_specs=self.generator.executor.ref_slot_specs, ) request_format = _resolve_tensor_only_format( request.format, request.extra_params, self.generator.extra_param_specs diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index be916bcc8ba0..458b655aed18 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -107,21 +107,107 @@ def _merge_extra_params( params.extra_params = None -def _read_reference_payload(reference) -> bytes: - """Read the ``input_reference`` payload (base64 JSON or multipart file). +def _read_reference_payload(reference: str) -> bytes: + """Decode one base64 (optionally ``data:`` URI) reference string to bytes. Payload size is deliberately not checked here: encoded size is not part of the request-validity contract, and body limits belong to the proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so malformed encodings — not sizes — are rejected. """ - if isinstance(reference, str): - try: - return base64.b64decode(reference, validate=True) - except ValueError as exc: - # binascii.Error subclasses ValueError. - raise ValueError("input_reference is not valid base64 data.") from exc - return reference.file.read() + data = reference + if data.startswith("data:"): + comma = data.find(",") + if comma == -1: + raise ValueError("reference data: URI is malformed (missing comma).") + data = data[comma + 1 :] + try: + return base64.b64decode(data, validate=True) + except ValueError as exc: + # binascii.Error subclasses ValueError. + raise ValueError("reference is not valid base64 data.") from exc + + +def _reference_payload_and_role(ref, data_field: str) -> tuple[bytes, Optional[str]]: + """Extract ``(payload_bytes, role)`` from one raw HTTP reference. + + ``ref`` is a base64/data-URI string, a multipart ``UploadFile`` (has + ``.file``), or a reference item exposing ``data_field`` ('image'/'video') + and, for images, an optional ``role``. + """ + role = getattr(ref, "role", None) + if isinstance(ref, str): + return _read_reference_payload(ref), role + if hasattr(ref, "file"): # multipart UploadFile + return ref.file.read(), role + data = getattr(ref, data_field, None) + if not isinstance(data, str): + raise ValueError(f"{data_field}_reference item must carry a base64 '{data_field}' string.") + return _read_reference_payload(data), role + + +def _materialize_reference( + payload: bytes, *, modality: str, ref_id: str, media_storage_path: Optional[str] +) -> str: + """Content-validate a reference payload and persist it, returning its path. + + HEIF/AVIF images are rejected on signature alone (Pillow support depends + on optional plugins the worker need not share). Video acceptance beyond the + container signature happens in the worker's NVDEC demux. + """ + if modality == "image": + if sniff_media_kind(payload) != "image": + raise ValueError( + "image_reference is not a recognized image; supported inputs are PNG/JPEG." + ) + if is_isobmff_image_bytes(payload): + raise ValueError( + "image_reference is a HEIF/AVIF image, which is not a supported " + "reference format; convert it to PNG or JPEG." + ) + elif modality == "video": + if sniff_media_kind(payload) != "video": + raise ValueError( + "video_reference is not a recognized media container; supported " + "inputs are MP4/AVI video." + ) + # audio: no signature sniffing (sniff_media_kind detects only image/video); + # the consuming pipeline validates the audio codec in its worker. + if media_storage_path is None: + raise ValueError(f"media_storage_path is required to store the {modality}_reference.") + ref_path = os.path.join(media_storage_path, ref_id) + with open(ref_path, "wb") as f: + f.write(payload) + return ref_path + + +def _build_reference_list( + value, *, modality: str, data_field: str, ref_cls, id: str, media_storage_path: Optional[str] +): + """Materialize an HTTP reference field into a list of typed ``*Ref`` objects. + + ``value`` is None, a base64/data-URI string, a multipart ``UploadFile``, a + reference item, or a list of any of those. Each entry is decoded, + content-validated for ``modality``, persisted to a per-index path, and + wrapped as ``ref_cls`` (carrying ``role`` for images). + """ + if value is None: + return None + raw_items = value if isinstance(value, list) else [value] + refs = [] + for i, item in enumerate(raw_items): + payload, role = _reference_payload_and_role(item, data_field) + ref_path = _materialize_reference( + payload, + modality=modality, + ref_id=f"{id}_{data_field}_ref_{i}", + media_storage_path=media_storage_path, + ) + kwargs = {data_field: ref_path} + if role is not None: + kwargs["role"] = role + refs.append(ref_cls(**kwargs)) + return refs def _decode_inline_media(extra_params: dict | None, specs) -> None: @@ -426,42 +512,43 @@ def parse_visual_gen_params( "directly." ) params.num_frames = derived - if request.input_reference is not None: - payload = _read_reference_payload(request.input_reference) - kind = sniff_media_kind(payload) - if kind == "image": - # Rejected on signature alone, not on a failed decode: - # whether Pillow reads HEIF/AVIF depends on optional plugins, - # and the worker process need not have the same ones. - if is_isobmff_image_bytes(payload): - raise ValueError( - "input_reference is a HEIF/AVIF image, which is not " - "a supported reference format; convert it to PNG or " - "JPEG." - ) - # I2V: the stored image file is the cross-model contract. - # every I2V pipeline reads ``params.image`` as a path. - if media_storage_path is None: - raise ValueError( - "media_storage_path is required when input_reference is an image" - ) - ref_path = os.path.join(media_storage_path, f"{id}_reference") - with open(ref_path, "wb") as f: - f.write(payload) - params.image = ref_path - elif kind == "video": - # V2V: encoded bytes pass through untouched; the worker - # demuxes and NVDEC-decodes them (acceptance happens there, - # so corrupt content behind a valid signature still fails as - # a client error). - if params.extra_params is None: - params.extra_params = {} - params.extra_params["video"] = payload - else: - raise ValueError( - "input_reference is not a recognized media container; " - "supported inputs are PNG/JPEG images and MP4/AVI video." - ) + # Reference inputs: materialize each transport (base64/data-URI/upload) + # to a stored file and hand the pipeline a typed ``*Ref`` carrying the + # local path. Decode stays model-specific in the worker. Local import: + # the visual_gen tree is already loaded in a VisualGen serving process, + # and this keeps it out of every plain-LLM process (see TYPE_CHECKING). + from tensorrt_llm.visual_gen.params import AudioRef, ImageRef, VideoRef + + image_refs = _build_reference_list( + request.image_reference, + modality="image", + data_field="image", + ref_cls=ImageRef, + id=id, + media_storage_path=media_storage_path, + ) + if image_refs: + params.image_reference = image_refs + video_refs = _build_reference_list( + request.video_reference, + modality="video", + data_field="video", + ref_cls=VideoRef, + id=id, + media_storage_path=media_storage_path, + ) + if video_refs: + params.video_reference = video_refs + audio_refs = _build_reference_list( + request.audio_reference, + modality="audio", + data_field="audio", + ref_cls=AudioRef, + id=id, + media_storage_path=media_storage_path, + ) + if audio_refs: + params.audio_reference = audio_refs _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _decode_inline_media(request.extra_params, generator.extra_param_specs) diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index 71d13d91bb17..32e4da6f3197 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -56,7 +56,7 @@ VisualGenArgs, ) from .output import VisualGenMetrics, VisualGenOutput - from .params import VisualGenParams + from .params import AudioRef, ImageRef, VideoRef, VisualGenParams from .visual_gen import ExtraParamSchema, VisualGen, VisualGenResult # Public name -> providing module. @@ -82,6 +82,9 @@ "VisualGenMetrics": "tensorrt_llm.visual_gen.output", "VisualGenOutput": "tensorrt_llm.visual_gen.output", "VisualGenParams": "tensorrt_llm.visual_gen.params", + "ImageRef": "tensorrt_llm.visual_gen.params", + "VideoRef": "tensorrt_llm.visual_gen.params", + "AudioRef": "tensorrt_llm.visual_gen.params", "QuantConfig": "tensorrt_llm.models.modeling_utils", } @@ -114,6 +117,9 @@ def __dir__(): "VisualGen", "VisualGenArgs", "VisualGenParams", + "ImageRef", + "VideoRef", + "AudioRef", "VisualGenResult", "VisualGenOutput", "VisualGenMetrics", diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index cd9d86ae19b8..cd2f834251c6 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -13,12 +13,61 @@ # See the License for the specific language governing permissions and # limitations under the License. import ast -from typing import Any, Dict, List, Optional, Union +from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import Field +from pydantic import Field, field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status +Role = Literal["reference", "first_frame", "last_frame"] + + +@set_api_status("prototype") +class ImageRef(StrictBaseModel): + """A single image reference carried by ``image_reference``. + + ``role`` is required only when the target model can accept the same + modality in more than one role (e.g. first + last frame); otherwise the + pipeline knows the image's meaning and ``role`` may be omitted. + """ + + image: Union[str, bytes] = Field( + description="Local path, ``http(s)``/``data:`` URL, or raw bytes." + ) + role: Optional[Role] = Field( + default=None, description="``reference`` | ``first_frame`` | ``last_frame``." + ) + + +@set_api_status("prototype") +class VideoRef(StrictBaseModel): + """A single video reference carried by ``video_reference`` (always ``reference`` role).""" + + video: Union[str, bytes] = Field( + description="Local path, ``http(s)``/``data:`` URL, or raw bytes." + ) + + +@set_api_status("prototype") +class AudioRef(StrictBaseModel): + """A single audio reference carried by ``audio_reference`` (always ``reference`` role).""" + + audio: Union[str, bytes] = Field( + description="Local path, ``http(s)``/``data:`` URL, or raw bytes." + ) + + +def _normalize_refs(value: Any, ref_cls: type, field: str) -> Optional[list]: + """Coerce a reference field to ``list[ref_cls]`` (or ``None``). + + Accepts a bare path/bytes, a single ref object, or a list mixing the two; + a bare path/bytes ``x`` becomes ``ref_cls(**{field: x})``. + """ + if value is None: + return None + items = value if isinstance(value, list) else [value] + return [x if isinstance(x, ref_cls) else ref_cls(**{field: x}) for x in items] + @set_api_status("prototype") class VisualGenParams(StrictBaseModel): @@ -78,9 +127,38 @@ class VisualGenParams(StrictBaseModel): # Conditioning inputs negative_prompt: Optional[str] = Field(default=None, description="Negative prompt for CFG.") - image: Optional[Union[str, bytes, List[Union[str, bytes]]]] = Field( - default=None, description="Reference image(s) for I2V/I2I." + # Per-modality reference inputs. A bare path/bytes, a single ref, or a + # list; normalized to ``list[*Ref]``. ``image_reference`` carries an + # optional per-item ``role`` (first_frame / last_frame / reference); + # video/audio references are always the single ``reference`` role. + image_reference: Optional[Union[str, bytes, ImageRef, List[Union[str, bytes, ImageRef]]]] = ( + Field( + default=None, + description="Reference image(s) for I2V/I2I; normalized to list[ImageRef].", + ) + ) + video_reference: Optional[Union[str, bytes, VideoRef, List[Union[str, bytes, VideoRef]]]] = ( + Field(default=None, description="Reference video(s) for V2V; normalized to list[VideoRef].") ) + audio_reference: Optional[Union[str, bytes, AudioRef, List[Union[str, bytes, AudioRef]]]] = ( + Field(default=None, description="Reference audio(s); normalized to list[AudioRef].") + ) + + @field_validator("image_reference", mode="after") + @classmethod + def _norm_image_reference(cls, v): + return _normalize_refs(v, ImageRef, "image") + + @field_validator("video_reference", mode="after") + @classmethod + def _norm_video_reference(cls, v): + return _normalize_refs(v, VideoRef, "video") + + @field_validator("audio_reference", mode="after") + @classmethod + def _norm_audio_reference(cls, v): + return _normalize_refs(v, AudioRef, "audio") + # Per-prompt multiplier num_images_per_prompt: int = Field(default=1, description="Number of images per prompt.") @@ -139,6 +217,7 @@ def validate_visual_gen_params( *, declared_defaults: Optional[Dict[str, Any]], extra_param_specs: Dict[str, Any], + ref_slot_specs: Optional[Dict[str, Any]] = None, ) -> None: """Validate *params* against pipeline-declared defaults and extra specs. @@ -228,6 +307,47 @@ def validate_visual_gen_params( f"extra_params['{key}'] value {value} is out of range [{lo}, {hi}]" ) + # --- reference role / arity checks (duck-typed RefSlotSpec) --- + # ``ref_slot_specs`` maps a reference field name to a spec exposing + # ``.roles`` (a list of role specs with ``.role`` / ``.min`` / ``.max``). + # role is required only when a modality declares more than one role; + # otherwise the single declared role is inferred. Reference fields are + # already normalized to ``list[*Ref]`` by the field validators. + if ref_slot_specs: + for field in ("image_reference", "video_reference", "audio_reference"): + refs = getattr(params, field, None) + if not refs: + continue + spec = ref_slot_specs.get(field) + if spec is None: + messages.append(f"'{field}' is not accepted by the loaded pipeline.") + continue + role_specs = list(spec.roles) + allowed = {rs.role for rs in role_specs} + role_required = len(role_specs) > 1 + counts: Dict[str, int] = {} + for r in refs: + role = getattr(r, "role", None) + if role is None: + if role_required: + messages.append( + f"{field}: 'role' is required for this model " + f"(one of {sorted(allowed)})." + ) + continue + role = role_specs[0].role + if role not in allowed: + messages.append( + f"{field}: role '{role}' not supported (allowed: {sorted(allowed)})." + ) + continue + counts[role] = counts.get(role, 0) + 1 + for rs in role_specs: + n = counts.get(rs.role, 0) + if n < rs.min or (rs.max is not None and n > rs.max): + bound = f"{rs.min}..{'inf' if rs.max is None else rs.max}" + messages.append(f"{field} role '{rs.role}': expected {bound}, got {n}.") + if not messages: return diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index b8f6a2b39e4d..5889076138f1 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -28,7 +28,7 @@ run_diffusion_worker, ) from tensorrt_llm._torch.visual_gen.output import split_visual_gen_output, to_visual_gen_output -from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema +from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema, RefSlotSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs from tensorrt_llm.visual_gen.output import VisualGenOutput @@ -38,6 +38,7 @@ "VisualGen", "VisualGenParams", "ExtraParamSchema", + "RefSlotSpec", "VisualGenResult", ] from tensorrt_llm.llmapi.utils import set_api_status @@ -291,6 +292,16 @@ def extra_param_specs(self) -> Dict[str, "ExtraParamSchema"]: """ return self.executor.extra_param_specs + @property + def ref_slot_specs(self) -> Dict[str, "RefSlotSpec"]: + """Reference slots the loaded pipeline accepts. + + Maps ``image_reference`` / ``video_reference`` / ``audio_reference`` to + a ``RefSlotSpec`` (accepted roles + per-role counts). Empty when the + pipeline takes no reference inputs. + """ + return self.executor.ref_slot_specs + @property def default_params(self) -> "VisualGenParams": """Returns a ``VisualGenParams`` with all defaults resolved for the loaded pipeline. @@ -413,6 +424,7 @@ def generate_async( resolved_params, declared_defaults=self.executor.default_generation_params, extra_param_specs=self.executor.extra_param_specs, + ref_slot_specs=self.executor.ref_slot_specs, ) else: resolved_params = self.default_params diff --git a/tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md b/tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md index 1c0808c41ea3..916b25b2ff39 100644 --- a/tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md +++ b/tests/integration/defs/perf/README_test_visual_gen_perf_sanity.md @@ -150,7 +150,7 @@ server_configs: `server_config.parallel_config` - `client_configs[].generation_mode` should be set explicitly for stable bucketing, especially for `i2v` and `t2v` -- `client_configs[].extra_body.input_reference` is the current way to express +- `client_configs[].extra_body.image_reference` is the current way to express `i2v` requests in checked-in YAMLs ## Test Case Formats diff --git a/tests/integration/defs/perf/visual_gen_perf_utils.py b/tests/integration/defs/perf/visual_gen_perf_utils.py index 08c7960e15e8..e840b7097b3d 100644 --- a/tests/integration/defs/perf/visual_gen_perf_utils.py +++ b/tests/integration/defs/perf/visual_gen_perf_utils.py @@ -105,7 +105,9 @@ def _infer_generation_mode(client_config: dict[str, Any]) -> str: except json.JSONDecodeError: extra_body = None - if isinstance(extra_body, dict) and "input_reference" in extra_body: + if isinstance(extra_body, dict) and ( + "image_reference" in extra_body or "video_reference" in extra_body + ): return "i2v" if backend == "openai-videos": diff --git a/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml b/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml index cfd8016072a6..c85816c826b1 100644 --- a/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml +++ b/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml @@ -38,7 +38,7 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - input_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= - name: ltx2_2stage_bf16_t2v_cfg2_ulysses4_compile_on model_name: ltx2_bf16 @@ -98,4 +98,4 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - input_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= diff --git a/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml b/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml index a57d9c235f86..4f12c02386ac 100644 --- a/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml +++ b/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml @@ -35,4 +35,4 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - input_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= diff --git a/tests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py b/tests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py index 4be47671a28a..2d890dd490ea 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_e2e.py @@ -381,7 +381,7 @@ def test_ti2v_sync(self, server, format_, expected_content_type): "format": format_, }, files={ - "input_reference": ("cat_piano.png", f, "image/png"), + "image_reference": ("cat_piano.png", f, "image/png"), }, ) assert resp.status_code == 200, resp.text @@ -418,7 +418,7 @@ def test_ti2v_async_lifecycle(self, server, format_, expected_content_type): "format": format_, }, files={ - "input_reference": ("cat_piano.png", f, "image/png"), + "image_reference": ("cat_piano.png", f, "image/png"), }, ) assert create_resp.status_code == 202, create_resp.text 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 27d658371451..ce76ef61ae52 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -230,6 +230,7 @@ def __init__( extra_param_specs=extra_param_specs or {"stg_scale": ExtraParamSchema(type="float", default=1.0)}, supports_image_edit=supports_image_edit, + ref_slot_specs={}, ) def _maybe_batch(self, tensor, n): @@ -1551,7 +1552,7 @@ def test_sync_video_generation_with_params(self, video_client): assert params.num_frames == int(2.0 * 8) def test_sync_video_generation_multipart(self, video_client, tmp_path): - """Multipart sync request with a real ``input_reference`` file.""" + """Multipart sync request with a real ``image_reference`` file.""" ref_path = tmp_path / "ref.png" Image.new("RGB", (4, 4), (64, 64, 64)).save(str(ref_path)) with open(ref_path, "rb") as f: @@ -1563,7 +1564,7 @@ def test_sync_video_generation_multipart(self, video_client, tmp_path): "seconds": "1.0", "fps": "8", }, - files={"input_reference": ("ref.png", f, "image/png")}, + files={"image_reference": ("ref.png", f, "image/png")}, ) assert resp.status_code == 200 assert len(resp.content) > 0 @@ -1582,25 +1583,23 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ "seconds": "1.0", "fps": "8", }, - files={"input_reference": ("ref.png", f, "image/png")}, + files={"image_reference": ("ref.png", f, "image/png")}, ) 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). + # image_reference is written to media storage and passed through as a + # typed ImageRef carrying the filesystem path. params = video_client.mock_gen.last_params - assert isinstance(params.image, str) - assert params.image.endswith("_reference") - assert os.path.exists(params.image) + ref_path = params.image_reference[0].image + assert isinstance(ref_path, str) + assert ref_path.endswith("_image_ref_0") + assert os.path.exists(ref_path) def test_sync_video_generation_multipart_with_video_reference(self, video_client): - """A video ``input_reference`` rides through as the encoded payload on - the model-specific ``video`` extra param (V2V), byte-identical — the - serve never decodes video; the worker demuxes/NVDEC-decodes it. - - Routed by container signature, so a checked-in H.264/MP4 fixture drives - the boundary directly. + """A ``video_reference`` upload is persisted byte-identical (V2V) — the + serve never decodes video; the worker demuxes/NVDEC-decodes the stored + file. A checked-in H.264/MP4 fixture drives the boundary directly. """ payload = _V2V_FIXTURE_MP4.read_bytes() with open(_V2V_FIXTURE_MP4, "rb") as f: @@ -1612,17 +1611,16 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client "seconds": "1.0", "fps": "8", }, - files={"input_reference": ("ref.mp4", f, "video/mp4")}, + files={"video_reference": ("ref.mp4", f, "video/mp4")}, ) assert resp.status_code == 200 assert len(resp.content) > 0 - # Video content must NOT land on params.image; it rides the - # model-specific ``video`` extra param as the untouched encoded bytes - # (the same intake the offline example's --video_path uses). + # Video conditioning arrives as a typed VideoRef holding a stored path; + # no image_reference is set, and the encoded bytes are byte-identical. params = video_client.mock_gen.last_params - assert params.image is None - assert params.extra_params["video"] == payload + assert params.image_reference is None + assert Path(params.video_reference[0].video).read_bytes() == payload def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content matching no image or video container signature is rejected @@ -1630,10 +1628,10 @@ def test_sync_video_generation_undecodable_reference_400(self, video_client): resp = video_client.post( "/v1/videos/sync", data={"prompt": "x"}, - files={"input_reference": ("doc.txt", BytesIO(b"not media"), "text/plain")}, + files={"image_reference": ("doc.txt", BytesIO(b"not media"), "text/plain")}, ) assert resp.status_code == 400 - assert "not a recognized media container" in resp.text + assert "not a recognized image" in resp.text def test_sync_video_failure(self, failing_client): resp = failing_client.post( @@ -1883,7 +1881,7 @@ def _blocking_save(self, *args, **kwargs): ) def test_async_video_multipart(self, video_client, tmp_path): - """Multipart async request with a real ``input_reference`` file.""" + """Multipart async request with a real ``image_reference`` file.""" ref_path = tmp_path / "ref.png" Image.new("RGB", (4, 4), (16, 16, 16)).save(str(ref_path)) with open(ref_path, "rb") as f: @@ -1895,7 +1893,7 @@ def test_async_video_multipart(self, video_client, tmp_path): "seconds": "1.0", "fps": "8", }, - files={"input_reference": ("ref.png", f, "image/png")}, + files={"image_reference": ("ref.png", f, "image/png")}, ) assert resp.status_code == 202 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 7dbf95922e8b..28dd267e286d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -54,7 +54,9 @@ def test_default_construction(self): assert params.num_frames is None assert params.frame_rate is None assert params.negative_prompt is None - assert params.image is None + assert params.image_reference is None + assert params.video_reference is None + assert params.audio_reference is None # ``image_cond_strength`` moved to per-pipeline ``extra_params`` # (only LTX-2 consumes it). It is no longer a top-level field. assert not hasattr(params, "image_cond_strength") @@ -99,23 +101,25 @@ def test_extra_params_accepted(self): assert params.extra_params["stg_scale"] == 0.5 assert params.extra_params["enhance_prompt"] is True - def test_image_accepts_str(self): + def test_image_reference_accepts_str(self): from tensorrt_llm.visual_gen import VisualGenParams - params = VisualGenParams(image="/path/to/image.png") - assert params.image == "/path/to/image.png" + params = VisualGenParams(image_reference="/path/to/image.png") + assert params.image_reference[0].image == "/path/to/image.png" + assert params.image_reference[0].role is None - def test_image_accepts_bytes(self): + def test_image_reference_accepts_bytes(self): from tensorrt_llm.visual_gen import VisualGenParams - params = VisualGenParams(image=b"\x89PNG") - assert params.image == b"\x89PNG" + params = VisualGenParams(image_reference=b"\x89PNG") + assert params.image_reference[0].image == b"\x89PNG" - def test_image_accepts_list(self): + def test_image_reference_accepts_list(self): from tensorrt_llm.visual_gen import VisualGenParams - params = VisualGenParams(image=["/path/a.png", b"\x89PNG"]) - assert len(params.image) == 2 + params = VisualGenParams(image_reference=["/path/a.png", b"\x89PNG"]) + assert len(params.image_reference) == 2 + assert params.image_reference[0].image == "/path/a.png" def test_model_dump(self): from tensorrt_llm.visual_gen import VisualGenParams @@ -299,9 +303,9 @@ def test_wan_i2v_extra_specs(self): specs = WanImageToVideoPipeline.extra_param_specs.fget( _wan_mock(is_wan22_14b=True, is_wan22_5b=False) ) - assert "last_image" in specs + # ``last_image`` moved to the typed image_reference 'last_frame' role. + assert "last_image" not in specs assert "guidance_scale_2" in specs - assert specs["last_image"].type == "str" def test_flux_no_extra_specs(self): from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux import FluxPipeline @@ -394,7 +398,7 @@ def test_flux2_reference_dimensions_remain_unset_for_pipeline_resolution(self): executor = self._make_mock_executor(Flux2Pipeline) executor.pipeline.derive_output_size_from_reference = True - req = self._make_request(image=b"encoded image") + req = self._make_request(image_reference=b"encoded image") self._merge(executor, req) @@ -407,7 +411,7 @@ def test_flux2_reference_dimensions_preserve_explicit_values(self): executor = self._make_mock_executor(Flux2Pipeline) executor.pipeline.derive_output_size_from_reference = True - req = self._make_request(image=b"encoded image", height=768, width=512) + req = self._make_request(image_reference=b"encoded image", height=768, width=512) self._merge(executor, req) @@ -818,8 +822,6 @@ def test_valid_extra_params_accepted(self): def test_spec_validator_runs_at_preflight(self): """Per-param validators turn deterministic client errors into 400s at the boundary instead of worker-side failures (Cosmos3 conditioning).""" - import torch - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params @@ -832,16 +834,6 @@ def _validate(extras): # Valid values pass. _validate({"condition_video_latent_indexes": [0, 1], "condition_video_keep": "last"}) - # ``video`` carries encoded MP4/AVI bytes: a video signature passes, - # empty / non-video bytes are client errors, and anything that is not - # bytes (e.g. a decoded tensor) fails the type check. - _validate({"video": b"\x00\x00\x00\x18ftypmp42\x00\x00\x00\x00mp42isom"}) - with pytest.raises(ValueError, match="empty"): - _validate({"video": b""}) - with pytest.raises(ValueError, match="not a recognized video container"): - _validate({"video": b"\x89PNG\r\n\x1a\n and not a video"}) - with pytest.raises(ValueError, match="expected type 'bytes'"): - _validate({"video": torch.zeros(3, 4, 4, 3, dtype=torch.uint8)}) with pytest.raises(ValueError, match="non-negative"): _validate({"condition_video_latent_indexes": [0, -1]}) @@ -887,8 +879,6 @@ def test_spec_validators_survive_pickling(self): specs = pickle.loads(pickle.dumps(COSMOS3_EXTRA_SPECS)) with pytest.raises(ValueError, match="first or last"): specs["condition_video_keep"].validator("middle") - with pytest.raises(ValueError, match="not a recognized video container"): - specs["video"].validator(b"garbage bytes") # --- unsupported universal fields --- @@ -929,13 +919,14 @@ def test_image_cond_strength_on_wan_via_extra_params_raises(self): with pytest.raises(ValueError, match="Unknown extra_params"): self._validate(executor, req) - def test_image_not_checked_by_validator(self): - """image is a conditioning input — validated at runtime by infer(), not here.""" + def test_image_reference_not_checked_without_ref_specs(self): + """Without ``ref_slot_specs``, image_reference is not role/arity checked + here — the pipeline's infer() consumes it at runtime.""" from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan import WanPipeline executor = self._make_mock_executor(WanPipeline, _wan_mock(num_heads=12)) - req = self._make_request(image="/path/to/img.png") - # Should not raise — image validation is the pipeline's responsibility + req = self._make_request(image_reference="/path/to/img.png") + # Should not raise — ``_validate`` here passes no ref_slot_specs. self._merge_and_validate(executor, req) def test_num_frames_on_video_pipeline_ok(self): @@ -946,14 +937,15 @@ def test_num_frames_on_video_pipeline_ok(self): req = self._make_request(num_frames=81) self._merge_and_validate(executor, req) - def test_image_on_i2v_pipeline_ok(self): - """image is declared by WanImageToVideoPipeline, should not raise.""" + def test_image_reference_on_i2v_pipeline_ok(self): + """image_reference is consumed by WanImageToVideoPipeline; validating + without ref_slot_specs should not raise.""" from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_i2v import ( WanImageToVideoPipeline, ) executor = self._make_mock_executor(WanImageToVideoPipeline, _wan_mock(num_heads=12)) - req = self._make_request(image="/path/to/img.png") + req = self._make_request(image_reference="/path/to/img.png") self._merge_and_validate(executor, req) def test_none_fields_not_flagged(self): @@ -1009,15 +1001,10 @@ def test_bool_rejected_for_float_spec(self): self._merge_and_validate(executor, req) def test_wrong_type_str_extra_param(self): - from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_i2v import ( - WanImageToVideoPipeline, - ) + from tensorrt_llm._torch.visual_gen.models.ltx2.pipeline_ltx2 import LTX2Pipeline - executor = self._make_mock_executor(WanImageToVideoPipeline, _wan_mock(num_heads=12)) - req = self._make_request( - image="/img.png", - extra_params={"last_image": 123}, - ) + executor = self._make_mock_executor(LTX2Pipeline) + req = self._make_request(extra_params={"output_type": 123}) with pytest.raises(ValueError, match="expected type 'str'"): self._merge_and_validate(executor, req) @@ -1375,7 +1362,7 @@ def request_warmup_cache_key(req): req = DiffusionRequest( request_id=8, prompt=["test"], - params=VisualGenParams(image=b"encoded image"), + params=VisualGenParams(image_reference=b"encoded image"), ) DiffusionExecutor.process_request(executor, req) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index de9d8956df4c..eaa971ec7c66 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -263,35 +263,51 @@ def test_video_does_not_carry_n(self): # ============================================================================= -# input_reference materialization +# reference materialization # ============================================================================= class TestInputReferenceMaterialization: - def test_base64_reference_written_to_disk(self, tmp_path): + def test_base64_image_reference_written_to_disk(self, tmp_path): generator = _StubVisualGen() img = Image.new("RGB", (4, 4), (10, 20, 30)) buf = BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) + request = VideoGenerationRequest(prompt="x", image_reference=b64) params = parse_visual_gen_params( request, "vid-1", generator, media_storage_path=str(tmp_path) ) - assert params.image is not None - assert str(params.image).endswith("vid-1_reference") + assert len(params.image_reference) == 1 + ref_path = params.image_reference[0].image + assert str(ref_path).endswith("vid-1_image_ref_0") # The decoded image is identical to what we passed in. - with open(params.image, "rb") as f: + with open(ref_path, "rb") as f: decoded = Image.open(f).convert("RGB") assert decoded.size == (4, 4) + def test_image_reference_role_and_list(self, tmp_path): + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + request = VideoGenerationRequest( + prompt="x", image_reference=[b64, {"image": b64, "role": "last_frame"}] + ) + params = parse_visual_gen_params( + request, "vid-r", generator, media_storage_path=str(tmp_path) + ) + assert [r.role for r in params.image_reference] == [None, "last_frame"] + paths = [r.image for r in params.image_reference] + assert len(set(paths)) == 2 # unique file per index + def test_missing_media_storage_path_raises(self): generator = _StubVisualGen() img = Image.new("RGB", (2, 2)) buf = BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) + request = VideoGenerationRequest(prompt="x", image_reference=b64) with pytest.raises(ValueError, match="media_storage_path"): parse_visual_gen_params(request, "vid-2", generator, media_storage_path=None) @@ -311,92 +327,116 @@ def _avi_bytes() -> bytes: TestInputReferenceMaterialization._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" ).read_bytes() - def test_multipart_avi_reference_routes_to_video(self, tmp_path): - # The AVI container signature must survive the real boundary: routed - # to the ``video`` extra param as untouched encoded bytes. + def test_multipart_avi_video_reference_written_to_disk(self, tmp_path): + # The AVI container survives the boundary and is persisted as untouched + # encoded bytes for the worker to demux. generator = _StubVisualGen() payload = self._avi_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.avi") - request = VideoGenerationRequest(prompt="x", input_reference=upload) + request = VideoGenerationRequest(prompt="x", video_reference=upload) params = parse_visual_gen_params( request, "vid-avi", generator, media_storage_path=str(tmp_path) ) - assert params.image is None - assert params.extra_params["video"] == payload + assert params.image_reference is None + assert Path(params.video_reference[0].video).read_bytes() == payload - def test_multipart_video_reference_routes_to_extra_params_bytes(self, tmp_path): + def test_multipart_mp4_video_reference_written_to_disk(self, tmp_path): generator = _StubVisualGen() payload = self._mp4_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.mp4") - request = VideoGenerationRequest(prompt="x", input_reference=upload) + request = VideoGenerationRequest(prompt="x", video_reference=upload) params = parse_visual_gen_params( request, "vid-3", generator, media_storage_path=str(tmp_path) ) - # Video content rides the model-specific ``video`` extra param as the - # encoded payload, byte-identical — the boundary never decodes video; - # the worker demuxes/NVDEC-decodes the conditioning window. - assert params.image is None - assert params.extra_params["video"] == payload - # Nothing lands in media storage for video references. - assert list(tmp_path.iterdir()) == [] - - def test_video_reference_needs_no_media_storage(self): - # Video bytes pass through in memory, so V2V works without a storage - # path at all (only image references persist a file for the worker). + # Encoded payload is persisted byte-identical — the boundary never + # decodes video; the worker demuxes/NVDEC-decodes the conditioning + # window from the stored file. + assert params.image_reference is None + vpath = params.video_reference[0].video + assert str(vpath).endswith("vid-3_video_ref_0") + assert Path(vpath).read_bytes() == payload + + def test_video_reference_needs_media_storage(self): + # Video references now persist to disk (the worker reads the path), so + # a storage path is required just like image references. generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) - params = parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) - assert params.extra_params["video"] == self._mp4_bytes() + request = VideoGenerationRequest(prompt="x", video_reference=b64) + with pytest.raises(ValueError, match="media_storage_path"): + parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) - def test_base64_video_reference_routes_to_extra_params_bytes(self, tmp_path): - # Routing is signature-based, so the JSON/base64 path can carry video - # even though it has no content-type or filename. + def test_base64_video_reference_written_to_disk(self, tmp_path): + # The JSON/base64 path carries video even though it has no content-type + # or filename; modality is declared by the field name. generator = _StubVisualGen() payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) + request = VideoGenerationRequest(prompt="x", video_reference=b64) params = parse_visual_gen_params( request, "vid-4", generator, media_storage_path=str(tmp_path) ) - assert params.image is None - assert params.extra_params["video"] == payload + assert params.image_reference is None + assert Path(params.video_reference[0].video).read_bytes() == payload - def test_video_reference_bytes_survive_real_specs(self): - """With the real cosmos3 specs loaded, parsing leaves the encoded - payload byte-identical in ``extra_params['video']`` — the boundary - never transforms video content; the worker decodes the window.""" + def test_video_reference_survives_real_specs(self, tmp_path): + """With the real cosmos3 specs loaded, the encoded payload is persisted + byte-identical — the boundary never transforms video content; the + worker decodes the conditioning window.""" from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS generator = _StubVisualGen(extra_param_specs=COSMOS3_EXTRA_SPECS) payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) - params = parse_visual_gen_params(request, "vid-10", generator, media_storage_path=None) - assert params.extra_params["video"] == payload - - def test_multipart_image_reference_routes_to_image(self, tmp_path): - # JPEG upload: content sniffing classifies it as an image and routes - # to params.image. The stored file has no type-suffix (PIL identifies - # by content, not name). + request = VideoGenerationRequest(prompt="x", video_reference=b64) + params = parse_visual_gen_params( + request, "vid-10", generator, media_storage_path=str(tmp_path) + ) + assert Path(params.video_reference[0].video).read_bytes() == payload + + def test_multipart_image_reference_written_to_disk(self, tmp_path): + # JPEG upload routed by field name to image_reference. The stored file + # has no type-suffix (PIL identifies by content, not name). generator = _StubVisualGen() img = Image.new("RGB", (4, 4), (10, 20, 30)) buf = BytesIO() img.save(buf, format="JPEG") buf.seek(0) upload = UploadFile(file=buf, filename="ref.jpg") - request = VideoGenerationRequest(prompt="x", input_reference=upload) + request = VideoGenerationRequest(prompt="x", image_reference=upload) params = parse_visual_gen_params( request, "vid-5", generator, media_storage_path=str(tmp_path) ) assert params.extra_params is None - assert str(params.image).endswith("vid-5_reference") + assert str(params.image_reference[0].image).endswith("vid-5_image_ref_0") + + def test_wrong_modality_content_raises(self, tmp_path): + # The field name declares modality; mismatched content is a client error. + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (2, 2)).save(buf, format="PNG") + img_b64 = base64.b64encode(buf.getvalue()).decode() + vid_b64 = base64.b64encode(self._mp4_bytes()).decode() + with pytest.raises(ValueError, match="video_reference is not a recognized"): + parse_visual_gen_params( + VideoGenerationRequest(prompt="x", video_reference=img_b64), + "vid-m1", + generator, + media_storage_path=str(tmp_path), + ) + with pytest.raises(ValueError, match="image_reference is not a recognized image"): + parse_visual_gen_params( + VideoGenerationRequest(prompt="x", image_reference=vid_b64), + "vid-m2", + generator, + media_storage_path=str(tmp_path), + ) + assert list(tmp_path.iterdir()) == [] - def test_undecodable_reference_raises_and_cleans_up(self, tmp_path): + def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() b64 = base64.b64encode(b"neither an image nor a video").decode() - request = VideoGenerationRequest(prompt="x", input_reference=b64) - with pytest.raises(ValueError, match="not a recognized media container"): + request = VideoGenerationRequest(prompt="x", image_reference=b64) + with pytest.raises(ValueError, match="not a recognized image"): parse_visual_gen_params(request, "vid-6", generator, media_storage_path=str(tmp_path)) # Classification runs on the bytes; rejected content never touches disk. assert list(tmp_path.iterdir()) == [] @@ -405,7 +445,7 @@ def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() # "ABC" survives the lenient alphabet filter but has an invalid # length, so b64decode raises. - request = VideoGenerationRequest(prompt="x", input_reference="ABC") + request = VideoGenerationRequest(prompt="x", image_reference="ABC") with pytest.raises(ValueError, match="not valid base64"): parse_visual_gen_params(request, "vid-7", generator, media_storage_path=str(tmp_path)) assert list(tmp_path.iterdir()) == [] @@ -418,11 +458,11 @@ def read(self, *args, **kwargs): raise OSError("client went away") upload = UploadFile(file=_BrokenStream(), filename="clip.mp4") - request = VideoGenerationRequest(prompt="x", input_reference=upload) + request = VideoGenerationRequest(prompt="x", video_reference=upload) # I/O failures keep their server-error semantics (no 400 masking) … with pytest.raises(OSError, match="client went away"): parse_visual_gen_params(request, "vid-8", generator, media_storage_path=str(tmp_path)) - # … but the partial materialization must not leak. + # … and the payload read fails before any file is written, so nothing leaks. assert list(tmp_path.iterdir()) == [] @@ -534,7 +574,7 @@ def test_heif_reference_rejected_with_actionable_message(self): generator = _StubVisualGen() heic = self._ftyp(b"heic", (b"mif1", b"heic")) + b"\x00" * 64 request = VideoGenerationRequest( - prompt="x", input_reference=base64.b64encode(heic).decode() + prompt="x", image_reference=base64.b64encode(heic).decode() ) with pytest.raises(ValueError, match="HEIF/AVIF"): parse_visual_gen_params(request, "vid-heic", generator, media_storage_path=None) @@ -557,12 +597,12 @@ def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): generator = _StubVisualGen() request = VideoGenerationRequest( - prompt="x", input_reference=base64.b64encode(truncated).decode() + prompt="x", image_reference=base64.b64encode(truncated).decode() ) params = parse_visual_gen_params( request, "vid-12", generator, media_storage_path=str(tmp_path) ) - assert Path(params.image).read_bytes() == truncated + assert Path(params.image_reference[0].image).read_bytes() == truncated # ============================================================================= diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index eae652bda660..e3ef313a287e 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1471,9 +1471,21 @@ models: default: null status: stable required: false - input_reference: + image_reference: kind: extension - type: Optional[Union[str, UploadFile]] + type: Optional[Union[str, UploadFile, ImageReferenceItem, List[Union[str, ImageReferenceItem]]]] + default: null + status: prototype + required: false + video_reference: + kind: extension + type: Optional[Union[str, UploadFile, VideoReferenceItem, List[Union[str, VideoReferenceItem]]]] + default: null + status: prototype + required: false + audio_reference: + kind: extension + type: Optional[Union[str, UploadFile, AudioReferenceItem, List[Union[str, AudioReferenceItem]]]] default: null status: prototype required: false From 484476a56902584e2b360c4f8801d5e8218d4b12 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 11 Aug 2026 21:29:45 -0700 Subject: [PATCH 02/61] [TRTLLM-15277][feat] Enforce required vs optional reference slots Make RoleSpec.min meaningful for whole-slot presence: validate_visual_gen_params now runs the per-role arity check on a declared reference slot even when it is empty, so a role with min>=1 is a required reference that fails as a clean 400 instead of crashing deep in the worker; min=0 leaves the slot optional. An absent undeclared slot stays silently accepted; only an unsolicited one is rejected. Set min=0 on the optional-image pipelines (LTX-2, Wan TI2V, Cosmos3 image+video, FLUX.2) whose infer() handles refs=None; keep min=1 on the required-image ones (Wan I2V first_frame, Qwen-edit, Qwen-Layered). Add test_ref_slot_required_vs_optional. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../models/cosmos3/pipeline_cosmos3.py | 6 ++-- .../visual_gen/models/flux/pipeline_flux2.py | 5 +-- .../visual_gen/models/ltx2/pipeline_ltx2.py | 3 +- .../visual_gen/models/wan/pipeline_wan.py | 2 +- tensorrt_llm/visual_gen/params.py | 12 ++++--- .../visual_gen/test_visual_gen_params.py | 36 +++++++++++++++++++ 6 files changed, 54 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 00fa4c70049c..3051b0b9466a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -606,13 +606,15 @@ def extra_param_specs(self): @property def ref_slot_specs(self): return { + # image (I2V) and video (V2V) are both optional; Cosmos3 also runs + # T2V with neither. "image_reference": RefSlotSpec( modality="image", - roles=[RoleSpec(role="first_frame", min=1, max=1)], + roles=[RoleSpec(role="first_frame", min=0, max=1)], ), "video_reference": RefSlotSpec( modality="video", - roles=[RoleSpec(role="reference", min=1, max=1)], + roles=[RoleSpec(role="reference", min=0, max=1)], ), } 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 0845beb6f070..a0ee93d2dcf4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -372,11 +372,12 @@ def default_generation_params(self): @property def ref_slot_specs(self): - # Reference image(s): single "reference" role, count 1..N (multi-subject). + # Optional reference image(s): "reference" role, count 0..N (multi-subject); + # FLUX.2 also runs plain text-to-image with none. return { "image_reference": RefSlotSpec( modality="image", - roles=[RoleSpec(role="reference", min=1, max=None)], + roles=[RoleSpec(role="reference", min=0, max=None)], ) } 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 27287bab5924..0673ece69952 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1384,7 +1384,8 @@ def ref_slot_specs(self): return { "image_reference": RefSlotSpec( modality="image", - roles=[RoleSpec(role="first_frame", min=1, max=1)], + # Optional first-frame conditioning (min=0); LTX-2 also runs T2V. + roles=[RoleSpec(role="first_frame", min=0, max=1)], ), } 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 04c9cda12e66..53c3c0c8e112 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -430,7 +430,7 @@ def ref_slot_specs(self): return { "image_reference": RefSlotSpec( modality="image", - roles=[RoleSpec(role="first_frame", min=1, max=1)], + roles=[RoleSpec(role="first_frame", min=0, max=1)], ) } diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index cd2f834251c6..20be51353088 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -315,12 +315,13 @@ def validate_visual_gen_params( # already normalized to ``list[*Ref]`` by the field validators. if ref_slot_specs: for field in ("image_reference", "video_reference", "audio_reference"): - refs = getattr(params, field, None) - if not refs: - continue + refs = getattr(params, field, None) or [] spec = ref_slot_specs.get(field) if spec is None: - messages.append(f"'{field}' is not accepted by the loaded pipeline.") + # An undeclared slot is only an error if the client actually + # sent one; an absent undeclared slot is fine. + if refs: + messages.append(f"'{field}' is not accepted by the loaded pipeline.") continue role_specs = list(spec.roles) allowed = {rs.role for rs in role_specs} @@ -342,6 +343,9 @@ def validate_visual_gen_params( ) continue counts[role] = counts.get(role, 0) + 1 + # Arity runs even for an absent slot: a role with ``min >= 1`` is a + # required reference, enforced here as a clean 400 instead of a deep + # worker crash. ``min == 0`` leaves the slot optional. for rs in role_specs: n = counts.get(rs.role, 0) if n < rs.min or (rs.max is not None and n > rs.max): 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 28dd267e286d..965120fafc5c 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -948,6 +948,42 @@ def test_image_reference_on_i2v_pipeline_ok(self): req = self._make_request(image_reference="/path/to/img.png") self._merge_and_validate(executor, req) + def test_ref_slot_required_vs_optional(self): + """``min >= 1`` marks a required reference (clean error when absent); + ``min == 0`` leaves the slot optional; an undeclared absent slot is + fine, but an unsolicited one is rejected.""" + from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + required = { + "image_reference": RefSlotSpec( + modality="image", roles=[RoleSpec(role="reference", min=1, max=1)] + ) + } + optional = { + "image_reference": RefSlotSpec( + modality="image", roles=[RoleSpec(role="first_frame", min=0, max=1)] + ) + } + + def run(params, spec): + validate_visual_gen_params( + params, declared_defaults=None, extra_param_specs={}, ref_slot_specs=spec + ) + + # Required slot, no image -> clean 400 here instead of a worker crash. + with pytest.raises(ValueError, match=r"expected 1\.\.1, got 0"): + run(VisualGenParams(), required) + # Optional slot, no image -> allowed (e.g. text-to-video). + run(VisualGenParams(), optional) + # Required slot with the image present -> allowed. + run(VisualGenParams(image_reference="a.png"), required) + # Undeclared slot left absent -> no spurious "not accepted". + run(VisualGenParams(), optional) + # Undeclared slot actually sent -> rejected. + with pytest.raises(ValueError, match=r"video_reference.*not accepted"): + run(VisualGenParams(video_reference="v.mp4"), optional) + def test_none_fields_not_flagged(self): """Fields left as None should never trigger unsupported-field errors.""" from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux import FluxPipeline From e4e4ff6cdd238f3e1d3eed4e5b2d4c84ac08ae6c Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 12 Aug 2026 01:04:38 -0700 Subject: [PATCH 03/61] [TRTLLM-15277][doc] Document reference inputs (image/video/audio_reference) Add a Reference inputs section to the VisualGen doc: the typed per-modality image_reference / video_reference / audio_reference fields shared by the Python API and serve, a simple no-role example (single-meaning image for I2V, Cosmos video_reference for V2V) and a complex role-required example (Wan I2V first_frame + last_frame). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 55 +++++++++++++++++++++++++ 1 file changed, 55 insertions(+) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 941e7df4b473..142ebf013148 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -119,6 +119,61 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue `response_format="path"` returns the generated file's server-side path (under `TRTLLM_MEDIA_STORAGE_PATH`) for co-located clients, enabled by default. Set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject such requests with HTTP 400. See the [serve examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/visual_gen/serve) for the full `response_format` reference. +### Reference Inputs + +Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, or uploaded file) to a local path before reaching the worker. + +Most models take a single reference whose role is unambiguous, so no `role` is specified: + +```python +from tensorrt_llm import VisualGen + +# The image conditions the generated video's first frame. +vg = VisualGen(model="Wan-AI/Wan2.2-TI2V-5B-Diffusers") +params = vg.default_params +params.image_reference = "start.png" +output = vg.generate(inputs="the scene comes alive with gentle motion", params=params) + +# Cosmos conditions generation on a reference video. +vg = VisualGen(model="nvidia/Cosmos3-Super") +params = vg.default_params +params.video_reference = "clip.mp4" +``` + +The equivalent serve request uploads the file, or sends a base64 string or `data:` URI in a JSON body: + +```bash +curl http://localhost:8000/v1/videos -F "prompt=the scene comes alive" -F "image_reference=@start.png" +curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_reference=@clip.mp4" +``` + +When a model accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame — the `role` is required to disambiguate: + +```python +from tensorrt_llm import VisualGen, ImageRef # VideoRef and AudioRef are also exported + +vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") +params = vg.default_params +params.image_reference = [ + ImageRef(image="start.png", role="first_frame"), + ImageRef(image="end.png", role="last_frame"), # optional +] +``` + +A JSON serve request carries the role and lists; a multipart upload is limited to a single file with no role: + +```bash +curl http://localhost:8000/v1/videos -H 'content-type: application/json' -d '{ + "prompt": "the subject comes alive", + "image_reference": [ + {"image": "", "role": "first_frame"}, + {"image": "", "role": "last_frame"} + ] +}' +``` + +FLUX.2 and Qwen-Image-Edit accept multiple reference images as a list on the same `image_reference` field through the Python API. + ## Optimizations ### Quantization From c872342d8cd966dee16992e635efd1c5f2c2d51a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:12:35 -0700 Subject: [PATCH 04/61] [TRTLLM-15277][refactor] Merge ImageRef/VideoRef/AudioRef into one MediaRef Per PR review: the three per-modality reference types were structurally identical (a Union[str, bytes] payload, plus role on ImageRef) and carried no modality/arity validation of their own -- that is owned by the pipeline via ref_slot_specs, keyed on the reference field name, not the item type. Collapse them into a single MediaRef(content, role); the serve items likewise become one MediaReferenceItem(content, role). The three outer fields (image_reference / video_reference / audio_reference) stay -- they carry the modality and drive ref_slot_specs -- and the pipelines read refs[i].content. _normalize_refs and _build_reference_list drop their ref_cls/data_field parameters. No change to ref_slot_specs, validation, or the outer API shape; the per-item payload key becomes content. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 10 +-- examples/visual_gen/models/wan_i2v.py | 4 +- tensorrt_llm/__init__.py | 14 ++-- .../models/cosmos3/pipeline_cosmos3.py | 4 +- .../visual_gen/models/flux/pipeline_flux2.py | 4 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 2 +- .../models/ltx2/pipeline_ltx2_two_stages.py | 2 +- .../qwen_image/pipeline_qwen_image_edit.py | 2 +- .../pipeline_qwen_image_layered.py | 2 +- .../visual_gen/models/wan/pipeline_wan.py | 2 +- .../visual_gen/models/wan/pipeline_wan_i2v.py | 4 +- tensorrt_llm/serve/openai_protocol.py | 59 ++++++-------- tensorrt_llm/serve/visual_gen_utils.py | 62 +++++--------- tensorrt_llm/visual_gen/__init__.py | 10 +-- tensorrt_llm/visual_gen/params.py | 80 +++++++------------ .../visual_gen/test_trtllm_serve_endpoints.py | 8 +- .../visual_gen/test_visual_gen_params.py | 6 +- .../visual_gen/test_visual_gen_utils.py | 18 ++--- .../references/trtllm_serve_api.yaml | 6 +- 19 files changed, 116 insertions(+), 183 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 142ebf013148..3103c4e48a4d 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -150,13 +150,13 @@ curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_re When a model accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame — the `role` is required to disambiguate: ```python -from tensorrt_llm import VisualGen, ImageRef # VideoRef and AudioRef are also exported +from tensorrt_llm import VisualGen, MediaRef vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") params = vg.default_params params.image_reference = [ - ImageRef(image="start.png", role="first_frame"), - ImageRef(image="end.png", role="last_frame"), # optional + MediaRef(content="start.png", role="first_frame"), + MediaRef(content="end.png", role="last_frame"), # optional ] ``` @@ -166,8 +166,8 @@ A JSON serve request carries the role and lists; a multipart upload is limited t curl http://localhost:8000/v1/videos -H 'content-type: application/json' -d '{ "prompt": "the subject comes alive", "image_reference": [ - {"image": "", "role": "first_frame"}, - {"image": "", "role": "last_frame"} + {"content": "", "role": "first_frame"}, + {"content": "", "role": "last_frame"} ] }' ``` diff --git a/examples/visual_gen/models/wan_i2v.py b/examples/visual_gen/models/wan_i2v.py index 9876b6ebaf45..b650fbdcd320 100644 --- a/examples/visual_gen/models/wan_i2v.py +++ b/examples/visual_gen/models/wan_i2v.py @@ -23,7 +23,7 @@ import argparse import os -from tensorrt_llm import ImageRef, VisualGen, VisualGenArgs +from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs _DEFAULT_IMAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cat_piano.png") @@ -66,7 +66,7 @@ def main(): # first-frame reference. Wan I2V also accepts a ``last_frame`` role, so the # role must be given to disambiguate. params = visual_gen.default_params - params.image_reference = [ImageRef(image=args.image, role="first_frame")] + params.image_reference = [MediaRef(content=args.image, role="first_frame")] output = visual_gen.generate( inputs="A cat presses the piano keys with its paws, soft notes filling the quiet room.", diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index aefaeaf503c6..e878bf7ad36b 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -62,9 +62,9 @@ from .mapping import Mapping from .models.automodel import AutoConfig, AutoModelForCausalLM from .sampling_params import SamplingParams - from .visual_gen import (AudioRef, ExtraParamSchema, ImageRef, VideoRef, - VisualGen, VisualGenArgs, VisualGenMetrics, - VisualGenOutput, VisualGenParams, VisualGenResult) + from .visual_gen import (ExtraParamSchema, MediaRef, VisualGen, + VisualGenArgs, VisualGenMetrics, VisualGenOutput, + VisualGenParams, VisualGenResult) # Public name -> (source module, attribute); attribute None = the module itself. _LAZY_ATTRS = { @@ -105,9 +105,7 @@ 'VisualGenMetrics': ('tensorrt_llm.visual_gen', 'VisualGenMetrics'), 'VisualGenOutput': ('tensorrt_llm.visual_gen', 'VisualGenOutput'), 'VisualGenParams': ('tensorrt_llm.visual_gen', 'VisualGenParams'), - 'ImageRef': ('tensorrt_llm.visual_gen', 'ImageRef'), - 'VideoRef': ('tensorrt_llm.visual_gen', 'VideoRef'), - 'AudioRef': ('tensorrt_llm.visual_gen', 'AudioRef'), + 'MediaRef': ('tensorrt_llm.visual_gen', 'MediaRef'), 'VisualGenResult': ('tensorrt_llm.visual_gen', 'VisualGenResult'), } @@ -176,9 +174,7 @@ def __dir__(): 'math_utils', 'VisualGen', 'VisualGenParams', - 'ImageRef', - 'VideoRef', - 'AudioRef', + 'MediaRef', '__version__', ] diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 3051b0b9466a..6e280fbfa26a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -759,7 +759,7 @@ def as_given(field_name): return value if field_name in specified else None refs_v = req.params.video_reference - video = refs_v[0].video if refs_v else None + video = refs_v[0].content if refs_v else None if isinstance(video, str): from pathlib import Path @@ -835,7 +835,7 @@ def as_given(field_name): return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, - image=refs_i[0].image if refs_i else None, + image=refs_i[0].content if refs_i else None, height=height, width=width, num_frames=req.params.num_frames, 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 a0ee93d2dcf4..57beb5b94f43 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -387,7 +387,7 @@ def prepare_request(self, req: Any) -> None: if not refs: return - reference_images = self._load_reference_images([r.image for r in refs]) + reference_images = self._load_reference_images([r.content for r in refs]) condition_images = self._preprocess_reference_images(reference_images) req.params.height, req.params.width = self._resolve_target_dimensions( req.params.height, @@ -408,7 +408,7 @@ def infer(self, req): seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, num_images_per_prompt=req.params.num_images_per_prompt, - image=[r.image for r in refs] if refs else None, + image=[r.content for r in refs] if refs else None, _condition_images=req.prepared_inputs.get("condition_images"), ) 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 0673ece69952..c36bee296819 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1406,7 +1406,7 @@ def infer(self, req): output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], max_sequence_length=req.params.max_sequence_length, - image=refs[0].image if refs else None, + image=refs[0].content if refs else None, image_cond_strength=extra["image_cond_strength"], stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], 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 0bc6c374e0be..72377d4113b0 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 @@ -1240,7 +1240,7 @@ def infer(self, req): output_type=extra["output_type"], guidance_rescale=extra["guidance_rescale"], max_sequence_length=req.params.max_sequence_length, - image=refs[0].image if refs else None, + image=refs[0].content if refs else None, image_cond_strength=extra["image_cond_strength"], stg_scale=extra["stg_scale"], stg_blocks=extra["stg_blocks"], 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 5f0d43cf2307..af329b32c272 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 @@ -352,7 +352,7 @@ def infer(self, req: Any) -> PipelineOutput: ) prompts = req.prompt if isinstance(req.prompt, list) else [req.prompt] refs = params.image_reference - pil_images = self._load_edit_images([r.image for r in refs] if refs else None) + pil_images = self._load_edit_images([r.content for r in refs] if refs else None) height = params.height width = params.width if height is None or width is None: 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 4edf20cb4eee..ae8489005c9c 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 @@ -774,7 +774,7 @@ def infer(self, req): refs = req.params.image_reference return self.forward( - image=refs[0].image if refs else None, + image=refs[0].content if refs else None, prompt=prompts, negative_prompt=negative, height=req.params.height, 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 53c3c0c8e112..ca945f27e80d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -439,7 +439,7 @@ def infer(self, req): extra = req.params.extra_params or {} # Wan 2.2 TI2V-5B takes one conditioning image if provided refs = req.params.image_reference - image = refs[0].image if refs else None + image = refs[0].content if refs else None return self.forward( prompt=req.prompt, 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 7eda46df2842..8c40493d193f 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 @@ -427,8 +427,8 @@ def infer(self, req): if first is None: raise ValueError("I2V pipeline requires a first_frame image_reference") last = by_role.get("last_frame") - image = first.image - last_image = last.image if last is not None else None + image = first.content + last_image = last.content if last is not None else None extra = req.params.extra_params or {} return self.forward( diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index b3f26c6df0ed..28452400db19 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2011,38 +2011,25 @@ class ImageGenerationResponse(OpenAIBaseModel): size: Optional[str] = None -class ImageReferenceItem(OpenAIBaseModel): - """One image reference for conditioning (mirrors ``ImageRef``). - - ``image`` carries base64-encoded bytes, optionally as a ``data:`` URI. - ``role`` is required only for models that accept more than one image - role (e.g. Wan first/last frame); single-role models infer it. +class MediaReferenceItem(OpenAIBaseModel): + """One media reference (image / video / audio) for conditioning (mirrors ``MediaRef``). + + ``content`` carries base64-encoded bytes, optionally as a ``data:`` URI. The + request field it sits in (``image_reference`` / ``video_reference`` / + ``audio_reference``) fixes the modality. ``role`` is required only for models + that accept more than one role for that modality (e.g. image first/last + frame); single-role models infer it. """ - image: str = Field( - description="Base64-encoded image bytes, optionally as a ``data:`` URI." + content: str = Field( + description="Base64-encoded media bytes, optionally as a ``data:`` URI." ) role: Optional[str] = Field( default=None, - description=( - "Reference role (e.g. 'reference', 'first_frame', 'last_frame'). " - "Required only when the model accepts multiple image roles."), - ) - - -class VideoReferenceItem(OpenAIBaseModel): - """One video reference for conditioning (mirrors ``VideoRef``).""" - - video: str = Field( - description="Base64-encoded video bytes, optionally as a ``data:`` URI." - ) - - -class AudioReferenceItem(OpenAIBaseModel): - """One audio reference for conditioning (mirrors ``AudioRef``).""" - - audio: str = Field( - description="Base64-encoded audio bytes, optionally as a ``data:`` URI." + description= + ("Reference role (e.g. 'reference', 'first_frame', 'last_frame'). " + "Required only when the model accepts multiple roles for the modality." + ), ) @@ -2070,32 +2057,32 @@ class VideoGenerationRequest(OpenAIBaseModel): ge=0, description="Random seed for reproducibility.") image_reference: Optional[Union[ - str, UploadFile, ImageReferenceItem, - List[Union[str, ImageReferenceItem]]]] = Field( + str, UploadFile, MediaReferenceItem, + List[Union[str, MediaReferenceItem]]]] = Field( default=None, description= ("Image reference(s) conditioning generation (e.g. image-to-video " - "first frame). JSON sends base64 bytes, an ``{image, role}`` " + "first frame). JSON sends base64 bytes, a ``{content, role}`` " "object, or a list of them; multipart uploads a single image file. " "PNG or JPEG only — HEIF/AVIF are not supported."), ) video_reference: Optional[Union[ - str, UploadFile, VideoReferenceItem, - List[Union[str, VideoReferenceItem]]]] = Field( + str, UploadFile, MediaReferenceItem, + List[Union[str, MediaReferenceItem]]]] = Field( default=None, description= ("Video reference(s) conditioning generation (video-to-video). JSON " - "sends base64 bytes, a ``{video}`` object, or a list of them; " + "sends base64 bytes, a ``{content}`` object, or a list of them; " "multipart uploads a single video file. MP4 or AVI, with H.264 the " "tested codec and others best-effort."), ) audio_reference: Optional[Union[ - str, UploadFile, AudioReferenceItem, - List[Union[str, AudioReferenceItem]]]] = Field( + str, UploadFile, MediaReferenceItem, + List[Union[str, MediaReferenceItem]]]] = Field( default=None, description= ("Audio reference(s) conditioning generation. JSON sends base64 " - "bytes, an ``{audio}`` object, or a list of them; multipart uploads " + "bytes, a ``{content}`` object, or a list of them; multipart uploads " "a single audio file. Accepted only by models that declare an audio " "reference slot."), ) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 458b655aed18..6974dc6455cd 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -128,21 +128,21 @@ def _read_reference_payload(reference: str) -> bytes: raise ValueError("reference is not valid base64 data.") from exc -def _reference_payload_and_role(ref, data_field: str) -> tuple[bytes, Optional[str]]: +def _reference_payload_and_role(ref) -> tuple[bytes, Optional[str]]: """Extract ``(payload_bytes, role)`` from one raw HTTP reference. ``ref`` is a base64/data-URI string, a multipart ``UploadFile`` (has - ``.file``), or a reference item exposing ``data_field`` ('image'/'video') - and, for images, an optional ``role``. + ``.file``), or a ``MediaReferenceItem`` exposing ``content`` and an optional + ``role``. """ role = getattr(ref, "role", None) if isinstance(ref, str): return _read_reference_payload(ref), role if hasattr(ref, "file"): # multipart UploadFile return ref.file.read(), role - data = getattr(ref, data_field, None) + data = getattr(ref, "content", None) if not isinstance(data, str): - raise ValueError(f"{data_field}_reference item must carry a base64 '{data_field}' string.") + raise ValueError("reference item must carry a base64 'content' string.") return _read_reference_payload(data), role @@ -181,32 +181,31 @@ def _materialize_reference( return ref_path -def _build_reference_list( - value, *, modality: str, data_field: str, ref_cls, id: str, media_storage_path: Optional[str] -): - """Materialize an HTTP reference field into a list of typed ``*Ref`` objects. +def _build_reference_list(value, *, modality: str, id: str, media_storage_path: Optional[str]): + """Materialize an HTTP reference field into a list of ``MediaRef`` objects. ``value`` is None, a base64/data-URI string, a multipart ``UploadFile``, a - reference item, or a list of any of those. Each entry is decoded, + ``MediaReferenceItem``, or a list of any of those. Each entry is decoded, content-validated for ``modality``, persisted to a per-index path, and - wrapped as ``ref_cls`` (carrying ``role`` for images). + wrapped as ``MediaRef`` (carrying ``role`` when present). """ if value is None: return None + # Local import: the visual_gen tree is already loaded in a VisualGen serving + # process, and this keeps it out of every plain-LLM process (see TYPE_CHECKING). + from tensorrt_llm.visual_gen.params import MediaRef + raw_items = value if isinstance(value, list) else [value] refs = [] for i, item in enumerate(raw_items): - payload, role = _reference_payload_and_role(item, data_field) + payload, role = _reference_payload_and_role(item) ref_path = _materialize_reference( payload, modality=modality, - ref_id=f"{id}_{data_field}_ref_{i}", + ref_id=f"{id}_{modality}_ref_{i}", media_storage_path=media_storage_path, ) - kwargs = {data_field: ref_path} - if role is not None: - kwargs["role"] = role - refs.append(ref_cls(**kwargs)) + refs.append(MediaRef(content=ref_path, role=role)) return refs @@ -513,39 +512,20 @@ def parse_visual_gen_params( ) params.num_frames = derived # Reference inputs: materialize each transport (base64/data-URI/upload) - # to a stored file and hand the pipeline a typed ``*Ref`` carrying the - # local path. Decode stays model-specific in the worker. Local import: - # the visual_gen tree is already loaded in a VisualGen serving process, - # and this keeps it out of every plain-LLM process (see TYPE_CHECKING). - from tensorrt_llm.visual_gen.params import AudioRef, ImageRef, VideoRef - + # to a stored file and hand the pipeline a ``MediaRef`` carrying the + # local path. Decode stays model-specific in the worker. image_refs = _build_reference_list( - request.image_reference, - modality="image", - data_field="image", - ref_cls=ImageRef, - id=id, - media_storage_path=media_storage_path, + request.image_reference, modality="image", id=id, media_storage_path=media_storage_path ) if image_refs: params.image_reference = image_refs video_refs = _build_reference_list( - request.video_reference, - modality="video", - data_field="video", - ref_cls=VideoRef, - id=id, - media_storage_path=media_storage_path, + request.video_reference, modality="video", id=id, media_storage_path=media_storage_path ) if video_refs: params.video_reference = video_refs audio_refs = _build_reference_list( - request.audio_reference, - modality="audio", - data_field="audio", - ref_cls=AudioRef, - id=id, - media_storage_path=media_storage_path, + request.audio_reference, modality="audio", id=id, media_storage_path=media_storage_path ) if audio_refs: params.audio_reference = audio_refs diff --git a/tensorrt_llm/visual_gen/__init__.py b/tensorrt_llm/visual_gen/__init__.py index 32e4da6f3197..98027abdf83d 100644 --- a/tensorrt_llm/visual_gen/__init__.py +++ b/tensorrt_llm/visual_gen/__init__.py @@ -56,7 +56,7 @@ VisualGenArgs, ) from .output import VisualGenMetrics, VisualGenOutput - from .params import AudioRef, ImageRef, VideoRef, VisualGenParams + from .params import MediaRef, VisualGenParams from .visual_gen import ExtraParamSchema, VisualGen, VisualGenResult # Public name -> providing module. @@ -82,9 +82,7 @@ "VisualGenMetrics": "tensorrt_llm.visual_gen.output", "VisualGenOutput": "tensorrt_llm.visual_gen.output", "VisualGenParams": "tensorrt_llm.visual_gen.params", - "ImageRef": "tensorrt_llm.visual_gen.params", - "VideoRef": "tensorrt_llm.visual_gen.params", - "AudioRef": "tensorrt_llm.visual_gen.params", + "MediaRef": "tensorrt_llm.visual_gen.params", "QuantConfig": "tensorrt_llm.models.modeling_utils", } @@ -117,9 +115,7 @@ def __dir__(): "VisualGen", "VisualGenArgs", "VisualGenParams", - "ImageRef", - "VideoRef", - "AudioRef", + "MediaRef", "VisualGenResult", "VisualGenOutput", "VisualGenMetrics", diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 20be51353088..5a59bb63eb25 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -23,15 +23,17 @@ @set_api_status("prototype") -class ImageRef(StrictBaseModel): - """A single image reference carried by ``image_reference``. - - ``role`` is required only when the target model can accept the same - modality in more than one role (e.g. first + last frame); otherwise the - pipeline knows the image's meaning and ``role`` may be omitted. +class MediaRef(StrictBaseModel): + """A single media reference (image / video / audio). + + Carried by ``image_reference`` / ``video_reference`` / ``audio_reference``; + the field it sits in fixes the modality. ``role`` is required only when the + target model accepts that modality in more than one role (e.g. image first + + last frame); otherwise the pipeline knows the reference's meaning and + ``role`` may be omitted (video/audio are always the single ``reference``). """ - image: Union[str, bytes] = Field( + content: Union[str, bytes] = Field( description="Local path, ``http(s)``/``data:`` URL, or raw bytes." ) role: Optional[Role] = Field( @@ -39,34 +41,16 @@ class ImageRef(StrictBaseModel): ) -@set_api_status("prototype") -class VideoRef(StrictBaseModel): - """A single video reference carried by ``video_reference`` (always ``reference`` role).""" - - video: Union[str, bytes] = Field( - description="Local path, ``http(s)``/``data:`` URL, or raw bytes." - ) +def _normalize_refs(value: Any) -> Optional[list]: + """Coerce a reference field to ``list[MediaRef]`` (or ``None``). - -@set_api_status("prototype") -class AudioRef(StrictBaseModel): - """A single audio reference carried by ``audio_reference`` (always ``reference`` role).""" - - audio: Union[str, bytes] = Field( - description="Local path, ``http(s)``/``data:`` URL, or raw bytes." - ) - - -def _normalize_refs(value: Any, ref_cls: type, field: str) -> Optional[list]: - """Coerce a reference field to ``list[ref_cls]`` (or ``None``). - - Accepts a bare path/bytes, a single ref object, or a list mixing the two; - a bare path/bytes ``x`` becomes ``ref_cls(**{field: x})``. + Accepts a bare path/bytes, a single ``MediaRef``, or a list mixing the two; + a bare path/bytes ``x`` becomes ``MediaRef(content=x)``. """ if value is None: return None items = value if isinstance(value, list) else [value] - return [x if isinstance(x, ref_cls) else ref_cls(**{field: x}) for x in items] + return [x if isinstance(x, MediaRef) else MediaRef(content=x) for x in items] @set_api_status("prototype") @@ -127,37 +111,27 @@ class VisualGenParams(StrictBaseModel): # Conditioning inputs negative_prompt: Optional[str] = Field(default=None, description="Negative prompt for CFG.") - # Per-modality reference inputs. A bare path/bytes, a single ref, or a - # list; normalized to ``list[*Ref]``. ``image_reference`` carries an - # optional per-item ``role`` (first_frame / last_frame / reference); - # video/audio references are always the single ``reference`` role. - image_reference: Optional[Union[str, bytes, ImageRef, List[Union[str, bytes, ImageRef]]]] = ( + # Per-modality reference inputs. A bare path/bytes, a single ``MediaRef``, + # or a list; normalized to ``list[MediaRef]``. The field fixes the modality; + # ``role`` is only meaningful where a model declares more than one role for + # it (e.g. image first_frame / last_frame). + image_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( Field( default=None, - description="Reference image(s) for I2V/I2I; normalized to list[ImageRef].", + description="Reference image(s) for I2V/I2I; normalized to list[MediaRef].", ) ) - video_reference: Optional[Union[str, bytes, VideoRef, List[Union[str, bytes, VideoRef]]]] = ( - Field(default=None, description="Reference video(s) for V2V; normalized to list[VideoRef].") + video_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( + Field(default=None, description="Reference video(s) for V2V; normalized to list[MediaRef].") ) - audio_reference: Optional[Union[str, bytes, AudioRef, List[Union[str, bytes, AudioRef]]]] = ( - Field(default=None, description="Reference audio(s); normalized to list[AudioRef].") + audio_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( + Field(default=None, description="Reference audio(s); normalized to list[MediaRef].") ) - @field_validator("image_reference", mode="after") - @classmethod - def _norm_image_reference(cls, v): - return _normalize_refs(v, ImageRef, "image") - - @field_validator("video_reference", mode="after") - @classmethod - def _norm_video_reference(cls, v): - return _normalize_refs(v, VideoRef, "video") - - @field_validator("audio_reference", mode="after") + @field_validator("image_reference", "video_reference", "audio_reference", mode="after") @classmethod - def _norm_audio_reference(cls, v): - return _normalize_refs(v, AudioRef, "audio") + def _norm_refs(cls, v): + return _normalize_refs(v) # Per-prompt multiplier num_images_per_prompt: int = Field(default=1, description="Number of images per prompt.") 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 ce76ef61ae52..983870b23330 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -1589,9 +1589,9 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert len(resp.content) > 0 # image_reference is written to media storage and passed through as a - # typed ImageRef carrying the filesystem path. + # MediaRef carrying the filesystem path. params = video_client.mock_gen.last_params - ref_path = params.image_reference[0].image + ref_path = params.image_reference[0].content assert isinstance(ref_path, str) assert ref_path.endswith("_image_ref_0") assert os.path.exists(ref_path) @@ -1616,11 +1616,11 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video conditioning arrives as a typed VideoRef holding a stored path; + # Video conditioning arrives as a MediaRef holding a stored path; # no image_reference is set, and the encoded bytes are byte-identical. params = video_client.mock_gen.last_params assert params.image_reference is None - assert Path(params.video_reference[0].video).read_bytes() == payload + assert Path(params.video_reference[0].content).read_bytes() == payload def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content matching no image or video container signature is rejected 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 965120fafc5c..b55c79425219 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -105,21 +105,21 @@ def test_image_reference_accepts_str(self): from tensorrt_llm.visual_gen import VisualGenParams params = VisualGenParams(image_reference="/path/to/image.png") - assert params.image_reference[0].image == "/path/to/image.png" + assert params.image_reference[0].content == "/path/to/image.png" assert params.image_reference[0].role is None def test_image_reference_accepts_bytes(self): from tensorrt_llm.visual_gen import VisualGenParams params = VisualGenParams(image_reference=b"\x89PNG") - assert params.image_reference[0].image == b"\x89PNG" + assert params.image_reference[0].content == b"\x89PNG" def test_image_reference_accepts_list(self): from tensorrt_llm.visual_gen import VisualGenParams params = VisualGenParams(image_reference=["/path/a.png", b"\x89PNG"]) assert len(params.image_reference) == 2 - assert params.image_reference[0].image == "/path/a.png" + assert params.image_reference[0].content == "/path/a.png" def test_model_dump(self): from tensorrt_llm.visual_gen import VisualGenParams diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index eaa971ec7c66..fdfb8240a749 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -279,7 +279,7 @@ def test_base64_image_reference_written_to_disk(self, tmp_path): request, "vid-1", generator, media_storage_path=str(tmp_path) ) assert len(params.image_reference) == 1 - ref_path = params.image_reference[0].image + ref_path = params.image_reference[0].content assert str(ref_path).endswith("vid-1_image_ref_0") # The decoded image is identical to what we passed in. with open(ref_path, "rb") as f: @@ -292,13 +292,13 @@ def test_image_reference_role_and_list(self, tmp_path): Image.new("RGB", (4, 4)).save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() request = VideoGenerationRequest( - prompt="x", image_reference=[b64, {"image": b64, "role": "last_frame"}] + prompt="x", image_reference=[b64, {"content": b64, "role": "last_frame"}] ) params = parse_visual_gen_params( request, "vid-r", generator, media_storage_path=str(tmp_path) ) assert [r.role for r in params.image_reference] == [None, "last_frame"] - paths = [r.image for r in params.image_reference] + paths = [r.content for r in params.image_reference] assert len(set(paths)) == 2 # unique file per index def test_missing_media_storage_path_raises(self): @@ -338,7 +338,7 @@ def test_multipart_avi_video_reference_written_to_disk(self, tmp_path): request, "vid-avi", generator, media_storage_path=str(tmp_path) ) assert params.image_reference is None - assert Path(params.video_reference[0].video).read_bytes() == payload + assert Path(params.video_reference[0].content).read_bytes() == payload def test_multipart_mp4_video_reference_written_to_disk(self, tmp_path): generator = _StubVisualGen() @@ -352,7 +352,7 @@ def test_multipart_mp4_video_reference_written_to_disk(self, tmp_path): # decodes video; the worker demuxes/NVDEC-decodes the conditioning # window from the stored file. assert params.image_reference is None - vpath = params.video_reference[0].video + vpath = params.video_reference[0].content assert str(vpath).endswith("vid-3_video_ref_0") assert Path(vpath).read_bytes() == payload @@ -376,7 +376,7 @@ def test_base64_video_reference_written_to_disk(self, tmp_path): request, "vid-4", generator, media_storage_path=str(tmp_path) ) assert params.image_reference is None - assert Path(params.video_reference[0].video).read_bytes() == payload + assert Path(params.video_reference[0].content).read_bytes() == payload def test_video_reference_survives_real_specs(self, tmp_path): """With the real cosmos3 specs loaded, the encoded payload is persisted @@ -391,7 +391,7 @@ def test_video_reference_survives_real_specs(self, tmp_path): params = parse_visual_gen_params( request, "vid-10", generator, media_storage_path=str(tmp_path) ) - assert Path(params.video_reference[0].video).read_bytes() == payload + assert Path(params.video_reference[0].content).read_bytes() == payload def test_multipart_image_reference_written_to_disk(self, tmp_path): # JPEG upload routed by field name to image_reference. The stored file @@ -407,7 +407,7 @@ def test_multipart_image_reference_written_to_disk(self, tmp_path): request, "vid-5", generator, media_storage_path=str(tmp_path) ) assert params.extra_params is None - assert str(params.image_reference[0].image).endswith("vid-5_image_ref_0") + assert str(params.image_reference[0].content).endswith("vid-5_image_ref_0") def test_wrong_modality_content_raises(self, tmp_path): # The field name declares modality; mismatched content is a client error. @@ -602,7 +602,7 @@ def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): params = parse_visual_gen_params( request, "vid-12", generator, media_storage_path=str(tmp_path) ) - assert Path(params.image_reference[0].image).read_bytes() == truncated + assert Path(params.image_reference[0].content).read_bytes() == truncated # ============================================================================= diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index e3ef313a287e..90f30a3a5448 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1473,19 +1473,19 @@ models: required: false image_reference: kind: extension - type: Optional[Union[str, UploadFile, ImageReferenceItem, List[Union[str, ImageReferenceItem]]]] + type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] default: null status: prototype required: false video_reference: kind: extension - type: Optional[Union[str, UploadFile, VideoReferenceItem, List[Union[str, VideoReferenceItem]]]] + type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] default: null status: prototype required: false audio_reference: kind: extension - type: Optional[Union[str, UploadFile, AudioReferenceItem, List[Union[str, AudioReferenceItem]]]] + type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] default: null status: prototype required: false From e1aaf831e39a627918e64d593bff8e5c47aff3a8 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 01:51:37 -0700 Subject: [PATCH 05/61] [TRTLLM-15277][feat] Keep input_reference as a deprecated back-compat alias Per PR review: re-add the single input_reference field on VideoGenerationRequest, marked deprecated, so existing clients keep working. It is sniff-routed to the typed image_reference / video_reference slot (image signature -> I2V, video container -> V2V, matching the pre-existing behavior where an input_reference video went to extra_params[video] for Cosmos), and is ignored when a typed image/video reference is also provided (typed fields take precedence). Examples use the recommended typed image_reference; docs and the api-stability reference note input_reference as deprecated. Adds tests for the routing and the precedence rule. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 + tensorrt_llm/serve/openai_protocol.py | 9 ++++ tensorrt_llm/serve/visual_gen_utils.py | 43 +++++++++++++++++++ .../visual_gen/test_visual_gen_utils.py | 40 +++++++++++++++++ .../references/trtllm_serve_api.yaml | 6 +++ 5 files changed, 100 insertions(+) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 3103c4e48a4d..0805bdac4973 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -174,6 +174,8 @@ curl http://localhost:8000/v1/videos -H 'content-type: application/json' -d '{ FLUX.2 and Qwen-Image-Edit accept multiple reference images as a list on the same `image_reference` field through the Python API. +A single `input_reference` field (deprecated) is still accepted on the serve video endpoints for backward compatibility; it is routed by content signature to image-to-video or video-to-video, and is ignored when a typed `image_reference` / `video_reference` is also provided. Prefer the typed fields. + ## Optimizations ### Quantization diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 28452400db19..c1966e6c6cbd 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2086,6 +2086,15 @@ class VideoGenerationRequest(OpenAIBaseModel): "a single audio file. Accepted only by models that declare an audio " "reference slot."), ) + input_reference: Optional[Union[str, UploadFile]] = Field( + default=None, + description=( + "Deprecated. A single image or video reference, routed by content " + "signature to image-to-video or video-to-video. Kept for backward " + "compatibility; prefer the typed ``image_reference`` / " + "``video_reference`` fields, which are used instead whenever the " + "same-modality typed field is also provided."), + ) # Resolution size: Optional[str] = Field(default=None, pattern=r"^(\d+x\d+|auto)$") diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 6974dc6455cd..b31ba1c0eaff 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -422,6 +422,46 @@ def cleanup_materialized_conditioning_inputs(value: Any) -> None: except OSError as exc: logger.warning("Failed to remove temporary image edit input %r: %s", path, exc) +def _apply_deprecated_input_reference( + input_reference, params, *, id: str, media_storage_path: Optional[str] +) -> None: + """Back-compat for the deprecated single ``input_reference``. + + Sniff-routes the payload to ``image_reference`` (image) or ``video_reference`` + (video), preserving the pre-typed-fields behavior. Ignored when a typed + image/video reference is already set — the typed fields take precedence. + """ + if input_reference is None: + return + logger.warning("'input_reference' is deprecated; use 'image_reference' / 'video_reference'.") + if params.image_reference or params.video_reference: + return + from tensorrt_llm.visual_gen.params import MediaRef + + payload, _ = _reference_payload_and_role(input_reference) + kind = sniff_media_kind(payload) + if kind == "image": + path = _materialize_reference( + payload, + modality="image", + ref_id=f"{id}_input_ref", + media_storage_path=media_storage_path, + ) + params.image_reference = [MediaRef(content=path)] + elif kind == "video": + path = _materialize_reference( + payload, + modality="video", + ref_id=f"{id}_input_ref", + media_storage_path=media_storage_path, + ) + params.video_reference = [MediaRef(content=path)] + else: + raise ValueError( + "input_reference is not a recognized media container; supported " + "inputs are PNG/JPEG images and MP4/AVI video." + ) + def parse_visual_gen_params( request: ImageGenerationRequest | ImageEditRequest | VideoGenerationRequest, @@ -529,6 +569,9 @@ def parse_visual_gen_params( ) if audio_refs: params.audio_reference = audio_refs + _apply_deprecated_input_reference( + request.input_reference, params, id=id, media_storage_path=media_storage_path + ) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _decode_inline_media(request.extra_params, generator.extra_param_specs) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index fdfb8240a749..27be0cb54e4d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -365,6 +365,46 @@ def test_video_reference_needs_media_storage(self): with pytest.raises(ValueError, match="media_storage_path"): parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) + def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): + # The deprecated single input_reference is sniff-routed to the typed slot. + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + img_b64 = base64.b64encode(buf.getvalue()).decode() + vid_b64 = base64.b64encode(self._mp4_bytes()).decode() + + p = parse_visual_gen_params( + VideoGenerationRequest(prompt="x", input_reference=img_b64), + "vid-i", + generator, + media_storage_path=str(tmp_path), + ) + assert len(p.image_reference) == 1 and p.video_reference is None + + p = parse_visual_gen_params( + VideoGenerationRequest(prompt="x", input_reference=vid_b64), + "vid-v", + generator, + media_storage_path=str(tmp_path), + ) + assert len(p.video_reference) == 1 and p.image_reference is None + + def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): + # A typed reference takes precedence; the deprecated input_reference is dropped. + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + img_b64 = base64.b64encode(buf.getvalue()).decode() + vid_b64 = base64.b64encode(self._mp4_bytes()).decode() + p = parse_visual_gen_params( + VideoGenerationRequest(prompt="x", image_reference=img_b64, input_reference=vid_b64), + "vid-x", + generator, + media_storage_path=str(tmp_path), + ) + assert len(p.image_reference) == 1 + assert p.video_reference is None # input_reference video dropped + def test_base64_video_reference_written_to_disk(self, tmp_path): # The JSON/base64 path carries video even though it has no content-type # or filename; modality is declared by the field name. diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 90f30a3a5448..118389dca435 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1489,6 +1489,12 @@ models: default: null status: prototype required: false + input_reference: + kind: extension + type: Optional[Union[str, UploadFile]] + default: null + status: deprecated + required: false size: kind: extension type: Optional[str] From 037edc79342f4d082e610d8370db7dc597b2dec5 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:05:08 -0700 Subject: [PATCH 06/61] [TRTLLM-15277][fix] Close reference file handles in serve examples The image_reference upload path in both serve examples opened the file without ever closing it (leaked handle). Wrap the open() in a context manager and keep the request (client.videos.create / requests.post) inside it, so the handle closes once the request completes. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- examples/visual_gen/serve/async_video_gen.py | 12 +++++----- examples/visual_gen/serve/sync_video_gen.py | 23 ++++++++++---------- 2 files changed, 19 insertions(+), 16 deletions(-) diff --git a/examples/visual_gen/serve/async_video_gen.py b/examples/visual_gen/serve/async_video_gen.py index 2c90ed305349..9dfd50925994 100755 --- a/examples/visual_gen/serve/async_video_gen.py +++ b/examples/visual_gen/serve/async_video_gen.py @@ -81,15 +81,17 @@ def test_async_video_generation( }, } - # Add input reference if provided (TI2V mode) + # Add input reference if provided (TI2V mode). Keep the create call + # inside the file's context so the handle closes once the request is sent. if image_reference: if not Path(image_reference).exists(): print(f"\n❌ Error: Input reference image not found: {image_reference}") return False - create_params["image_reference"] = open(image_reference, "rb") - - # Create video generation job - job = client.videos.create(**create_params) + with open(image_reference, "rb") as ref_file: + create_params["image_reference"] = ref_file + job = client.videos.create(**create_params) + else: + job = client.videos.create(**create_params) print("Video generation started: \n", job.model_dump_json(indent=2)) diff --git a/examples/visual_gen/serve/sync_video_gen.py b/examples/visual_gen/serve/sync_video_gen.py index 52100a27ea2c..d4ee70abd18a 100755 --- a/examples/visual_gen/serve/sync_video_gen.py +++ b/examples/visual_gen/serve/sync_video_gen.py @@ -83,18 +83,19 @@ def test_sync_video_generation( "format": format, } - # Add the file + # Add the file. Keep the request inside the file's context so the + # handle closes once the upload completes. ## Note: The content-type must be multipart/form-data. - files = { - "image_reference": ( - Path(image_reference).name, - open(image_reference, "rb"), - "multipart/form-data", - ) - } - - print("\n Uploading reference image and generating video...") - response_video = requests.post(endpoint, data=form_data, files=files) + with open(image_reference, "rb") as ref_file: + files = { + "image_reference": ( + Path(image_reference).name, + ref_file, + "multipart/form-data", + ) + } + print("\n Uploading reference image and generating video...") + response_video = requests.post(endpoint, data=form_data, files=files) else: # T2V mode - Use JSON response_video = requests.post( From 69896cc7bd614848d21d83628eb99900b5c3b154 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:09:44 -0700 Subject: [PATCH 07/61] [TRTLLM-15277][chore] Annotate ref_slot_specs property return type Add the dict[str, RefSlotSpec] return annotation to each pipeline's ref_slot_specs override (Cosmos3, FLUX.2, LTX-2, Qwen-edit, Qwen-Layered, plus Wan TI2V and Wan I2V), matching the base BasePipeline.ref_slot_specs signature. Annotation-only; contents unchanged. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 2 +- tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py | 2 +- tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py | 2 +- .../visual_gen/models/qwen_image/pipeline_qwen_image_edit.py | 2 +- .../models/qwen_image_layered/pipeline_qwen_image_layered.py | 2 +- tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py | 2 +- tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan_i2v.py | 2 +- 7 files changed, 7 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 6e280fbfa26a..69b0693877f8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -604,7 +604,7 @@ def extra_param_specs(self): return dict(COSMOS3_EXTRA_SPECS) @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { # image (I2V) and video (V2V) are both optional; Cosmos3 also runs # T2V with neither. 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 57beb5b94f43..f5fdfa91d390 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -371,7 +371,7 @@ def default_generation_params(self): } @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: # Optional reference image(s): "reference" role, count 0..N (multi-subject); # FLUX.2 also runs plain text-to-image with none. return { 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 c36bee296819..6bf0b82e5ba4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1380,7 +1380,7 @@ def extra_param_specs(self): } @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { "image_reference": RefSlotSpec( modality="image", 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 af329b32c272..19ec7704cbc4 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 @@ -133,7 +133,7 @@ def default_warmup_resolutions(self) -> list[tuple[int, int]]: return [(1024, 1024)] @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { "image_reference": RefSlotSpec( modality="image", 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 ae8489005c9c..e6771dc8b928 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 @@ -254,7 +254,7 @@ def extra_param_specs(self) -> dict: } @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { "image_reference": RefSlotSpec( modality="image", 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 ca945f27e80d..05eb06b3d863 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -425,7 +425,7 @@ def extra_param_specs(self): return get_wan_extra_param_specs(self.is_wan22_14b) @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: # Optional single conditioning image (first frame); T2V when absent. return { "image_reference": RefSlotSpec( 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 8c40493d193f..7bd717878728 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 @@ -405,7 +405,7 @@ def extra_param_specs(self): return get_wan_extra_param_specs(self.is_wan22_14b) @property - def ref_slot_specs(self): + def ref_slot_specs(self) -> dict[str, RefSlotSpec]: # I2V first frame (required) + optional last frame for interpolation. return { "image_reference": RefSlotSpec( From be219d19298ea85934d266b0ee9641b0df6af8a5 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:17:27 -0700 Subject: [PATCH 08/61] [TRTLLM-15277][fix] Tighten reference preflight validation WanPipeline.ref_slot_specs declares the image_reference slot only for the Wan 2.2 TI2V-5B variant (is_wan22_5b); the T2V variants accept no image (forward() already raises otherwise), so an image request to them now fails cleanly at preflight instead of deep in the worker. validate_visual_gen_params enters the reference-check block whenever ref_slot_specs is not None (including an empty mapping), so a pipeline that declares no slots rejects any reference the client sent; only None skips. Update the endpoint mock to declare image/video slots and add a test for the empty-mapping rejection. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/models/wan/pipeline_wan.py | 5 +++++ tensorrt_llm/visual_gen/params.py | 6 ++++-- .../visual_gen/test_trtllm_serve_endpoints.py | 11 ++++++++++- .../_torch/visual_gen/test_visual_gen_params.py | 15 +++++++++++++++ 4 files changed, 34 insertions(+), 3 deletions(-) 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 05eb06b3d863..360ebcdd063d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -426,6 +426,11 @@ def extra_param_specs(self): @property def ref_slot_specs(self) -> dict[str, RefSlotSpec]: + # Only Wan 2.2 TI2V-5B conditions on a first frame; the T2V variants + # accept no reference (forward() rejects an image otherwise), so they + # declare no slot and unsupported image requests fail at preflight. + if not self.is_wan22_5b: + return {} # Optional single conditioning image (first frame); T2V when absent. return { "image_reference": RefSlotSpec( diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 5a59bb63eb25..a2115d659029 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -286,8 +286,10 @@ def validate_visual_gen_params( # ``.roles`` (a list of role specs with ``.role`` / ``.min`` / ``.max``). # role is required only when a modality declares more than one role; # otherwise the single declared role is inferred. Reference fields are - # already normalized to ``list[*Ref]`` by the field validators. - if ref_slot_specs: + # already normalized to ``list[*Ref]`` by the field validators. An empty + # (but non-None) mapping means the pipeline declares no slots, so any + # reference the client sent is rejected; only ``None`` skips validation. + if ref_slot_specs is not None: for field in ("image_reference", "video_reference", "audio_reference"): refs = getattr(params, field, None) or [] spec = ref_slot_specs.get(field) 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 983870b23330..3e89238279b8 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -217,6 +217,8 @@ 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 RefSlotSpec, RoleSpec + self.executor = SimpleNamespace( default_generation_params={ "height": 64, @@ -230,7 +232,14 @@ def __init__( extra_param_specs=extra_param_specs or {"stg_scale": ExtraParamSchema(type="float", default=1.0)}, supports_image_edit=supports_image_edit, - ref_slot_specs={}, + ref_slot_specs={ + "image_reference": RefSlotSpec( + modality="image", roles=[RoleSpec(role="first_frame", min=0, max=1)] + ), + "video_reference": RefSlotSpec( + modality="video", roles=[RoleSpec(role="reference", min=0, max=1)] + ), + }, ) def _maybe_batch(self, tensor, n): 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 b55c79425219..9ecd254c3a64 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -984,6 +984,21 @@ def run(params, spec): with pytest.raises(ValueError, match=r"video_reference.*not accepted"): run(VisualGenParams(video_reference="v.mp4"), optional) + def test_empty_ref_slot_specs_rejects_references(self): + """An empty (non-None) ref_slot_specs means the pipeline declares no + slots, so a reference is rejected; only ``None`` skips validation.""" + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + def run(params, spec): + validate_visual_gen_params( + params, declared_defaults=None, extra_param_specs={}, ref_slot_specs=spec + ) + + with pytest.raises(ValueError, match="not accepted"): + run(VisualGenParams(image_reference="a.png"), {}) + run(VisualGenParams(image_reference="a.png"), None) # None -> skipped + run(VisualGenParams(), {}) # no reference -> allowed + def test_none_fields_not_flagged(self): """Fields left as None should never trigger unsupported-field errors.""" from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux import FluxPipeline From 0789d4dc3bd13950c92bfa9c120d01d646963696 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 02:31:33 -0700 Subject: [PATCH 09/61] [TRTLLM-15277][fix] Clean up partial reference materialization and align input_reference docs _build_reference_list now removes the files earlier items already wrote when a later item in a multi-reference request is rejected, so a 400 leaves nothing on disk. Correct the deprecated input_reference field description to match the code: the typed image_reference / video_reference fields take precedence globally (input_reference is ignored whenever either is set), not only for the same modality. Add type annotations to _apply_deprecated_input_reference. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol.py | 13 ++++--- tensorrt_llm/serve/visual_gen_utils.py | 38 ++++++++++++++----- .../visual_gen/test_visual_gen_utils.py | 13 +++++++ 3 files changed, 48 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index c1966e6c6cbd..84512e5734f0 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2088,12 +2088,13 @@ class VideoGenerationRequest(OpenAIBaseModel): ) input_reference: Optional[Union[str, UploadFile]] = Field( default=None, - description=( - "Deprecated. A single image or video reference, routed by content " - "signature to image-to-video or video-to-video. Kept for backward " - "compatibility; prefer the typed ``image_reference`` / " - "``video_reference`` fields, which are used instead whenever the " - "same-modality typed field is also provided."), + description= + ("Deprecated. A single image or video reference, routed by content " + "signature to image-to-video or video-to-video. Kept for backward " + "compatibility; prefer the typed ``image_reference`` / " + "``video_reference`` fields, which take precedence — this field is " + "ignored whenever a typed ``image_reference`` or ``video_reference`` " + "is provided."), ) # Resolution diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index b31ba1c0eaff..bbb12f56ccb9 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -20,6 +20,8 @@ ) if TYPE_CHECKING: + from fastapi import UploadFile + # 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 @@ -197,15 +199,27 @@ def _build_reference_list(value, *, modality: str, id: str, media_storage_path: raw_items = value if isinstance(value, list) else [value] refs = [] - for i, item in enumerate(raw_items): - payload, role = _reference_payload_and_role(item) - ref_path = _materialize_reference( - payload, - modality=modality, - ref_id=f"{id}_{modality}_ref_{i}", - media_storage_path=media_storage_path, - ) - refs.append(MediaRef(content=ref_path, role=role)) + created_paths: list[str] = [] + try: + for i, item in enumerate(raw_items): + payload, role = _reference_payload_and_role(item) + ref_path = _materialize_reference( + payload, + modality=modality, + ref_id=f"{id}_{modality}_ref_{i}", + media_storage_path=media_storage_path, + ) + created_paths.append(ref_path) + refs.append(MediaRef(content=ref_path, role=role)) + except Exception: + # A later item failed; remove the files earlier items already wrote so + # a rejected multi-reference request leaves nothing on disk. + for path in created_paths: + try: + os.remove(path) + except OSError: + pass + raise return refs @@ -423,7 +437,11 @@ def cleanup_materialized_conditioning_inputs(value: Any) -> None: logger.warning("Failed to remove temporary image edit input %r: %s", path, exc) def _apply_deprecated_input_reference( - input_reference, params, *, id: str, media_storage_path: Optional[str] + input_reference: str | UploadFile | None, + params: VisualGenParams, + *, + id: str, + media_storage_path: str | None, ) -> None: """Back-compat for the deprecated single ``input_reference``. diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 27be0cb54e4d..a5bb6040e850 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -505,6 +505,19 @@ def read(self, *args, **kwargs): # … and the payload read fails before any file is written, so nothing leaks. assert list(tmp_path.iterdir()) == [] + def test_multi_reference_partial_failure_cleans_up(self, tmp_path): + # A later item's rejection removes the files earlier items already wrote, + # so a rejected multi-reference request leaves nothing on disk. + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + good = base64.b64encode(buf.getvalue()).decode() + bad = base64.b64encode(b"neither an image nor a video").decode() + request = VideoGenerationRequest(prompt="x", image_reference=[good, bad]) + with pytest.raises(ValueError, match="not a recognized image"): + parse_visual_gen_params(request, "vid-11", generator, media_storage_path=str(tmp_path)) + assert list(tmp_path.iterdir()) == [] + class TestMediaBytesProbes: """The in-memory signature probes the serve boundary routes on.""" From f8016f28fd360cef6521ed9d0258f4e34838f211 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:40:23 -0700 Subject: [PATCH 10/61] [TRTLLM-15277][feat] Accept http(s) URL and file:// reference inputs on serve Serve reference strings now dispatch on URL scheme: http(s) URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, size cap), file:// paths are read from local disk, and data: / bare strings continue to decode as base64. The resolved bytes flow into the existing materialize-then-write unchanged, bringing serve references to parity with the LLM multimodal input forms. Fetch / read failures surface as ValueError (HTTP 400) rather than a server 500. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- tensorrt_llm/serve/openai_protocol.py | 28 +++++---- tensorrt_llm/serve/visual_gen_utils.py | 43 +++++++++++--- .../visual_gen/test_visual_gen_utils.py | 59 +++++++++++++++++++ 4 files changed, 112 insertions(+), 20 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 0805bdac4973..06f5fde5b986 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -121,7 +121,7 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue ### Reference Inputs -Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, or uploaded file) to a local path before reaching the worker. +Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, `http(s)` URL, `file://` path, or uploaded file) to a local path before reaching the worker. `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). Most models take a single reference whose role is unambiguous, so no `role` is specified: diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 84512e5734f0..c5dd5262858d 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2062,29 +2062,33 @@ class VideoGenerationRequest(OpenAIBaseModel): default=None, description= ("Image reference(s) conditioning generation (e.g. image-to-video " - "first frame). JSON sends base64 bytes, a ``{content, role}`` " - "object, or a list of them; multipart uploads a single image file. " - "PNG or JPEG only — HEIF/AVIF are not supported."), + "first frame). A JSON string is base64 bytes (raw or ``data:`` " + "URI), an ``http(s)`` URL, or a ``file://`` path; or send a " + "``{content, role}`` object or a list of them; multipart uploads a " + "single image file. PNG or JPEG only — HEIF/AVIF are not supported." + ), ) video_reference: Optional[Union[ str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] = Field( default=None, description= - ("Video reference(s) conditioning generation (video-to-video). JSON " - "sends base64 bytes, a ``{content}`` object, or a list of them; " - "multipart uploads a single video file. MP4 or AVI, with H.264 the " - "tested codec and others best-effort."), + ("Video reference(s) conditioning generation (video-to-video). A " + "JSON string is base64 bytes (raw or ``data:`` URI), an ``http(s)`` " + "URL, or a ``file://`` path; or send a ``{content}`` object or a " + "list of them; multipart uploads a single video file. MP4 or AVI, " + "with H.264 the tested codec and others best-effort."), ) audio_reference: Optional[Union[ str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] = Field( default=None, - description= - ("Audio reference(s) conditioning generation. JSON sends base64 " - "bytes, a ``{content}`` object, or a list of them; multipart uploads " - "a single audio file. Accepted only by models that declare an audio " - "reference slot."), + description=( + "Audio reference(s) conditioning generation. A JSON string is " + "base64 bytes (raw or ``data:`` URI), an ``http(s)`` URL, or a " + "``file://`` path; or send a ``{content}`` object or a list of " + "them; multipart uploads a single audio file. Accepted only by " + "models that declare an audio reference slot."), ) input_reference: Optional[Union[str, UploadFile]] = Field( default=None, diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index bbb12f56ccb9..eaef9053c3da 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -6,12 +6,18 @@ import os from collections.abc import Mapping from io import BytesIO +from pathlib import Path 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.inputs.media_io import ( + _normalize_file_uri, + _safe_request_get, + is_isobmff_image_bytes, + sniff_media_kind, +) from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ( ImageEditRequest, @@ -130,22 +136,45 @@ def _read_reference_payload(reference: str) -> bytes: raise ValueError("reference is not valid base64 data.") from exc +def _resolve_reference_string(reference: str) -> bytes: + """Resolve one reference string to raw bytes, dispatching on URL scheme. + + Mirrors the LLM multimodal loader so serve references accept the same forms: + ``http(s)`` fetches through the SSRF-guarded loader (private-address block, + redirect re-validation, timeout, size cap); ``file://`` reads a local file; + ``data:`` and bare strings decode as base64. Fetch/read failures become + ``ValueError`` so a bad URL or path is a client 400, not a server 500. + """ + scheme = urlparse(reference).scheme + if scheme in ("http", "https"): + try: + return _safe_request_get(reference).content + except Exception as exc: + raise ValueError(f"reference URL could not be fetched: {exc}") from exc + if scheme == "file": + try: + return Path(_normalize_file_uri(reference)).read_bytes() + except OSError as exc: + raise ValueError(f"reference file could not be read: {exc}") from exc + return _read_reference_payload(reference) + + def _reference_payload_and_role(ref) -> tuple[bytes, Optional[str]]: """Extract ``(payload_bytes, role)`` from one raw HTTP reference. - ``ref`` is a base64/data-URI string, a multipart ``UploadFile`` (has - ``.file``), or a ``MediaReferenceItem`` exposing ``content`` and an optional - ``role``. + ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or ``file://`` + path), a multipart ``UploadFile`` (has ``.file``), or a ``MediaReferenceItem`` + exposing ``content`` and an optional ``role``. """ role = getattr(ref, "role", None) if isinstance(ref, str): - return _read_reference_payload(ref), role + return _resolve_reference_string(ref), role if hasattr(ref, "file"): # multipart UploadFile return ref.file.read(), role data = getattr(ref, "content", None) if not isinstance(data, str): - raise ValueError("reference item must carry a base64 'content' string.") - return _read_reference_payload(data), role + raise ValueError("reference item must carry a 'content' string.") + return _resolve_reference_string(data), role def _materialize_reference( diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index a5bb6040e850..0e3c21884bbe 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -518,6 +518,65 @@ def test_multi_reference_partial_failure_cleans_up(self, tmp_path): parse_visual_gen_params(request, "vid-11", generator, media_storage_path=str(tmp_path)) assert list(tmp_path.iterdir()) == [] + def test_file_uri_image_reference_read_and_materialized(self, tmp_path): + # A file:// reference is read from local disk and persisted like any other. + generator = _StubVisualGen() + src = tmp_path / "ref.png" + Image.new("RGB", (4, 4), (7, 8, 9)).save(src, format="PNG") + store = tmp_path / "store" + store.mkdir() + request = VideoGenerationRequest(prompt="x", image_reference=src.as_uri()) + params = parse_visual_gen_params( + request, "vid-file", generator, media_storage_path=str(store) + ) + assert Path(params.image_reference[0].content).read_bytes() == src.read_bytes() + + def test_http_url_image_reference_fetched_and_materialized(self, tmp_path, monkeypatch): + # An http(s) reference is fetched through the guarded loader, then stored. + generator = _StubVisualGen() + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + png = buf.getvalue() + + class _FakeResp: + def __init__(self, content): + self.content = content + + monkeypatch.setattr( + "tensorrt_llm.serve.visual_gen_utils._safe_request_get", + lambda url, **kwargs: _FakeResp(png), + ) + request = VideoGenerationRequest(prompt="x", image_reference="https://example.com/a.png") + params = parse_visual_gen_params( + request, "vid-url", generator, media_storage_path=str(tmp_path) + ) + assert Path(params.image_reference[0].content).read_bytes() == png + + def test_http_url_fetch_failure_is_client_error(self, tmp_path, monkeypatch): + # A blocked/failed fetch (e.g. SSRF guard) is a client 400, not a 500, + # and leaves nothing on disk. + generator = _StubVisualGen() + + def _blocked(url, **kwargs): + raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") + + monkeypatch.setattr("tensorrt_llm.serve.visual_gen_utils._safe_request_get", _blocked) + request = VideoGenerationRequest(prompt="x", image_reference="http://10.0.0.1/a.png") + with pytest.raises(ValueError, match="reference URL could not be fetched"): + parse_visual_gen_params( + request, "vid-ssrf", generator, media_storage_path=str(tmp_path) + ) + assert list(tmp_path.iterdir()) == [] + + def test_missing_file_uri_is_client_error(self, tmp_path): + # A file:// path that does not exist is a client 400, not a server 500. + generator = _StubVisualGen() + missing = (tmp_path / "does_not_exist.png").as_uri() + request = VideoGenerationRequest(prompt="x", image_reference=missing) + with pytest.raises(ValueError, match="reference file could not be read"): + parse_visual_gen_params(request, "vid-nf", generator, media_storage_path=str(tmp_path)) + assert list(tmp_path.iterdir()) == [] + class TestMediaBytesProbes: """The in-memory signature probes the serve boundary routes on.""" From 977400d26be93dae7ac78e9f2e1f695a6fd6efc9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 21:58:17 -0700 Subject: [PATCH 11/61] [TRTLLM-15277][doc] Note http(s) URL and file:// serve reference forms in examples Make the Reference Inputs serve example consistent with the materialization description: a JSON reference string may be base64, a data: URI, an http(s) URL, or a file:// path, and multipart uploads raw bytes. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 06f5fde5b986..d2ee9fa50523 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -140,11 +140,16 @@ params = vg.default_params params.video_reference = "clip.mp4" ``` -The equivalent serve request uploads the file, or sends a base64 string or `data:` URI in a JSON body: +The equivalent serve request uploads the file, or sends a base64 string, `data:` URI, `http(s)` URL, or `file://` path as the field value in a JSON body: ```bash +# multipart file upload (raw bytes, no base64) curl http://localhost:8000/v1/videos -F "prompt=the scene comes alive" -F "image_reference=@start.png" curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_reference=@clip.mp4" + +# JSON body: the reference string may be base64, a data: URI, an http(s) URL, or a file:// path +curl http://localhost:8000/v1/videos -H 'content-type: application/json' \ + -d '{"prompt": "the scene comes alive", "image_reference": "https://example.com/start.png"}' ``` When a model accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame — the `role` is required to disambiguate: From f5fdb45189da6cdbcba8a2613428f8b2450e5b91 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 22:17:38 -0700 Subject: [PATCH 12/61] [TRTLLM-15277][feat] Accept bare local file paths as serve reference input MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A bare reference string is now decoded as base64 first and, if that fails, read as a local file path, so a plain path works without the file:// scheme — matching the LLM multimodal loader's empty-scheme handling. file:// URIs and http(s) URLs are unchanged. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 6 ++--- tensorrt_llm/serve/openai_protocol.py | 6 ++--- tensorrt_llm/serve/visual_gen_utils.py | 24 +++++++++++++++---- .../visual_gen/test_visual_gen_utils.py | 14 +++++++++++ 4 files changed, 39 insertions(+), 11 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index d2ee9fa50523..bbdaba35aa61 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -121,7 +121,7 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue ### Reference Inputs -Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, `http(s)` URL, `file://` path, or uploaded file) to a local path before reaching the worker. `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). +Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, `http(s)` URL, local file path, or uploaded file) to a local path before reaching the worker. A local path may be given bare or as a `file://` URI; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). Most models take a single reference whose role is unambiguous, so no `role` is specified: @@ -140,14 +140,14 @@ params = vg.default_params params.video_reference = "clip.mp4" ``` -The equivalent serve request uploads the file, or sends a base64 string, `data:` URI, `http(s)` URL, or `file://` path as the field value in a JSON body: +The equivalent serve request uploads the file, or sends a base64 string, `data:` URI, `http(s)` URL, or local file path as the field value in a JSON body: ```bash # multipart file upload (raw bytes, no base64) curl http://localhost:8000/v1/videos -F "prompt=the scene comes alive" -F "image_reference=@start.png" curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_reference=@clip.mp4" -# JSON body: the reference string may be base64, a data: URI, an http(s) URL, or a file:// path +# JSON body: the reference string may be base64, a data: URI, an http(s) URL, or a local file path curl http://localhost:8000/v1/videos -H 'content-type: application/json' \ -d '{"prompt": "the scene comes alive", "image_reference": "https://example.com/start.png"}' ``` diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index c5dd5262858d..211f48e8ace7 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2063,7 +2063,7 @@ class VideoGenerationRequest(OpenAIBaseModel): description= ("Image reference(s) conditioning generation (e.g. image-to-video " "first frame). A JSON string is base64 bytes (raw or ``data:`` " - "URI), an ``http(s)`` URL, or a ``file://`` path; or send a " + "URI), an ``http(s)`` URL, or a local file path; or send a " "``{content, role}`` object or a list of them; multipart uploads a " "single image file. PNG or JPEG only — HEIF/AVIF are not supported." ), @@ -2075,7 +2075,7 @@ class VideoGenerationRequest(OpenAIBaseModel): description= ("Video reference(s) conditioning generation (video-to-video). A " "JSON string is base64 bytes (raw or ``data:`` URI), an ``http(s)`` " - "URL, or a ``file://`` path; or send a ``{content}`` object or a " + "URL, or a local file path; or send a ``{content}`` object or a " "list of them; multipart uploads a single video file. MP4 or AVI, " "with H.264 the tested codec and others best-effort."), ) @@ -2086,7 +2086,7 @@ class VideoGenerationRequest(OpenAIBaseModel): description=( "Audio reference(s) conditioning generation. A JSON string is " "base64 bytes (raw or ``data:`` URI), an ``http(s)`` URL, or a " - "``file://`` path; or send a ``{content}`` object or a list of " + "local file path; or send a ``{content}`` object or a list of " "them; multipart uploads a single audio file. Accepted only by " "models that declare an audio reference slot."), ) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index eaef9053c3da..e87173da2c7e 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -141,9 +141,11 @@ def _resolve_reference_string(reference: str) -> bytes: Mirrors the LLM multimodal loader so serve references accept the same forms: ``http(s)`` fetches through the SSRF-guarded loader (private-address block, - redirect re-validation, timeout, size cap); ``file://`` reads a local file; - ``data:`` and bare strings decode as base64. Fetch/read failures become - ``ValueError`` so a bad URL or path is a client 400, not a server 500. + redirect re-validation, timeout, size cap); ``file://`` and bare local paths + read from disk; ``data:`` and base64 strings decode inline. A bare string is + decoded as base64 first and, failing that, read as a local file path. + Fetch/read failures become ``ValueError`` so a bad URL or path is a client + 400, not a server 500. """ scheme = urlparse(reference).scheme if scheme in ("http", "https"): @@ -156,13 +158,25 @@ def _resolve_reference_string(reference: str) -> bytes: return Path(_normalize_file_uri(reference)).read_bytes() except OSError as exc: raise ValueError(f"reference file could not be read: {exc}") from exc - return _read_reference_payload(reference) + if scheme == "data": + return _read_reference_payload(reference) + # Bare string: base64 first (the established default), else a local file path + # so a plain path works without the file:// scheme. + try: + return _read_reference_payload(reference) + except ValueError: + try: + return Path(reference).read_bytes() + except OSError as exc: + raise ValueError( + f"reference is not valid base64 data, and not a readable local file: {exc}" + ) from exc def _reference_payload_and_role(ref) -> tuple[bytes, Optional[str]]: """Extract ``(payload_bytes, role)`` from one raw HTTP reference. - ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or ``file://`` + ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or a local file path), a multipart ``UploadFile`` (has ``.file``), or a ``MediaReferenceItem`` exposing ``content`` and an optional ``role``. """ diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 0e3c21884bbe..8da3d18b579e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -531,6 +531,20 @@ def test_file_uri_image_reference_read_and_materialized(self, tmp_path): ) assert Path(params.image_reference[0].content).read_bytes() == src.read_bytes() + def test_bare_path_image_reference_read_and_materialized(self, tmp_path): + # A bare local path (no file:// scheme) is read from disk after the + # base64 decode attempt fails. + generator = _StubVisualGen() + src = tmp_path / "ref.png" + Image.new("RGB", (4, 4), (11, 22, 33)).save(src, format="PNG") + store = tmp_path / "store" + store.mkdir() + request = VideoGenerationRequest(prompt="x", image_reference=str(src)) + params = parse_visual_gen_params( + request, "vid-bare", generator, media_storage_path=str(store) + ) + assert Path(params.image_reference[0].content).read_bytes() == src.read_bytes() + def test_http_url_image_reference_fetched_and_materialized(self, tmp_path, monkeypatch): # An http(s) reference is fetched through the guarded loader, then stored. generator = _StubVisualGen() From 14b84726ad1a4337ec71487fc85b64b3b1ceb845 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 13 Aug 2026 23:55:22 -0700 Subject: [PATCH 13/61] [TRTLLM-15277][fix] Reclaim materialized reference files on request completion/failure Reference inputs materialized to media storage (`{id}_{modality}_ref_*`, and the deprecated `{id}_input_ref`) were never removed, so every conditioned request left a copy in TRTLLM_MEDIA_STORAGE_PATH for the job's lifetime, and a later failure in a multi-reference request leaked the earlier inputs. Add cleanup_reference_files (glob by the request-id prefix; output files carry no `ref` and are untouched) and call it at every request exit: the sync route `finally`, the async background task `finally` (which runs on CancelledError too, so a delete that cancels an in-flight job still reclaims the inputs), and the async enqueue-failure path. delete_video needs nothing extra since completion/failure cleanup runs first. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_video_routes.py | 32 +++++++++++++++-- tensorrt_llm/serve/visual_gen_utils.py | 20 +++++++++++ .../visual_gen/test_trtllm_serve_endpoints.py | 34 +++++++++++++++---- .../visual_gen/test_visual_gen_utils.py | 24 +++++++++++++ 4 files changed, 101 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index a0338e459f19..5f906252bafe 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -35,7 +35,11 @@ build_visual_gen_server_timings, build_visual_gen_timing_headers, ) -from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE, parse_visual_gen_params +from tensorrt_llm.serve.visual_gen_utils import ( + VIDEO_STORE, + cleanup_reference_files, + parse_visual_gen_params, +) if TYPE_CHECKING: # Type-only: importing tensorrt_llm.visual_gen at runtime would pull the @@ -144,6 +148,9 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time + # Assigned before the try so the ``finally`` can always clean up any + # reference inputs materialized for this request id. + video_id = f"video_{uuid.uuid4().hex}" try: # Client-side ValueErrors from content-type parsing, request # translation, encoder-format preflight, parameter validation, @@ -156,7 +163,6 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: path_error = self._reject_disabled_path(request.response_format) if path_error is not None: return path_error - video_id = f"video_{uuid.uuid4().hex}" params = parse_visual_gen_params( request, video_id, @@ -278,6 +284,11 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: err_type="InternalServerError", status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) + finally: + # References are input-only; the pipeline has consumed them by now, + # so remove them (success or failure) — conditioned requests must not + # accumulate materialized inputs in media storage. + cleanup_reference_files(str(self.media_storage_path), video_id) async def _parse_video_generation_request( self, @@ -376,6 +387,10 @@ async def openai_video_generation_async( - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time + # Assigned before the try so the ``finally`` can clean up references + # materialized for this request if no background task takes ownership. + video_id = f"video_{uuid.uuid4().hex}" + task_started = False try: # Parse request based on content-type request = await self._parse_video_generation_request(raw_request) @@ -383,7 +398,6 @@ async def openai_video_generation_async( if path_error is not None: return path_error - 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) ) @@ -434,6 +448,8 @@ async def openai_video_generation_async( ) ) self.video_gen_tasks[video_id] = task + # The background task now owns reference cleanup (its ``finally``). + task_started = True task.add_done_callback(lambda t, vid=video_id: self._on_video_task_done(vid, t)) return JSONResponse(content=video_job.model_dump(), status_code=202) @@ -450,6 +466,11 @@ async def openai_video_generation_async( err_type="InternalServerError", status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) + finally: + if not task_started: + # Failed before the background task was scheduled — nothing else + # will clean these up. + cleanup_reference_files(str(self.media_storage_path), video_id) async def _generate_video_background( self, @@ -554,6 +575,11 @@ async def _generate_video_background( job.completed_at = int(time.time()) job.error = str(e) await VIDEO_STORE.upsert(video_id, job) + finally: + # References are input-only; remove them once generation has run, + # failed, or been cancelled. Runs on CancelledError too, so a delete + # that cancels an in-flight job still reclaims the inputs. + cleanup_reference_files(str(self.media_storage_path), video_id) async def list_videos(self, raw_request: Request) -> Response: """List all generated videos. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index e87173da2c7e..8f1f25773c46 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -226,6 +226,26 @@ def _materialize_reference( return ref_path +def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) -> None: + """Remove the materialized reference inputs for one request. + + References are materialized as ``{request_id}_{modality}_ref_{i}`` (and the + deprecated ``{request_id}_input_ref``) under ``media_storage_path``. They are + input-only — unneeded once the pipeline has consumed them — so the request + owner removes them by the ``request_id`` prefix, covering image/video/audio + and the deprecated single reference regardless of count. Output files + (``{request_id}_{i}.``) carry no ``ref`` and are left untouched. + Best-effort: already-removed files are ignored. + """ + if media_storage_path is None: + return + for path in Path(media_storage_path).glob(f"{request_id}_*ref*"): + try: + path.unlink() + except OSError: + pass + + def _build_reference_list(value, *, modality: str, id: str, media_storage_path: Optional[str]): """Materialize an HTTP reference field into a list of ``MediaRef`` objects. 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 3e89238279b8..2623c6bbc282 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -209,6 +209,9 @@ def __init__( # used by tests to assert forwarded VisualGenParams fields. self.last_inputs = None self.last_params = None + # Snapshot of materialized reference-file contents at generation time, + # captured before the route cleans them up. Keyed by stored path. + self.last_ref_bytes = {} # Stand-in for the coordinator-side executor proxy. The async video # route reads ``default_generation_params`` / ``extra_param_specs`` # directly off this attribute when running synchronous pre-flight @@ -250,9 +253,21 @@ def _maybe_batch(self, tensor, n): # --- VisualGen interface --- + def _snapshot_refs(self, params) -> None: + # Capture materialized reference bytes before the route cleans them up, + # so tests can still assert byte-identity after the request finishes. + self.last_ref_bytes = {} + for field in ("image_reference", "video_reference", "audio_reference"): + for ref in getattr(params, field, None) or []: + path = getattr(ref, "content", None) + if isinstance(path, str) and os.path.exists(path): + with open(path, "rb") as fh: + self.last_ref_bytes[path] = fh.read() + def generate(self, inputs=None, params=None) -> VisualGenOutput: self.last_inputs = inputs self.last_params = params + self._snapshot_refs(params) if self._validation_error is not None: raise self._validation_error if self._generate_error is not None: @@ -271,6 +286,7 @@ def generate(self, inputs=None, params=None) -> VisualGenOutput: def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": self.last_inputs = inputs self.last_params = params + self._snapshot_refs(params) if self._validation_error is not None: raise self._validation_error n = getattr(params, "num_images_per_prompt", 1) if params else 1 @@ -1597,13 +1613,15 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert resp.status_code == 200 assert len(resp.content) > 0 - # image_reference is written to media storage and passed through as a - # MediaRef carrying the filesystem path. + # image_reference is materialized to media storage and passed through as + # a MediaRef carrying the filesystem path. params = video_client.mock_gen.last_params ref_path = params.image_reference[0].content assert isinstance(ref_path, str) assert ref_path.endswith("_image_ref_0") - assert os.path.exists(ref_path) + # The materialized reference is input-only and is cleaned up once the + # request finishes, so it must not linger in media storage. + assert not os.path.exists(ref_path) def test_sync_video_generation_multipart_with_video_reference(self, video_client): """A ``video_reference`` upload is persisted byte-identical (V2V) — the @@ -1625,11 +1643,15 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video conditioning arrives as a MediaRef holding a stored path; - # no image_reference is set, and the encoded bytes are byte-identical. + # Video conditioning arrives as a MediaRef holding a stored path; no + # image_reference is set, and the encoded bytes were persisted + # byte-identical (snapshotted at generation time, since the route cleans + # the reference up once the request finishes). params = video_client.mock_gen.last_params assert params.image_reference is None - assert Path(params.video_reference[0].content).read_bytes() == payload + ref_path = params.video_reference[0].content + assert video_client.mock_gen.last_ref_bytes[ref_path] == payload + assert not os.path.exists(ref_path) def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content matching no image or video container signature is rejected diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 8da3d18b579e..7d29d4b7681a 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -25,6 +25,7 @@ from tensorrt_llm.serve.visual_gen_utils import ( _merge_extra_params, _warn_if_set_with_no_semantic, + cleanup_reference_files, parse_visual_gen_params, ) from tensorrt_llm.visual_gen import VisualGenParams @@ -835,3 +836,26 @@ def test_malformed_base64_is_a_client_error(self): request = VideoGenerationRequest(prompt="storm", extra_params={"video": "not!b64!"}) with pytest.raises(ValueError, match="not valid base64"): parse_visual_gen_params(request, "id-bad", self._generator()) + + +class TestCleanupReferenceFiles: + """The reference-file reclaim helper keyed on the request id prefix.""" + + def test_removes_only_this_request_ref_files(self, tmp_path): + vid = "video_abc123" + (tmp_path / f"{vid}_image_ref_0").write_bytes(b"a") + (tmp_path / f"{vid}_video_ref_1").write_bytes(b"b") + (tmp_path / f"{vid}_input_ref").write_bytes(b"c") # deprecated alias + (tmp_path / f"{vid}_0.mp4").write_bytes(b"out") # output — keep + (tmp_path / "video_other_image_ref_0").write_bytes(b"d") # other id — keep + cleanup_reference_files(str(tmp_path), vid) + assert sorted(p.name for p in tmp_path.iterdir()) == [ + f"{vid}_0.mp4", + "video_other_image_ref_0", + ] + + def test_none_storage_is_noop(self): + cleanup_reference_files(None, "video_x") # no raise + + def test_missing_files_are_ignored(self, tmp_path): + cleanup_reference_files(str(tmp_path), "video_absent") # no raise From 52c29dd2ab30c2c933c9444c1da1fa428bfcabde Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Sun, 16 Aug 2026 22:46:52 -0700 Subject: [PATCH 14/61] [TRTLLM-15277][refactor] Use shared Literal types for reference role/modality MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit RoleSpec.role and MediaReferenceItem.role now use the shared Role literal (tensorrt_llm.visual_gen.params, already used by MediaRef.role); RefSlotSpec.modality uses the core-lib MediaModality literal (tensorrt_llm.inputs.media_io). This replaces bare str + a description that only listed the allowed values, so pydantic rejects an unknown role/modality with a 422 at the request boundary instead of accepting any string. One source of truth per concept — no duplicate literals. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/pipeline.py | 6 ++++-- tensorrt_llm/serve/openai_protocol.py | 9 ++++----- 2 files changed, 8 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 7cb0fd815a81..fd14dbd30610 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -27,9 +27,11 @@ from tensorrt_llm._torch.autotuner import autotune from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.inputs.media_io import MediaModality from tensorrt_llm.llmapi.utils import StrictBaseModel from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping +from tensorrt_llm.visual_gen.params import Role from .cache import CacheDiTAccelerator, TeaCacheAccelerator from .checkpoints import WeightLoader @@ -73,7 +75,7 @@ class ExtraParamSchema(StrictBaseModel): class RoleSpec(StrictBaseModel): """One accepted role for a reference modality, with its count bounds.""" - role: str = Field(description="'reference' | 'first_frame' | 'last_frame'.") + role: Role = Field(description="Role of the reference input.") min: int = Field(default=1, description="Minimum count for this role.") max: Optional[int] = Field( default=1, description="Maximum count for this role (None = unbounded)." @@ -90,7 +92,7 @@ class RefSlotSpec(StrictBaseModel): Pickled to the coordinator in the READY handshake, so keep it plain data. """ - modality: str = Field(description="'image' | 'video' | 'audio'.") + modality: MediaModality = Field(description="Reference modality.") roles: List[RoleSpec] = Field(description="Accepted roles + counts for this modality.") diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 211f48e8ace7..af05bc168a93 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -61,6 +61,7 @@ from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy +from tensorrt_llm.visual_gen.params import Role _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 @@ -2024,12 +2025,10 @@ class MediaReferenceItem(OpenAIBaseModel): content: str = Field( description="Base64-encoded media bytes, optionally as a ``data:`` URI." ) - role: Optional[str] = Field( + role: Optional[Role] = Field( default=None, - description= - ("Reference role (e.g. 'reference', 'first_frame', 'last_frame'). " - "Required only when the model accepts multiple roles for the modality." - ), + description="Reference role. Required only when the model accepts " + "multiple roles for the modality.", ) From 14971b39d9fe728ce9b7a7cb0f45e27d1ade88df Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 17 Aug 2026 00:29:34 -0700 Subject: [PATCH 15/61] [TRTLLM-15277][fix] Add type hints to reference-materialization helpers Annotate the untyped params/returns on _reference_payload_and_role and _build_reference_list so the serve reference helpers carry full type hints. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/visual_gen_utils.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 8f1f25773c46..62e17fe9141d 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -173,7 +173,7 @@ def _resolve_reference_string(reference: str) -> bytes: ) from exc -def _reference_payload_and_role(ref) -> tuple[bytes, Optional[str]]: +def _reference_payload_and_role(ref: Any) -> tuple[bytes, Optional[str]]: """Extract ``(payload_bytes, role)`` from one raw HTTP reference. ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or a local file @@ -246,7 +246,9 @@ def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) pass -def _build_reference_list(value, *, modality: str, id: str, media_storage_path: Optional[str]): +def _build_reference_list( + value: Any, *, modality: str, id: str, media_storage_path: Optional[str] +) -> Optional[list]: """Materialize an HTTP reference field into a list of ``MediaRef`` objects. ``value`` is None, a base64/data-URI string, a multipart ``UploadFile``, a From b7da2fad0f9bc6a9f189cd8d174ac9b3005475a8 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 17 Aug 2026 03:03:28 -0700 Subject: [PATCH 16/61] [TRTLLM-15277][refactor] Extract shared media_refs module for reference resolve/materialize/cleanup Move _read_reference_payload / _resolve_reference_string / _materialize_reference / cleanup_reference_files out of tensorrt_llm/serve/visual_gen_utils.py into tensorrt_llm/visual_gen/media_refs.py so both the serve boundary and the engine frontend can use them without an engine->serve import. Behavior-preserving move; serve imports them back. Prep for pushing reference materialization down into the engine. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_video_routes.py | 7 +- tensorrt_llm/serve/visual_gen_utils.py | 122 +------------- tensorrt_llm/visual_gen/media_refs.py | 149 ++++++++++++++++++ .../visual_gen/test_visual_gen_utils.py | 6 +- 4 files changed, 156 insertions(+), 128 deletions(-) create mode 100644 tensorrt_llm/visual_gen/media_refs.py diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 5f906252bafe..d26b321beab3 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -35,11 +35,8 @@ build_visual_gen_server_timings, build_visual_gen_timing_headers, ) -from tensorrt_llm.serve.visual_gen_utils import ( - VIDEO_STORE, - cleanup_reference_files, - parse_visual_gen_params, -) +from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE, parse_visual_gen_params +from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files if TYPE_CHECKING: # Type-only: importing tensorrt_llm.visual_gen at runtime would pull the diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 62e17fe9141d..a09103a5150c 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -6,24 +6,19 @@ import os from collections.abc import Mapping from io import BytesIO -from pathlib import Path 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 ( - _normalize_file_uri, - _safe_request_get, - is_isobmff_image_bytes, - sniff_media_kind, -) +from tensorrt_llm.inputs.media_io import sniff_media_kind from tensorrt_llm.logger import logger from tensorrt_llm.serve.openai_protocol import ( ImageEditRequest, ImageGenerationRequest, VideoGenerationRequest, ) +from tensorrt_llm.visual_gen.media_refs import _materialize_reference, _resolve_reference_string if TYPE_CHECKING: from fastapi import UploadFile @@ -115,64 +110,6 @@ def _merge_extra_params( params.extra_params = None -def _read_reference_payload(reference: str) -> bytes: - """Decode one base64 (optionally ``data:`` URI) reference string to bytes. - - Payload size is deliberately not checked here: encoded size is not part - of the request-validity contract, and body limits belong to the - proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so - malformed encodings — not sizes — are rejected. - """ - data = reference - if data.startswith("data:"): - comma = data.find(",") - if comma == -1: - raise ValueError("reference data: URI is malformed (missing comma).") - data = data[comma + 1 :] - try: - return base64.b64decode(data, validate=True) - except ValueError as exc: - # binascii.Error subclasses ValueError. - raise ValueError("reference is not valid base64 data.") from exc - - -def _resolve_reference_string(reference: str) -> bytes: - """Resolve one reference string to raw bytes, dispatching on URL scheme. - - Mirrors the LLM multimodal loader so serve references accept the same forms: - ``http(s)`` fetches through the SSRF-guarded loader (private-address block, - redirect re-validation, timeout, size cap); ``file://`` and bare local paths - read from disk; ``data:`` and base64 strings decode inline. A bare string is - decoded as base64 first and, failing that, read as a local file path. - Fetch/read failures become ``ValueError`` so a bad URL or path is a client - 400, not a server 500. - """ - scheme = urlparse(reference).scheme - if scheme in ("http", "https"): - try: - return _safe_request_get(reference).content - except Exception as exc: - raise ValueError(f"reference URL could not be fetched: {exc}") from exc - if scheme == "file": - try: - return Path(_normalize_file_uri(reference)).read_bytes() - except OSError as exc: - raise ValueError(f"reference file could not be read: {exc}") from exc - if scheme == "data": - return _read_reference_payload(reference) - # Bare string: base64 first (the established default), else a local file path - # so a plain path works without the file:// scheme. - try: - return _read_reference_payload(reference) - except ValueError: - try: - return Path(reference).read_bytes() - except OSError as exc: - raise ValueError( - f"reference is not valid base64 data, and not a readable local file: {exc}" - ) from exc - - def _reference_payload_and_role(ref: Any) -> tuple[bytes, Optional[str]]: """Extract ``(payload_bytes, role)`` from one raw HTTP reference. @@ -191,61 +128,6 @@ def _reference_payload_and_role(ref: Any) -> tuple[bytes, Optional[str]]: return _resolve_reference_string(data), role -def _materialize_reference( - payload: bytes, *, modality: str, ref_id: str, media_storage_path: Optional[str] -) -> str: - """Content-validate a reference payload and persist it, returning its path. - - HEIF/AVIF images are rejected on signature alone (Pillow support depends - on optional plugins the worker need not share). Video acceptance beyond the - container signature happens in the worker's NVDEC demux. - """ - if modality == "image": - if sniff_media_kind(payload) != "image": - raise ValueError( - "image_reference is not a recognized image; supported inputs are PNG/JPEG." - ) - if is_isobmff_image_bytes(payload): - raise ValueError( - "image_reference is a HEIF/AVIF image, which is not a supported " - "reference format; convert it to PNG or JPEG." - ) - elif modality == "video": - if sniff_media_kind(payload) != "video": - raise ValueError( - "video_reference is not a recognized media container; supported " - "inputs are MP4/AVI video." - ) - # audio: no signature sniffing (sniff_media_kind detects only image/video); - # the consuming pipeline validates the audio codec in its worker. - if media_storage_path is None: - raise ValueError(f"media_storage_path is required to store the {modality}_reference.") - ref_path = os.path.join(media_storage_path, ref_id) - with open(ref_path, "wb") as f: - f.write(payload) - return ref_path - - -def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) -> None: - """Remove the materialized reference inputs for one request. - - References are materialized as ``{request_id}_{modality}_ref_{i}`` (and the - deprecated ``{request_id}_input_ref``) under ``media_storage_path``. They are - input-only — unneeded once the pipeline has consumed them — so the request - owner removes them by the ``request_id`` prefix, covering image/video/audio - and the deprecated single reference regardless of count. Output files - (``{request_id}_{i}.``) carry no ``ref`` and are left untouched. - Best-effort: already-removed files are ignored. - """ - if media_storage_path is None: - return - for path in Path(media_storage_path).glob(f"{request_id}_*ref*"): - try: - path.unlink() - except OSError: - pass - - def _build_reference_list( value: Any, *, modality: str, id: str, media_storage_path: Optional[str] ) -> Optional[list]: diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py new file mode 100644 index 000000000000..6ad4fd1a43cd --- /dev/null +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -0,0 +1,149 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Reference-media resolve / materialize / cleanup, shared by serve and engine. + +Verb convention: ``resolve`` -> bytes, ``materialize`` -> path. These are used +by both the serve boundary (``tensorrt_llm/serve``) and the engine frontend +(``VisualGen.generate_async``), so they live here rather than under ``serve`` to +avoid an engine -> serve import. +""" + +from __future__ import annotations + +import base64 +import os +from pathlib import Path +from typing import Optional +from urllib.parse import urlparse + +from tensorrt_llm.inputs.media_io import ( + _normalize_file_uri, + _safe_request_get, + is_isobmff_image_bytes, + sniff_media_kind, +) + + +def _read_reference_payload(reference: str) -> bytes: + """Decode one base64 (optionally ``data:`` URI) reference string to bytes. + + Payload size is deliberately not checked here: encoded size is not part + of the request-validity contract, and body limits belong to the + proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so + malformed encodings — not sizes — are rejected. + """ + data = reference + if data.startswith("data:"): + comma = data.find(",") + if comma == -1: + raise ValueError("reference data: URI is malformed (missing comma).") + data = data[comma + 1 :] + try: + return base64.b64decode(data, validate=True) + except ValueError as exc: + # binascii.Error subclasses ValueError. + raise ValueError("reference is not valid base64 data.") from exc + + +def _resolve_reference_string(reference: str) -> bytes: + """Resolve one reference string to raw bytes, dispatching on URL scheme. + + Mirrors the LLM multimodal loader so serve references accept the same forms: + ``http(s)`` fetches through the SSRF-guarded loader (private-address block, + redirect re-validation, timeout, size cap); ``file://`` and bare local paths + read from disk; ``data:`` and base64 strings decode inline. A bare string is + decoded as base64 first and, failing that, read as a local file path. + Fetch/read failures become ``ValueError`` so a bad URL or path is a client + 400, not a server 500. + """ + scheme = urlparse(reference).scheme + if scheme in ("http", "https"): + try: + return _safe_request_get(reference).content + except Exception as exc: + raise ValueError(f"reference URL could not be fetched: {exc}") from exc + if scheme == "file": + try: + return Path(_normalize_file_uri(reference)).read_bytes() + except OSError as exc: + raise ValueError(f"reference file could not be read: {exc}") from exc + if scheme == "data": + return _read_reference_payload(reference) + # Bare string: base64 first (the established default), else a local file path + # so a plain path works without the file:// scheme. + try: + return _read_reference_payload(reference) + except ValueError: + try: + return Path(reference).read_bytes() + except OSError as exc: + raise ValueError( + f"reference is not valid base64 data, and not a readable local file: {exc}" + ) from exc + + +def _materialize_reference( + payload: bytes, *, modality: str, ref_id: str, media_storage_path: Optional[str] +) -> str: + """Content-validate a reference payload and persist it, returning its path. + + HEIF/AVIF images are rejected on signature alone (Pillow support depends + on optional plugins the worker need not share). Video acceptance beyond the + container signature happens in the worker's NVDEC demux. + """ + if modality == "image": + if sniff_media_kind(payload) != "image": + raise ValueError( + "image_reference is not a recognized image; supported inputs are PNG/JPEG." + ) + if is_isobmff_image_bytes(payload): + raise ValueError( + "image_reference is a HEIF/AVIF image, which is not a supported " + "reference format; convert it to PNG or JPEG." + ) + elif modality == "video": + if sniff_media_kind(payload) != "video": + raise ValueError( + "video_reference is not a recognized media container; supported " + "inputs are MP4/AVI video." + ) + # audio: no signature sniffing (sniff_media_kind detects only image/video); + # the consuming pipeline validates the audio codec in its worker. + if media_storage_path is None: + raise ValueError(f"media_storage_path is required to store the {modality}_reference.") + ref_path = os.path.join(media_storage_path, ref_id) + with open(ref_path, "wb") as f: + f.write(payload) + return ref_path + + +def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) -> None: + """Remove the materialized reference inputs for one request. + + References are materialized as ``{request_id}_{modality}_ref_{i}`` (and the + deprecated ``{request_id}_input_ref``) under ``media_storage_path``. They are + input-only — unneeded once the pipeline has consumed them — so the request + owner removes them by the ``request_id`` prefix, covering image/video/audio + and the deprecated single reference regardless of count. Output files + (``{request_id}_{i}.``) carry no ``ref`` and are left untouched. + Best-effort: already-removed files are ignored. + """ + if media_storage_path is None: + return + for path in Path(media_storage_path).glob(f"{request_id}_*ref*"): + try: + path.unlink() + except OSError: + pass diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 7d29d4b7681a..8c412aff57e2 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -25,10 +25,10 @@ from tensorrt_llm.serve.visual_gen_utils import ( _merge_extra_params, _warn_if_set_with_no_semantic, - cleanup_reference_files, parse_visual_gen_params, ) from tensorrt_llm.visual_gen import VisualGenParams +from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files pytestmark = pytest.mark.cpu_only @@ -558,7 +558,7 @@ def __init__(self, content): self.content = content monkeypatch.setattr( - "tensorrt_llm.serve.visual_gen_utils._safe_request_get", + "tensorrt_llm.visual_gen.media_refs._safe_request_get", lambda url, **kwargs: _FakeResp(png), ) request = VideoGenerationRequest(prompt="x", image_reference="https://example.com/a.png") @@ -575,7 +575,7 @@ def test_http_url_fetch_failure_is_client_error(self, tmp_path, monkeypatch): def _blocked(url, **kwargs): raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") - monkeypatch.setattr("tensorrt_llm.serve.visual_gen_utils._safe_request_get", _blocked) + monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) request = VideoGenerationRequest(prompt="x", image_reference="http://10.0.0.1/a.png") with pytest.raises(ValueError, match="reference URL could not be fetched"): parse_visual_gen_params( From c890b30ebabbe7fa3cd3142ec5f8b1d20d00fd8e Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 17 Aug 2026 03:11:55 -0700 Subject: [PATCH 17/61] [TRTLLM-15277][feat] Materialize reference inputs in the engine (generate_async) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit generate_async now resolves + materializes reference inputs on the coordinator before broadcast, via prepare_reference_slots: a trusted local path passes through untouched (not copied, not cleaned), and everything else (bytes / http(s) / data: / base64) is resolved to bytes and materialized to TRTLLM_MEDIA_STORAGE_PATH. Those files are reclaimed by cleanup_reference_files when the request reaches terminal state (VisualGenResult on_finish), keyed on request_id so user paths are never touched. This gives the standalone Python API the same materialize + cleanup the serve routes already had — fixing multi-GPU broadcasting raw bytes N-ways and fetching a URL once per rank. It is a no-op for serve today (serve still materializes; the engine sees a local path and passes it through); the serve slim-down that makes the engine the sole materializer is a follow-up. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/visual_gen/media_refs.py | 58 ++++++++++++++++++- tensorrt_llm/visual_gen/visual_gen.py | 41 ++++++++++++- .../visual_gen/test_visual_gen_utils.py | 55 ++++++++++++++++++ 3 files changed, 151 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index 6ad4fd1a43cd..3de2492e2f98 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -25,7 +25,7 @@ import base64 import os from pathlib import Path -from typing import Optional +from typing import Any, Optional from urllib.parse import urlparse from tensorrt_llm.inputs.media_io import ( @@ -147,3 +147,59 @@ def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) path.unlink() except OSError: pass + + +def resolve_media_storage_path() -> Path: + """Resolve the media storage directory, creating it if needed. + + Reads ``TRTLLM_MEDIA_STORAGE_PATH`` (default ``/tmp/trtllm_generated``), + shared by the serve boundary and the engine so both write materialized + references to the same place. + """ + path = Path(os.getenv("TRTLLM_MEDIA_STORAGE_PATH", "/tmp/trtllm_generated")) # nosec B108 + path.mkdir(parents=True, exist_ok=True) + return path + + +def _is_local_path(content: Any) -> bool: + """True if ``content`` is a trusted local path to pass through untouched. + + A ``file://`` URI, or a bare string naming an existing file. Everything else + (bytes, ``http(s)`` / ``data:`` URLs, base64) is materialized. + """ + if not isinstance(content, str): + return False + scheme = urlparse(content).scheme + if scheme == "file": + return True + return scheme == "" and os.path.exists(content) + + +def prepare_reference_slots( + params: Any, *, request_id: str, media_storage_path: Optional[str] +) -> None: + """Resolve + materialize each reference to a local path, in place. + + The single reference choke point, used by the engine (``generate_async``) + so serve and the standalone Python API share one path. A trusted local path + (``file://`` / existing bare path) passes through unchanged — not + materialized, not cleaned up (it is the caller's file). Everything else + (bytes / ``http(s)`` / ``data:`` / base64) resolves to bytes and materializes + to ``media_storage_path``; those files are reclaimed by + :func:`cleanup_reference_files` keyed on ``request_id``. Runs before the + coordinator broadcasts the request, so bad-media ``ValueError`` surfaces to + the caller synchronously (serve keeps its immediate 400). + """ + for slot in ("image_reference", "video_reference", "audio_reference"): + modality = slot.split("_", 1)[0] + for i, ref in enumerate(getattr(params, slot, None) or []): + content = ref.content + if _is_local_path(content): + continue + data = content if isinstance(content, bytes) else _resolve_reference_string(content) + ref.content = _materialize_reference( + data, + modality=modality, + ref_id=f"{request_id}_{modality}_ref_{i}", + media_storage_path=media_storage_path, + ) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 5889076138f1..44e589f8d989 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -19,7 +19,7 @@ import sys import weakref from pathlib import Path -from typing import Any, AsyncIterator, Dict, List, Literal, Optional, Union +from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional, Union from tensorrt_llm._torch.visual_gen import DiffusionRequest, DiffusionResponse from tensorrt_llm._torch.visual_gen.executor import ( @@ -31,6 +31,11 @@ from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema, RefSlotSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs +from tensorrt_llm.visual_gen.media_refs import ( + cleanup_reference_files, + prepare_reference_slots, + resolve_media_storage_path, +) from tensorrt_llm.visual_gen.output import VisualGenOutput from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params @@ -71,6 +76,7 @@ def __init__( request_id: int, executor: "DiffusionRemoteClient", batch_size: Optional[int] = None, + on_finish: Optional[Callable[[], None]] = None, ): self.request_id = request_id self.executor = executor @@ -79,6 +85,10 @@ def __init__( self._batch_size = batch_size self._resolved = None self._finished = False + # Run once at terminal state (success/error/timeout) to reclaim the + # engine-materialized reference files for this request; None otherwise. + self._on_finish = on_finish + self._cleaned = False @property def done(self) -> bool: @@ -128,10 +138,12 @@ async def aresult(self, timeout: Optional[float] = None): for _ in range(self._batch_size) ] self._finished = True + self._run_finish() return self._resolved_value() self._resolved = self._build_resolved(response) self._finished = True + self._run_finish() return self._resolved_value() def result(self, timeout: Optional[float] = None): @@ -156,6 +168,16 @@ def cancel(self): # ----- internals ----- + def _run_finish(self) -> None: + """Run the terminal cleanup callback once (idempotent, best-effort).""" + if self._cleaned or self._on_finish is None: + return + self._cleaned = True + try: + self._on_finish() + except Exception: + pass + def _build_resolved(self, response: "DiffusionResponse"): # Failure class travels on the result object, not on the public # ``VisualGenOutput`` — no new public field for an error taxonomy. @@ -437,6 +459,16 @@ def generate_async( if resolved_params.seed is None: resolved_params.seed = secrets.randbits(63) + # Resolve/materialize references to local paths here, on the coordinator, + # before the request is broadcast — the single choke point shared by serve + # and the standalone Python API. Runs synchronously so bad-media + # ``ValueError`` reaches the caller before dispatch (serve keeps its 400); + # a trusted local path passes through untouched (not materialized/cleaned). + media_storage_path = str(resolve_media_storage_path()) + prepare_reference_slots( + resolved_params, request_id=str(req_id), media_storage_path=media_storage_path + ) + request = DiffusionRequest( request_id=req_id, prompt=prompt, @@ -444,7 +476,12 @@ def generate_async( ) self.executor.enqueue_requests([request]) - return VisualGenResult(req_id, self.executor, batch_size=batch_size) + return VisualGenResult( + req_id, + self.executor, + batch_size=batch_size, + on_finish=lambda: cleanup_reference_files(media_storage_path, str(req_id)), + ) @staticmethod def _atexit_shutdown(self_ref): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 8c412aff57e2..c13dcb156be5 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -859,3 +859,58 @@ def test_none_storage_is_noop(self): def test_missing_files_are_ignored(self, tmp_path): cleanup_reference_files(str(tmp_path), "video_absent") # no raise + + +class TestPrepareReferenceSlots: + """Engine-side reference choke point: passthrough local paths, materialize the rest.""" + + def test_local_path_passthrough_not_materialized_or_cleaned(self, tmp_path): + from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + + src = tmp_path / "user.png" + Image.new("RGB", (4, 4)).save(src, format="PNG") + store = tmp_path / "store" + store.mkdir() + params = VisualGenParams(image_reference=str(src)) # bare local path -> passthrough + prepare_reference_slots(params, request_id="req1", media_storage_path=str(store)) + assert params.image_reference[0].content == str(src) # unchanged + assert list(store.iterdir()) == [] # nothing materialized + cleanup_reference_files(str(store), "req1") + assert src.exists() # user file untouched by cleanup + + def test_bytes_materialized_to_storage_and_cleaned(self, tmp_path): + from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + png = buf.getvalue() + store = tmp_path / "store" + store.mkdir() + params = VisualGenParams(image_reference=png) # bytes -> materialize + prepare_reference_slots(params, request_id="req2", media_storage_path=str(store)) + path = params.image_reference[0].content + assert path == str(store / "req2_image_ref_0") + assert Path(path).read_bytes() == png + cleanup_reference_files(str(store), "req2") + assert not Path(path).exists() + + def test_is_local_path(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _is_local_path + + f = tmp_path / "x" + f.write_bytes(b"a") + assert _is_local_path(str(f)) is True + assert _is_local_path(f.as_uri()) is True # file:// + assert _is_local_path("iVBORw0KGgo=") is False # base64-ish, no such file + assert _is_local_path("https://example.com/a.png") is False + assert _is_local_path(b"raw bytes") is False + + def test_resolve_media_storage_path(self, tmp_path, monkeypatch): + from tensorrt_llm.visual_gen.media_refs import resolve_media_storage_path + + target = tmp_path / "ms" + monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(target)) + resolved = resolve_media_storage_path() + assert resolved == target and target.is_dir() From 970552b74abb4ada9c35a6086508c087b72440b9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 17 Aug 2026 20:33:49 -0700 Subject: [PATCH 18/61] [TRTLLM-15277][refactor] Sink reference materialize/cleanup into the engine; serve becomes transport-only MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The engine's generate_async is now the single choke point for reference inputs: it resolves (http(s)/file://, data:/base64), validates, materializes to media storage, and reclaims the files via VisualGenResult's terminal hook keyed on the engine request id. Serve no longer materializes — parse_visual_gen_params only normalizes transport (multipart upload to bytes, strings pass through) and hands the pipeline a MediaRef carrying the raw content, so the standalone Python API and serve share one path. Both video routes and the image route offload the blocking generate_async via asyncio.to_thread and await it, then await generation through aresult (async, non-blocking). Bad media / unknown extra_params still surface as an immediate 400 before enqueue (sync) or before the 202 (async), so the async route drops its duplicate validate_visual_gen_params call and neither route keeps serve-side cleanup wiring. A trusted local path (file:// or an existing bare path) passes through untouched — not materialized, not cleaned up — with file:// normalized to a plain path for the pipeline; a missing path falls through to resolve so it becomes a client 400 rather than a silent passthrough. Partial-materialize failures reclaim earlier-written files at the choke point, and aresult runs cleanup on cancellation too so deleting an in-flight job reclaims its references. Tests: the serve util tests drive the real parse -> prepare_reference_slots flow and assert file:// / bare paths pass through; the endpoint mock mirrors the engine (validate + materialize + terminal cleanup). Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_server.py | 29 ++-- tensorrt_llm/serve/openai_video_routes.py | 69 ++++----- tensorrt_llm/serve/visual_gen_utils.py | 109 ++++--------- tensorrt_llm/visual_gen/media_refs.py | 51 ++++--- tensorrt_llm/visual_gen/visual_gen.py | 9 +- .../visual_gen/test_trtllm_serve_endpoints.py | 94 ++++++++---- .../visual_gen/test_visual_gen_utils.py | 144 ++++++++---------- 7 files changed, 240 insertions(+), 265 deletions(-) diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index dd5669fac0e9..ed501d164ec2 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -107,8 +107,7 @@ from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, build_visual_gen_timing_headers) -from tensorrt_llm.serve.visual_gen_utils import ( - cleanup_materialized_conditioning_inputs, parse_visual_gen_params) +from tensorrt_llm.serve.visual_gen_utils import parse_visual_gen_params from tensorrt_llm.usage import TerminalOutcome, record_termination_observation from tensorrt_llm.version import __version__ as VERSION @@ -3070,14 +3069,16 @@ async def openai_image_generation(self, request: ImageGenerationRequest, # through to the outer ``except Exception`` → 500 so the # client doesn't get blamed for a server-internal failure. try: - params = parse_visual_gen_params(request, image_id, - self.generator) + params = parse_visual_gen_params(request, self.generator) logger.info( f"Generating image: {image_id} with params: {params} and prompt: {request.prompt}" ) image_gen_start = time.perf_counter() - output = self.generator.generate(inputs=request.prompt, - params=params) + # Offload the blocking materialize/enqueue off the event loop but + # await it (bad params → 400 here); then await generation. + handle = await asyncio.to_thread(self.generator.generate_async, + request.prompt, params) + output = await handle.aresult() except ValueError as exc: logger.error(f"Image request error: {exc}") return self.create_error_response( @@ -3328,29 +3329,19 @@ async def openai_image_edit(self, raw_request: Request) -> Response: try: image_id = f"image_{uuid.uuid4().hex}" - input_paths = None try: request = await self._parse_image_edit_request(raw_request) path_error = self._reject_disabled_path(request.response_format) if path_error is not None: return path_error - params = parse_visual_gen_params( - request, - image_id, - self.generator, - media_storage_path=str(self.media_storage_path), - ) - input_paths = params.image + params = parse_visual_gen_params(request, self.generator) 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) + output = self.generator.generate(inputs=request.prompt, + params=params) except ValidationError as exc: return self._render_pydantic_validation_error(exc) except ValueError as exc: diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index d26b321beab3..91c80b46b74d 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -36,12 +36,12 @@ build_visual_gen_timing_headers, ) from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE, parse_visual_gen_params -from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files 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.params import VisualGenParams + from tensorrt_llm.visual_gen.visual_gen import VisualGenResult def _video_content_type(suffix: str) -> str: @@ -145,27 +145,22 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time - # Assigned before the try so the ``finally`` can always clean up any - # reference inputs materialized for this request id. + # Names this request's output files (``{video_id}_{i}``) and the b64 + # response id; references are keyed and reclaimed by the engine instead. video_id = f"video_{uuid.uuid4().hex}" try: # Client-side ValueErrors from content-type parsing, request # translation, encoder-format preflight, parameter validation, - # and the synchronous engine call return 400. Serialization / - # encoder failures further down (server-side) fall through to - # the outer ``except Exception`` → 500. + # and the engine call return 400. Serialization / encoder failures + # further down (server-side) fall through to the outer + # ``except Exception`` → 500. try: # Parse request based on content-type request = await self._parse_video_generation_request(raw_request) path_error = self._reject_disabled_path(request.response_format) if path_error is not None: return path_error - params = parse_visual_gen_params( - request, - video_id, - self.generator, - media_storage_path=str(self.media_storage_path), - ) + params = parse_visual_gen_params(request, self.generator) request_format = _resolve_tensor_only_format( request.format, request.extra_params, self.generator.extra_param_specs ) @@ -174,7 +169,13 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) sync_video_start = time.perf_counter() - output = self.generator.generate(inputs=request.prompt, params=params) + # Offload the blocking resolve/materialize/enqueue off the event + # loop but await it, so bad media / bad params still surface as + # 400 here; then await generation on the executor's loop. + handle = await asyncio.to_thread( + self.generator.generate_async, request.prompt, params + ) + output = await handle.aresult() except ValidationError as exc: return self._render_pydantic_validation_error(exc) except ValueError as exc: @@ -281,11 +282,6 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: err_type="InternalServerError", status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) - finally: - # References are input-only; the pipeline has consumed them by now, - # so remove them (success or failure) — conditioned requests must not - # accumulate materialized inputs in media storage. - cleanup_reference_files(str(self.media_storage_path), video_id) async def _parse_video_generation_request( self, @@ -384,10 +380,9 @@ async def openai_video_generation_async( - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time - # Assigned before the try so the ``finally`` can clean up references - # materialized for this request if no background task takes ownership. + # Names this request's output files and VIDEO_STORE entry; references + # are keyed and reclaimed by the engine when the task awaits the handle. video_id = f"video_{uuid.uuid4().hex}" - task_started = False try: # Parse request based on content-type request = await self._parse_video_generation_request(raw_request) @@ -395,9 +390,7 @@ async def openai_video_generation_async( if path_error is not None: return path_error - params = parse_visual_gen_params( - request, video_id, self.generator, media_storage_path=str(self.media_storage_path) - ) + params = parse_visual_gen_params(request, self.generator) # Synchronously validate the resolved params against the # loaded pipeline's extra-param specs / declared defaults # so unknown ``extra_params`` keys and similar engine-side @@ -419,6 +412,13 @@ async def openai_video_generation_async( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) + # Resolve/materialize references, validate params, and enqueue in the + # foreground (offloaded but awaited) so bad media / unknown + # extra_params surface as 400 here, before the 202 — not as a queued + # job that later fails. The engine reclaims the references via its + # terminal hook once the background task awaits the handle. + handle = await asyncio.to_thread(self.generator.generate_async, request.prompt, params) + # Persist the queued job before scheduling the background task so # that a fast-completing task can always look it up in VIDEO_STORE. video_job = VideoJob( @@ -435,18 +435,17 @@ async def openai_video_generation_async( ) await VIDEO_STORE.upsert(video_id, video_job) - # Start background generation task + # Start background task to await generation and save the result. task = asyncio.create_task( self._generate_video_background( video_id=video_id, request=request, params=params, request_format=request_format, + handle=handle, ) ) self.video_gen_tasks[video_id] = task - # The background task now owns reference cleanup (its ``finally``). - task_started = True task.add_done_callback(lambda t, vid=video_id: self._on_video_task_done(vid, t)) return JSONResponse(content=video_job.model_dump(), status_code=202) @@ -463,11 +462,6 @@ async def openai_video_generation_async( err_type="InternalServerError", status_code=HTTPStatus.INTERNAL_SERVER_ERROR, ) - finally: - if not task_started: - # Failed before the background task was scheduled — nothing else - # will clean these up. - cleanup_reference_files(str(self.media_storage_path), video_id) async def _generate_video_background( self, @@ -475,8 +469,9 @@ async def _generate_video_background( request: VideoGenerationRequest, params: VisualGenParams, request_format: str, + handle: "VisualGenResult", ): - """Background task to generate video and save to storage. + """Background task to await generation and save to storage. ``request_format`` is the format already resolved by the route (see :func:`_resolve_tensor_only_format`), not ``request.format``: the @@ -489,8 +484,7 @@ async def _generate_video_background( if job: job.status = "generating" await VIDEO_STORE.upsert(video_id, job) - future = self.generator.generate_async(inputs=request.prompt, params=params) - output = await future + output = await handle if output.video is None: # Update job status to failed since we're in a background task @@ -572,11 +566,6 @@ async def _generate_video_background( job.completed_at = int(time.time()) job.error = str(e) await VIDEO_STORE.upsert(video_id, job) - finally: - # References are input-only; remove them once generation has run, - # failed, or been cancelled. Runs on CancelledError too, so a delete - # that cancels an in-flight job still reclaims the inputs. - cleanup_reference_files(str(self.media_storage_path), video_id) async def list_videos(self, raw_request: Request) -> Response: """List all generated videos. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index a09103a5150c..9ce3f29c0f66 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -18,7 +18,7 @@ ImageGenerationRequest, VideoGenerationRequest, ) -from tensorrt_llm.visual_gen.media_refs import _materialize_reference, _resolve_reference_string +from tensorrt_llm.visual_gen.media_refs import _resolve_reference_string if TYPE_CHECKING: from fastapi import UploadFile @@ -110,33 +110,34 @@ def _merge_extra_params( params.extra_params = None -def _reference_payload_and_role(ref: Any) -> tuple[bytes, Optional[str]]: - """Extract ``(payload_bytes, role)`` from one raw HTTP reference. +def _reference_payload_and_role(ref: Any) -> tuple[Any, Optional[str]]: + """Extract ``(content, role)`` from one raw HTTP reference for transport. ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or a local file path), a multipart ``UploadFile`` (has ``.file``), or a ``MediaReferenceItem`` - exposing ``content`` and an optional ``role``. + exposing ``content`` and an optional ``role``. An upload is read to ``bytes`` + here — the only decode the boundary owns; strings pass through untouched for + the engine to resolve and materialize. """ role = getattr(ref, "role", None) if isinstance(ref, str): - return _resolve_reference_string(ref), role + return ref, role if hasattr(ref, "file"): # multipart UploadFile return ref.file.read(), role data = getattr(ref, "content", None) if not isinstance(data, str): raise ValueError("reference item must carry a 'content' string.") - return _resolve_reference_string(data), role + return data, role -def _build_reference_list( - value: Any, *, modality: str, id: str, media_storage_path: Optional[str] -) -> Optional[list]: - """Materialize an HTTP reference field into a list of ``MediaRef`` objects. +def _build_reference_list(value: Any) -> Optional[list]: + """Normalize one HTTP reference field into a list of ``MediaRef`` objects. - ``value`` is None, a base64/data-URI string, a multipart ``UploadFile``, a - ``MediaReferenceItem``, or a list of any of those. Each entry is decoded, - content-validated for ``modality``, persisted to a per-index path, and - wrapped as ``MediaRef`` (carrying ``role`` when present). + ``value`` is None, a base64/data-URI/URL/path string, a multipart + ``UploadFile``, a ``MediaReferenceItem``, or a list of any of those. Each + entry becomes a ``MediaRef`` carrying its transport content — ``bytes`` for + an upload, the string otherwise — plus its ``role``. Resolution and + materialization happen later at the engine choke point. """ if value is None: return None @@ -146,27 +147,9 @@ def _build_reference_list( raw_items = value if isinstance(value, list) else [value] refs = [] - created_paths: list[str] = [] - try: - for i, item in enumerate(raw_items): - payload, role = _reference_payload_and_role(item) - ref_path = _materialize_reference( - payload, - modality=modality, - ref_id=f"{id}_{modality}_ref_{i}", - media_storage_path=media_storage_path, - ) - created_paths.append(ref_path) - refs.append(MediaRef(content=ref_path, role=role)) - except Exception: - # A later item failed; remove the files earlier items already wrote so - # a rejected multi-reference request leaves nothing on disk. - for path in created_paths: - try: - os.remove(path) - except OSError: - pass - raise + for item in raw_items: + content, role = _reference_payload_and_role(item) + refs.append(MediaRef(content=content, role=role)) return refs @@ -386,15 +369,14 @@ def cleanup_materialized_conditioning_inputs(value: Any) -> None: def _apply_deprecated_input_reference( input_reference: str | UploadFile | None, params: VisualGenParams, - *, - id: str, - media_storage_path: str | None, ) -> None: """Back-compat for the deprecated single ``input_reference``. Sniff-routes the payload to ``image_reference`` (image) or ``video_reference`` (video), preserving the pre-typed-fields behavior. Ignored when a typed image/video reference is already set — the typed fields take precedence. + Routing needs the bytes, so a string payload is resolved here; the engine + materializes the resulting ``MediaRef`` like any other reference. """ if input_reference is None: return @@ -404,23 +386,13 @@ def _apply_deprecated_input_reference( from tensorrt_llm.visual_gen.params import MediaRef payload, _ = _reference_payload_and_role(input_reference) + if isinstance(payload, str): + payload = _resolve_reference_string(payload) kind = sniff_media_kind(payload) if kind == "image": - path = _materialize_reference( - payload, - modality="image", - ref_id=f"{id}_input_ref", - media_storage_path=media_storage_path, - ) - params.image_reference = [MediaRef(content=path)] + params.image_reference = [MediaRef(content=payload)] elif kind == "video": - path = _materialize_reference( - payload, - modality="video", - ref_id=f"{id}_input_ref", - media_storage_path=media_storage_path, - ) - params.video_reference = [MediaRef(content=path)] + params.video_reference = [MediaRef(content=payload)] else: raise ValueError( "input_reference is not a recognized media container; supported " @@ -430,9 +402,7 @@ def _apply_deprecated_input_reference( def parse_visual_gen_params( request: ImageGenerationRequest | ImageEditRequest | VideoGenerationRequest, - id: str, generator: VisualGen, - media_storage_path: Optional[str] = None, ) -> VisualGenParams: """Translate an HTTP request into :class:`VisualGenParams`. @@ -481,15 +451,8 @@ def parse_visual_gen_params( 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, - ) + params.image_reference = _build_reference_list(request.image) elif isinstance(request, VideoGenerationRequest): if request.frame_rate is not None: @@ -516,27 +479,19 @@ def parse_visual_gen_params( "directly." ) params.num_frames = derived - # Reference inputs: materialize each transport (base64/data-URI/upload) - # to a stored file and hand the pipeline a ``MediaRef`` carrying the - # local path. Decode stays model-specific in the worker. - image_refs = _build_reference_list( - request.image_reference, modality="image", id=id, media_storage_path=media_storage_path - ) + # Reference inputs: hand the pipeline a ``MediaRef`` carrying the + # transport content (``bytes`` for an upload, the string otherwise). + # The engine resolves and materializes; the boundary never touches disk. + image_refs = _build_reference_list(request.image_reference) if image_refs: params.image_reference = image_refs - video_refs = _build_reference_list( - request.video_reference, modality="video", id=id, media_storage_path=media_storage_path - ) + video_refs = _build_reference_list(request.video_reference) if video_refs: params.video_reference = video_refs - audio_refs = _build_reference_list( - request.audio_reference, modality="audio", id=id, media_storage_path=media_storage_path - ) + audio_refs = _build_reference_list(request.audio_reference) if audio_refs: params.audio_reference = audio_refs - _apply_deprecated_input_reference( - request.input_reference, params, id=id, media_storage_path=media_storage_path - ) + _apply_deprecated_input_reference(request.input_reference, params) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _decode_inline_media(request.extra_params, generator.extra_param_specs) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index 3de2492e2f98..efa3b5068dfa 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -164,14 +164,16 @@ def resolve_media_storage_path() -> Path: def _is_local_path(content: Any) -> bool: """True if ``content`` is a trusted local path to pass through untouched. - A ``file://`` URI, or a bare string naming an existing file. Everything else - (bytes, ``http(s)`` / ``data:`` URLs, base64) is materialized. + A ``file://`` URI or bare string naming an *existing* file. A missing path + (either form) returns False so it falls through to the resolve step, whose + read raises ``ValueError`` — a client 400, not a silent passthrough. + Everything else (bytes, ``http(s)`` / ``data:`` URLs, base64) materializes. """ if not isinstance(content, str): return False scheme = urlparse(content).scheme if scheme == "file": - return True + return os.path.exists(_normalize_file_uri(content)) return scheme == "" and os.path.exists(content) @@ -182,24 +184,35 @@ def prepare_reference_slots( The single reference choke point, used by the engine (``generate_async``) so serve and the standalone Python API share one path. A trusted local path - (``file://`` / existing bare path) passes through unchanged — not - materialized, not cleaned up (it is the caller's file). Everything else + (``file://`` / existing bare path) passes through — not materialized, not + cleaned up (it is the caller's file); a ``file://`` URI is normalized to a + plain path so the pipeline, which opens paths, can read it. Everything else (bytes / ``http(s)`` / ``data:`` / base64) resolves to bytes and materializes to ``media_storage_path``; those files are reclaimed by :func:`cleanup_reference_files` keyed on ``request_id``. Runs before the coordinator broadcasts the request, so bad-media ``ValueError`` surfaces to - the caller synchronously (serve keeps its immediate 400). + the caller synchronously (serve keeps its immediate 400). If a later slot + fails mid-materialize, the files earlier slots wrote are reclaimed here so a + rejected request leaves nothing on disk. """ - for slot in ("image_reference", "video_reference", "audio_reference"): - modality = slot.split("_", 1)[0] - for i, ref in enumerate(getattr(params, slot, None) or []): - content = ref.content - if _is_local_path(content): - continue - data = content if isinstance(content, bytes) else _resolve_reference_string(content) - ref.content = _materialize_reference( - data, - modality=modality, - ref_id=f"{request_id}_{modality}_ref_{i}", - media_storage_path=media_storage_path, - ) + try: + for slot in ("image_reference", "video_reference", "audio_reference"): + modality = slot.split("_", 1)[0] + for i, ref in enumerate(getattr(params, slot, None) or []): + content = ref.content + if _is_local_path(content): + if urlparse(content).scheme == "file": + ref.content = _normalize_file_uri(content) + continue + data = content if isinstance(content, bytes) else _resolve_reference_string(content) + ref.content = _materialize_reference( + data, + modality=modality, + ref_id=f"{request_id}_{modality}_ref_{i}", + media_storage_path=media_storage_path, + ) + except Exception: + # The terminal on_finish hook is not wired yet (the request is never + # enqueued on failure), so reclaim any files earlier slots wrote here. + cleanup_reference_files(media_storage_path, request_id) + raise diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 44e589f8d989..b30d180eef08 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -114,7 +114,14 @@ async def aresult(self, timeout: Optional[float] = None): self.executor.await_responses(self.request_id, timeout=timeout), self.executor._event_loop, ) - response = await asyncio.wrap_future(future) + try: + response = await asyncio.wrap_future(future) + except asyncio.CancelledError: + # A caller (e.g. serve's delete-in-flight) dropped the wait. The + # worker keeps running but its output is discarded, so run the + # terminal cleanup here to reclaim the materialized references. + self._run_finish() + raise if response is None: # Timeout before any response. Tell the executor to drop any 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 2623c6bbc282..5f218637bcda 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -35,7 +35,13 @@ from tensorrt_llm.serve.openai_server import _normalize_image_output from tensorrt_llm.serve.visual_gen_metrics import SERVER_TIMING_HEADER from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE +from tensorrt_llm.visual_gen.media_refs import ( + cleanup_reference_files, + prepare_reference_slots, + resolve_media_storage_path, +) from tensorrt_llm.visual_gen.output import VisualGenMetrics, VisualGenOutput +from tensorrt_llm.visual_gen.params import validate_visual_gen_params pytestmark = pytest.mark.cpu_only @@ -265,37 +271,40 @@ def _snapshot_refs(self, params) -> None: self.last_ref_bytes[path] = fh.read() def generate(self, inputs=None, params=None) -> VisualGenOutput: - self.last_inputs = inputs - self.last_params = params - self._snapshot_refs(params) - if self._validation_error is not None: - raise self._validation_error - if self._generate_error is not None: - raise self._generate_error - if self._should_fail: - raise RuntimeError("Generation intentionally failed") - n = getattr(params, "num_images_per_prompt", 1) if params else 1 - return VisualGenOutput( - request_id=self._next_request_id(), - image=self._maybe_batch(self._image, n), - video=self._maybe_batch(self._video, n), - audio=self._audio, - metrics=_make_dummy_metrics(), - ) + return self.generate_async(inputs=inputs, params=params).result() def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": self.last_inputs = inputs self.last_params = params - self._snapshot_refs(params) if self._validation_error is not None: raise self._validation_error + # Mirror the real engine entry: validate against the pipeline metadata + # (unknown extra_params / undeclared fields / ref arity) before doing any + # work, so the route's synchronous 400 path is exercised end-to-end. + if params is not None: + validate_visual_gen_params( + params, + declared_defaults=self.executor.default_generation_params, + extra_param_specs=self.executor.extra_param_specs, + ref_slot_specs=self.executor.ref_slot_specs, + ) + # Materialize references at the coordinator, then hand the result a + # terminal cleanup keyed on the same request id. + req_id = self._next_request_id() + media_storage_path = str(resolve_media_storage_path()) + prepare_reference_slots( + params, request_id=str(req_id), media_storage_path=media_storage_path + ) + self._snapshot_refs(params) n = getattr(params, "num_images_per_prompt", 1) if params else 1 return MockVisualGenResult( - request_id=self._next_request_id(), + request_id=req_id, image=self._maybe_batch(self._image, n), video=self._maybe_batch(self._video, n), audio=self._audio, should_fail=self._should_fail, + generate_error=self._generate_error, + on_finish=lambda: cleanup_reference_files(media_storage_path, str(req_id)), ) def _next_request_id(self) -> int: @@ -349,17 +358,34 @@ def __init__( video: Optional[torch.Tensor] = None, audio: Optional[torch.Tensor] = None, should_fail: bool = False, + generate_error: Optional[BaseException] = None, + on_finish=None, ): self.request_id = request_id self._image = image self._video = video self._audio = audio self._should_fail = should_fail + # Engine-side failure surfaced through the result (capacity/client), + # distinct from a coordinator preflight rejection. + self._generate_error = generate_error + self._on_finish = on_finish + self._cleaned = False + + def _run_finish(self): + # Terminal reference cleanup, run once (idempotent), mirroring the real + # VisualGenResult so it fires on success and failure alike. + if self._cleaned or self._on_finish is None: + return + self._cleaned = True + try: + self._on_finish() + except Exception: + pass - def __await__(self): - return self.aresult().__await__() - - async def aresult(self, timeout=None): + def _resolve(self) -> VisualGenOutput: + if self._generate_error is not None: + raise self._generate_error if self._should_fail: raise RuntimeError("Async generation intentionally failed") return VisualGenOutput( @@ -370,16 +396,20 @@ async def aresult(self, timeout=None): metrics=_make_dummy_metrics(), ) + def __await__(self): + return self.aresult().__await__() + + async def aresult(self, timeout=None): + try: + return self._resolve() + finally: + self._run_finish() + def result(self, timeout=None): - if self._should_fail: - raise RuntimeError("Async generation intentionally failed") - return VisualGenOutput( - request_id=self.request_id, - image=self._image, - video=self._video, - audio=self._audio, - metrics=_make_dummy_metrics(), - ) + try: + return self._resolve() + finally: + self._run_finish() # --------------------------------------------------------------------------- diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index c13dcb156be5..08b6fac602ac 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -28,7 +28,19 @@ parse_visual_gen_params, ) from tensorrt_llm.visual_gen import VisualGenParams -from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files +from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files, prepare_reference_slots + + +def _parse_and_prepare(request, generator, request_id, media_storage_path): + """Run the production reference flow: serve transport, then engine materialize. + + ``parse_visual_gen_params`` only normalizes transport (upload -> bytes, + strings pass through); resolution + materialization happen at the engine + choke point, so tests that assert on stored paths drive both stages. + """ + params = parse_visual_gen_params(request, generator) + prepare_reference_slots(params, request_id=request_id, media_storage_path=media_storage_path) + return params pytestmark = pytest.mark.cpu_only @@ -83,7 +95,7 @@ def test_all_none_request_keeps_pipeline_defaults(self, image_request_defaults): generator = _StubVisualGen( defaults={"width": 1024, "height": 1024, "num_inference_steps": 30}, ) - params = parse_visual_gen_params(image_request_defaults, "id-1", generator) + params = parse_visual_gen_params(image_request_defaults, generator) assert params.width == 1024 assert params.height == 1024 assert params.num_inference_steps == 30 @@ -103,7 +115,7 @@ def test_image_explicit_fields_override_defaults(self): n=4, negative_prompt="blurry", ) - params = parse_visual_gen_params(request, "id-2", generator) + params = parse_visual_gen_params(request, generator) assert (params.width, params.height) == (512, 512) assert params.num_inference_steps == 10 assert params.guidance_scale == 4.0 @@ -115,19 +127,19 @@ def test_image_explicit_fields_override_defaults(self): def test_size_string_used_when_width_height_absent(self): generator = _StubVisualGen() request = ImageGenerationRequest(prompt="cat", size="768x256") - params = parse_visual_gen_params(request, "id-3", generator) + params = parse_visual_gen_params(request, generator) assert (params.width, params.height) == (768, 256) def test_width_height_pair_wins_over_size(self): generator = _StubVisualGen() request = ImageGenerationRequest(prompt="cat", size="768x256", width=128, height=64) - params = parse_visual_gen_params(request, "id-4", generator) + params = parse_visual_gen_params(request, generator) assert (params.width, params.height) == (128, 64) def test_image_seed_propagates(self): generator = _StubVisualGen() request = ImageGenerationRequest(prompt="cat", seed=12345) - params = parse_visual_gen_params(request, "id-seed", generator) + params = parse_visual_gen_params(request, generator) assert params.seed == 12345 @@ -194,7 +206,7 @@ def _fake_warning(msg: str, *args: object, **kwargs: object) -> None: def test_quality_hd_does_not_override_steps(self): generator = _StubVisualGen(defaults={"num_inference_steps": 25}) request = ImageGenerationRequest(prompt="cat", quality="hd") - params = parse_visual_gen_params(request, "id-q", generator) + params = parse_visual_gen_params(request, generator) # ``quality`` is an OpenAI-shape no-semantic field. The pipeline # default for ``num_inference_steps`` must reach the engine # unchanged. @@ -234,21 +246,21 @@ class TestVideoFrameBudget: def test_num_frames_wins_over_seconds_times_frame_rate(self): generator = _StubVisualGen(defaults={"frame_rate": 24.0}) request = VideoGenerationRequest(prompt="x", num_frames=33, seconds=10.0) - params = parse_visual_gen_params(request, "id-v1", generator) + params = parse_visual_gen_params(request, generator) assert params.num_frames == 33 def test_seconds_and_frame_rate_derive_num_frames(self): generator = _StubVisualGen(defaults={"frame_rate": 12.0}) # fps alias resolves to frame_rate via populate_by_name=True request = VideoGenerationRequest(prompt="x", seconds=2.5, fps=24) - params = parse_visual_gen_params(request, "id-v2", generator) + params = parse_visual_gen_params(request, generator) assert params.frame_rate == 24.0 assert params.num_frames == int(2.5 * 24.0) def test_seconds_alone_uses_pipeline_frame_rate(self): generator = _StubVisualGen(defaults={"frame_rate": 16.0}) request = VideoGenerationRequest(prompt="x", seconds=4.0) - params = parse_visual_gen_params(request, "id-v3", generator) + params = parse_visual_gen_params(request, generator) assert params.frame_rate == 16.0 assert params.num_frames == int(4.0 * 16.0) @@ -259,7 +271,7 @@ def test_video_does_not_carry_n(self): # leave ``num_images_per_prompt`` unchanged from the pipeline # default. request = VideoGenerationRequest(prompt="x") - params = parse_visual_gen_params(request, "id-v4", generator) + params = parse_visual_gen_params(request, generator) assert params.num_images_per_prompt == 1 @@ -276,9 +288,7 @@ def test_base64_image_reference_written_to_disk(self, tmp_path): img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() request = VideoGenerationRequest(prompt="x", image_reference=b64) - params = parse_visual_gen_params( - request, "vid-1", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-1", str(tmp_path)) assert len(params.image_reference) == 1 ref_path = params.image_reference[0].content assert str(ref_path).endswith("vid-1_image_ref_0") @@ -295,9 +305,7 @@ def test_image_reference_role_and_list(self, tmp_path): request = VideoGenerationRequest( prompt="x", image_reference=[b64, {"content": b64, "role": "last_frame"}] ) - params = parse_visual_gen_params( - request, "vid-r", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-r", str(tmp_path)) assert [r.role for r in params.image_reference] == [None, "last_frame"] paths = [r.content for r in params.image_reference] assert len(set(paths)) == 2 # unique file per index @@ -310,7 +318,7 @@ def test_missing_media_storage_path_raises(self): b64 = base64.b64encode(buf.getvalue()).decode() request = VideoGenerationRequest(prompt="x", image_reference=b64) with pytest.raises(ValueError, match="media_storage_path"): - parse_visual_gen_params(request, "vid-2", generator, media_storage_path=None) + _parse_and_prepare(request, generator, "vid-2", None) _TEST_DATA = Path(__file__).parent / "test_data" @@ -335,9 +343,7 @@ def test_multipart_avi_video_reference_written_to_disk(self, tmp_path): payload = self._avi_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.avi") request = VideoGenerationRequest(prompt="x", video_reference=upload) - params = parse_visual_gen_params( - request, "vid-avi", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-avi", str(tmp_path)) assert params.image_reference is None assert Path(params.video_reference[0].content).read_bytes() == payload @@ -346,9 +352,7 @@ def test_multipart_mp4_video_reference_written_to_disk(self, tmp_path): payload = self._mp4_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.mp4") request = VideoGenerationRequest(prompt="x", video_reference=upload) - params = parse_visual_gen_params( - request, "vid-3", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-3", str(tmp_path)) # Encoded payload is persisted byte-identical — the boundary never # decodes video; the worker demuxes/NVDEC-decodes the conditioning # window from the stored file. @@ -364,7 +368,7 @@ def test_video_reference_needs_media_storage(self): b64 = base64.b64encode(self._mp4_bytes()).decode() request = VideoGenerationRequest(prompt="x", video_reference=b64) with pytest.raises(ValueError, match="media_storage_path"): - parse_visual_gen_params(request, "vid-9", generator, media_storage_path=None) + _parse_and_prepare(request, generator, "vid-9", None) def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): # The deprecated single input_reference is sniff-routed to the typed slot. @@ -374,19 +378,19 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): img_b64 = base64.b64encode(buf.getvalue()).decode() vid_b64 = base64.b64encode(self._mp4_bytes()).decode() - p = parse_visual_gen_params( + p = _parse_and_prepare( VideoGenerationRequest(prompt="x", input_reference=img_b64), - "vid-i", generator, - media_storage_path=str(tmp_path), + "vid-i", + str(tmp_path), ) assert len(p.image_reference) == 1 and p.video_reference is None - p = parse_visual_gen_params( + p = _parse_and_prepare( VideoGenerationRequest(prompt="x", input_reference=vid_b64), - "vid-v", generator, - media_storage_path=str(tmp_path), + "vid-v", + str(tmp_path), ) assert len(p.video_reference) == 1 and p.image_reference is None @@ -397,11 +401,11 @@ def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): Image.new("RGB", (4, 4)).save(buf, format="PNG") img_b64 = base64.b64encode(buf.getvalue()).decode() vid_b64 = base64.b64encode(self._mp4_bytes()).decode() - p = parse_visual_gen_params( + p = _parse_and_prepare( VideoGenerationRequest(prompt="x", image_reference=img_b64, input_reference=vid_b64), - "vid-x", generator, - media_storage_path=str(tmp_path), + "vid-x", + str(tmp_path), ) assert len(p.image_reference) == 1 assert p.video_reference is None # input_reference video dropped @@ -413,9 +417,7 @@ def test_base64_video_reference_written_to_disk(self, tmp_path): payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() request = VideoGenerationRequest(prompt="x", video_reference=b64) - params = parse_visual_gen_params( - request, "vid-4", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-4", str(tmp_path)) assert params.image_reference is None assert Path(params.video_reference[0].content).read_bytes() == payload @@ -429,9 +431,7 @@ def test_video_reference_survives_real_specs(self, tmp_path): payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() request = VideoGenerationRequest(prompt="x", video_reference=b64) - params = parse_visual_gen_params( - request, "vid-10", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-10", str(tmp_path)) assert Path(params.video_reference[0].content).read_bytes() == payload def test_multipart_image_reference_written_to_disk(self, tmp_path): @@ -444,9 +444,7 @@ def test_multipart_image_reference_written_to_disk(self, tmp_path): buf.seek(0) upload = UploadFile(file=buf, filename="ref.jpg") request = VideoGenerationRequest(prompt="x", image_reference=upload) - params = parse_visual_gen_params( - request, "vid-5", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-5", str(tmp_path)) assert params.extra_params is None assert str(params.image_reference[0].content).endswith("vid-5_image_ref_0") @@ -458,18 +456,18 @@ def test_wrong_modality_content_raises(self, tmp_path): img_b64 = base64.b64encode(buf.getvalue()).decode() vid_b64 = base64.b64encode(self._mp4_bytes()).decode() with pytest.raises(ValueError, match="video_reference is not a recognized"): - parse_visual_gen_params( + _parse_and_prepare( VideoGenerationRequest(prompt="x", video_reference=img_b64), - "vid-m1", generator, - media_storage_path=str(tmp_path), + "vid-m1", + str(tmp_path), ) with pytest.raises(ValueError, match="image_reference is not a recognized image"): - parse_visual_gen_params( + _parse_and_prepare( VideoGenerationRequest(prompt="x", image_reference=vid_b64), - "vid-m2", generator, - media_storage_path=str(tmp_path), + "vid-m2", + str(tmp_path), ) assert list(tmp_path.iterdir()) == [] @@ -478,7 +476,7 @@ def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): b64 = base64.b64encode(b"neither an image nor a video").decode() request = VideoGenerationRequest(prompt="x", image_reference=b64) with pytest.raises(ValueError, match="not a recognized image"): - parse_visual_gen_params(request, "vid-6", generator, media_storage_path=str(tmp_path)) + _parse_and_prepare(request, generator, "vid-6", str(tmp_path)) # Classification runs on the bytes; rejected content never touches disk. assert list(tmp_path.iterdir()) == [] @@ -488,7 +486,7 @@ def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): # length, so b64decode raises. request = VideoGenerationRequest(prompt="x", image_reference="ABC") with pytest.raises(ValueError, match="not valid base64"): - parse_visual_gen_params(request, "vid-7", generator, media_storage_path=str(tmp_path)) + _parse_and_prepare(request, generator, "vid-7", str(tmp_path)) assert list(tmp_path.iterdir()) == [] def test_upload_stream_failure_cleans_up_tmp(self, tmp_path): @@ -502,7 +500,7 @@ def read(self, *args, **kwargs): request = VideoGenerationRequest(prompt="x", video_reference=upload) # I/O failures keep their server-error semantics (no 400 masking) … with pytest.raises(OSError, match="client went away"): - parse_visual_gen_params(request, "vid-8", generator, media_storage_path=str(tmp_path)) + _parse_and_prepare(request, generator, "vid-8", str(tmp_path)) # … and the payload read fails before any file is written, so nothing leaks. assert list(tmp_path.iterdir()) == [] @@ -516,35 +514,33 @@ def test_multi_reference_partial_failure_cleans_up(self, tmp_path): bad = base64.b64encode(b"neither an image nor a video").decode() request = VideoGenerationRequest(prompt="x", image_reference=[good, bad]) with pytest.raises(ValueError, match="not a recognized image"): - parse_visual_gen_params(request, "vid-11", generator, media_storage_path=str(tmp_path)) + _parse_and_prepare(request, generator, "vid-11", str(tmp_path)) assert list(tmp_path.iterdir()) == [] - def test_file_uri_image_reference_read_and_materialized(self, tmp_path): - # A file:// reference is read from local disk and persisted like any other. + def test_file_uri_image_reference_passthrough(self, tmp_path): + # A file:// reference is a trusted local path: normalized to a plain path + # and passed through untouched — not copied into media storage. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (7, 8, 9)).save(src, format="PNG") store = tmp_path / "store" store.mkdir() request = VideoGenerationRequest(prompt="x", image_reference=src.as_uri()) - params = parse_visual_gen_params( - request, "vid-file", generator, media_storage_path=str(store) - ) - assert Path(params.image_reference[0].content).read_bytes() == src.read_bytes() + params = _parse_and_prepare(request, generator, "vid-file", str(store)) + assert params.image_reference[0].content == str(src) + assert list(store.iterdir()) == [] - def test_bare_path_image_reference_read_and_materialized(self, tmp_path): - # A bare local path (no file:// scheme) is read from disk after the - # base64 decode attempt fails. + def test_bare_path_image_reference_passthrough(self, tmp_path): + # A bare local path is passed through unchanged — not materialized. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (11, 22, 33)).save(src, format="PNG") store = tmp_path / "store" store.mkdir() request = VideoGenerationRequest(prompt="x", image_reference=str(src)) - params = parse_visual_gen_params( - request, "vid-bare", generator, media_storage_path=str(store) - ) - assert Path(params.image_reference[0].content).read_bytes() == src.read_bytes() + params = _parse_and_prepare(request, generator, "vid-bare", str(store)) + assert params.image_reference[0].content == str(src) + assert list(store.iterdir()) == [] def test_http_url_image_reference_fetched_and_materialized(self, tmp_path, monkeypatch): # An http(s) reference is fetched through the guarded loader, then stored. @@ -562,9 +558,7 @@ def __init__(self, content): lambda url, **kwargs: _FakeResp(png), ) request = VideoGenerationRequest(prompt="x", image_reference="https://example.com/a.png") - params = parse_visual_gen_params( - request, "vid-url", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-url", str(tmp_path)) assert Path(params.image_reference[0].content).read_bytes() == png def test_http_url_fetch_failure_is_client_error(self, tmp_path, monkeypatch): @@ -578,9 +572,7 @@ def _blocked(url, **kwargs): monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) request = VideoGenerationRequest(prompt="x", image_reference="http://10.0.0.1/a.png") with pytest.raises(ValueError, match="reference URL could not be fetched"): - parse_visual_gen_params( - request, "vid-ssrf", generator, media_storage_path=str(tmp_path) - ) + _parse_and_prepare(request, generator, "vid-ssrf", str(tmp_path)) assert list(tmp_path.iterdir()) == [] def test_missing_file_uri_is_client_error(self, tmp_path): @@ -589,7 +581,7 @@ def test_missing_file_uri_is_client_error(self, tmp_path): missing = (tmp_path / "does_not_exist.png").as_uri() request = VideoGenerationRequest(prompt="x", image_reference=missing) with pytest.raises(ValueError, match="reference file could not be read"): - parse_visual_gen_params(request, "vid-nf", generator, media_storage_path=str(tmp_path)) + _parse_and_prepare(request, generator, "vid-nf", str(tmp_path)) assert list(tmp_path.iterdir()) == [] @@ -704,7 +696,7 @@ def test_heif_reference_rejected_with_actionable_message(self): prompt="x", image_reference=base64.b64encode(heic).decode() ) with pytest.raises(ValueError, match="HEIF/AVIF"): - parse_visual_gen_params(request, "vid-heic", generator, media_storage_path=None) + _parse_and_prepare(request, generator, "vid-heic", None) def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): """The boundary routes on signature and never decodes. @@ -726,9 +718,7 @@ def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): request = VideoGenerationRequest( prompt="x", image_reference=base64.b64encode(truncated).decode() ) - params = parse_visual_gen_params( - request, "vid-12", generator, media_storage_path=str(tmp_path) - ) + params = _parse_and_prepare(request, generator, "vid-12", str(tmp_path)) assert Path(params.image_reference[0].content).read_bytes() == truncated From f1e6913e7af91526b4b0cdc763cbc33a953843c5 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 17 Aug 2026 22:55:38 -0700 Subject: [PATCH 19/61] [TRTLLM-15277][fix] Infer a reference slot's single required role; update stale Flux2 reference tests MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit validate_visual_gen_params infers a role-less reference to a multi-role slot's unique required role (min >= 1) when exactly one exists — e.g. i2v's first_frame — matching pipeline_wan_i2v's own default (by_role.get("first_frame") or by_role.get(None)). This lets a single multipart image-to-video upload validate without an explicit role, which the multipart transport cannot carry, while a genuinely ambiguous slot (two required roles) still requires one. MediaReferenceItem docstrings updated to match. The Flux2 unit tests still supplied the pre-Scheme-C params.image field; they now build params.image_reference (a list of refs exposing .content), matching the pipeline's reference API. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol.py | 10 ++--- tensorrt_llm/visual_gen/params.py | 18 ++++++--- .../test_flux2_image_conditioning.py | 2 +- .../_torch/visual_gen/test_flux_infer.py | 2 +- .../visual_gen/test_visual_gen_params.py | 39 +++++++++++++++++++ 5 files changed, 59 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index af05bc168a93..f334b3e147e9 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -2017,9 +2017,9 @@ class MediaReferenceItem(OpenAIBaseModel): ``content`` carries base64-encoded bytes, optionally as a ``data:`` URI. The request field it sits in (``image_reference`` / ``video_reference`` / - ``audio_reference``) fixes the modality. ``role`` is required only for models - that accept more than one role for that modality (e.g. image first/last - frame); single-role models infer it. + ``audio_reference``) fixes the modality. ``role`` is required only when it is + ambiguous — a model with more than one required role for that modality; a + single role, or a single required role (e.g. i2v first_frame), is inferred. """ content: str = Field( @@ -2027,8 +2027,8 @@ class MediaReferenceItem(OpenAIBaseModel): ) role: Optional[Role] = Field( default=None, - description="Reference role. Required only when the model accepts " - "multiple roles for the modality.", + description="Reference role. Required only when the model has more than " + "one required role for the modality; otherwise inferred.", ) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index a2115d659029..82301d1d3106 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -284,8 +284,9 @@ def validate_visual_gen_params( # --- reference role / arity checks (duck-typed RefSlotSpec) --- # ``ref_slot_specs`` maps a reference field name to a spec exposing # ``.roles`` (a list of role specs with ``.role`` / ``.min`` / ``.max``). - # role is required only when a modality declares more than one role; - # otherwise the single declared role is inferred. Reference fields are + # role must be explicit only when the assignment is ambiguous (a multi-role + # slot with more than one required role); a single-role slot or a single + # required role is inferred. Reference fields are # already normalized to ``list[*Ref]`` by the field validators. An empty # (but non-None) mapping means the pipeline declares no slots, so any # reference the client sent is rejected; only ``None`` skips validation. @@ -301,18 +302,25 @@ def validate_visual_gen_params( continue role_specs = list(spec.roles) allowed = {rs.role for rs in role_specs} - role_required = len(role_specs) > 1 + # A role-less ref is inferred when unambiguous: a single-role slot, + # or a multi-role slot with exactly one required role (min >= 1) — + # e.g. i2v's first_frame — matching the pipeline's own default. Only + # a genuinely ambiguous slot (multiple required roles) demands one. + required_roles = [rs.role for rs in role_specs if rs.min >= 1] counts: Dict[str, int] = {} for r in refs: role = getattr(r, "role", None) if role is None: - if role_required: + if len(role_specs) == 1: + role = role_specs[0].role + elif len(required_roles) == 1: + role = required_roles[0] + else: messages.append( f"{field}: 'role' is required for this model " f"(one of {sorted(allowed)})." ) continue - role = role_specs[0].role if role not in allowed: messages.append( f"{field}: role '{role}' not supported (allowed: {sorted(allowed)})." diff --git a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py index ac4a5ee84b14..729c86292dd3 100644 --- a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py +++ b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py @@ -71,7 +71,7 @@ def test_prepare_request_resolves_dimensions_and_infer_reuses_images() -> None: req = SimpleNamespace( prompt=["edit this image"], params=SimpleNamespace( - image=_png_bytes(), + image_reference=[SimpleNamespace(content=_png_bytes())], height=None, width=None, num_inference_steps=1, diff --git a/tests/unittest/_torch/visual_gen/test_flux_infer.py b/tests/unittest/_torch/visual_gen/test_flux_infer.py index a50435a095ea..1f4f4956cd51 100644 --- a/tests/unittest/_torch/visual_gen/test_flux_infer.py +++ b/tests/unittest/_torch/visual_gen/test_flux_infer.py @@ -30,7 +30,7 @@ def test_infer_forwards_num_images_per_prompt( seed=42, max_sequence_length=512, num_images_per_prompt=2, - image=[b"reference"], + image_reference=[SimpleNamespace(content=b"reference")], ), ) 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 9ecd254c3a64..d105e391608a 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -984,6 +984,45 @@ def run(params, spec): with pytest.raises(ValueError, match=r"video_reference.*not accepted"): run(VisualGenParams(video_reference="v.mp4"), optional) + def test_multi_role_slot_infers_single_required_role(self): + """A role-less ref against a multi-role slot is inferred when only one + role is required (e.g. i2v first_frame required, last_frame optional), + matching the pipeline's own default; a genuinely ambiguous slot (two + required roles) still demands an explicit role.""" + from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec + from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + + def run(params, spec): + validate_visual_gen_params( + params, declared_defaults=None, extra_param_specs={}, ref_slot_specs=spec + ) + + # i2v shape: first_frame required, last_frame optional. A single + # role-less upload fills first_frame -> allowed (no explicit role). + i2v = { + "image_reference": RefSlotSpec( + modality="image", + roles=[ + RoleSpec(role="first_frame", min=1, max=1), + RoleSpec(role="last_frame", min=0, max=1), + ], + ) + } + run(VisualGenParams(image_reference="a.png"), i2v) + + # Two required roles -> ambiguous, role stays mandatory. + ambiguous = { + "image_reference": RefSlotSpec( + modality="image", + roles=[ + RoleSpec(role="first_frame", min=1, max=1), + RoleSpec(role="last_frame", min=1, max=1), + ], + ) + } + with pytest.raises(ValueError, match="'role' is required"): + run(VisualGenParams(image_reference="a.png"), ambiguous) + def test_empty_ref_slot_specs_rejects_references(self): """An empty (non-None) ref_slot_specs means the pipeline declares no slots, so a reference is rejected; only ``None`` skips validation.""" From a60ae6c06fb8e4929f67fa1d5a0f4970cfe29b11 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 18 Aug 2026 02:56:21 -0700 Subject: [PATCH 20/61] [TRTLLM-15277][feat] BREAKING: require an explicit format on media references MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit A media reference now declares the wire form of its content instead of having it guessed. MediaRef and MediaReferenceItem gain a required 'format' field — path, url, base64 or bytes — and the bare path/bytes shorthand is gone, so image_reference="start.png" becomes MediaRef(content="start.png", format="path"). Callers no longer need the file:// or data: prefixes to disambiguate; both are still accepted, but format is what decides. This removes two stacked guesses in the resolver. A bare string was decoded as base64 and, failing that, read as a local file, so a malformed base64 string could silently reach the filesystem; and _is_local_path decided passthrough-vs-materialize by os.path.exists, so a mistyped path silently reclassified as base64. Both are deleted in favour of one explicit four-way dispatch, and a declared path that does not exist is now a clean client error instead of a reclassification. Closing the accidental filesystem-read vector also makes the intentional one auditable, and easy to gate later. prepare_reference_slots rewrites 'format' alongside 'content' when it materializes a reference: the mutated params object is what gets broadcast to the workers, so a stale format would hand a worker a filesystem path while claiming base64. Multipart uploads are unchanged for clients. A file part carries its form in the transport, so serve supplies format="bytes" itself and every existing curl, SDK call, example and test keeps working untouched; a reference sent as a multipart text part is now parsed as the JSON object form. format="bytes" over JSON is rejected with a message pointing at multipart, since JSON cannot carry raw bytes. The deprecated input_reference, having no wrapper object to hold a format, gains the sibling input_reference_format, required when it is a string and implied for an upload. MediaReferenceItem is promoted to a gated api-stability model. It was previously invisible to the gate — the nested type string is documentation-only — so a required field on it would have landed with green CI and an empty api_stability diff, which also silences the api-breaking label workflow and the api-review-committee gate. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 33 ++- examples/visual_gen/models/cosmos3/cosmos3.py | 8 +- examples/visual_gen/models/flux2.py | 6 +- examples/visual_gen/models/qwen_image_edit.py | 6 +- .../visual_gen/models/qwen_image_layered.py | 4 +- examples/visual_gen/models/wan_i2v.py | 2 +- examples/visual_gen/serve/README.md | 8 +- tensorrt_llm/serve/openai_protocol.py | 99 +++++-- tensorrt_llm/serve/openai_video_routes.py | 16 + tensorrt_llm/serve/visual_gen_utils.py | 88 ++++-- tensorrt_llm/visual_gen/media_refs.py | 117 ++++---- tensorrt_llm/visual_gen/params.py | 68 +++-- .../visual_gen/test_visual_gen_params.py | 140 +++++++-- .../visual_gen/test_visual_gen_utils.py | 275 +++++++++++++++--- .../references/trtllm_serve_api.yaml | 32 +- 15 files changed, 673 insertions(+), 229 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index bbdaba35aa61..e18b34f00602 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -121,35 +121,44 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue ### Reference Inputs -Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a path, raw bytes, a single reference, or a list. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized (base64, `data:` URI, `http(s)` URL, local file path, or uploaded file) to a local path before reaching the worker. A local path may be given bare or as a `file://` URI; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). +Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a single reference or a list. A reference always declares the wire form of its content — `MediaRef(content=..., format=...)` in Python, `{"content": ..., "format": ...}` in JSON. `format` is **required** and nothing is guessed: a bare string or bare bytes is rejected, so a mistyped path can never be silently read as base64. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized to a local path before reaching the worker; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). + +| `format` | Content | Notes | +|---|---|---| +| `path` | A local file readable by the process running generation | Bare path or `file://` URI. The file must exist, and passes through in place — it is neither copied nor deleted. | +| `url` | An `http(s)` URL | Fetched through the SSRF-guarded loader, then materialized to a local path. | +| `base64` | Base64 text | A `data:` URI is also accepted. | +| `bytes` | Raw `bytes` | Python API only. Rejected over JSON (HTTP 422) — send `base64` or upload the file via multipart. | + +The `file://` and `data:` prefixes are still accepted, but are no longer needed to disambiguate: `format` already states which form the content is in. A multipart file upload carries no `format` — the server implies `bytes` from the transport. A reference item's `format` is distinct from the request's top-level `format` field, which selects the *output* encoding (`mp4`, `png`, …). Most models take a single reference whose role is unambiguous, so no `role` is specified: ```python -from tensorrt_llm import VisualGen +from tensorrt_llm import VisualGen, MediaRef # The image conditions the generated video's first frame. vg = VisualGen(model="Wan-AI/Wan2.2-TI2V-5B-Diffusers") params = vg.default_params -params.image_reference = "start.png" +params.image_reference = MediaRef(content="start.png", format="path") output = vg.generate(inputs="the scene comes alive with gentle motion", params=params) # Cosmos conditions generation on a reference video. vg = VisualGen(model="nvidia/Cosmos3-Super") params = vg.default_params -params.video_reference = "clip.mp4" +params.video_reference = MediaRef(content="clip.mp4", format="path") ``` -The equivalent serve request uploads the file, or sends a base64 string, `data:` URI, `http(s)` URL, or local file path as the field value in a JSON body: +The equivalent serve request uploads the file, or sends a `{content, format}` object as the field value in a JSON body: ```bash # multipart file upload (raw bytes, no base64) curl http://localhost:8000/v1/videos -F "prompt=the scene comes alive" -F "image_reference=@start.png" curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_reference=@clip.mp4" -# JSON body: the reference string may be base64, a data: URI, an http(s) URL, or a local file path +# JSON body: the content plus the format it is in (path, url, or base64) curl http://localhost:8000/v1/videos -H 'content-type: application/json' \ - -d '{"prompt": "the scene comes alive", "image_reference": "https://example.com/start.png"}' + -d '{"prompt": "the scene comes alive", "image_reference": {"content": "https://example.com/start.png", "format": "url"}}' ``` When a model accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame — the `role` is required to disambiguate: @@ -160,8 +169,8 @@ from tensorrt_llm import VisualGen, MediaRef vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") params = vg.default_params params.image_reference = [ - MediaRef(content="start.png", role="first_frame"), - MediaRef(content="end.png", role="last_frame"), # optional + MediaRef(content="start.png", format="path", role="first_frame"), + MediaRef(content="end.png", format="path", role="last_frame"), # optional ] ``` @@ -171,15 +180,15 @@ A JSON serve request carries the role and lists; a multipart upload is limited t curl http://localhost:8000/v1/videos -H 'content-type: application/json' -d '{ "prompt": "the subject comes alive", "image_reference": [ - {"content": "", "role": "first_frame"}, - {"content": "", "role": "last_frame"} + {"content": "", "format": "base64", "role": "first_frame"}, + {"content": "", "format": "base64", "role": "last_frame"} ] }' ``` FLUX.2 and Qwen-Image-Edit accept multiple reference images as a list on the same `image_reference` field through the Python API. -A single `input_reference` field (deprecated) is still accepted on the serve video endpoints for backward compatibility; it is routed by content signature to image-to-video or video-to-video, and is ignored when a typed `image_reference` / `video_reference` is also provided. Prefer the typed fields. +A single `input_reference` field (deprecated) is still accepted on the serve video endpoints for backward compatibility; it is routed by content signature to image-to-video or video-to-video, and is ignored when a typed `image_reference` / `video_reference` is also provided. Being a bare value, it declares its wire form through the sibling `input_reference_format` field (`path` / `url` / `base64`), which is required when `input_reference` is a string and implied for a multipart upload. Prefer the typed fields. ## Optimizations diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 4f35e67c1ca8..402239944b23 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -27,7 +27,7 @@ from pathlib import Path from typing import Any, Dict, Optional -from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS _SCRIPT_DIR = Path(__file__).resolve().parent @@ -336,7 +336,7 @@ def main(): "--image_path", type=str, default=None, - help="Optional conditioning image path or URL for I2V/TI2V", + help="Optional conditioning image path for I2V/TI2V", ) parser.add_argument( "--output_path", @@ -477,7 +477,7 @@ def main(): # Query per-model defaults (resolution, steps, guidance, seed, etc.). params = visual_gen.default_params if image_path is not None: - params.image_reference = image_path + params.image_reference = [MediaRef(content=image_path, format="path")] negative_prompt = resolve_negative_prompt( negative_prompt=args.negative_prompt, @@ -514,7 +514,7 @@ def main(): with open(args.action_json, encoding="utf-8") as f: params.extra_params["action"] = json.load(f) if args.video_path is not None: - params.video_reference = args.video_path + params.video_reference = [MediaRef(content=args.video_path, format="path")] if args.extra_params: # Merged last: explicit JSON wins over flag-derived values. params.extra_params.update(args.extra_params) diff --git a/examples/visual_gen/models/flux2.py b/examples/visual_gen/models/flux2.py index ad983d72b321..adf50af3bf8f 100644 --- a/examples/visual_gen/models/flux2.py +++ b/examples/visual_gen/models/flux2.py @@ -25,7 +25,7 @@ import argparse from pathlib import Path -from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs def _output_paths(output_path: str, num_images: int) -> str | list[str]: @@ -118,7 +118,9 @@ def main(): # Start from per-model defaults and override only user-provided request fields. params = visual_gen.default_params params.num_images_per_prompt = args.num_images_per_prompt - params.image_reference = args.image + params.image_reference = ( + [MediaRef(content=path, format="path") for path in args.image] if args.image else None + ) if args.image: # Let FLUX.2 derive omitted dimensions from the first processed reference. params.height = args.height diff --git a/examples/visual_gen/models/qwen_image_edit.py b/examples/visual_gen/models/qwen_image_edit.py index 4d1740f340da..d48904f06b6d 100644 --- a/examples/visual_gen/models/qwen_image_edit.py +++ b/examples/visual_gen/models/qwen_image_edit.py @@ -23,7 +23,7 @@ import argparse -from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs def parse_args() -> argparse.Namespace: @@ -44,7 +44,7 @@ def parse_args() -> argparse.Namespace: "--image", nargs="+", required=True, - help="One or more input image paths or URLs.", + help="One or more input image paths.", ) parser.add_argument( "--prompt", @@ -64,7 +64,7 @@ def main() -> None: extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None visual_gen = VisualGen(model=args.model, args=extra_args) params = visual_gen.default_params - params.image_reference = args.image + params.image_reference = [MediaRef(content=path, format="path") for path in args.image] output = visual_gen.generate(inputs=args.prompt, params=params) saved = output.save(args.output_path) print(f"Saved edited image to {saved}") diff --git a/examples/visual_gen/models/qwen_image_layered.py b/examples/visual_gen/models/qwen_image_layered.py index f5fc47d2be82..46115c7c15e1 100644 --- a/examples/visual_gen/models/qwen_image_layered.py +++ b/examples/visual_gen/models/qwen_image_layered.py @@ -23,7 +23,7 @@ import argparse from pathlib import Path -from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs def parse_args() -> argparse.Namespace: @@ -63,7 +63,7 @@ def main() -> None: visual_gen = VisualGen(model=args.model, args=extra_args) params = visual_gen.default_params - params.image_reference = args.image + params.image_reference = [MediaRef(content=args.image, format="path")] output = visual_gen.generate(inputs=args.prompt, params=params) if output.image is not None and output.image.shape[0] > 1: diff --git a/examples/visual_gen/models/wan_i2v.py b/examples/visual_gen/models/wan_i2v.py index b650fbdcd320..2438680157ca 100644 --- a/examples/visual_gen/models/wan_i2v.py +++ b/examples/visual_gen/models/wan_i2v.py @@ -66,7 +66,7 @@ def main(): # first-frame reference. Wan I2V also accepts a ``last_frame`` role, so the # role must be given to disambiguate. params = visual_gen.default_params - params.image_reference = [MediaRef(content=args.image, role="first_frame")] + params.image_reference = [MediaRef(content=args.image, format="path", role="first_frame")] output = visual_gen.generate( inputs="A cat presses the piano keys with its paws, soft notes filling the quiet room.", diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 604942e4d86f..5e605f800568 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,7 +286,13 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `image_reference`: Reference image(s) for I2V/TI2V. `video_reference`: reference video(s) for V2V. `audio_reference`: reference audio(s). Each accepts a base64-encoded string (optionally a `data:` URI), a `{image|video|audio, role}` object, or a list of them in JSON; or a single uploaded file in multipart form-data. +- `image_reference`: Reference image(s) for I2V/TI2V. `video_reference`: reference video(s) for V2V. `audio_reference`: reference audio(s). In JSON each accepts a `{content, format, role}` object or a list of them; `format` is required and declares how to read `content` — `"path"` (a file readable by the server; a `file://` URI is also accepted), `"url"` (`http(s)`), or `"base64"` (a `data:` URI is also accepted). Nothing is guessed, so a bare string is rejected, and `"bytes"` is rejected over JSON — upload the file instead. A multipart file upload needs no `format`: the transport implies it. + + ```json + {"image_reference": {"content": "iVBORw0KGgoAAAANSUhEUg...", "format": "base64"}} + ``` + + - `format` here is the *input* wire form; the top-level `format` selects the *output* encoding. - **Supported formats**: PNG and JPEG images; MP4 and AVI video, with H.264 the tested codec and others best-effort. HEIF/AVIF are not supported. - `extra_params`: model-specific overflow (see below) - `response_format`: `"file"` (default; `FileResponse` byte download) or `"path"` (server-side output path JSON, for co-located clients) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index f334b3e147e9..b60920c5283d 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -61,7 +61,7 @@ from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy -from tensorrt_llm.visual_gen.params import Role +from tensorrt_llm.visual_gen.params import ContentFormat, Role _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 @@ -2015,22 +2015,39 @@ class ImageGenerationResponse(OpenAIBaseModel): class MediaReferenceItem(OpenAIBaseModel): """One media reference (image / video / audio) for conditioning (mirrors ``MediaRef``). - ``content`` carries base64-encoded bytes, optionally as a ``data:`` URI. The - request field it sits in (``image_reference`` / ``video_reference`` / - ``audio_reference``) fixes the modality. ``role`` is required only when it is - ambiguous — a model with more than one required role for that modality; a - single role, or a single required role (e.g. i2v first_frame), is inferred. + ``format`` declares how to read ``content``; it is required, so no wire form + is ever guessed. The request field it sits in (``image_reference`` / + ``video_reference`` / ``audio_reference``) fixes the modality. ``role`` is + required only when it is ambiguous — a model with more than one required role + for that modality; a single role, or a single required role (e.g. i2v + first_frame), is inferred. """ content: str = Field( - description="Base64-encoded media bytes, optionally as a ``data:`` URI." + description="The reference payload, in the form declared by ``format``." ) + format: ContentFormat = Field(description=( + "Wire form of ``content``: ``path`` (a file readable by the server; a " + "``file://`` URI is also accepted), ``url`` (``http(s)``, fetched " + "through the SSRF-guarded loader), or ``base64`` (a ``data:`` URI is " + "also accepted). ``bytes`` cannot be carried in JSON — upload the file " + "as multipart/form-data instead. Distinct from the top-level " + "``format``, which selects the *output* encoding.")) role: Optional[Role] = Field( default=None, description="Reference role. Required only when the model has more than " "one required role for the modality; otherwise inferred.", ) + @field_validator("format") + @classmethod + def _reject_bytes_over_json(cls, v: str) -> str: + if v == "bytes": + raise ValueError( + "format='bytes' cannot be carried in JSON; upload the file as " + "multipart/form-data, or send format='base64'.") + return v + class VideoGenerationRequest(OpenAIBaseModel): """Video generation request (extended API). @@ -2056,38 +2073,34 @@ class VideoGenerationRequest(OpenAIBaseModel): ge=0, description="Random seed for reproducibility.") image_reference: Optional[Union[ - str, UploadFile, MediaReferenceItem, - List[Union[str, MediaReferenceItem]]]] = Field( + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( default=None, description= ("Image reference(s) conditioning generation (e.g. image-to-video " - "first frame). A JSON string is base64 bytes (raw or ``data:`` " - "URI), an ``http(s)`` URL, or a local file path; or send a " - "``{content, role}`` object or a list of them; multipart uploads a " - "single image file. PNG or JPEG only — HEIF/AVIF are not supported." - ), + "first frame). Send a ``{content, format, role}`` object or a list " + "of them, where ``format`` is ``path`` / ``url`` / ``base64``; or " + "upload a single image file via multipart, whose form is implied. " + "PNG or JPEG only — HEIF/AVIF are not supported."), ) video_reference: Optional[Union[ - str, UploadFile, MediaReferenceItem, - List[Union[str, MediaReferenceItem]]]] = Field( + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( default=None, description= - ("Video reference(s) conditioning generation (video-to-video). A " - "JSON string is base64 bytes (raw or ``data:`` URI), an ``http(s)`` " - "URL, or a local file path; or send a ``{content}`` object or a " - "list of them; multipart uploads a single video file. MP4 or AVI, " - "with H.264 the tested codec and others best-effort."), + ("Video reference(s) conditioning generation (video-to-video). Send " + "a ``{content, format}`` object or a list of them, where ``format`` " + "is ``path`` / ``url`` / ``base64``; or upload a single video file " + "via multipart, whose form is implied. MP4 or AVI, with H.264 the " + "tested codec and others best-effort."), ) audio_reference: Optional[Union[ - str, UploadFile, MediaReferenceItem, - List[Union[str, MediaReferenceItem]]]] = Field( + UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] = Field( default=None, - description=( - "Audio reference(s) conditioning generation. A JSON string is " - "base64 bytes (raw or ``data:`` URI), an ``http(s)`` URL, or a " - "local file path; or send a ``{content}`` object or a list of " - "them; multipart uploads a single audio file. Accepted only by " - "models that declare an audio reference slot."), + description= + ("Audio reference(s) conditioning generation. Send a " + "``{content, format}`` object or a list of them, where ``format`` " + "is ``path`` / ``url`` / ``base64``; or upload a single audio " + "file via multipart, whose form is implied. Accepted only by " + "models that declare an audio reference slot."), ) input_reference: Optional[Union[str, UploadFile]] = Field( default=None, @@ -2097,7 +2110,14 @@ class VideoGenerationRequest(OpenAIBaseModel): "compatibility; prefer the typed ``image_reference`` / " "``video_reference`` fields, which take precedence — this field is " "ignored whenever a typed ``image_reference`` or ``video_reference`` " - "is provided."), + "is provided. A string form requires ``input_reference_format``."), + ) + input_reference_format: Optional[ContentFormat] = Field( + default=None, + description=( + "Deprecated, alongside ``input_reference``: the wire form of that " + "field's value (``path`` / ``url`` / ``base64``). Required when " + "``input_reference`` is a string; implied for a multipart upload."), ) # Resolution @@ -2171,6 +2191,25 @@ def _reject_removed_response_format(cls, value): raise ValueError(removed[value]) return value + @model_validator(mode="after") + def _check_input_reference_format(self): + """Require the deprecated ``input_reference``'s wire form when it is a string. + + A multipart upload carries its own form, so the sibling is only needed + for the string spelling. + """ + if isinstance(self.input_reference, str): + if self.input_reference_format is None: + raise ValueError( + "'input_reference_format' is required when 'input_reference' is a " + "string; send 'path', 'url' or 'base64' (or upload the file via " + "multipart/form-data)") + if self.input_reference_format == "bytes": + raise ValueError( + "input_reference_format='bytes' cannot be carried in JSON; upload " + "the file as multipart/form-data, or send 'base64'") + return self + class VideoJob(OpenAIBaseModel): """Metadata for an asynchronous video generation job. diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 91c80b46b74d..d2b6b2535584 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -57,6 +57,9 @@ def _video_content_type(suffix: str) -> str: # /v1/videos/{id} routes try when the stored output_path is missing. _KNOWN_VIDEO_OUTPUT_SUFFIXES = (".mp4", ".avi", ".safetensors", ".pt") +# Reference fields whose multipart text form is a JSON object, not a scalar. +_REFERENCE_FIELDS = ("image_reference", "video_reference", "audio_reference") + def _resolve_tensor_only_format(fmt, extra_params, extra_param_specs): """Resolve ``format`` for a request whose result an encoder cannot carry. @@ -327,6 +330,19 @@ async def _parse_video_generation_request( f"'extra_params' must be a JSON object string; {exc}" ) from exc continue + if key in _REFERENCE_FIELDS: + # A reference sent as a text part is the object form; a file + # part was already routed above and carries its own format. + if value == "": + continue + try: + data[key] = json.loads(value) + except json.JSONDecodeError as exc: + raise ValueError( + f"'{key}' must be an uploaded file or a JSON object string " + f'like {{"content": ..., "format": "url"}}; {exc}' + ) from exc + continue if value == "": continue data[key] = value diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 9ce3f29c0f66..7ca39ffdb6f1 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -18,7 +18,7 @@ ImageGenerationRequest, VideoGenerationRequest, ) -from tensorrt_llm.visual_gen.media_refs import _resolve_reference_string +from tensorrt_llm.visual_gen.media_refs import _resolve_reference if TYPE_CHECKING: from fastapi import UploadFile @@ -110,34 +110,30 @@ def _merge_extra_params( params.extra_params = None -def _reference_payload_and_role(ref: Any) -> tuple[Any, Optional[str]]: - """Extract ``(content, role)`` from one raw HTTP reference for transport. +def _reference_transport(ref: Any) -> tuple[Any, str, Optional[str]]: + """Extract ``(content, format, role)`` from one raw HTTP reference. - ``ref`` is a string (base64/``data:`` URI, ``http(s)`` URL, or a local file - path), a multipart ``UploadFile`` (has ``.file``), or a ``MediaReferenceItem`` - exposing ``content`` and an optional ``role``. An upload is read to ``bytes`` - here — the only decode the boundary owns; strings pass through untouched for - the engine to resolve and materialize. + ``ref`` is a multipart ``UploadFile`` (has ``.file``) or a + ``MediaReferenceItem`` exposing ``content`` / ``format`` / ``role``. An + upload is read to ``bytes`` here — the only decode the boundary owns — and + its format is implied by the transport rather than declared by the client; + everything else passes through for the engine to resolve. """ - role = getattr(ref, "role", None) - if isinstance(ref, str): - return ref, role if hasattr(ref, "file"): # multipart UploadFile - return ref.file.read(), role - data = getattr(ref, "content", None) - if not isinstance(data, str): + return ref.file.read(), "bytes", getattr(ref, "role", None) + content = getattr(ref, "content", None) + if not isinstance(content, str): raise ValueError("reference item must carry a 'content' string.") - return data, role + return content, ref.format, getattr(ref, "role", None) def _build_reference_list(value: Any) -> Optional[list]: """Normalize one HTTP reference field into a list of ``MediaRef`` objects. - ``value`` is None, a base64/data-URI/URL/path string, a multipart - ``UploadFile``, a ``MediaReferenceItem``, or a list of any of those. Each - entry becomes a ``MediaRef`` carrying its transport content — ``bytes`` for - an upload, the string otherwise — plus its ``role``. Resolution and - materialization happen later at the engine choke point. + ``value`` is None, a multipart ``UploadFile``, a ``MediaReferenceItem``, or + a list of those. Each entry becomes a ``MediaRef`` carrying its transport + content plus the declared (or, for an upload, implied) wire format. + Resolution and materialization happen later at the engine choke point. """ if value is None: return None @@ -148,8 +144,33 @@ def _build_reference_list(value: Any) -> Optional[list]: raw_items = value if isinstance(value, list) else [value] refs = [] for item in raw_items: - content, role = _reference_payload_and_role(item) - refs.append(MediaRef(content=content, role=role)) + content, content_format, role = _reference_transport(item) + refs.append(MediaRef(content=content, format=content_format, role=role)) + return refs + + +def _build_image_edit_reference_list(value: Any) -> Optional[list]: + """Build references from an image-edit request's OpenAI-shaped ``image``. + + That field follows OpenAI's schema, which has no place to declare a wire + form: an entry is a bare base64 string or a multipart upload, so the format + is implied by the transport instead of read off the item. + """ + if value is None: + return None + from tensorrt_llm.visual_gen.params import MediaRef + + refs = [] + for item in value if isinstance(value, list) else [value]: + if hasattr(item, "file"): # multipart UploadFile + refs.append(MediaRef(content=item.file.read(), format="bytes")) + elif isinstance(item, str): + refs.append(MediaRef(content=item, format="base64")) + else: + raise ValueError( + "image edit inputs must be base64-encoded images or uploaded files, " + f"got {type(item).__name__}." + ) return refs @@ -369,14 +390,16 @@ def cleanup_materialized_conditioning_inputs(value: Any) -> None: def _apply_deprecated_input_reference( input_reference: str | UploadFile | None, params: VisualGenParams, + input_reference_format: Optional[str] = None, ) -> None: """Back-compat for the deprecated single ``input_reference``. Sniff-routes the payload to ``image_reference`` (image) or ``video_reference`` (video), preserving the pre-typed-fields behavior. Ignored when a typed image/video reference is already set — the typed fields take precedence. - Routing needs the bytes, so a string payload is resolved here; the engine - materializes the resulting ``MediaRef`` like any other reference. + Routing needs the bytes, so the payload is resolved here using the wire form + from the sibling ``input_reference_format`` (implied for an upload); the + resolved bytes are then handed to the engine like any other reference. """ if input_reference is None: return @@ -385,14 +408,15 @@ def _apply_deprecated_input_reference( return from tensorrt_llm.visual_gen.params import MediaRef - payload, _ = _reference_payload_and_role(input_reference) - if isinstance(payload, str): - payload = _resolve_reference_string(payload) + if hasattr(input_reference, "file"): # multipart upload — form implied + payload = input_reference.file.read() + else: + payload = _resolve_reference(input_reference, input_reference_format) kind = sniff_media_kind(payload) if kind == "image": - params.image_reference = [MediaRef(content=payload)] + params.image_reference = [MediaRef(content=payload, format="bytes")] elif kind == "video": - params.video_reference = [MediaRef(content=payload)] + params.video_reference = [MediaRef(content=payload, format="bytes")] else: raise ValueError( "input_reference is not a recognized media container; supported " @@ -452,7 +476,7 @@ def parse_visual_gen_params( if request.n is not None: params.num_images_per_prompt = request.n _validate_image_edit_request_limits(request, generator) - params.image_reference = _build_reference_list(request.image) + params.image_reference = _build_image_edit_reference_list(request.image) elif isinstance(request, VideoGenerationRequest): if request.frame_rate is not None: @@ -491,7 +515,9 @@ def parse_visual_gen_params( audio_refs = _build_reference_list(request.audio_reference) if audio_refs: params.audio_reference = audio_refs - _apply_deprecated_input_reference(request.input_reference, params) + _apply_deprecated_input_reference( + request.input_reference, params, request.input_reference_format + ) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _decode_inline_media(request.extra_params, generator.extra_param_specs) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index efa3b5068dfa..edb7a0ae58cd 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -26,7 +26,6 @@ import os from pathlib import Path from typing import Any, Optional -from urllib.parse import urlparse from tensorrt_llm.inputs.media_io import ( _normalize_file_uri, @@ -49,6 +48,10 @@ def _read_reference_payload(reference: str) -> bytes: comma = data.find(",") if comma == -1: raise ValueError("reference data: URI is malformed (missing comma).") + # Match the LLM loader: only base64 payloads are supported, and saying so + # beats letting a percent-encoded body fail as "not valid base64". + if "base64" not in data[:comma].split(";")[1:]: + raise ValueError("only base64 data: URIs are supported for references.") data = data[comma + 1 :] try: return base64.b64decode(data, validate=True) @@ -57,41 +60,43 @@ def _read_reference_payload(reference: str) -> bytes: raise ValueError("reference is not valid base64 data.") from exc -def _resolve_reference_string(reference: str) -> bytes: - """Resolve one reference string to raw bytes, dispatching on URL scheme. +def _local_path(reference: str) -> Path: + """Normalize a ``path`` reference (bare or ``file://``) to a ``Path``.""" + return Path(_normalize_file_uri(reference)) - Mirrors the LLM multimodal loader so serve references accept the same forms: - ``http(s)`` fetches through the SSRF-guarded loader (private-address block, - redirect re-validation, timeout, size cap); ``file://`` and bare local paths - read from disk; ``data:`` and base64 strings decode inline. A bare string is - decoded as base64 first and, failing that, read as a local file path. - Fetch/read failures become ``ValueError`` so a bad URL or path is a client - 400, not a server 500. + +def _resolve_reference(content: Any, content_format: str) -> bytes: + """Resolve one reference to raw bytes using its declared wire form. + + Dispatch is on the caller-declared ``format``, never on the shape of the + value: a bare string is otherwise ambiguous between a local path and + base64, and guessing lets a mistyped path become base64 (or a malformed + base64 become a filesystem read). Fetch/read/decode failures become + ``ValueError`` so a bad reference is a client 400, not a server 500. """ - scheme = urlparse(reference).scheme - if scheme in ("http", "https"): + if content_format == "bytes": + if not isinstance(content, bytes): + raise ValueError( + f"format='bytes' requires bytes content, got {type(content).__name__}." + ) + return content + if not isinstance(content, str): + raise ValueError( + f"format={content_format!r} requires string content, got {type(content).__name__}." + ) + if content_format == "url": try: - return _safe_request_get(reference).content + return _safe_request_get(content).content except Exception as exc: raise ValueError(f"reference URL could not be fetched: {exc}") from exc - if scheme == "file": + if content_format == "path": try: - return Path(_normalize_file_uri(reference)).read_bytes() + return _local_path(content).read_bytes() except OSError as exc: raise ValueError(f"reference file could not be read: {exc}") from exc - if scheme == "data": - return _read_reference_payload(reference) - # Bare string: base64 first (the established default), else a local file path - # so a plain path works without the file:// scheme. - try: - return _read_reference_payload(reference) - except ValueError: - try: - return Path(reference).read_bytes() - except OSError as exc: - raise ValueError( - f"reference is not valid base64 data, and not a readable local file: {exc}" - ) from exc + if content_format == "base64": + return _read_reference_payload(content) + raise ValueError(f"unsupported reference format: {content_format!r}") def _materialize_reference( @@ -161,56 +166,48 @@ def resolve_media_storage_path() -> Path: return path -def _is_local_path(content: Any) -> bool: - """True if ``content`` is a trusted local path to pass through untouched. - - A ``file://`` URI or bare string naming an *existing* file. A missing path - (either form) returns False so it falls through to the resolve step, whose - read raises ``ValueError`` — a client 400, not a silent passthrough. - Everything else (bytes, ``http(s)`` / ``data:`` URLs, base64) materializes. - """ - if not isinstance(content, str): - return False - scheme = urlparse(content).scheme - if scheme == "file": - return os.path.exists(_normalize_file_uri(content)) - return scheme == "" and os.path.exists(content) - - def prepare_reference_slots( params: Any, *, request_id: str, media_storage_path: Optional[str] ) -> None: """Resolve + materialize each reference to a local path, in place. The single reference choke point, used by the engine (``generate_async``) - so serve and the standalone Python API share one path. A trusted local path - (``file://`` / existing bare path) passes through — not materialized, not - cleaned up (it is the caller's file); a ``file://`` URI is normalized to a - plain path so the pipeline, which opens paths, can read it. Everything else - (bytes / ``http(s)`` / ``data:`` / base64) resolves to bytes and materializes - to ``media_storage_path``; those files are reclaimed by - :func:`cleanup_reference_files` keyed on ``request_id``. Runs before the - coordinator broadcasts the request, so bad-media ``ValueError`` surfaces to - the caller synchronously (serve keeps its immediate 400). If a later slot - fails mid-materialize, the files earlier slots wrote are reclaimed here so a - rejected request leaves nothing on disk. + so serve and the standalone Python API share one path. Dispatch is on each + reference's declared ``format``, never on the shape of its content. A + ``path`` reference is the caller's own file: it passes through — not + materialized, not cleaned up — with a ``file://`` URI normalized to a plain + path so the pipeline, which opens paths, can read it. Every other form + (``url`` / ``base64`` / ``bytes``) resolves to bytes and materializes to + ``media_storage_path``; those files are reclaimed by + :func:`cleanup_reference_files` keyed on ``request_id``. + + ``format`` is rewritten alongside ``content``: the mutated params object is + what gets broadcast to the workers, so a stale format would send e.g. + ``base64`` to a worker holding a filesystem path. + + Runs before the coordinator broadcasts the request, so a bad reference + raises ``ValueError`` synchronously (serve keeps its immediate 400). If a + later slot fails mid-materialize, the files earlier slots wrote are + reclaimed here so a rejected request leaves nothing on disk. """ try: for slot in ("image_reference", "video_reference", "audio_reference"): modality = slot.split("_", 1)[0] for i, ref in enumerate(getattr(params, slot, None) or []): - content = ref.content - if _is_local_path(content): - if urlparse(content).scheme == "file": - ref.content = _normalize_file_uri(content) + if ref.format == "path": + path = _local_path(ref.content) + if not path.exists(): + raise ValueError(f"reference file does not exist: {ref.content}") + ref.content = str(path) continue - data = content if isinstance(content, bytes) else _resolve_reference_string(content) + data = _resolve_reference(ref.content, ref.format) ref.content = _materialize_reference( data, modality=modality, ref_id=f"{request_id}_{modality}_ref_{i}", media_storage_path=media_storage_path, ) + ref.format = "path" except Exception: # The terminal on_finish hook is not wired yet (the request is never # enqueued on failure), so reclaim any files earlier slots wrote here. diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 82301d1d3106..915234d96267 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -21,6 +21,12 @@ Role = Literal["reference", "first_frame", "last_frame"] +# Wire form of a reference's ``content``. Declared explicitly rather than +# sniffed: a bare string is otherwise ambiguous between a local path and +# base64, and guessing lets a mistyped path silently become base64 (or a +# malformed base64 silently become a filesystem read). +ContentFormat = Literal["path", "url", "base64", "bytes"] + @set_api_status("prototype") class MediaRef(StrictBaseModel): @@ -34,23 +40,42 @@ class MediaRef(StrictBaseModel): """ content: Union[str, bytes] = Field( - description="Local path, ``http(s)``/``data:`` URL, or raw bytes." + description="The reference payload, in the form declared by ``format``." + ) + format: ContentFormat = Field( + description=( + "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " + "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " + "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." + ) ) role: Optional[Role] = Field( default=None, description="``reference`` | ``first_frame`` | ``last_frame``." ) -def _normalize_refs(value: Any) -> Optional[list]: - """Coerce a reference field to ``list[MediaRef]`` (or ``None``). +def _reject_bare_refs(value: Any) -> Any: + """Reject the bare path/bytes shorthand with an actionable message. - Accepts a bare path/bytes, a single ``MediaRef``, or a list mixing the two; - a bare path/bytes ``x`` becomes ``MediaRef(content=x)``. + Runs before coercion, so the caller sees what to do instead of a union + mismatch reported against an inner model. A bare string has nowhere to + declare its wire form, and guessing is what ``format`` exists to prevent. """ + for x in value if isinstance(value, list) else [value]: + if isinstance(x, (str, bytes)): + raise ValueError( + "a reference must declare its wire form; a bare " + f"{type(x).__name__} is no longer accepted. Pass " + 'MediaRef(content=..., format="path"|"url"|"base64"|"bytes").' + ) + return value + + +def _normalize_refs(value: Any) -> Optional[list]: + """Coerce a reference field to ``list[MediaRef]`` (or ``None``).""" if value is None: return None - items = value if isinstance(value, list) else [value] - return [x if isinstance(x, MediaRef) else MediaRef(content=x) for x in items] + return value if isinstance(value, list) else [value] @set_api_status("prototype") @@ -111,23 +136,26 @@ class VisualGenParams(StrictBaseModel): # Conditioning inputs negative_prompt: Optional[str] = Field(default=None, description="Negative prompt for CFG.") - # Per-modality reference inputs. A bare path/bytes, a single ``MediaRef``, - # or a list; normalized to ``list[MediaRef]``. The field fixes the modality; - # ``role`` is only meaningful where a model declares more than one role for - # it (e.g. image first_frame / last_frame). - image_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( - Field( - default=None, - description="Reference image(s) for I2V/I2I; normalized to list[MediaRef].", - ) + # Per-modality reference inputs. A single ``MediaRef`` or a list; normalized + # to ``list[MediaRef]``. The field fixes the modality; ``role`` is only + # meaningful where a model declares more than one role for it (e.g. image + # first_frame / last_frame), and each ref declares its own ``format``. + image_reference: Optional[Union[MediaRef, List[MediaRef]]] = Field( + default=None, + description="Reference image(s) for I2V/I2I; normalized to list[MediaRef].", ) - video_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( - Field(default=None, description="Reference video(s) for V2V; normalized to list[MediaRef].") + video_reference: Optional[Union[MediaRef, List[MediaRef]]] = Field( + default=None, description="Reference video(s) for V2V; normalized to list[MediaRef]." ) - audio_reference: Optional[Union[str, bytes, MediaRef, List[Union[str, bytes, MediaRef]]]] = ( - Field(default=None, description="Reference audio(s); normalized to list[MediaRef].") + audio_reference: Optional[Union[MediaRef, List[MediaRef]]] = Field( + default=None, description="Reference audio(s); normalized to list[MediaRef]." ) + @field_validator("image_reference", "video_reference", "audio_reference", mode="before") + @classmethod + def _reject_bare(cls, v): + return v if v is None else _reject_bare_refs(v) + @field_validator("image_reference", "video_reference", "audio_reference", mode="after") @classmethod def _norm_refs(cls, v): 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 d105e391608a..780090b6ac19 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -102,22 +102,31 @@ def test_extra_params_accepted(self): assert params.extra_params["enhance_prompt"] is True def test_image_reference_accepts_str(self): - from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - params = VisualGenParams(image_reference="/path/to/image.png") + params = VisualGenParams( + image_reference=MediaRef(content="/path/to/image.png", format="path") + ) assert params.image_reference[0].content == "/path/to/image.png" + assert params.image_reference[0].format == "path" assert params.image_reference[0].role is None def test_image_reference_accepts_bytes(self): - from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - params = VisualGenParams(image_reference=b"\x89PNG") + params = VisualGenParams(image_reference=MediaRef(content=b"\x89PNG", format="bytes")) assert params.image_reference[0].content == b"\x89PNG" + assert params.image_reference[0].format == "bytes" def test_image_reference_accepts_list(self): - from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - params = VisualGenParams(image_reference=["/path/a.png", b"\x89PNG"]) + params = VisualGenParams( + image_reference=[ + MediaRef(content="/path/a.png", format="path"), + MediaRef(content=b"\x89PNG", format="bytes"), + ] + ) assert len(params.image_reference) == 2 assert params.image_reference[0].content == "/path/a.png" @@ -148,6 +157,66 @@ def test_seed_accepts_int64_range(self): assert VisualGenParams(seed=2**40).seed == 2**40 +# ============================================================================= +# MediaRef — wire form is declared, never guessed +# ============================================================================= + + +class TestMediaRefValidation: + """Every reference declares its ``format``; the bare shorthand is gone.""" + + @pytest.mark.parametrize("bare", ["a.png", b"\x89PNG"]) + def test_bare_reference_rejected_with_actionable_message(self, bare): + """A bare str/bytes has nowhere to declare its wire form, and the + rejection must say what to pass instead.""" + from pydantic import ValidationError + + from tensorrt_llm.visual_gen import VisualGenParams + + with pytest.raises(ValidationError, match="must declare its wire form"): + VisualGenParams(image_reference=bare) + + def test_bare_reference_in_list_rejected(self): + from pydantic import ValidationError + + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + + with pytest.raises(ValidationError, match="must declare its wire form"): + VisualGenParams( + image_reference=[MediaRef(content="a.png", format="path"), "b.png"], + ) + + def test_format_is_required(self): + from pydantic import ValidationError + + from tensorrt_llm.visual_gen import MediaRef + + with pytest.raises(ValidationError, match=r"format\s+Field required"): + MediaRef(content="a.png") + + def test_unknown_format_rejected(self): + from pydantic import ValidationError + + from tensorrt_llm.visual_gen import MediaRef + + with pytest.raises( + ValidationError, match="Input should be 'path', 'url', 'base64' or 'bytes'" + ): + MediaRef(content="a.png", format="filepath") + + def test_single_ref_normalized_to_list_and_list_preserved(self): + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + + single = VisualGenParams(video_reference=MediaRef(content="v.mp4", format="path")) + assert single.video_reference == [MediaRef(content="v.mp4", format="path")] + + refs = [ + MediaRef(content="https://example.com/a.png", format="url"), + MediaRef(content="Zm9v", format="base64"), + ] + assert VisualGenParams(image_reference=refs).image_reference == refs + + # ============================================================================= # ExtraParamSchema # ============================================================================= @@ -395,10 +464,11 @@ def test_user_values_not_overwritten(self): def test_flux2_reference_dimensions_remain_unset_for_pipeline_resolution(self): from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + from tensorrt_llm.visual_gen.params import MediaRef executor = self._make_mock_executor(Flux2Pipeline) executor.pipeline.derive_output_size_from_reference = True - req = self._make_request(image_reference=b"encoded image") + req = self._make_request(image_reference=MediaRef(content=b"encoded image", format="bytes")) self._merge(executor, req) @@ -408,10 +478,15 @@ def test_flux2_reference_dimensions_remain_unset_for_pipeline_resolution(self): def test_flux2_reference_dimensions_preserve_explicit_values(self): from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline + from tensorrt_llm.visual_gen.params import MediaRef executor = self._make_mock_executor(Flux2Pipeline) executor.pipeline.derive_output_size_from_reference = True - req = self._make_request(image_reference=b"encoded image", height=768, width=512) + req = self._make_request( + image_reference=MediaRef(content=b"encoded image", format="bytes"), + height=768, + width=512, + ) self._merge(executor, req) @@ -923,9 +998,12 @@ def test_image_reference_not_checked_without_ref_specs(self): """Without ``ref_slot_specs``, image_reference is not role/arity checked here — the pipeline's infer() consumes it at runtime.""" from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan import WanPipeline + from tensorrt_llm.visual_gen.params import MediaRef executor = self._make_mock_executor(WanPipeline, _wan_mock(num_heads=12)) - req = self._make_request(image_reference="/path/to/img.png") + req = self._make_request( + image_reference=MediaRef(content="/path/to/img.png", format="path") + ) # Should not raise — ``_validate`` here passes no ref_slot_specs. self._merge_and_validate(executor, req) @@ -943,9 +1021,12 @@ def test_image_reference_on_i2v_pipeline_ok(self): from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_i2v import ( WanImageToVideoPipeline, ) + from tensorrt_llm.visual_gen.params import MediaRef executor = self._make_mock_executor(WanImageToVideoPipeline, _wan_mock(num_heads=12)) - req = self._make_request(image_reference="/path/to/img.png") + req = self._make_request( + image_reference=MediaRef(content="/path/to/img.png", format="path") + ) self._merge_and_validate(executor, req) def test_ref_slot_required_vs_optional(self): @@ -953,7 +1034,11 @@ def test_ref_slot_required_vs_optional(self): ``min == 0`` leaves the slot optional; an undeclared absent slot is fine, but an unsolicited one is rejected.""" from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec - from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + from tensorrt_llm.visual_gen.params import ( + MediaRef, + VisualGenParams, + validate_visual_gen_params, + ) required = { "image_reference": RefSlotSpec( @@ -977,12 +1062,12 @@ def run(params, spec): # Optional slot, no image -> allowed (e.g. text-to-video). run(VisualGenParams(), optional) # Required slot with the image present -> allowed. - run(VisualGenParams(image_reference="a.png"), required) + run(VisualGenParams(image_reference=MediaRef(content="a.png", format="path")), required) # Undeclared slot left absent -> no spurious "not accepted". run(VisualGenParams(), optional) # Undeclared slot actually sent -> rejected. with pytest.raises(ValueError, match=r"video_reference.*not accepted"): - run(VisualGenParams(video_reference="v.mp4"), optional) + run(VisualGenParams(video_reference=MediaRef(content="v.mp4", format="path")), optional) def test_multi_role_slot_infers_single_required_role(self): """A role-less ref against a multi-role slot is inferred when only one @@ -990,7 +1075,11 @@ def test_multi_role_slot_infers_single_required_role(self): matching the pipeline's own default; a genuinely ambiguous slot (two required roles) still demands an explicit role.""" from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec - from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + from tensorrt_llm.visual_gen.params import ( + MediaRef, + VisualGenParams, + validate_visual_gen_params, + ) def run(params, spec): validate_visual_gen_params( @@ -1008,7 +1097,7 @@ def run(params, spec): ], ) } - run(VisualGenParams(image_reference="a.png"), i2v) + run(VisualGenParams(image_reference=MediaRef(content="a.png", format="path")), i2v) # Two required roles -> ambiguous, role stays mandatory. ambiguous = { @@ -1021,21 +1110,28 @@ def run(params, spec): ) } with pytest.raises(ValueError, match="'role' is required"): - run(VisualGenParams(image_reference="a.png"), ambiguous) + run( + VisualGenParams(image_reference=MediaRef(content="a.png", format="path")), ambiguous + ) def test_empty_ref_slot_specs_rejects_references(self): """An empty (non-None) ref_slot_specs means the pipeline declares no slots, so a reference is rejected; only ``None`` skips validation.""" - from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params + from tensorrt_llm.visual_gen.params import ( + MediaRef, + VisualGenParams, + validate_visual_gen_params, + ) def run(params, spec): validate_visual_gen_params( params, declared_defaults=None, extra_param_specs={}, ref_slot_specs=spec ) + ref = MediaRef(content="a.png", format="path") with pytest.raises(ValueError, match="not accepted"): - run(VisualGenParams(image_reference="a.png"), {}) - run(VisualGenParams(image_reference="a.png"), None) # None -> skipped + run(VisualGenParams(image_reference=ref), {}) + run(VisualGenParams(image_reference=ref), None) # None -> skipped run(VisualGenParams(), {}) # no reference -> allowed def test_none_fields_not_flagged(self): @@ -1426,7 +1522,7 @@ def test_runtime_error_carried_on_response(self): def test_reference_size_is_prepared_before_warmup_lookup(self): from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor, DiffusionRequest from tensorrt_llm._torch.visual_gen.models.flux.pipeline_flux2 import Flux2Pipeline - from tensorrt_llm.visual_gen.params import VisualGenParams + from tensorrt_llm.visual_gen.params import MediaRef, VisualGenParams events = [] executor = self._make_executor(Flux2Pipeline) @@ -1452,7 +1548,9 @@ def request_warmup_cache_key(req): req = DiffusionRequest( request_id=8, prompt=["test"], - params=VisualGenParams(image_reference=b"encoded image"), + params=VisualGenParams( + image_reference=MediaRef(content=b"encoded image", format="bytes") + ), ) DiffusionExecutor.process_request(executor, req) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 08b6fac602ac..23920b33cd50 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -287,7 +287,9 @@ def test_base64_image_reference_written_to_disk(self, tmp_path): buf = BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest(prompt="x", image_reference=b64) + request = VideoGenerationRequest( + prompt="x", image_reference={"content": b64, "format": "base64"} + ) params = _parse_and_prepare(request, generator, "vid-1", str(tmp_path)) assert len(params.image_reference) == 1 ref_path = params.image_reference[0].content @@ -303,7 +305,11 @@ def test_image_reference_role_and_list(self, tmp_path): Image.new("RGB", (4, 4)).save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() request = VideoGenerationRequest( - prompt="x", image_reference=[b64, {"content": b64, "role": "last_frame"}] + prompt="x", + image_reference=[ + {"content": b64, "format": "base64"}, + {"content": b64, "format": "base64", "role": "last_frame"}, + ], ) params = _parse_and_prepare(request, generator, "vid-r", str(tmp_path)) assert [r.role for r in params.image_reference] == [None, "last_frame"] @@ -316,7 +322,9 @@ def test_missing_media_storage_path_raises(self): buf = BytesIO() img.save(buf, format="PNG") b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest(prompt="x", image_reference=b64) + request = VideoGenerationRequest( + prompt="x", image_reference={"content": b64, "format": "base64"} + ) with pytest.raises(ValueError, match="media_storage_path"): _parse_and_prepare(request, generator, "vid-2", None) @@ -366,7 +374,9 @@ def test_video_reference_needs_media_storage(self): # a storage path is required just like image references. generator = _StubVisualGen() b64 = base64.b64encode(self._mp4_bytes()).decode() - request = VideoGenerationRequest(prompt="x", video_reference=b64) + request = VideoGenerationRequest( + prompt="x", video_reference={"content": b64, "format": "base64"} + ) with pytest.raises(ValueError, match="media_storage_path"): _parse_and_prepare(request, generator, "vid-9", None) @@ -379,7 +389,9 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): vid_b64 = base64.b64encode(self._mp4_bytes()).decode() p = _parse_and_prepare( - VideoGenerationRequest(prompt="x", input_reference=img_b64), + VideoGenerationRequest( + prompt="x", input_reference=img_b64, input_reference_format="base64" + ), generator, "vid-i", str(tmp_path), @@ -387,7 +399,9 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): assert len(p.image_reference) == 1 and p.video_reference is None p = _parse_and_prepare( - VideoGenerationRequest(prompt="x", input_reference=vid_b64), + VideoGenerationRequest( + prompt="x", input_reference=vid_b64, input_reference_format="base64" + ), generator, "vid-v", str(tmp_path), @@ -402,7 +416,12 @@ def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): img_b64 = base64.b64encode(buf.getvalue()).decode() vid_b64 = base64.b64encode(self._mp4_bytes()).decode() p = _parse_and_prepare( - VideoGenerationRequest(prompt="x", image_reference=img_b64, input_reference=vid_b64), + VideoGenerationRequest( + prompt="x", + image_reference={"content": img_b64, "format": "base64"}, + input_reference=vid_b64, + input_reference_format="base64", + ), generator, "vid-x", str(tmp_path), @@ -416,7 +435,9 @@ def test_base64_video_reference_written_to_disk(self, tmp_path): generator = _StubVisualGen() payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest(prompt="x", video_reference=b64) + request = VideoGenerationRequest( + prompt="x", video_reference={"content": b64, "format": "base64"} + ) params = _parse_and_prepare(request, generator, "vid-4", str(tmp_path)) assert params.image_reference is None assert Path(params.video_reference[0].content).read_bytes() == payload @@ -430,7 +451,9 @@ def test_video_reference_survives_real_specs(self, tmp_path): generator = _StubVisualGen(extra_param_specs=COSMOS3_EXTRA_SPECS) payload = self._mp4_bytes() b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest(prompt="x", video_reference=b64) + request = VideoGenerationRequest( + prompt="x", video_reference={"content": b64, "format": "base64"} + ) params = _parse_and_prepare(request, generator, "vid-10", str(tmp_path)) assert Path(params.video_reference[0].content).read_bytes() == payload @@ -457,14 +480,18 @@ def test_wrong_modality_content_raises(self, tmp_path): vid_b64 = base64.b64encode(self._mp4_bytes()).decode() with pytest.raises(ValueError, match="video_reference is not a recognized"): _parse_and_prepare( - VideoGenerationRequest(prompt="x", video_reference=img_b64), + VideoGenerationRequest( + prompt="x", video_reference={"content": img_b64, "format": "base64"} + ), generator, "vid-m1", str(tmp_path), ) with pytest.raises(ValueError, match="image_reference is not a recognized image"): _parse_and_prepare( - VideoGenerationRequest(prompt="x", image_reference=vid_b64), + VideoGenerationRequest( + prompt="x", image_reference={"content": vid_b64, "format": "base64"} + ), generator, "vid-m2", str(tmp_path), @@ -474,7 +501,9 @@ def test_wrong_modality_content_raises(self, tmp_path): def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() b64 = base64.b64encode(b"neither an image nor a video").decode() - request = VideoGenerationRequest(prompt="x", image_reference=b64) + request = VideoGenerationRequest( + prompt="x", image_reference={"content": b64, "format": "base64"} + ) with pytest.raises(ValueError, match="not a recognized image"): _parse_and_prepare(request, generator, "vid-6", str(tmp_path)) # Classification runs on the bytes; rejected content never touches disk. @@ -482,9 +511,11 @@ def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() - # "ABC" survives the lenient alphabet filter but has an invalid - # length, so b64decode raises. - request = VideoGenerationRequest(prompt="x", image_reference="ABC") + # "ABC" has an invalid base64 length. The declared format is honored, + # so this is a decode error rather than a fallback to a filesystem read. + request = VideoGenerationRequest( + prompt="x", image_reference={"content": "ABC", "format": "base64"} + ) with pytest.raises(ValueError, match="not valid base64"): _parse_and_prepare(request, generator, "vid-7", str(tmp_path)) assert list(tmp_path.iterdir()) == [] @@ -512,32 +543,42 @@ def test_multi_reference_partial_failure_cleans_up(self, tmp_path): Image.new("RGB", (4, 4)).save(buf, format="PNG") good = base64.b64encode(buf.getvalue()).decode() bad = base64.b64encode(b"neither an image nor a video").decode() - request = VideoGenerationRequest(prompt="x", image_reference=[good, bad]) + request = VideoGenerationRequest( + prompt="x", + image_reference=[ + {"content": good, "format": "base64"}, + {"content": bad, "format": "base64"}, + ], + ) with pytest.raises(ValueError, match="not a recognized image"): _parse_and_prepare(request, generator, "vid-11", str(tmp_path)) assert list(tmp_path.iterdir()) == [] def test_file_uri_image_reference_passthrough(self, tmp_path): - # A file:// reference is a trusted local path: normalized to a plain path + # format="path" also accepts a file:// URI: normalized to a plain path # and passed through untouched — not copied into media storage. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (7, 8, 9)).save(src, format="PNG") store = tmp_path / "store" store.mkdir() - request = VideoGenerationRequest(prompt="x", image_reference=src.as_uri()) + request = VideoGenerationRequest( + prompt="x", image_reference={"content": src.as_uri(), "format": "path"} + ) params = _parse_and_prepare(request, generator, "vid-file", str(store)) assert params.image_reference[0].content == str(src) assert list(store.iterdir()) == [] def test_bare_path_image_reference_passthrough(self, tmp_path): - # A bare local path is passed through unchanged — not materialized. + # A bare local path under format="path" is passed through unchanged. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (11, 22, 33)).save(src, format="PNG") store = tmp_path / "store" store.mkdir() - request = VideoGenerationRequest(prompt="x", image_reference=str(src)) + request = VideoGenerationRequest( + prompt="x", image_reference={"content": str(src), "format": "path"} + ) params = _parse_and_prepare(request, generator, "vid-bare", str(store)) assert params.image_reference[0].content == str(src) assert list(store.iterdir()) == [] @@ -557,7 +598,10 @@ def __init__(self, content): "tensorrt_llm.visual_gen.media_refs._safe_request_get", lambda url, **kwargs: _FakeResp(png), ) - request = VideoGenerationRequest(prompt="x", image_reference="https://example.com/a.png") + request = VideoGenerationRequest( + prompt="x", + image_reference={"content": "https://example.com/a.png", "format": "url"}, + ) params = _parse_and_prepare(request, generator, "vid-url", str(tmp_path)) assert Path(params.image_reference[0].content).read_bytes() == png @@ -570,7 +614,9 @@ def _blocked(url, **kwargs): raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) - request = VideoGenerationRequest(prompt="x", image_reference="http://10.0.0.1/a.png") + request = VideoGenerationRequest( + prompt="x", image_reference={"content": "http://10.0.0.1/a.png", "format": "url"} + ) with pytest.raises(ValueError, match="reference URL could not be fetched"): _parse_and_prepare(request, generator, "vid-ssrf", str(tmp_path)) assert list(tmp_path.iterdir()) == [] @@ -579,11 +625,36 @@ def test_missing_file_uri_is_client_error(self, tmp_path): # A file:// path that does not exist is a client 400, not a server 500. generator = _StubVisualGen() missing = (tmp_path / "does_not_exist.png").as_uri() - request = VideoGenerationRequest(prompt="x", image_reference=missing) - with pytest.raises(ValueError, match="reference file could not be read"): + request = VideoGenerationRequest( + prompt="x", image_reference={"content": missing, "format": "path"} + ) + with pytest.raises(ValueError, match="reference file does not exist"): _parse_and_prepare(request, generator, "vid-nf", str(tmp_path)) assert list(tmp_path.iterdir()) == [] + def test_bare_reference_string_is_rejected(self): + # The bare-string shorthand is gone: a reference must declare its wire + # form rather than have it guessed from the shape of the value. + from pydantic import ValidationError + + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + b64 = base64.b64encode(buf.getvalue()).decode() + with pytest.raises(ValidationError): + VideoGenerationRequest(prompt="x", image_reference=b64) + with pytest.raises(ValidationError, match="a bare str is no longer accepted"): + VisualGenParams(image_reference=b64) + + def test_json_reference_cannot_declare_bytes(self): + # JSON cannot carry raw bytes; the HTTP schema says so instead of + # letting a str reach the engine claiming to be bytes. + from pydantic import ValidationError + + with pytest.raises(ValidationError, match="multipart/form-data"): + VideoGenerationRequest( + prompt="x", image_reference={"content": "abc", "format": "bytes"} + ) + class TestMediaBytesProbes: """The in-memory signature probes the serve boundary routes on.""" @@ -693,7 +764,8 @@ def test_heif_reference_rejected_with_actionable_message(self): generator = _StubVisualGen() heic = self._ftyp(b"heic", (b"mif1", b"heic")) + b"\x00" * 64 request = VideoGenerationRequest( - prompt="x", image_reference=base64.b64encode(heic).decode() + prompt="x", + image_reference={"content": base64.b64encode(heic).decode(), "format": "base64"}, ) with pytest.raises(ValueError, match="HEIF/AVIF"): _parse_and_prepare(request, generator, "vid-heic", None) @@ -716,7 +788,8 @@ def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): generator = _StubVisualGen() request = VideoGenerationRequest( - prompt="x", image_reference=base64.b64encode(truncated).decode() + prompt="x", + image_reference={"content": base64.b64encode(truncated).decode(), "format": "base64"}, ) params = _parse_and_prepare(request, generator, "vid-12", str(tmp_path)) assert Path(params.image_reference[0].content).read_bytes() == truncated @@ -855,14 +928,14 @@ class TestPrepareReferenceSlots: """Engine-side reference choke point: passthrough local paths, materialize the rest.""" def test_local_path_passthrough_not_materialized_or_cleaned(self, tmp_path): - from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots src = tmp_path / "user.png" Image.new("RGB", (4, 4)).save(src, format="PNG") store = tmp_path / "store" store.mkdir() - params = VisualGenParams(image_reference=str(src)) # bare local path -> passthrough + params = VisualGenParams(image_reference=MediaRef(content=str(src), format="path")) prepare_reference_slots(params, request_id="req1", media_storage_path=str(store)) assert params.image_reference[0].content == str(src) # unchanged assert list(store.iterdir()) == [] # nothing materialized @@ -870,7 +943,7 @@ def test_local_path_passthrough_not_materialized_or_cleaned(self, tmp_path): assert src.exists() # user file untouched by cleanup def test_bytes_materialized_to_storage_and_cleaned(self, tmp_path): - from tensorrt_llm.visual_gen import VisualGenParams + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots buf = BytesIO() @@ -878,7 +951,7 @@ def test_bytes_materialized_to_storage_and_cleaned(self, tmp_path): png = buf.getvalue() store = tmp_path / "store" store.mkdir() - params = VisualGenParams(image_reference=png) # bytes -> materialize + params = VisualGenParams(image_reference=MediaRef(content=png, format="bytes")) prepare_reference_slots(params, request_id="req2", media_storage_path=str(store)) path = params.image_reference[0].content assert path == str(store / "req2_image_ref_0") @@ -886,16 +959,60 @@ def test_bytes_materialized_to_storage_and_cleaned(self, tmp_path): cleanup_reference_files(str(store), "req2") assert not Path(path).exists() - def test_is_local_path(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _is_local_path + def test_materialized_reference_format_is_rewritten_to_path(self, tmp_path): + """``content`` and ``format`` are rewritten together. - f = tmp_path / "x" - f.write_bytes(b"a") - assert _is_local_path(str(f)) is True - assert _is_local_path(f.as_uri()) is True # file:// - assert _is_local_path("iVBORw0KGgo=") is False # base64-ish, no such file - assert _is_local_path("https://example.com/a.png") is False - assert _is_local_path(b"raw bytes") is False + The mutated params object is what gets broadcast to the workers, so a + stale ``base64`` format would reach a worker holding a filesystem path. + """ + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + + buf = BytesIO() + Image.new("RGB", (4, 4)).save(buf, format="PNG") + png = buf.getvalue() + src = tmp_path / "user.png" + src.write_bytes(png) + store = tmp_path / "store" + store.mkdir() + params = VisualGenParams( + image_reference=[ + MediaRef(content=base64.b64encode(png).decode(), format="base64"), + MediaRef(content=png, format="bytes"), + MediaRef(content=str(src), format="path"), + ] + ) + prepare_reference_slots(params, request_id="req3", media_storage_path=str(store)) + assert [r.format for r in params.image_reference] == ["path", "path", "path"] + assert all(Path(r.content).read_bytes() == png for r in params.image_reference) + + def test_path_format_on_missing_file_raises(self, tmp_path): + """A mistyped path stays a path error — it is never retried as base64.""" + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + + store = tmp_path / "store" + store.mkdir() + params = VisualGenParams( + image_reference=MediaRef(content=str(tmp_path / "absent.png"), format="path") + ) + with pytest.raises(ValueError, match="reference file does not exist"): + prepare_reference_slots(params, request_id="req4", media_storage_path=str(store)) + assert list(store.iterdir()) == [] + + def test_base64_format_on_a_filesystem_path_raises(self, tmp_path): + """A path sent as ``base64`` is a decode error — it is never read from disk.""" + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + + src = tmp_path / "user.png" + Image.new("RGB", (4, 4)).save(src, format="PNG") + store = tmp_path / "store" + store.mkdir() + params = VisualGenParams(image_reference=MediaRef(content=str(src), format="base64")) + with pytest.raises(ValueError, match="not valid base64"): + prepare_reference_slots(params, request_id="req5", media_storage_path=str(store)) + assert list(store.iterdir()) == [] def test_resolve_media_storage_path(self, tmp_path, monkeypatch): from tensorrt_llm.visual_gen.media_refs import resolve_media_storage_path @@ -904,3 +1021,83 @@ def test_resolve_media_storage_path(self, tmp_path, monkeypatch): monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(target)) resolved = resolve_media_storage_path() assert resolved == target and target.is_dir() + + +class TestResolveReference: + """``_resolve_reference`` dispatches on the declared format, never on the value.""" + + @staticmethod + def _png() -> bytes: + buf = BytesIO() + Image.new("RGB", (4, 4), (1, 2, 3)).save(buf, format="PNG") + return buf.getvalue() + + def test_every_format_resolves_to_the_same_bytes(self, tmp_path, monkeypatch): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + png = self._png() + src = tmp_path / "ref.png" + src.write_bytes(png) + b64 = base64.b64encode(png).decode() + + class _FakeResp: + def __init__(self, content): + self.content = content + + monkeypatch.setattr( + "tensorrt_llm.visual_gen.media_refs._safe_request_get", + lambda url, **kwargs: _FakeResp(png), + ) + assert _resolve_reference(str(src), "path") == png + assert _resolve_reference(src.as_uri(), "path") == png + assert _resolve_reference("https://example.com/ref.png", "url") == png + assert _resolve_reference(b64, "base64") == png + assert _resolve_reference(f"data:image/png;base64,{b64}", "base64") == png + assert _resolve_reference(png, "bytes") == png + + def test_missing_path_is_a_read_error(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + with pytest.raises(ValueError, match="reference file could not be read"): + _resolve_reference(str(tmp_path / "absent.png"), "path") + + def test_base64_does_not_fall_back_to_a_disk_read(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + src = tmp_path / "ref.png" + src.write_bytes(self._png()) + with pytest.raises(ValueError, match="not valid base64"): + _resolve_reference(str(src), "base64") + + def test_url_fetch_failure_is_a_client_error(self, monkeypatch): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + def _blocked(url, **kwargs): + raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") + + monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) + with pytest.raises(ValueError, match="reference URL could not be fetched"): + _resolve_reference("http://10.0.0.1/a.png", "url") + + def test_non_base64_data_uri_is_rejected(self): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + with pytest.raises(ValueError, match="only base64 data: URIs"): + _resolve_reference("data:image/png,%89PNG", "base64") + with pytest.raises(ValueError, match="data: URI is malformed"): + _resolve_reference("data:image/png;base64", "base64") + + def test_content_type_must_match_the_declared_format(self): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + with pytest.raises(ValueError, match="requires bytes content"): + _resolve_reference("not bytes", "bytes") + for content_format in ("path", "url", "base64"): + with pytest.raises(ValueError, match="requires string content"): + _resolve_reference(b"raw bytes", content_format) + + def test_unknown_format_is_rejected(self): + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + + with pytest.raises(ValueError, match="unsupported reference format"): + _resolve_reference("a.png", "filepath") diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 118389dca435..b6fd5a84c555 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1473,19 +1473,19 @@ models: required: false image_reference: kind: extension - type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] + type: Optional[Union[UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] default: null status: prototype required: false video_reference: kind: extension - type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] + type: Optional[Union[UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] default: null status: prototype required: false audio_reference: kind: extension - type: Optional[Union[str, UploadFile, MediaReferenceItem, List[Union[str, MediaReferenceItem]]]] + type: Optional[Union[UploadFile, MediaReferenceItem, List[MediaReferenceItem]]] default: null status: prototype required: false @@ -1495,6 +1495,12 @@ models: default: null status: deprecated required: false + input_reference_format: + kind: extension + type: Optional[ContentFormat] + default: null + status: deprecated + required: false size: kind: extension type: Optional[str] @@ -1567,3 +1573,23 @@ models: default: null status: stable required: false + MediaReferenceItem: + fields: + content: + kind: extension + type: str + default: null + status: prototype + required: true + format: + kind: extension + type: ContentFormat + default: null + status: prototype + required: true + role: + kind: extension + type: Optional[Role] + default: null + status: prototype + required: false From 283a8d7d3aac569f2d54071e2299aa12f08a6b52 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 18 Aug 2026 05:18:41 -0700 Subject: [PATCH 21/61] [TRTLLM-15277][fix] Enforce the reference content/format pairing at construction 'bytes' is the only format carrying a binary payload; 'path', 'url' and 'base64' name a location or an encoding and are therefore strings. That pairing was only checked when the engine resolved the reference, so MediaRef(content='oops', format='bytes') constructed cleanly and failed deep in generate_async. A model_validator moves the check to construction, where it surfaces as an immediate ValueError or an HTTP 422. The check runs on construction only. StrictBaseModel does not enable validate_assignment, which prepare_reference_slots depends on: it rewrites content and then format, and the statement between the two assignments holds a state that contradicts the pairing. A regression test pins that. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/visual_gen/params.py | 23 +++++++++++++- .../visual_gen/test_visual_gen_params.py | 31 +++++++++++++++++++ 2 files changed, 53 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 915234d96267..fd1b67d1f5c4 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -15,7 +15,7 @@ import ast from typing import Any, Dict, List, Literal, Optional, Union -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status @@ -53,6 +53,27 @@ class MediaRef(StrictBaseModel): default=None, description="``reference`` | ``first_frame`` | ``last_frame``." ) + @model_validator(mode="after") + def _check_content_matches_format(self): + """Reject a ``content`` whose Python type contradicts ``format``. + + ``bytes`` is the only format carrying a binary payload; the other three + name a location or an encoding and are therefore strings. Checking the + pairing here fails at construction — an HTTP 422 or an immediate + ``ValueError`` — instead of deep in the engine's resolve step. + """ + if self.format == "bytes": + if not isinstance(self.content, bytes): + raise ValueError( + f"format='bytes' requires bytes content, got {type(self.content).__name__}." + ) + elif not isinstance(self.content, str): + raise ValueError( + f"format={self.format!r} requires string content, got " + f"{type(self.content).__name__}." + ) + return self + def _reject_bare_refs(value: Any) -> Any: """Reject the bare path/bytes shorthand with an actionable message. 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 780090b6ac19..378c6aeecc57 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -194,6 +194,37 @@ def test_format_is_required(self): with pytest.raises(ValidationError, match=r"format\s+Field required"): MediaRef(content="a.png") + @pytest.mark.parametrize( + "content,content_format", + [("not-bytes", "bytes"), (b"\x89PNG", "base64"), (b"\x89PNG", "path"), (b"x", "url")], + ) + def test_content_type_must_match_format(self, content, content_format): + """The pairing is enforced at construction, not deep in the engine.""" + from pydantic import ValidationError + + from tensorrt_llm.visual_gen import MediaRef + + with pytest.raises(ValidationError, match="requires (bytes|string) content"): + MediaRef(content=content, format=content_format) + + def test_content_type_matching_format_accepted(self): + from tensorrt_llm.visual_gen import MediaRef + + assert MediaRef(content=b"raw", format="bytes").content == b"raw" + for fmt in ("path", "url", "base64"): + assert MediaRef(content="x", format=fmt).format == fmt + + def test_engine_rewrite_is_not_blocked_by_the_pairing_check(self): + """``prepare_reference_slots`` rewrites content then format, so the + intermediate state contradicts the pairing; assignment must not + re-validate or that rewrite would be impossible.""" + from tensorrt_llm.visual_gen import MediaRef + + ref = MediaRef(content="aGk=", format="base64") + ref.content = "/tmp/materialized" # contradicts format for one statement + ref.format = "path" + assert (ref.content, ref.format) == ("/tmp/materialized", "path") + def test_unknown_format_rejected(self): from pydantic import ValidationError From 30efd2bff9e6e768cca0fdd80a3803ca022206c9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 18 Aug 2026 06:18:25 -0700 Subject: [PATCH 22/61] [TRTLLM-15277][feat] Let ImageMediaIO callers choose the alpha and target mode MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit convert_image_mode, _load_and_convert_image and ImageMediaIO gain 'mode' and 'drop_alpha'. Both default to today's behavior — RGB with the alpha composited onto white — so every existing caller is byte-identical; the five call sites all pass (image, "RGB") and are unaffected. The knobs exist because RGBA-to-RGB has two defensible semantics and the right one depends on the consumer. Compositing onto white changes every pixel with alpha < 255, not just the fully transparent ones. PIL's own convert("RGB") and diffusers' load_image instead drop the channel and keep the stored RGB, so a pipeline ported from diffusers needs drop_alpha=True to stay numerically aligned with upstream. A consumer that composites layers itself needs mode="RGBA" to keep the channel at all. Tests pin all three properties: drop_alpha=True is pixel-identical to PIL/diffusers, the default is pixel-identical to the previous behavior, and the two semantics coincide for opaque media — which is what every committed golden uses. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/inputs/media_io.py | 43 +++++++--- tests/unittest/llmapi/apps/test_media_io.py | 88 ++++++++++++++++++++- 2 files changed, 121 insertions(+), 10 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 19c10697416f..9bf36b1dda13 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -60,11 +60,19 @@ def rgba_to_rgb( return converted -def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: - """Convert image to specified mode with proper handling of RGBA to RGB conversion.""" +def convert_image_mode(image: Image.Image, to_mode: str, drop_alpha: bool = False) -> Image.Image: + """Convert image to specified mode with proper handling of RGBA to RGB conversion. + + ``drop_alpha`` selects how an RGBA source sheds its alpha channel on the way + to RGB: ``True`` discards the channel and keeps the stored RGB, which is what + PIL's own ``convert("RGB")`` and diffusers' ``load_image`` do; ``False`` + composites onto a white background, which changes every pixel with + ``alpha < 255``. It only applies to that one direction, and defaults to + compositing so existing callers are unaffected. + """ if image.mode == to_mode: return image - elif image.mode == "RGBA" and to_mode == "RGB": + elif image.mode == "RGBA" and to_mode == "RGB" and not drop_alpha: return rgba_to_rgb(image) else: return image.convert(to_mode) @@ -242,10 +250,10 @@ async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: return await _fetch(owned_session) -def _load_and_convert_image(image): +def _load_and_convert_image(image, mode: str = "RGB", drop_alpha: bool = False): image = Image.open(image) image.load() - return convert_image_mode(image, "RGB") + return convert_image_mode(image, mode, drop_alpha) def _audio_frame_to_array(frame, mono: bool) -> np.ndarray: @@ -825,11 +833,22 @@ async def _run_in_executor(fn, *args, **kwargs): class ImageMediaIO(BaseMediaIO[Union[Image.Image, torch.Tensor, np.ndarray]]): """I/O for the image modality.""" - def __init__(self, format: str = "pt", device: str = "cpu") -> None: + def __init__( + self, + format: str = "pt", + device: str = "cpu", + mode: str = "RGB", + drop_alpha: bool = False, + ) -> None: if format not in _SUPPORTED_IMAGE_FORMATS: raise ValueError(f"format must be one of {_SUPPORTED_IMAGE_FORMATS}, got {format!r}") self._format = format self._device = device + # Target PIL mode, plus how RGBA sheds its alpha en route to RGB. A + # consumer that composites layers asks for mode="RGBA"; one that must + # match diffusers' preprocessing asks for drop_alpha=True. + self._mode = mode + self._drop_alpha = drop_alpha def _postprocess(self, image: Image.Image) -> Union[Image.Image, torch.Tensor, np.ndarray]: if self._format == "pt": @@ -842,15 +861,21 @@ def _postprocess(self, image: Image.Image) -> Union[Image.Image, torch.Tensor, n return image def load_bytes(self, data: bytes) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess(_load_and_convert_image(BytesIO(data))) + return self._postprocess( + _load_and_convert_image(BytesIO(data), self._mode, self._drop_alpha) + ) def load_base64( self, media_type: str, data: str ) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess(_load_and_convert_image(BytesIO(base64.b64decode(data)))) + return self._postprocess( + _load_and_convert_image(BytesIO(base64.b64decode(data)), self._mode, self._drop_alpha) + ) def load_file(self, url: str) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess(_load_and_convert_image(Path(_normalize_file_uri(url)))) + return self._postprocess( + _load_and_convert_image(Path(_normalize_file_uri(url)), self._mode, self._drop_alpha) + ) class AudioMediaIO(BaseMediaIO[Tuple[np.ndarray, int]]): diff --git a/tests/unittest/llmapi/apps/test_media_io.py b/tests/unittest/llmapi/apps/test_media_io.py index d139ea2cd4a6..f34cdc475a08 100644 --- a/tests/unittest/llmapi/apps/test_media_io.py +++ b/tests/unittest/llmapi/apps/test_media_io.py @@ -5,7 +5,13 @@ import pytest from tensorrt_llm.inputs import MultimodalDataTracker -from tensorrt_llm.inputs.media_io import AudioMediaIO, BaseMediaIO, ImageMediaIO, VideoMediaIO +from tensorrt_llm.inputs.media_io import ( + AudioMediaIO, + BaseMediaIO, + ImageMediaIO, + VideoMediaIO, + convert_image_mode, +) from tensorrt_llm.serve.chat_utils import parse_chat_message_content_part pytestmark = pytest.mark.cpu_only @@ -91,3 +97,83 @@ def test_non_video_classes_use_plain_shallow_merge(self, media_io_cls): {"num_frames": 32}, ) assert merged == {"num_frames": 32, "fps": 1} + + +class TestImageAlphaHandling: + """RGBA -> RGB has two defensible semantics. + + The caller picks, and the default must stay what every existing caller + already gets. + """ + + @staticmethod + def _rgba_png() -> bytes: + """Build a 3-pixel RGBA fixture. + + One opaque, one half-transparent and one fully transparent pixel, all + sharing the same stored RGB so the two semantics are separable. + """ + from io import BytesIO + + from PIL import Image + + im = Image.new("RGBA", (3, 1)) + im.putpixel((0, 0), (200, 30, 30, 255)) + im.putpixel((1, 0), (200, 30, 30, 128)) + im.putpixel((2, 0), (200, 30, 30, 0)) + buf = BytesIO() + im.save(buf, format="PNG") + return buf.getvalue() + + def test_drop_alpha_matches_pil_and_diffusers(self): + """Match diffusers. + + Its load_image defaults to image.convert("RGB"), so pipelines ported + from diffusers need that exact behavior to stay aligned. + """ + from io import BytesIO + + from PIL import Image + + png = self._rgba_png() + reference = list(Image.open(BytesIO(png)).convert("RGB").getdata()) + got = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(png) + assert list(got.getdata()) == reference + + def test_default_composites_and_is_unchanged(self): + """Keep compositing onto white by default. + + No existing LLM or VLM caller may shift. + """ + got = ImageMediaIO(format="pil").load_bytes(self._rgba_png()) + assert list(got.getdata()) == [(200, 30, 30), (227, 142, 142), (255, 255, 255)] + + def test_semantics_coincide_for_opaque_images(self): + """Coincide on opaque media. + + Every committed golden uses opaque media, so the two semantics must be + bit-identical there. + """ + from io import BytesIO + + from PIL import Image + + buf = BytesIO() + Image.new("RGBA", (2, 1), (10, 20, 30, 255)).save(buf, format="PNG") + png = buf.getvalue() + composited = ImageMediaIO(format="pil").load_bytes(png) + dropped = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(png) + assert list(composited.getdata()) == list(dropped.getdata()) + + def test_mode_rgba_preserves_alpha(self): + got = ImageMediaIO(format="pil", mode="RGBA").load_bytes(self._rgba_png()) + assert got.mode == "RGBA" + assert list(got.getdata())[2] == (200, 30, 30, 0) + + def test_convert_image_mode_default_is_unchanged(self): + """The shared helper is publicly exported; its default must not move.""" + from PIL import Image + + im = Image.new("RGBA", (1, 1), (200, 30, 30, 0)) + assert convert_image_mode(im, "RGB").getpixel((0, 0)) == (255, 255, 255) + assert convert_image_mode(im, "RGB", drop_alpha=True).getpixel((0, 0)) == (200, 30, 30) From 384ce701e3bf65e556e93fb94d539cca8f2bf283 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 18 Aug 2026 22:57:19 -0700 Subject: [PATCH 23/61] [TRTLLM-15277][refactor] Split NVDEC decoding into mechanism, selector and MediaIO leaf MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit decode_video_reference_window was one function fusing three concerns: the NVDEC demux/decode loop, the choice of which frames to keep, and a mandatory resize. They are now separable — _nvdec_decode is the mechanism, FrameSelector names the retention policy, and target_hw is optional — so a consumer that wants NVDEC with a different frame policy, or at the source resolution, no longer has to fork the loop. Selection is explicit by construction: FrameSelector has no default and an unimplemented strategy raises rather than silently degrading to a window, so a caller cannot inherit another model's frame convention by accident. WindowSelector validates its own range, moving the check off the decode path and onto the value. NvdecVideoMediaIO exposes this through the shared BaseMediaIO contract. It is a sibling of VideoMediaIO rather than a subclass because that one decodes with cv2 on the CPU and returns VideoData, which its VLM callers depend on; this one returns a device tensor. It lives under tensorrt_llm/media rather than tensorrt_llm/inputs so the NVDEC body stays with its owners — the dependency runs VG -> inputs, which is the legal direction. decode_video_reference_window survives as a thin spelling over WindowSelector, so pipeline_cosmos3 and the existing tests are untouched; the 23 pre-existing cases pass unchanged as the equivalence baseline. Resize still happens before retention, which is what bounds retained memory to the target rather than to the source. The one behavioral subtlety is that a ring sized from the first frame would return None when the range matches no frame, so the ring is still pre-allocated whenever a target is known; a regression test pins the empty result. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/test_media_decode.py | 73 +++++++++++++++++++ 1 file changed, 73 insertions(+) diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index 0c88f47ff86b..f157731c65a8 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -27,7 +27,11 @@ from tensorrt_llm._torch.visual_gen.utils import synchronize_media_prepare_status from tensorrt_llm.media.decoding import ( + FrameSelector, + NvdecVideoMediaIO, + WindowSelector, _lanczos_taps, + _nvdec_decode, decode_video_reference_window, resize_center_crop_uint8, resize_fit_pad_uint8, @@ -379,3 +383,72 @@ def test_resize_perf_representative(self): torch.cuda.synchronize() per_frame = (time.perf_counter() - start) / 10 assert per_frame < 0.25, f"resize took {per_frame * 1e3:.1f} ms/frame" + + +class TestSelectorAndMediaIO: + """The decode mechanism, its explicit frame selector, and the MediaIO leaf.""" + + _DEVICE = torch.device("cuda:0") + + def test_window_selector_rejects_mixed_and_reversed_ranges(self): + """The range is validated when the selector is built, not at decode.""" + with pytest.raises(ValueError, match="both count from"): + WindowSelector(0, -1) + with pytest.raises(ValueError, match="must not exceed"): + WindowSelector(3, 1) + + def test_unimplemented_selector_is_rejected(self): + """Selection has no default: an unknown strategy must not silently + fall back to a window.""" + + class FpsSelector(FrameSelector): + pass + + with pytest.raises(NotImplementedError, match="FpsSelector"): + _nvdec_decode(_MP4.read_bytes(), selector=FpsSelector(), device=self._DEVICE) + + @pytest.mark.parametrize("fixture", [_MP4, _AVI], ids=["mp4", "avi"]) + def test_media_io_matches_the_window_function(self, fixture): + """The MediaIO leaf and the legacy free function are the same decode.""" + data = fixture.read_bytes() + reference = decode_video_reference_window( + data, first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE + ) + got = NvdecVideoMediaIO( + selector=WindowSelector(0, 4), device=self._DEVICE, target_hw=(64, 64) + ).load_bytes(data) + assert torch.equal(got, reference) + + def test_media_io_load_file_matches_load_bytes(self, tmp_path): + io = NvdecVideoMediaIO( + selector=WindowSelector(0, 4), device=self._DEVICE, target_hw=(64, 64) + ) + assert torch.equal(io.load_file(str(_MP4)), io.load_bytes(_MP4.read_bytes())) + + def test_no_target_keeps_the_source_resolution(self): + """Resize only happens when a target is given. + + The fixture is natively 64x64, so the target has to be a different + size for the two paths to be distinguishable at all. + """ + data = _MP4.read_bytes() + native = _nvdec_decode(data, selector=WindowSelector(0, 0), device=self._DEVICE) + resized = _nvdec_decode( + data, selector=WindowSelector(0, 0), device=self._DEVICE, target_hw=(32, 32) + ) + assert native.shape == (1, 64, 64, 3) + assert resized.shape == (1, 32, 32, 3) + assert native.dtype == torch.uint8 + + def test_range_past_the_clip_returns_empty_not_error(self): + """A window beyond the clip yields what exists — here, nothing. + + The target's spatial dims are kept so the caller can still pad. + """ + window = _nvdec_decode( + _MP4.read_bytes(), + selector=WindowSelector(100, 104), + device=self._DEVICE, + target_hw=(64, 64), + ) + assert window.shape == (0, 64, 64, 3) From 0ed7818fd1cec011877f6ae34347057632ce4b5f Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 19 Aug 2026 01:49:28 -0700 Subject: [PATCH 24/61] [TRTLLM-15277][feat] VisualGen: carry references as bytes, decode through MediaIO References no longer round-trip through the filesystem. `prepare_reference_slots` resolves every declared wire form (path/url/base64/bytes) to raw bytes on the coordinator and rewrites `format` to `"bytes"`, so bytes are the canonical form all the way to the pipeline. A worker therefore needs no filesystem shared with the client, and there is nothing to clean up when a request ends: the terminal `on_finish` hook, `cleanup_reference_files`, and `resolve_media_storage_path` are gone, along with the `VisualGenResult` cancellation branch that existed only to fire that hook. With paths off the wire every worker-side decode has a single input type, so the hand-rolled `PIL.Image.open(...).convert(...)` / `load_image(...)` sites across Cosmos3, FLUX.2, LTX-2, Qwen-Image-Edit, Qwen-Image-Layered and Wan collapse onto `ImageMediaIO.load_bytes`, and the `str` type guards and hints on the same paths become `bytes`. Each site keeps the alpha semantics it had: `drop_alpha=True` where the pipeline follows diffusers' `convert("RGB")`, the compositing default where it went through `load_image`, and `mode="RGBA"` for layer decomposition. No direct `Image.open` / `load_image` call is left under `_torch/visual_gen/`. Cosmos3 also loses its `Path(video).read_bytes()` re-read: the video reference already arrives as the encoded bytes NVDEC demuxes from memory. Verified on real weights that `format="path"` and `format="base64"` naming the same file produce bitwise-identical output: Wan2.2-I2V-A14B, Cosmos3-Super I2V, Cosmos3-Super V2V, and FLUX.2-dev all pass. LTX-2 decodes its reference and generates a non-degenerate video, but is nondeterministic run-to-run (a path-vs-path control with a fixed seed also differs), so bitwise equivalence is not measurable there. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- .../models/cosmos3/pipeline_cosmos3.py | 18 +- .../visual_gen/models/flux/pipeline_flux2.py | 23 +- .../visual_gen/models/ltx2/pipeline_ltx2.py | 5 +- .../qwen_image/pipeline_qwen_image_edit.py | 14 +- .../pipeline_qwen_image_layered.py | 9 +- .../visual_gen/models/wan/pipeline_wan.py | 9 +- .../visual_gen/models/wan/pipeline_wan_i2v.py | 36 +- tensorrt_llm/serve/openai_video_routes.py | 5 +- tensorrt_llm/visual_gen/media_refs.py | 120 ++----- tensorrt_llm/visual_gen/visual_gen.py | 55 +-- .../visual_gen/test_cosmos3_pipeline.py | 25 +- .../test_flux2_image_conditioning.py | 31 +- .../visual_gen/test_trtllm_serve_endpoints.py | 80 ++--- .../visual_gen/test_visual_gen_utils.py | 318 ++++++------------ 15 files changed, 258 insertions(+), 492 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index e18b34f00602..c65a6b42fda1 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -121,7 +121,7 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue ### Reference Inputs -Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a single reference or a list. A reference always declares the wire form of its content — `MediaRef(content=..., format=...)` in Python, `{"content": ..., "format": ...}` in JSON. `format` is **required** and nothing is guessed: a bare string or bare bytes is rejected, so a mistyped path can never be silently read as base64. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. When served, references are carried on the video endpoints and are materialized to a local path before reaching the worker; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). +Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a single reference or a list. A reference always declares the wire form of its content — `MediaRef(content=..., format=...)` in Python, `{"content": ..., "format": ...}` in JSON. `format` is **required** and nothing is guessed: a bare string or bare bytes is rejected, so a mistyped path can never be silently read as base64. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. Whatever form a reference is declared in, it is resolved to raw bytes on the coordinator before the request is broadcast, so a worker never needs a filesystem shared with the client; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). | `format` | Content | Notes | |---|---|---| diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 69b0693877f8..fd2953487b0a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -45,7 +45,7 @@ def tqdm(iterable, **kwargs): synchronize_media_prepare_status, ) from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.inputs.utils import load_image +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from tensorrt_llm.media.decoding import decode_video_reference_window, video_stream_info @@ -237,7 +237,7 @@ def _condition_pixel_frame_count( return max(condition_video_latent_indexes) * int(temporal_compression) + 1 -def _load_reference_image(path: str): +def _load_reference_image(data: bytes): """Load an I2V reference, reporting unreadable content as a client error. The worker's load is the acceptance check — the serve boundary only @@ -247,7 +247,7 @@ def _load_reference_image(path: str): upload would be reported as a server fault. """ try: - return load_image(path, format="pil") + return ImageMediaIO(format="pil").load_bytes(data) except OSError as exc: raise ValueError( f"Image reference could not be decoded; it may be truncated, " @@ -760,10 +760,6 @@ def as_given(field_name): refs_v = req.params.video_reference video = refs_v[0].content if refs_v else None - if isinstance(video, str): - from pathlib import Path - - video = Path(video).read_bytes() # forward() NVDEC-demuxes from memory if video is not None: _validate_video_reference(video) is_action = extra_params.get("action_mode") is not None @@ -1448,7 +1444,7 @@ def forward( self, prompt: Union[str, List[str]], negative_prompt: Optional[str] = None, - image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, height: Optional[int] = None, width: Optional[int] = None, num_frames: Optional[int] = None, @@ -1756,9 +1752,9 @@ def forward( raise ValueError("Batch generation is not supported for Cosmos3") # Validate image input — only single image is supported for batch generation - if image is not None and not isinstance(image, (PIL.Image.Image, torch.Tensor, str)): + if image is not None and not isinstance(image, (PIL.Image.Image, torch.Tensor, bytes)): raise ValueError( - f"`image` must be a PIL.Image, torch.Tensor, or file path string, " + f"`image` must be a PIL.Image, torch.Tensor, or encoded bytes, " f"got {type(image)}. Batch of different images is not supported; " f"use a single image with multiple prompts instead." ) @@ -2012,7 +2008,7 @@ def forward( elif image is not None: prepare_error: Optional[Exception] = None try: - if isinstance(image, str): + if isinstance(image, bytes): image = _load_reference_image(image) if isinstance(image, PIL.Image.Image): 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 f5fdfa91d390..b1465590d925 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -18,7 +18,6 @@ - 4-axis RoPE: (32, 32, 32, 32) instead of 3-axis """ -import io import json import os import time @@ -46,6 +45,7 @@ from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm.inputs.media_io import ImageMediaIO, convert_image_mode from tensorrt_llm.logger import logger from .transformer_flux2 import Flux2Transformer2DModel @@ -426,9 +426,8 @@ def forward( image: Optional[ Union[ PIL.Image.Image, - str, bytes, - List[Union[PIL.Image.Image, str, bytes]], + List[Union[PIL.Image.Image, bytes]], ] ] = None, _condition_images: Optional[List[torch.Tensor]] = None, @@ -754,30 +753,28 @@ def _prepare_latent_ids(self, height: int, width: int) -> torch.Tensor: def _load_reference_images( image: Union[ PIL.Image.Image, - str, bytes, - List[Union[PIL.Image.Image, str, bytes]], + List[Union[PIL.Image.Image, bytes]], ], ) -> List[PIL.Image.Image]: - """Normalize supported reference-image inputs to materialized RGB images.""" + """Normalize supported reference-image inputs to decoded RGB images.""" inputs = image if isinstance(image, list) else [image] if not inputs: raise ValueError("`image` must contain at least one reference image.") + # drop_alpha keeps diffusers' semantics: convert("RGB") drops the alpha + # channel rather than compositing it onto white. + media_io = ImageMediaIO(format="pil", drop_alpha=True) images = [] for index, item in enumerate(inputs): try: if isinstance(item, PIL.Image.Image): - images.append(item.convert("RGB")) - elif isinstance(item, str): - with PIL.Image.open(item) as loaded: - images.append(loaded.convert("RGB")) + images.append(convert_image_mode(item, "RGB", drop_alpha=True)) elif isinstance(item, bytes): - with PIL.Image.open(io.BytesIO(item)) as loaded: - images.append(loaded.convert("RGB")) + images.append(media_io.load_bytes(item)) else: raise ValueError( - "Reference images must be PIL images, file paths, or encoded bytes; " + "Reference images must be PIL images or encoded bytes; " f"item {index} has type {type(item).__name__}." ) except OSError as exc: 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 6bf0b82e5ba4..b2a25cc0a432 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -28,6 +28,7 @@ ) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .ltx2_core.audio_vae import AudioDecoderConfigurator, VocoderConfigurator, decode_audio @@ -1233,10 +1234,10 @@ def _load_and_preprocess_image( Returns: Tensor of shape ``(1, 3, 1, H, W)`` in ``[-1, 1]``. """ - if isinstance(image, str): + if isinstance(image, bytes): from PIL import Image - pil_img = Image.open(image).convert("RGB") + pil_img = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) pil_img = pil_img.resize((width, height), Image.LANCZOS) import numpy as np 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 19ec7704cbc4..3e2d32d3f6a0 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 @@ -10,7 +10,6 @@ import math import time -from io import BytesIO from typing import Any import numpy as np @@ -20,7 +19,7 @@ from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline -from tensorrt_llm.inputs.utils import load_image +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .pipeline_qwen_image import QwenImagePipeline, _calculate_shift @@ -164,15 +163,8 @@ def _load_edit_images(image: Any) -> list[Any]: if image is None: raise ValueError("Qwen-Image-Edit requires image_reference.") images = image if isinstance(image, list) else [image] - pil_images = [] - for item in images: - if isinstance(item, bytes): - from PIL import Image - - pil_images.append(Image.open(BytesIO(item)).convert("RGB")) - else: - pil_images.append(load_image(item, format="pil")) - return pil_images + media_io = ImageMediaIO(format="pil") + return [media_io.load_bytes(item) for item in images] def _preprocess_edit_images( self, 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 e6771dc8b928..35dd70734b52 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 @@ -14,7 +14,6 @@ # limitations under the License. """Qwen-Image-Layered image decomposition pipeline.""" -import io import math import time from typing import List, Optional, Tuple, Union @@ -30,6 +29,7 @@ RoleSpec, ) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .transformer_qwen_image_layered import QwenImageLayeredTransformer2DModel @@ -361,14 +361,11 @@ def load_weights(self, weights: dict) -> None: @staticmethod def _load_image_input(image): - from PIL import Image - if isinstance(image, list): return [QwenImageLayeredPipeline._load_image_input(item) for item in image] - if isinstance(image, str): - return Image.open(image).convert("RGBA") if isinstance(image, bytes): - return Image.open(io.BytesIO(image)).convert("RGBA") + # Layer decomposition needs the alpha channel, not a flattened RGB. + return ImageMediaIO(format="pil", mode="RGBA").load_bytes(image) if hasattr(image, "convert") and getattr(image, "mode", None) != "RGBA": return image.convert("RGBA") return image 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 360ebcdd063d..f5ce9dee270d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -43,6 +43,7 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .transformer_wan import WanTransformer3DModel @@ -476,7 +477,7 @@ def forward( guidance_scale_2: Optional[float] = None, boundary_ratio: Optional[float] = None, max_sequence_length: int = 512, - image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -793,7 +794,7 @@ def _prepare_latents( def _prepare_latents_wan22_5B_i2v( self, batch_size: int, - image: Union[PIL.Image.Image, torch.Tensor, str], + image: Union[PIL.Image.Image, torch.Tensor, bytes], height: int, width: int, num_frames: int, @@ -814,8 +815,8 @@ def _prepare_latents_wan22_5B_i2v( latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) # Load and preprocess image - if isinstance(image, str): - image = PIL.Image.open(image).convert("RGB") + if isinstance(image, bytes): + image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) image = ( self.video_processor.preprocess(image, height=height, width=width) .to(self.device, dtype=self.vae.dtype) 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 7bd717878728..b3e9460590e9 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 @@ -39,6 +39,7 @@ from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor +from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger # Supported Wan I2V 14B models: @@ -450,7 +451,7 @@ def infer(self, req): @torch.no_grad() def forward( self, - image: Union[PIL.Image.Image, torch.Tensor, str], + image: Union[PIL.Image.Image, torch.Tensor, bytes], prompt: Union[str, List[str]], seed: int, negative_prompt: Optional[str] = None, @@ -462,16 +463,16 @@ def forward( guidance_scale_2: Optional[float] = None, boundary_ratio: Optional[float] = None, max_sequence_length: int = 512, - last_image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + last_image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, ): pipeline_start = time.time() timer = CudaPhaseTimer() timer.mark_pre_start() # Validate image input — only single image is supported for batch generation - if not isinstance(image, (PIL.Image.Image, torch.Tensor, str)): + if not isinstance(image, (PIL.Image.Image, torch.Tensor, bytes)): raise ValueError( - f"`image` must be a PIL.Image, torch.Tensor, or file path string, " + f"`image` must be a PIL.Image, torch.Tensor, or encoded bytes, " f"got {type(image)}. Batch of different images is not supported; " f"use a single image with multiple prompts instead." ) @@ -723,14 +724,17 @@ def get_embeds(texts): def _encode_image( self, - image: Union[PIL.Image.Image, torch.Tensor, str], - last_image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + image: Union[PIL.Image.Image, torch.Tensor, bytes], + last_image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, ) -> torch.Tensor: """Encode image(s) using CLIP image encoder (Wan 2.1 I2V only).""" - if isinstance(image, str): - image = PIL.Image.open(image).convert("RGB") - if isinstance(last_image, str): - last_image = PIL.Image.open(last_image).convert("RGB") + # drop_alpha keeps diffusers' semantics: convert("RGB") drops the alpha + # channel rather than compositing it onto white. + media_io = ImageMediaIO(format="pil", drop_alpha=True) + if isinstance(image, bytes): + image = media_io.load_bytes(image) + if isinstance(last_image, bytes): + last_image = media_io.load_bytes(last_image) images_to_encode = [image] if last_image is None else [image, last_image] @@ -744,12 +748,12 @@ def _encode_image( def _prepare_latents( self, batch_size: int, - image: Union[PIL.Image.Image, torch.Tensor, str], + image: Union[PIL.Image.Image, torch.Tensor, bytes], height: int, width: int, num_frames: int, generator: torch.Generator, - last_image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + last_image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: """Prepare latents with image conditioning for I2V generation.""" num_channels_latents = 16 @@ -761,15 +765,15 @@ def _prepare_latents( latents = randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) # Load and preprocess image(s) - if isinstance(image, str): - image = PIL.Image.open(image).convert("RGB") + if isinstance(image, bytes): + image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) image = self.video_processor.preprocess(image, height=height, width=width).to( self.device, dtype=torch.float32 ) if last_image is not None: - if isinstance(last_image, str): - last_image = PIL.Image.open(last_image).convert("RGB") + if isinstance(last_image, bytes): + last_image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(last_image) last_image = self.video_processor.preprocess(last_image, height=height, width=width).to( self.device, dtype=torch.float32 ) diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index d2b6b2535584..61c847ea997b 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -428,11 +428,10 @@ async def openai_video_generation_async( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) - # Resolve/materialize references, validate params, and enqueue in the + # Resolve references, validate params, and enqueue in the # foreground (offloaded but awaited) so bad media / unknown # extra_params surface as 400 here, before the 202 — not as a queued - # job that later fails. The engine reclaims the references via its - # terminal hook once the background task awaits the handle. + # job that later fails. handle = await asyncio.to_thread(self.generator.generate_async, request.prompt, params) # Persist the queued job before scheduling the background task so diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index edb7a0ae58cd..86041266b134 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -12,20 +12,20 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. -"""Reference-media resolve / materialize / cleanup, shared by serve and engine. +"""Reference-media resolution, shared by serve and engine. -Verb convention: ``resolve`` -> bytes, ``materialize`` -> path. These are used -by both the serve boundary (``tensorrt_llm/serve``) and the engine frontend -(``VisualGen.generate_async``), so they live here rather than under ``serve`` to -avoid an engine -> serve import. +Every declared wire form resolves to raw bytes, the canonical form carried all +the way to the pipeline. Used by both the serve boundary +(``tensorrt_llm/serve``) and the engine frontend (``VisualGen.generate_async``), +so this lives here rather than under ``serve`` to avoid an engine -> serve +import. """ from __future__ import annotations import base64 -import os from pathlib import Path -from typing import Any, Optional +from typing import Any from tensorrt_llm.inputs.media_io import ( _normalize_file_uri, @@ -99,10 +99,8 @@ def _resolve_reference(content: Any, content_format: str) -> bytes: raise ValueError(f"unsupported reference format: {content_format!r}") -def _materialize_reference( - payload: bytes, *, modality: str, ref_id: str, media_storage_path: Optional[str] -) -> str: - """Content-validate a reference payload and persist it, returning its path. +def _validate_reference_payload(payload: bytes, *, modality: str) -> None: + """Reject a payload whose container does not match the declared modality. HEIF/AVIF images are rejected on signature alone (Pillow support depends on optional plugins the worker need not share). Video acceptance beyond the @@ -126,90 +124,32 @@ def _materialize_reference( ) # audio: no signature sniffing (sniff_media_kind detects only image/video); # the consuming pipeline validates the audio codec in its worker. - if media_storage_path is None: - raise ValueError(f"media_storage_path is required to store the {modality}_reference.") - ref_path = os.path.join(media_storage_path, ref_id) - with open(ref_path, "wb") as f: - f.write(payload) - return ref_path - - -def cleanup_reference_files(media_storage_path: Optional[str], request_id: str) -> None: - """Remove the materialized reference inputs for one request. - - References are materialized as ``{request_id}_{modality}_ref_{i}`` (and the - deprecated ``{request_id}_input_ref``) under ``media_storage_path``. They are - input-only — unneeded once the pipeline has consumed them — so the request - owner removes them by the ``request_id`` prefix, covering image/video/audio - and the deprecated single reference regardless of count. Output files - (``{request_id}_{i}.``) carry no ``ref`` and are left untouched. - Best-effort: already-removed files are ignored. - """ - if media_storage_path is None: - return - for path in Path(media_storage_path).glob(f"{request_id}_*ref*"): - try: - path.unlink() - except OSError: - pass - -def resolve_media_storage_path() -> Path: - """Resolve the media storage directory, creating it if needed. - - Reads ``TRTLLM_MEDIA_STORAGE_PATH`` (default ``/tmp/trtllm_generated``), - shared by the serve boundary and the engine so both write materialized - references to the same place. - """ - path = Path(os.getenv("TRTLLM_MEDIA_STORAGE_PATH", "/tmp/trtllm_generated")) # nosec B108 - path.mkdir(parents=True, exist_ok=True) - return path - -def prepare_reference_slots( - params: Any, *, request_id: str, media_storage_path: Optional[str] -) -> None: - """Resolve + materialize each reference to a local path, in place. +def prepare_reference_slots(params: Any) -> None: + """Resolve every reference to raw bytes, in place. The single reference choke point, used by the engine (``generate_async``) so serve and the standalone Python API share one path. Dispatch is on each - reference's declared ``format``, never on the shape of its content. A - ``path`` reference is the caller's own file: it passes through — not - materialized, not cleaned up — with a ``file://`` URI normalized to a plain - path so the pipeline, which opens paths, can read it. Every other form - (``url`` / ``base64`` / ``bytes``) resolves to bytes and materializes to - ``media_storage_path``; those files are reclaimed by - :func:`cleanup_reference_files` keyed on ``request_id``. - - ``format`` is rewritten alongside ``content``: the mutated params object is - what gets broadcast to the workers, so a stale format would send e.g. - ``base64`` to a worker holding a filesystem path. + reference's declared ``format``, never on the shape of its content: the + declared form is resolved to bytes, content-validated against the slot's + modality, and written back with ``format`` set to ``"bytes"``. + + ``format`` is rewritten alongside ``content`` because the mutated params + object is what gets broadcast to the workers; a stale format would tell a + worker it is holding base64 when it is holding raw bytes. + + Bytes are the canonical form all the way to the pipeline, so a reference + never touches the filesystem: there is nothing to clean up afterwards, and + a worker needs no shared filesystem to read what the coordinator resolved. Runs before the coordinator broadcasts the request, so a bad reference - raises ``ValueError`` synchronously (serve keeps its immediate 400). If a - later slot fails mid-materialize, the files earlier slots wrote are - reclaimed here so a rejected request leaves nothing on disk. + raises ``ValueError`` synchronously and serve keeps its immediate 400. """ - try: - for slot in ("image_reference", "video_reference", "audio_reference"): - modality = slot.split("_", 1)[0] - for i, ref in enumerate(getattr(params, slot, None) or []): - if ref.format == "path": - path = _local_path(ref.content) - if not path.exists(): - raise ValueError(f"reference file does not exist: {ref.content}") - ref.content = str(path) - continue - data = _resolve_reference(ref.content, ref.format) - ref.content = _materialize_reference( - data, - modality=modality, - ref_id=f"{request_id}_{modality}_ref_{i}", - media_storage_path=media_storage_path, - ) - ref.format = "path" - except Exception: - # The terminal on_finish hook is not wired yet (the request is never - # enqueued on failure), so reclaim any files earlier slots wrote here. - cleanup_reference_files(media_storage_path, request_id) - raise + for slot in ("image_reference", "video_reference", "audio_reference"): + modality = slot.split("_", 1)[0] + for ref in getattr(params, slot, None) or []: + data = _resolve_reference(ref.content, ref.format) + _validate_reference_payload(data, modality=modality) + ref.content = data + ref.format = "bytes" diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index b30d180eef08..ee908e59f050 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -19,7 +19,7 @@ import sys import weakref from pathlib import Path -from typing import Any, AsyncIterator, Callable, Dict, List, Literal, Optional, Union +from typing import Any, AsyncIterator, Dict, List, Literal, Optional, Union from tensorrt_llm._torch.visual_gen import DiffusionRequest, DiffusionResponse from tensorrt_llm._torch.visual_gen.executor import ( @@ -31,11 +31,7 @@ from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema, RefSlotSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs -from tensorrt_llm.visual_gen.media_refs import ( - cleanup_reference_files, - prepare_reference_slots, - resolve_media_storage_path, -) +from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots from tensorrt_llm.visual_gen.output import VisualGenOutput from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params @@ -76,7 +72,6 @@ def __init__( request_id: int, executor: "DiffusionRemoteClient", batch_size: Optional[int] = None, - on_finish: Optional[Callable[[], None]] = None, ): self.request_id = request_id self.executor = executor @@ -85,10 +80,6 @@ def __init__( self._batch_size = batch_size self._resolved = None self._finished = False - # Run once at terminal state (success/error/timeout) to reclaim the - # engine-materialized reference files for this request; None otherwise. - self._on_finish = on_finish - self._cleaned = False @property def done(self) -> bool: @@ -114,14 +105,7 @@ async def aresult(self, timeout: Optional[float] = None): self.executor.await_responses(self.request_id, timeout=timeout), self.executor._event_loop, ) - try: - response = await asyncio.wrap_future(future) - except asyncio.CancelledError: - # A caller (e.g. serve's delete-in-flight) dropped the wait. The - # worker keeps running but its output is discarded, so run the - # terminal cleanup here to reclaim the materialized references. - self._run_finish() - raise + response = await asyncio.wrap_future(future) if response is None: # Timeout before any response. Tell the executor to drop any @@ -145,12 +129,10 @@ async def aresult(self, timeout: Optional[float] = None): for _ in range(self._batch_size) ] self._finished = True - self._run_finish() return self._resolved_value() self._resolved = self._build_resolved(response) self._finished = True - self._run_finish() return self._resolved_value() def result(self, timeout: Optional[float] = None): @@ -175,16 +157,6 @@ def cancel(self): # ----- internals ----- - def _run_finish(self) -> None: - """Run the terminal cleanup callback once (idempotent, best-effort).""" - if self._cleaned or self._on_finish is None: - return - self._cleaned = True - try: - self._on_finish() - except Exception: - pass - def _build_resolved(self, response: "DiffusionResponse"): # Failure class travels on the result object, not on the public # ``VisualGenOutput`` — no new public field for an error taxonomy. @@ -466,15 +438,11 @@ def generate_async( if resolved_params.seed is None: resolved_params.seed = secrets.randbits(63) - # Resolve/materialize references to local paths here, on the coordinator, - # before the request is broadcast — the single choke point shared by serve - # and the standalone Python API. Runs synchronously so bad-media - # ``ValueError`` reaches the caller before dispatch (serve keeps its 400); - # a trusted local path passes through untouched (not materialized/cleaned). - media_storage_path = str(resolve_media_storage_path()) - prepare_reference_slots( - resolved_params, request_id=str(req_id), media_storage_path=media_storage_path - ) + # Resolve references to bytes here, on the coordinator, before the + # request is broadcast — the single choke point shared by serve and the + # standalone Python API. Runs synchronously so bad-media ``ValueError`` + # reaches the caller before dispatch (serve keeps its 400). + prepare_reference_slots(resolved_params) request = DiffusionRequest( request_id=req_id, @@ -483,12 +451,7 @@ def generate_async( ) self.executor.enqueue_requests([request]) - return VisualGenResult( - req_id, - self.executor, - batch_size=batch_size, - on_finish=lambda: cleanup_reference_files(media_storage_path, str(req_id)), - ) + return VisualGenResult(req_id, self.executor, batch_size=batch_size) @staticmethod def _atexit_shutdown(self_ref): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 149bc3be2c7c..5f0e2f96a1f9 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -25,6 +25,7 @@ """ import gc +import io import json import os from collections.abc import Generator @@ -794,31 +795,23 @@ class TestReferenceImageLoad: """ def test_truncated_image_is_a_client_error(self, tmp_path): - # Incompressible content, so half the file is genuinely half the image. + # Incompressible content, so half the payload is genuinely half the image. noise = PIL.Image.frombytes("RGB", (64, 64), os.urandom(64 * 64 * 3)) whole = tmp_path / "whole.png" noise.save(whole, format="PNG") data = whole.read_bytes() - path = tmp_path / "truncated.png" - path.write_bytes(data[: len(data) // 2]) with pytest.raises(ValueError, match="could not be decoded"): - _load_reference_image(str(path)) + _load_reference_image(data[: len(data) // 2]) - def test_unidentifiable_content_is_a_client_error(self, tmp_path): - path = tmp_path / "notreally.png" - path.write_bytes(b"not an image at all") + def test_unidentifiable_content_is_a_client_error(self): with pytest.raises(ValueError, match="could not be decoded"): - _load_reference_image(str(path)) + _load_reference_image(b"not an image at all") - def test_missing_file_is_a_client_error(self, tmp_path): - with pytest.raises(ValueError, match="could not be decoded"): - _load_reference_image(str(tmp_path / "nope.png")) - - def test_valid_image_loads(self, tmp_path): - path = tmp_path / "ok.png" - PIL.Image.new("RGB", (8, 8), (1, 2, 3)).save(path, format="PNG") - assert _load_reference_image(str(path)).size == (8, 8) + def test_valid_image_loads(self): + buffer = io.BytesIO() + PIL.Image.new("RGB", (8, 8), (1, 2, 3)).save(buffer, format="PNG") + assert _load_reference_image(buffer.getvalue()).size == (8, 8) _V2V_FIXTURE_MP4 = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" diff --git a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py index 729c86292dd3..38b3656ebe11 100644 --- a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py +++ b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py @@ -22,16 +22,35 @@ def _png_bytes() -> bytes: return buffer.getvalue() -def test_load_reference_images_accepts_pil_path_and_bytes(tmp_path) -> None: +def test_load_reference_images_accepts_pil_and_bytes(tmp_path) -> None: pil_image = PIL.Image.new("L", (64, 64), color=128) - image_path = tmp_path / "reference.png" - pil_image.save(image_path) - images = Flux2Pipeline._load_reference_images([pil_image, str(image_path), _png_bytes()]) + images = Flux2Pipeline._load_reference_images([pil_image, _png_bytes()]) - assert len(images) == 3 + assert len(images) == 2 assert all(image.mode == "RGB" for image in images) - assert [image.size for image in images] == [(64, 64), (64, 64), (64, 64)] + assert [image.size for image in images] == [(64, 64), (64, 64)] + + +def test_load_reference_images_rejects_a_path() -> None: + """References reach the worker as bytes, so a path is a type error here, + not a filesystem read.""" + with pytest.raises(ValueError, match="PIL images or encoded bytes"): + Flux2Pipeline._load_reference_images(["/tmp/nope.png"]) + + +def test_load_reference_images_drops_alpha_without_compositing() -> None: + """RGBA is converted the way diffusers does it: the alpha channel is + dropped, so a fully-transparent pixel keeps its RGB value instead of + turning white.""" + rgba = PIL.Image.new("RGBA", (8, 8), color=(10, 20, 30, 0)) + buffer = io.BytesIO() + rgba.save(buffer, format="PNG") + + (image,) = Flux2Pipeline._load_reference_images([buffer.getvalue()]) + + assert image.mode == "RGB" + assert image.getpixel((0, 0)) == (10, 20, 30) @pytest.mark.parametrize("image", [[], [object()]]) 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 5f218637bcda..91376c00f7db 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -35,11 +35,7 @@ from tensorrt_llm.serve.openai_server import _normalize_image_output from tensorrt_llm.serve.visual_gen_metrics import SERVER_TIMING_HEADER from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE -from tensorrt_llm.visual_gen.media_refs import ( - cleanup_reference_files, - prepare_reference_slots, - resolve_media_storage_path, -) +from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots from tensorrt_llm.visual_gen.output import VisualGenMetrics, VisualGenOutput from tensorrt_llm.visual_gen.params import validate_visual_gen_params @@ -215,7 +211,7 @@ def __init__( # used by tests to assert forwarded VisualGenParams fields. self.last_inputs = None self.last_params = None - # Snapshot of materialized reference-file contents at generation time, + # Snapshot of the resolved reference bytes at generation time, # captured before the route cleans them up. Keyed by stored path. self.last_ref_bytes = {} # Stand-in for the coordinator-side executor proxy. The async video @@ -260,15 +256,12 @@ def _maybe_batch(self, tensor, n): # --- VisualGen interface --- def _snapshot_refs(self, params) -> None: - # Capture materialized reference bytes before the route cleans them up, - # so tests can still assert byte-identity after the request finishes. + # Record what each slot resolved to, so tests can assert byte-identity + # against the payload the client sent. self.last_ref_bytes = {} for field in ("image_reference", "video_reference", "audio_reference"): - for ref in getattr(params, field, None) or []: - path = getattr(ref, "content", None) - if isinstance(path, str) and os.path.exists(path): - with open(path, "rb") as fh: - self.last_ref_bytes[path] = fh.read() + refs = getattr(params, field, None) or [] + self.last_ref_bytes[field] = [ref.content for ref in refs] def generate(self, inputs=None, params=None) -> VisualGenOutput: return self.generate_async(inputs=inputs, params=params).result() @@ -288,13 +281,9 @@ def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": extra_param_specs=self.executor.extra_param_specs, ref_slot_specs=self.executor.ref_slot_specs, ) - # Materialize references at the coordinator, then hand the result a - # terminal cleanup keyed on the same request id. + # Mirror the engine: resolve every reference to bytes at the coordinator. req_id = self._next_request_id() - media_storage_path = str(resolve_media_storage_path()) - prepare_reference_slots( - params, request_id=str(req_id), media_storage_path=media_storage_path - ) + prepare_reference_slots(params) self._snapshot_refs(params) n = getattr(params, "num_images_per_prompt", 1) if params else 1 return MockVisualGenResult( @@ -304,7 +293,6 @@ def generate_async(self, inputs=None, params=None) -> "MockVisualGenResult": audio=self._audio, should_fail=self._should_fail, generate_error=self._generate_error, - on_finish=lambda: cleanup_reference_files(media_storage_path, str(req_id)), ) def _next_request_id(self) -> int: @@ -359,7 +347,6 @@ def __init__( audio: Optional[torch.Tensor] = None, should_fail: bool = False, generate_error: Optional[BaseException] = None, - on_finish=None, ): self.request_id = request_id self._image = image @@ -369,19 +356,6 @@ def __init__( # Engine-side failure surfaced through the result (capacity/client), # distinct from a coordinator preflight rejection. self._generate_error = generate_error - self._on_finish = on_finish - self._cleaned = False - - def _run_finish(self): - # Terminal reference cleanup, run once (idempotent), mirroring the real - # VisualGenResult so it fires on success and failure alike. - if self._cleaned or self._on_finish is None: - return - self._cleaned = True - try: - self._on_finish() - except Exception: - pass def _resolve(self) -> VisualGenOutput: if self._generate_error is not None: @@ -400,16 +374,10 @@ def __await__(self): return self.aresult().__await__() async def aresult(self, timeout=None): - try: - return self._resolve() - finally: - self._run_finish() + return self._resolve() def result(self, timeout=None): - try: - return self._resolve() - finally: - self._run_finish() + return self._resolve() # --------------------------------------------------------------------------- @@ -1643,20 +1611,15 @@ def test_sync_video_generation_multipart_with_reference(self, video_client, tmp_ assert resp.status_code == 200 assert len(resp.content) > 0 - # image_reference is materialized to media storage and passed through as - # a MediaRef carrying the filesystem path. + # image_reference reaches the engine as a MediaRef carrying raw bytes. params = video_client.mock_gen.last_params - ref_path = params.image_reference[0].content - assert isinstance(ref_path, str) - assert ref_path.endswith("_image_ref_0") - # The materialized reference is input-only and is cleaned up once the - # request finishes, so it must not linger in media storage. - assert not os.path.exists(ref_path) + assert params.image_reference[0].content == ref_path.read_bytes() + assert params.image_reference[0].format == "bytes" def test_sync_video_generation_multipart_with_video_reference(self, video_client): - """A ``video_reference`` upload is persisted byte-identical (V2V) — the - serve never decodes video; the worker demuxes/NVDEC-decodes the stored - file. A checked-in H.264/MP4 fixture drives the boundary directly. + """A ``video_reference`` upload reaches the engine byte-identical (V2V) — + serve never decodes video; the worker demuxes/NVDEC-decodes it. A + checked-in H.264/MP4 fixture drives the boundary directly. """ payload = _V2V_FIXTURE_MP4.read_bytes() with open(_V2V_FIXTURE_MP4, "rb") as f: @@ -1673,15 +1636,12 @@ def test_sync_video_generation_multipart_with_video_reference(self, video_client assert resp.status_code == 200 assert len(resp.content) > 0 - # Video conditioning arrives as a MediaRef holding a stored path; no - # image_reference is set, and the encoded bytes were persisted - # byte-identical (snapshotted at generation time, since the route cleans - # the reference up once the request finishes). + # Video conditioning reaches the engine as raw bytes, byte-identical to + # the upload, and no image_reference is set. params = video_client.mock_gen.last_params assert params.image_reference is None - ref_path = params.video_reference[0].content - assert video_client.mock_gen.last_ref_bytes[ref_path] == payload - assert not os.path.exists(ref_path) + assert params.video_reference[0].format == "bytes" + assert video_client.mock_gen.last_ref_bytes["video_reference"] == [payload] def test_sync_video_generation_undecodable_reference_400(self, video_client): """Content matching no image or video container signature is rejected diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 23920b33cd50..ae0f59c8c6c7 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -28,20 +28,22 @@ parse_visual_gen_params, ) from tensorrt_llm.visual_gen import VisualGenParams -from tensorrt_llm.visual_gen.media_refs import cleanup_reference_files, prepare_reference_slots +from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots -def _parse_and_prepare(request, generator, request_id, media_storage_path): - """Run the production reference flow: serve transport, then engine materialize. +def _parse_and_prepare(request, generator): + """Run the production reference flow: serve transport, then engine resolve. ``parse_visual_gen_params`` only normalizes transport (upload -> bytes, - strings pass through); resolution + materialization happen at the engine - choke point, so tests that assert on stored paths drive both stages. + strings pass through with their declared format); resolution to bytes + happens at the engine choke point, so tests that assert on the resolved + payload drive both stages. """ params = parse_visual_gen_params(request, generator) - prepare_reference_slots(params, request_id=request_id, media_storage_path=media_storage_path) + prepare_reference_slots(params) return params + pytestmark = pytest.mark.cpu_only @@ -276,12 +278,12 @@ def test_video_does_not_carry_n(self): # ============================================================================= -# reference materialization +# reference resolution # ============================================================================= -class TestInputReferenceMaterialization: - def test_base64_image_reference_written_to_disk(self, tmp_path): +class TestInputReferenceResolution: + def test_base64_image_reference_resolves_to_bytes(self, tmp_path): generator = _StubVisualGen() img = Image.new("RGB", (4, 4), (10, 20, 30)) buf = BytesIO() @@ -290,14 +292,11 @@ def test_base64_image_reference_written_to_disk(self, tmp_path): request = VideoGenerationRequest( prompt="x", image_reference={"content": b64, "format": "base64"} ) - params = _parse_and_prepare(request, generator, "vid-1", str(tmp_path)) + params = _parse_and_prepare(request, generator) assert len(params.image_reference) == 1 ref_path = params.image_reference[0].content - assert str(ref_path).endswith("vid-1_image_ref_0") - # The decoded image is identical to what we passed in. - with open(ref_path, "rb") as f: - decoded = Image.open(f).convert("RGB") - assert decoded.size == (4, 4) + assert ref_path == buf.getvalue() + assert params.image_reference[0].format == "bytes" def test_image_reference_role_and_list(self, tmp_path): generator = _StubVisualGen() @@ -311,22 +310,9 @@ def test_image_reference_role_and_list(self, tmp_path): {"content": b64, "format": "base64", "role": "last_frame"}, ], ) - params = _parse_and_prepare(request, generator, "vid-r", str(tmp_path)) + params = _parse_and_prepare(request, generator) assert [r.role for r in params.image_reference] == [None, "last_frame"] - paths = [r.content for r in params.image_reference] - assert len(set(paths)) == 2 # unique file per index - - def test_missing_media_storage_path_raises(self): - generator = _StubVisualGen() - img = Image.new("RGB", (2, 2)) - buf = BytesIO() - img.save(buf, format="PNG") - b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest( - prompt="x", image_reference={"content": b64, "format": "base64"} - ) - with pytest.raises(ValueError, match="media_storage_path"): - _parse_and_prepare(request, generator, "vid-2", None) + assert [r.content for r in params.image_reference] == [buf.getvalue()] * 2 _TEST_DATA = Path(__file__).parent / "test_data" @@ -334,51 +320,38 @@ def test_missing_media_storage_path_raises(self): def _mp4_bytes() -> bytes: """9-frame H.264-in-MP4 fixture (provenance: test_data/README.md).""" return ( - TestInputReferenceMaterialization._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.mp4" + TestInputReferenceResolution._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.mp4" ).read_bytes() @staticmethod def _avi_bytes() -> bytes: """Same 9 frames as H.264-in-AVI (provenance: test_data/README.md).""" return ( - TestInputReferenceMaterialization._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" + TestInputReferenceResolution._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" ).read_bytes() - def test_multipart_avi_video_reference_written_to_disk(self, tmp_path): + def test_multipart_avi_video_reference_resolves_to_bytes(self, tmp_path): # The AVI container survives the boundary and is persisted as untouched # encoded bytes for the worker to demux. generator = _StubVisualGen() payload = self._avi_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.avi") request = VideoGenerationRequest(prompt="x", video_reference=upload) - params = _parse_and_prepare(request, generator, "vid-avi", str(tmp_path)) + params = _parse_and_prepare(request, generator) assert params.image_reference is None - assert Path(params.video_reference[0].content).read_bytes() == payload + assert params.video_reference[0].content == payload - def test_multipart_mp4_video_reference_written_to_disk(self, tmp_path): + def test_multipart_mp4_video_reference_resolves_to_bytes(self, tmp_path): generator = _StubVisualGen() payload = self._mp4_bytes() upload = UploadFile(file=BytesIO(payload), filename="clip.mp4") request = VideoGenerationRequest(prompt="x", video_reference=upload) - params = _parse_and_prepare(request, generator, "vid-3", str(tmp_path)) + params = _parse_and_prepare(request, generator) # Encoded payload is persisted byte-identical — the boundary never # decodes video; the worker demuxes/NVDEC-decodes the conditioning # window from the stored file. assert params.image_reference is None - vpath = params.video_reference[0].content - assert str(vpath).endswith("vid-3_video_ref_0") - assert Path(vpath).read_bytes() == payload - - def test_video_reference_needs_media_storage(self): - # Video references now persist to disk (the worker reads the path), so - # a storage path is required just like image references. - generator = _StubVisualGen() - b64 = base64.b64encode(self._mp4_bytes()).decode() - request = VideoGenerationRequest( - prompt="x", video_reference={"content": b64, "format": "base64"} - ) - with pytest.raises(ValueError, match="media_storage_path"): - _parse_and_prepare(request, generator, "vid-9", None) + assert params.video_reference[0].content == payload def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): # The deprecated single input_reference is sniff-routed to the typed slot. @@ -393,8 +366,6 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): prompt="x", input_reference=img_b64, input_reference_format="base64" ), generator, - "vid-i", - str(tmp_path), ) assert len(p.image_reference) == 1 and p.video_reference is None @@ -403,8 +374,6 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): prompt="x", input_reference=vid_b64, input_reference_format="base64" ), generator, - "vid-v", - str(tmp_path), ) assert len(p.video_reference) == 1 and p.image_reference is None @@ -423,13 +392,11 @@ def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): input_reference_format="base64", ), generator, - "vid-x", - str(tmp_path), ) assert len(p.image_reference) == 1 assert p.video_reference is None # input_reference video dropped - def test_base64_video_reference_written_to_disk(self, tmp_path): + def test_base64_video_reference_resolves_to_bytes(self, tmp_path): # The JSON/base64 path carries video even though it has no content-type # or filename; modality is declared by the field name. generator = _StubVisualGen() @@ -438,9 +405,9 @@ def test_base64_video_reference_written_to_disk(self, tmp_path): request = VideoGenerationRequest( prompt="x", video_reference={"content": b64, "format": "base64"} ) - params = _parse_and_prepare(request, generator, "vid-4", str(tmp_path)) + params = _parse_and_prepare(request, generator) assert params.image_reference is None - assert Path(params.video_reference[0].content).read_bytes() == payload + assert params.video_reference[0].content == payload def test_video_reference_survives_real_specs(self, tmp_path): """With the real cosmos3 specs loaded, the encoded payload is persisted @@ -454,10 +421,10 @@ def test_video_reference_survives_real_specs(self, tmp_path): request = VideoGenerationRequest( prompt="x", video_reference={"content": b64, "format": "base64"} ) - params = _parse_and_prepare(request, generator, "vid-10", str(tmp_path)) - assert Path(params.video_reference[0].content).read_bytes() == payload + params = _parse_and_prepare(request, generator) + assert params.video_reference[0].content == payload - def test_multipart_image_reference_written_to_disk(self, tmp_path): + def test_multipart_image_reference_resolves_to_bytes(self, tmp_path): # JPEG upload routed by field name to image_reference. The stored file # has no type-suffix (PIL identifies by content, not name). generator = _StubVisualGen() @@ -467,9 +434,9 @@ def test_multipart_image_reference_written_to_disk(self, tmp_path): buf.seek(0) upload = UploadFile(file=buf, filename="ref.jpg") request = VideoGenerationRequest(prompt="x", image_reference=upload) - params = _parse_and_prepare(request, generator, "vid-5", str(tmp_path)) + params = _parse_and_prepare(request, generator) assert params.extra_params is None - assert str(params.image_reference[0].content).endswith("vid-5_image_ref_0") + assert isinstance(params.image_reference[0].content, bytes) def test_wrong_modality_content_raises(self, tmp_path): # The field name declares modality; mismatched content is a client error. @@ -484,8 +451,6 @@ def test_wrong_modality_content_raises(self, tmp_path): prompt="x", video_reference={"content": img_b64, "format": "base64"} ), generator, - "vid-m1", - str(tmp_path), ) with pytest.raises(ValueError, match="image_reference is not a recognized image"): _parse_and_prepare( @@ -493,10 +458,8 @@ def test_wrong_modality_content_raises(self, tmp_path): prompt="x", image_reference={"content": vid_b64, "format": "base64"} ), generator, - "vid-m2", - str(tmp_path), ) - assert list(tmp_path.iterdir()) == [] + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() @@ -505,9 +468,9 @@ def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): prompt="x", image_reference={"content": b64, "format": "base64"} ) with pytest.raises(ValueError, match="not a recognized image"): - _parse_and_prepare(request, generator, "vid-6", str(tmp_path)) + _parse_and_prepare(request, generator) # Classification runs on the bytes; rejected content never touches disk. - assert list(tmp_path.iterdir()) == [] + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): generator = _StubVisualGen() @@ -517,8 +480,8 @@ def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): prompt="x", image_reference={"content": "ABC", "format": "base64"} ) with pytest.raises(ValueError, match="not valid base64"): - _parse_and_prepare(request, generator, "vid-7", str(tmp_path)) - assert list(tmp_path.iterdir()) == [] + _parse_and_prepare(request, generator) + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_upload_stream_failure_cleans_up_tmp(self, tmp_path): generator = _StubVisualGen() @@ -531,9 +494,9 @@ def read(self, *args, **kwargs): request = VideoGenerationRequest(prompt="x", video_reference=upload) # I/O failures keep their server-error semantics (no 400 masking) … with pytest.raises(OSError, match="client went away"): - _parse_and_prepare(request, generator, "vid-8", str(tmp_path)) + _parse_and_prepare(request, generator) # … and the payload read fails before any file is written, so nothing leaks. - assert list(tmp_path.iterdir()) == [] + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_multi_reference_partial_failure_cleans_up(self, tmp_path): # A later item's rejection removes the files earlier items already wrote, @@ -551,40 +514,36 @@ def test_multi_reference_partial_failure_cleans_up(self, tmp_path): ], ) with pytest.raises(ValueError, match="not a recognized image"): - _parse_and_prepare(request, generator, "vid-11", str(tmp_path)) - assert list(tmp_path.iterdir()) == [] + _parse_and_prepare(request, generator) + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - def test_file_uri_image_reference_passthrough(self, tmp_path): - # format="path" also accepts a file:// URI: normalized to a plain path - # and passed through untouched — not copied into media storage. + def test_file_uri_image_reference_is_read(self, tmp_path): + # format="path" also accepts a file:// URI, normalized before the read. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (7, 8, 9)).save(src, format="PNG") - store = tmp_path / "store" - store.mkdir() request = VideoGenerationRequest( prompt="x", image_reference={"content": src.as_uri(), "format": "path"} ) - params = _parse_and_prepare(request, generator, "vid-file", str(store)) - assert params.image_reference[0].content == str(src) - assert list(store.iterdir()) == [] + params = _parse_and_prepare(request, generator) + assert params.image_reference[0].content == src.read_bytes() + assert params.image_reference[0].format == "bytes" - def test_bare_path_image_reference_passthrough(self, tmp_path): - # A bare local path under format="path" is passed through unchanged. + def test_bare_path_image_reference_is_read(self, tmp_path): + # A path is read at the coordinator, so the worker needs no shared + # filesystem to see what the client named. generator = _StubVisualGen() src = tmp_path / "ref.png" Image.new("RGB", (4, 4), (11, 22, 33)).save(src, format="PNG") - store = tmp_path / "store" - store.mkdir() request = VideoGenerationRequest( prompt="x", image_reference={"content": str(src), "format": "path"} ) - params = _parse_and_prepare(request, generator, "vid-bare", str(store)) - assert params.image_reference[0].content == str(src) - assert list(store.iterdir()) == [] + params = _parse_and_prepare(request, generator) + assert params.image_reference[0].content == src.read_bytes() + assert params.image_reference[0].format == "bytes" - def test_http_url_image_reference_fetched_and_materialized(self, tmp_path, monkeypatch): - # An http(s) reference is fetched through the guarded loader, then stored. + def test_http_url_image_reference_is_fetched(self, tmp_path, monkeypatch): + # An http(s) reference is fetched through the guarded loader. generator = _StubVisualGen() buf = BytesIO() Image.new("RGB", (4, 4)).save(buf, format="PNG") @@ -602,8 +561,8 @@ def __init__(self, content): prompt="x", image_reference={"content": "https://example.com/a.png", "format": "url"}, ) - params = _parse_and_prepare(request, generator, "vid-url", str(tmp_path)) - assert Path(params.image_reference[0].content).read_bytes() == png + params = _parse_and_prepare(request, generator) + assert params.image_reference[0].content == png def test_http_url_fetch_failure_is_client_error(self, tmp_path, monkeypatch): # A blocked/failed fetch (e.g. SSRF guard) is a client 400, not a 500, @@ -618,8 +577,8 @@ def _blocked(url, **kwargs): prompt="x", image_reference={"content": "http://10.0.0.1/a.png", "format": "url"} ) with pytest.raises(ValueError, match="reference URL could not be fetched"): - _parse_and_prepare(request, generator, "vid-ssrf", str(tmp_path)) - assert list(tmp_path.iterdir()) == [] + _parse_and_prepare(request, generator) + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_missing_file_uri_is_client_error(self, tmp_path): # A file:// path that does not exist is a client 400, not a server 500. @@ -628,9 +587,9 @@ def test_missing_file_uri_is_client_error(self, tmp_path): request = VideoGenerationRequest( prompt="x", image_reference={"content": missing, "format": "path"} ) - with pytest.raises(ValueError, match="reference file does not exist"): - _parse_and_prepare(request, generator, "vid-nf", str(tmp_path)) - assert list(tmp_path.iterdir()) == [] + with pytest.raises(ValueError, match="reference file could not be read"): + _parse_and_prepare(request, generator) + assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk def test_bare_reference_string_is_rejected(self): # The bare-string shorthand is gone: a reference must declare its wire @@ -668,8 +627,8 @@ def test_sniff_media_kind(self): Image.new("RGB", (2, 2)).save(jpg, format="JPEG") assert sniff_media_kind(png.getvalue()) == "image" assert sniff_media_kind(jpg.getvalue()) == "image" - assert sniff_media_kind(TestInputReferenceMaterialization._mp4_bytes()) == "video" - assert sniff_media_kind(TestInputReferenceMaterialization._avi_bytes()) == "video" + assert sniff_media_kind(TestInputReferenceResolution._mp4_bytes()) == "video" + assert sniff_media_kind(TestInputReferenceResolution._avi_bytes()) == "video" assert sniff_media_kind(b"plain text, not media") is None assert sniff_media_kind(b"") is None # RIFF alone is not AVI (e.g. WAV audio is RIFF too). @@ -768,7 +727,7 @@ def test_heif_reference_rejected_with_actionable_message(self): image_reference={"content": base64.b64encode(heic).decode(), "format": "base64"}, ) with pytest.raises(ValueError, match="HEIF/AVIF"): - _parse_and_prepare(request, generator, "vid-heic", None) + _parse_and_prepare(request, generator) def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): """The boundary routes on signature and never decodes. @@ -791,8 +750,8 @@ def test_truncated_image_reference_is_routed_not_decoded(self, tmp_path): prompt="x", image_reference={"content": base64.b64encode(truncated).decode(), "format": "base64"}, ) - params = _parse_and_prepare(request, generator, "vid-12", str(tmp_path)) - assert Path(params.image_reference[0].content).read_bytes() == truncated + params = _parse_and_prepare(request, generator) + assert params.image_reference[0].content == truncated # ============================================================================= @@ -876,7 +835,7 @@ def test_base64_extra_param_reaches_the_pipeline_as_bytes(self): prompt="storm", extra_params={"video": base64.b64encode(b"\x00mp4").decode()}, ) - params = parse_visual_gen_params(request, "id-b64", self._generator()) + params = parse_visual_gen_params(request, self._generator()) assert params.extra_params["video"] == b"\x00mp4" def test_nested_control_reaches_the_pipeline_as_bytes(self): @@ -884,143 +843,88 @@ def test_nested_control_reaches_the_pipeline_as_bytes(self): prompt="storm", extra_params={"edge": {"control": base64.b64encode(b"\x00ctrl").decode()}}, ) - params = parse_visual_gen_params(request, "id-nested", self._generator()) + params = parse_visual_gen_params(request, self._generator()) assert params.extra_params["edge"]["control"] == b"\x00ctrl" def test_non_media_params_are_untouched(self): request = VideoGenerationRequest( prompt="storm", extra_params={"resolution": "720", "edge": True} ) - params = parse_visual_gen_params(request, "id-plain", self._generator()) + params = parse_visual_gen_params(request, self._generator()) assert params.extra_params["resolution"] == "720" assert params.extra_params["edge"] is True def test_malformed_base64_is_a_client_error(self): request = VideoGenerationRequest(prompt="storm", extra_params={"video": "not!b64!"}) with pytest.raises(ValueError, match="not valid base64"): - parse_visual_gen_params(request, "id-bad", self._generator()) - - -class TestCleanupReferenceFiles: - """The reference-file reclaim helper keyed on the request id prefix.""" - - def test_removes_only_this_request_ref_files(self, tmp_path): - vid = "video_abc123" - (tmp_path / f"{vid}_image_ref_0").write_bytes(b"a") - (tmp_path / f"{vid}_video_ref_1").write_bytes(b"b") - (tmp_path / f"{vid}_input_ref").write_bytes(b"c") # deprecated alias - (tmp_path / f"{vid}_0.mp4").write_bytes(b"out") # output — keep - (tmp_path / "video_other_image_ref_0").write_bytes(b"d") # other id — keep - cleanup_reference_files(str(tmp_path), vid) - assert sorted(p.name for p in tmp_path.iterdir()) == [ - f"{vid}_0.mp4", - "video_other_image_ref_0", - ] - - def test_none_storage_is_noop(self): - cleanup_reference_files(None, "video_x") # no raise - - def test_missing_files_are_ignored(self, tmp_path): - cleanup_reference_files(str(tmp_path), "video_absent") # no raise + parse_visual_gen_params(request, self._generator()) class TestPrepareReferenceSlots: - """Engine-side reference choke point: passthrough local paths, materialize the rest.""" - - def test_local_path_passthrough_not_materialized_or_cleaned(self, tmp_path): - from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + """The engine choke point: every declared form resolves to raw bytes.""" - src = tmp_path / "user.png" - Image.new("RGB", (4, 4)).save(src, format="PNG") - store = tmp_path / "store" - store.mkdir() - params = VisualGenParams(image_reference=MediaRef(content=str(src), format="path")) - prepare_reference_slots(params, request_id="req1", media_storage_path=str(store)) - assert params.image_reference[0].content == str(src) # unchanged - assert list(store.iterdir()) == [] # nothing materialized - cleanup_reference_files(str(store), "req1") - assert src.exists() # user file untouched by cleanup - - def test_bytes_materialized_to_storage_and_cleaned(self, tmp_path): + def test_every_format_resolves_to_the_same_bytes(self, tmp_path): + """path / base64 / bytes all name the same payload, so all three must + land on byte-identical content with format rewritten to "bytes".""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots buf = BytesIO() - Image.new("RGB", (4, 4)).save(buf, format="PNG") + Image.new("RGB", (4, 4), (7, 8, 9)).save(buf, format="PNG") png = buf.getvalue() - store = tmp_path / "store" - store.mkdir() - params = VisualGenParams(image_reference=MediaRef(content=png, format="bytes")) - prepare_reference_slots(params, request_id="req2", media_storage_path=str(store)) - path = params.image_reference[0].content - assert path == str(store / "req2_image_ref_0") - assert Path(path).read_bytes() == png - cleanup_reference_files(str(store), "req2") - assert not Path(path).exists() - - def test_materialized_reference_format_is_rewritten_to_path(self, tmp_path): - """``content`` and ``format`` are rewritten together. - - The mutated params object is what gets broadcast to the workers, so a - stale ``base64`` format would reach a worker holding a filesystem path. - """ + src = tmp_path / "ref.png" + src.write_bytes(png) + + for ref in ( + MediaRef(content=str(src), format="path"), + MediaRef(content=base64.b64encode(png).decode(), format="base64"), + MediaRef(content=png, format="bytes"), + ): + params = VisualGenParams(image_reference=ref) + prepare_reference_slots(params) + assert params.image_reference[0].content == png + assert params.image_reference[0].format == "bytes" + + def test_nothing_is_resolves_to_bytes(self, tmp_path, monkeypatch): + """References never touch the filesystem, so a worker needs no shared + filesystem to read what the coordinator resolved.""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) buf = BytesIO() Image.new("RGB", (4, 4)).save(buf, format="PNG") - png = buf.getvalue() - src = tmp_path / "user.png" - src.write_bytes(png) - store = tmp_path / "store" - store.mkdir() params = VisualGenParams( - image_reference=[ - MediaRef(content=base64.b64encode(png).decode(), format="base64"), - MediaRef(content=png, format="bytes"), - MediaRef(content=str(src), format="path"), - ] + image_reference=MediaRef( + content=base64.b64encode(buf.getvalue()).decode(), format="base64" + ) ) - prepare_reference_slots(params, request_id="req3", media_storage_path=str(store)) - assert [r.format for r in params.image_reference] == ["path", "path", "path"] - assert all(Path(r.content).read_bytes() == png for r in params.image_reference) + prepare_reference_slots(params) + assert list(tmp_path.iterdir()) == [] - def test_path_format_on_missing_file_raises(self, tmp_path): - """A mistyped path stays a path error — it is never retried as base64.""" + def test_wrong_modality_is_rejected(self, tmp_path): + """Content is validated against the slot's modality before dispatch.""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots - store = tmp_path / "store" - store.mkdir() + buf = BytesIO() + Image.new("RGB", (2, 2)).save(buf, format="PNG") params = VisualGenParams( - image_reference=MediaRef(content=str(tmp_path / "absent.png"), format="path") + image_reference=MediaRef(content=buf.getvalue(), format="bytes"), + video_reference=MediaRef(content=buf.getvalue(), format="bytes"), ) - with pytest.raises(ValueError, match="reference file does not exist"): - prepare_reference_slots(params, request_id="req4", media_storage_path=str(store)) - assert list(store.iterdir()) == [] + with pytest.raises(ValueError, match="video_reference is not a recognized"): + prepare_reference_slots(params) - def test_base64_format_on_a_filesystem_path_raises(self, tmp_path): - """A path sent as ``base64`` is a decode error — it is never read from disk.""" + def test_missing_path_is_a_client_error(self, tmp_path): from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots - src = tmp_path / "user.png" - Image.new("RGB", (4, 4)).save(src, format="PNG") - store = tmp_path / "store" - store.mkdir() - params = VisualGenParams(image_reference=MediaRef(content=str(src), format="base64")) - with pytest.raises(ValueError, match="not valid base64"): - prepare_reference_slots(params, request_id="req5", media_storage_path=str(store)) - assert list(store.iterdir()) == [] - - def test_resolve_media_storage_path(self, tmp_path, monkeypatch): - from tensorrt_llm.visual_gen.media_refs import resolve_media_storage_path - - target = tmp_path / "ms" - monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(target)) - resolved = resolve_media_storage_path() - assert resolved == target and target.is_dir() + params = VisualGenParams( + image_reference=MediaRef(content=str(tmp_path / "nope.png"), format="path") + ) + with pytest.raises(ValueError, match="could not be read"): + prepare_reference_slots(params) class TestResolveReference: From bb7897cdfa1f9831aab7825b700fd738aa043017 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 19 Aug 2026 07:52:04 -0700 Subject: [PATCH 25/61] [TRTLLM-15277][perf] VisualGen: hand reference payloads to rank0 through shared memory MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Reference bytes used to ride inside the request pickle across the coordinator -> rank0 hop, so every byte was copied to serialize, copied into the socket, copied out, and copied again to deserialize. `DiffusionRequest.refs_to_handles()` now moves each payload into a shared-memory block on the coordinator and leaves only the handle dict in the request; rank0 calls `refs_to_bytes()` the moment it takes the request off the queue. This mirrors what the LLM path already does for multimodal tensors (`llm.py` `to_handle` -> `base_worker.py` `to_tensor`) and what VisualGen already does in the response direction (`PipelineOutput.to_handle`/`to_tensor`). The restore happens before `broadcast_object_list`, so the rank0 -> N-rank hop keeps broadcasting raw bytes. That is a correctness constraint, not an omission: a shared-tensor handle is consumed exactly once, so minting one handle for N ranks would free the block N-1 times. The LLM path draws the same line — it has the full `RequestBroadcaster` machinery available and still broadcasts materialized tensors. `VisualGenParams` and `MediaRef` are untouched; the handle lives on the internal `DiffusionRequest`, so this is invisible to the public API and costs nothing for text-only requests, which carry no handle at all. Measured end to end over the real `ZeroMqQueue` across a process boundary, including the consumer-side restore back to bytes: 0.5 MB 1.58 -> 0.86 ms (1.85x), 2 MB 4.10 -> 1.79 ms (2.30x), 8 MB 21.6 -> 6.45 ms (3.35x), 32 MB 135.9 -> 39.4 ms (3.45x). Verified on real weights that output is unchanged and `format="path"` still equals `format="base64"` bitwise: Wan2.2-I2V-A14B, Cosmos3-Super I2V, Cosmos3-Super V2V and FLUX.2-dev on one GPU, plus Cosmos3-Super V2V and Wan2.2-I2V-A14B at `ulysses_size=2` to exercise the broadcast hop. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 38 +++++++++ tensorrt_llm/visual_gen/visual_gen.py | 4 + .../visual_gen/test_visual_gen_utils.py | 79 +++++++++++++++++++ 3 files changed, 121 insertions(+) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 6dbdce267008..ec65ab0cda4e 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -15,6 +15,7 @@ import torch.multiprocessing as mp import zmq +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.executor.ipc import ZeroMqQueue @@ -257,6 +258,41 @@ class DiffusionRequest: prompt: List[str] params: Optional["VisualGenParams"] = None prepared_inputs: Dict[str, Any] = field(default_factory=dict, repr=False) + # Set only between the two ends of the coordinator -> rank0 hop; see + # ``refs_to_handles``. + ref_handles: Optional[List[Dict[str, Any]]] = field(default=None, repr=False) + + def refs_to_handles(self) -> None: + """Move reference payloads into shared memory, in place (producer side). + + Only the coordinator -> rank0 hop travels as handles: rank0 restores the + bytes before broadcasting, because a shared-tensor handle is consumed + exactly once and minting one for N ranks would free the block N-1 times. + """ + if self.params is None: + return + handles = [] + for slot in ("image_reference", "video_reference", "audio_reference"): + for index, ref in enumerate(getattr(self.params, slot, None) or []): + # bytearray() because the shared storage must be writable. + buffer = torch.frombuffer(bytearray(ref.content), dtype=torch.uint8) + handles.append( + { + "slot": slot, + "index": index, + "handle": SharedTensorContainer.from_tensor(buffer).dump_to_dict(), + } + ) + ref.content = b"" + self.ref_handles = handles or None + + def refs_to_bytes(self) -> None: + """Restore reference payloads from shared memory, in place (consumer side).""" + for entry in self.ref_handles or []: + ref = getattr(self.params, entry["slot"])[entry["index"]] + container = SharedTensorContainer.from_dict(entry["handle"]) + ref.content = container.get_local_view().numpy().tobytes() + self.ref_handles = None @dataclass @@ -396,6 +432,8 @@ def serve_forever(self): req = None if self.rank == 0: req = self.requests_ipc.get() + if req is not None: + req.refs_to_bytes() logger.info(f"Worker {self.device_id}: Request available") # Broadcast to all ranks. ``req.params.seed`` is already a diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ee908e59f050..e4c64a749a5e 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -449,6 +449,10 @@ def generate_async( prompt=prompt, params=resolved_params, ) + # Hand the reference payloads to rank0 through shared memory instead of + # through the request pickle, which copies every reference byte to + # cross one process boundary. + request.refs_to_handles() self.executor.enqueue_requests([request]) return VisualGenResult(req_id, self.executor, batch_size=batch_size) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index ae0f59c8c6c7..37dde47aa6ac 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -12,6 +12,8 @@ from __future__ import annotations import base64 +import os +import pickle from io import BytesIO from pathlib import Path from typing import Any, Dict, Optional @@ -1005,3 +1007,80 @@ def test_unknown_format_is_rejected(self): with pytest.raises(ValueError, match="unsupported reference format"): _resolve_reference("a.png", "filepath") + + +# ============================================================================= +# reference transport (coordinator -> rank0) +# ============================================================================= + + +class TestReferenceHandleTransport: + """References cross the coordinator -> rank0 hop as shared-memory handles. + + The payload must survive the round trip byte-identically and must not ride + the request pickle, which is the whole point of the handle. + """ + + @staticmethod + def _request(*payloads: bytes): + from tensorrt_llm._torch.visual_gen import DiffusionRequest + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + + params = VisualGenParams( + image_reference=[MediaRef(content=p, format="bytes") for p in payloads] + ) + return DiffusionRequest(request_id=1, prompt=["x"], params=params) + + def test_round_trip_is_byte_identical(self): + payloads = (b"\x89PNG\r\n\x1a\n" + os.urandom(4096), os.urandom(1024)) + req = self._request(*payloads) + + req.refs_to_handles() + req.refs_to_bytes() + + assert tuple(r.content for r in req.params.image_reference) == payloads + assert all(r.format == "bytes" for r in req.params.image_reference) + + def test_payload_leaves_the_request_pickle(self): + """The handle is the transport, so the bytes must not also be pickled — + otherwise the hop still pays for a full copy of every reference.""" + payload = os.urandom(256 * 1024) + req = self._request(payload) + before = len(pickle.dumps(req)) + + req.refs_to_handles() + after = len(pickle.dumps(req)) + + assert after < before - len(payload) // 2 + assert payload not in pickle.dumps(req) + + def test_survives_a_real_pickle_round_trip(self): + """A handle is only useful if it still resolves after being serialized + and rebuilt, which is what the IPC queue does to it.""" + payload = os.urandom(8192) + req = self._request(payload) + req.refs_to_handles() + + received = pickle.loads(pickle.dumps(req)) + received.refs_to_bytes() + + assert received.params.image_reference[0].content == payload + + def test_no_references_costs_nothing(self): + """T2V/T2I requests carry no handle, so the hop is untouched for them.""" + from tensorrt_llm._torch.visual_gen import DiffusionRequest + from tensorrt_llm.visual_gen import VisualGenParams + + req = DiffusionRequest(request_id=1, prompt=["x"], params=VisualGenParams()) + req.refs_to_handles() + assert req.ref_handles is None + + def test_restore_is_idempotent(self): + """rank0 restores unconditionally; a second call must not re-consume a + handle that has already been resolved.""" + payload = os.urandom(2048) + req = self._request(payload) + req.refs_to_handles() + req.refs_to_bytes() + req.refs_to_bytes() + assert req.params.image_reference[0].content == payload From a1e33c98789c40e072d9098cd1ca0c3ee8976f16 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 19 Aug 2026 11:01:58 -0700 Subject: [PATCH 26/61] [TRTLLM-15277][perf] VisualGen: keep reference bytes out of the request broadcast MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `broadcast_object_list` serializes whatever it is handed into a uint8 tensor before the collective runs, so carrying reference payloads inside the request object meant every reference byte was copied on rank0 for each request. That copy, not the collective, dominated the hop: for a 32 MB reference at 8 ranks it was 131 ms against 2.9 ms for the pre-bytes path, which only ever broadcast a filename. Two changes. Single-rank runs skip the hop entirely — there is no peer, and `broadcast_object_list` has no `world_size == 1` early exit, so it would serialize the whole request before discovering that. For multi-rank runs, `refs_detach()` lifts the payloads out of the request and records their sizes; the now-small object goes over `broadcast_object_list` and each payload follows as a raw uint8 tensor that every rank sizes from `ref_sizes`. Measured on the broadcast phase, same run, same machine state: | ranks | reference | before | after | |---|---|---|---| | 1 | 8 MB | 6.6 ms | 0 (skipped) | | 1 | 32 MB | 37.8 ms | 0 (skipped) | | 2 | 32 MB | 79.3 ms | 15.0 ms | | 4 | 32 MB | 99.4 ms | 26.1 ms | | 8 | 8 MB | 37.6 ms | 20.4 ms | | 8 | 32 MB | 130.9 ms | 71.4 ms | A raw CUDA/NCCL broadcast measured faster still at 32 MB (57.8 ms at 8 ranks), but it needs the payload staged on and off the device and is slower than CPU below 8 ranks, so it is not worth the extra path here. Verified on real weights that output is unchanged: Wan2.2-I2V-A14B and Cosmos3-Super V2V pass both at one rank (the skipped hop) and at `ulysses_size=2` (the split broadcast), and `test_trtllm_serve_e2e.py::TestWanImageToVideo` passes. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 65 +++++++++++++++++-- .../visual_gen/test_visual_gen_utils.py | 53 +++++++++++++++ 2 files changed, 114 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index ec65ab0cda4e..5f42d504cb45 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -261,6 +261,9 @@ class DiffusionRequest: # Set only between the two ends of the coordinator -> rank0 hop; see # ``refs_to_handles``. ref_handles: Optional[List[Dict[str, Any]]] = field(default=None, repr=False) + # Set only while the request is in flight on the rank0 -> N-rank hop; see + # ``refs_detach``. + ref_sizes: Optional[List[int]] = field(default=None, repr=False) def refs_to_handles(self) -> None: """Move reference payloads into shared memory, in place (producer side). @@ -294,6 +297,31 @@ def refs_to_bytes(self) -> None: ref.content = container.get_local_view().numpy().tobytes() self.ref_handles = None + def _refs(self): + for slot in ("image_reference", "video_reference", "audio_reference"): + yield from getattr(self.params, slot, None) or [] + + def refs_detach(self) -> List[bytes]: + """Take the reference payloads out of the request and record their sizes. + + Broadcasting them inside the request object would make + ``broadcast_object_list`` serialize every reference byte into a tensor + first, which costs more than the collective that follows. + """ + if self.params is None: + return [] + payloads = [ref.content for ref in self._refs()] + for ref in self._refs(): + ref.content = b"" + self.ref_sizes = [len(p) for p in payloads] + return payloads + + def refs_attach(self, payloads: List[bytes]) -> None: + """Put the separately broadcast payloads back, in place.""" + for ref, payload in zip(self._refs(), payloads): + ref.content = payload + self.ref_sizes = None + @dataclass class DiffusionResponse: @@ -347,6 +375,7 @@ def __init__( self.pipeline = None # initialized in _load_pipeline self.requests_ipc = None self.rank = dist.get_rank() + self.world_size = dist.get_world_size() self.response_queue = queue.Queue() self.sender_thread = None @@ -426,6 +455,33 @@ def _load_pipeline(self): ) ) + def _broadcast_request(self, req: Optional[DiffusionRequest]) -> Optional[DiffusionRequest]: + """Send one request from rank0 to every rank. + + Reference payloads ride as raw uint8 tensors alongside the request + rather than inside it, because ``broadcast_object_list`` pickles the + object into a tensor first and that copy dominates the collective. + """ + payloads = req.refs_detach() if self.rank == 0 and req is not None else [] + obj_list = [req] + dist.broadcast_object_list(obj_list, src=0) + req = obj_list[0] + if req is None: + return None + + if self.rank == 0: + # bytearray() because a tensor over an immutable buffer is read-only. + buffers = [torch.frombuffer(bytearray(p), dtype=torch.uint8) for p in payloads] + else: + buffers = [torch.empty(n, dtype=torch.uint8) for n in req.ref_sizes or []] + for buffer in buffers: + dist.broadcast(buffer, src=0) + + if self.rank != 0: + payloads = [b.numpy().tobytes() for b in buffers] + req.refs_attach(payloads) + return req + def serve_forever(self): """Main execution loop.""" while True: @@ -439,10 +495,11 @@ def serve_forever(self): # Broadcast to all ranks. ``req.params.seed`` is already a # concrete int — resolved once on the coordinator process at # :meth:`VisualGen.generate_async` entry — so the broadcast - # propagates the same value to every rank. - obj_list = [req] - dist.broadcast_object_list(obj_list, src=0) - req = obj_list[0] + # propagates the same value to every rank. Single-rank runs skip + # it: there is no peer, and the object broadcast would still + # serialize the whole request to a tensor before discovering that. + if self.world_size > 1: + req = self._broadcast_request(req) if req is None: logger.info(f"Worker {self.device_id}: Shutdown signal received") diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 37dde47aa6ac..eae440569aec 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1084,3 +1084,56 @@ def test_restore_is_idempotent(self): req.refs_to_bytes() req.refs_to_bytes() assert req.params.image_reference[0].content == payload + + +class TestReferenceBroadcastSplit: + """Reference payloads leave the object before the rank0 -> N-rank hop. + + ``broadcast_object_list`` serializes whatever it is handed into a tensor, + so leaving the bytes inside the request would copy every reference byte + before the collective even starts. + """ + + @staticmethod + def _request(*payloads: bytes): + from tensorrt_llm._torch.visual_gen import DiffusionRequest + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + + params = VisualGenParams( + image_reference=[MediaRef(content=p, format="bytes") for p in payloads], + video_reference=MediaRef(content=b"\x00\x00\x00\x18ftypmp42", format="bytes"), + ) + return DiffusionRequest(request_id=1, prompt=["x"], params=params) + + def test_detach_empties_the_object_and_records_sizes(self): + payloads = (os.urandom(4096), os.urandom(64)) + req = self._request(*payloads) + + detached = req.refs_detach() + + assert detached[:2] == list(payloads) + assert req.ref_sizes == [len(p) for p in detached] + assert all(r.content == b"" for r in req.params.image_reference) + assert payloads[0] not in pickle.dumps(req) + + def test_attach_restores_every_slot_in_order(self): + payloads = (os.urandom(2048), os.urandom(128)) + req = self._request(*payloads) + + detached = req.refs_detach() + req.refs_attach(detached) + + assert tuple(r.content for r in req.params.image_reference) == payloads + assert req.params.video_reference[0].content == b"\x00\x00\x00\x18ftypmp42" + assert req.ref_sizes is None + + def test_sizes_let_a_peer_size_its_buffers(self): + """Non-source ranks allocate from ``ref_sizes`` alone, so it has to + survive the object hop and match the payloads exactly.""" + payloads = (os.urandom(1024), os.urandom(7)) + req = self._request(*payloads) + req.refs_detach() + + peer = pickle.loads(pickle.dumps(req)) + assert peer.ref_sizes == req.ref_sizes + assert peer.ref_sizes[:2] == [1024, 7] From c5b0474670b3060e2137ef350135f134db7fb5b0 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:14:05 -0700 Subject: [PATCH 27/61] [TRTLLM-15277][fix] VisualGen: release a dropped request's shared memory, refresh stale docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit An unconsumed shared-tensor handle keeps its block mapped until the producing process exits — measured: five handles created and dropped leave five `/dev/shm` entries, while a handle consumed once leaves none. The sender thread logs and swallows a failed `requests_ipc.put`, so before this change a request that never reached rank0 leaked one block per reference for the lifetime of a `trtllm-serve` process. It now consumes its own handles on that path. Docs and comments still described the old materialize-to-disk transport. Two were user-facing and wrong, not merely stale: the reference-format table said a `path` reference "passes through in place" and had to be readable by "the process running generation", when the coordinator now reads it and the worker never touches the filesystem; and a `url` reference is no longer "materialized to a local path". Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 4 ++-- tensorrt_llm/_torch/visual_gen/executor.py | 6 +++++ tensorrt_llm/serve/openai_server.py | 2 +- tensorrt_llm/serve/openai_video_routes.py | 6 ++--- tensorrt_llm/serve/visual_gen_utils.py | 4 ++-- .../visual_gen/test_trtllm_serve_endpoints.py | 3 +-- .../visual_gen/test_visual_gen_params.py | 4 ++-- .../visual_gen/test_visual_gen_utils.py | 22 +++++++++++++++++++ 8 files changed, 39 insertions(+), 12 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index c65a6b42fda1..7b065ac2b739 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -125,8 +125,8 @@ Conditioning references are supplied through the typed, per-modality fields `ima | `format` | Content | Notes | |---|---|---| -| `path` | A local file readable by the process running generation | Bare path or `file://` URI. The file must exist, and passes through in place — it is neither copied nor deleted. | -| `url` | An `http(s)` URL | Fetched through the SSRF-guarded loader, then materialized to a local path. | +| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must exist; it is read once on the coordinator and is never modified or deleted. | +| `url` | An `http(s)` URL | Fetched on the coordinator through the SSRF-guarded loader. | | `base64` | Base64 text | A `data:` URI is also accepted. | | `bytes` | Raw `bytes` | Python API only. Rejected over JSON (HTTP 422) — send `base64` or upload the file via multipart. | diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 5f42d504cb45..02a9576dfd37 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -977,6 +977,7 @@ def _send_shutdown(self): def _process_requests(self): """Process pending requests.""" + req = None try: req = self.pending_requests.get(timeout=POLL_TIMEOUT) if req is None: @@ -994,6 +995,11 @@ def _process_requests(self): except Exception as e: logger.error(f"DiffusionClient: Error sending request: {e}") logger.error(traceback.format_exc()) + if req is not None and req.ref_handles: + # The request never reached rank0, so nothing downstream will + # consume its handles, and an unconsumed handle keeps its + # shared-memory block mapped until this process exits. + req.refs_to_bytes() def _process_responses(self): """Poll and process responses.""" diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index ed501d164ec2..7fc7f392b153 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -3074,7 +3074,7 @@ async def openai_image_generation(self, request: ImageGenerationRequest, f"Generating image: {image_id} with params: {params} and prompt: {request.prompt}" ) image_gen_start = time.perf_counter() - # Offload the blocking materialize/enqueue off the event loop but + # Offload the blocking resolve/enqueue off the event loop but # await it (bad params → 400 here); then await generation. handle = await asyncio.to_thread(self.generator.generate_async, request.prompt, params) diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 61c847ea997b..2d61ce2a7ae1 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -172,9 +172,9 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) sync_video_start = time.perf_counter() - # Offload the blocking resolve/materialize/enqueue off the event - # loop but await it, so bad media / bad params still surface as - # 400 here; then await generation on the executor's loop. + # Offload the blocking resolve/enqueue off the event loop but + # await it, so bad media / bad params still surface as 400 + # here; then await generation on the executor's loop. handle = await asyncio.to_thread( self.generator.generate_async, request.prompt, params ) diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 7ca39ffdb6f1..a98f77f41acc 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -133,7 +133,7 @@ def _build_reference_list(value: Any) -> Optional[list]: ``value`` is None, a multipart ``UploadFile``, a ``MediaReferenceItem``, or a list of those. Each entry becomes a ``MediaRef`` carrying its transport content plus the declared (or, for an upload, implied) wire format. - Resolution and materialization happen later at the engine choke point. + Resolution to bytes happens later, at the engine choke point. """ if value is None: return None @@ -505,7 +505,7 @@ def parse_visual_gen_params( params.num_frames = derived # Reference inputs: hand the pipeline a ``MediaRef`` carrying the # transport content (``bytes`` for an upload, the string otherwise). - # The engine resolves and materializes; the boundary never touches disk. + # The engine resolves it to bytes; the boundary never touches disk. image_refs = _build_reference_list(request.image_reference) if image_refs: params.image_reference = image_refs 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 91376c00f7db..44fe4e7e5803 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -211,8 +211,7 @@ def __init__( # used by tests to assert forwarded VisualGenParams fields. self.last_inputs = None self.last_params = None - # Snapshot of the resolved reference bytes at generation time, - # captured before the route cleans them up. Keyed by stored path. + # Resolved reference bytes per slot, recorded at generation time. self.last_ref_bytes = {} # Stand-in for the coordinator-side executor proxy. The async video # route reads ``default_generation_params`` / ``extra_param_specs`` 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 378c6aeecc57..fda5948de713 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -221,9 +221,9 @@ def test_engine_rewrite_is_not_blocked_by_the_pairing_check(self): from tensorrt_llm.visual_gen import MediaRef ref = MediaRef(content="aGk=", format="base64") - ref.content = "/tmp/materialized" # contradicts format for one statement + ref.content = "/tmp/ref.png" # contradicts format for one statement ref.format = "path" - assert (ref.content, ref.format) == ("/tmp/materialized", "path") + assert (ref.content, ref.format) == ("/tmp/ref.png", "path") def test_unknown_format_rejected(self): from pydantic import ValidationError diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index eae440569aec..88aa5bb9790e 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1137,3 +1137,25 @@ def test_sizes_let_a_peer_size_its_buffers(self): peer = pickle.loads(pickle.dumps(req)) assert peer.ref_sizes == req.ref_sizes assert peer.ref_sizes[:2] == [1024, 7] + + def test_dropped_request_releases_its_shared_memory(self): + """A handle that is never consumed keeps its shared-memory block mapped + until the process exits, so a request that fails on its way to rank0 + has to consume its own handles on the way out.""" + import gc + import glob + + def blocks(): + return set(glob.glob("/dev/shm/torch_*")) + + before = blocks() + req = self._request(os.urandom(1024 * 1024)) + req.refs_to_handles() + gc.collect() + assert len(blocks() - before) == len(req.ref_handles) + + # What the sender thread does when the request never reaches rank0. + req.refs_to_bytes() + del req + gc.collect() + assert blocks() - before == set() From 304ccf1c4d040df6fd5a536e5a2e29e414441a93 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 19 Aug 2026 19:33:43 -0700 Subject: [PATCH 28/61] [TRTLLM-15277][fix] VisualGen: make reference handle bookkeeping fail loudly instead of leaking or hanging MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Follow-ups from an independent review of the reference transport. None of these is reachable through a normal request today, but each one fails in a way that is expensive to diagnose. `refs_to_handles` accumulated handles in a local list and published it only on success, so a reference that failed to reach shared memory left every block taken before it unreachable — including by the reclaim path added for a dropped request. The list is now published before it is filled. `refs_attach` zipped references against payloads, so a short list quietly left the tail of the references empty and a long one was discarded, after which clearing `ref_sizes` erased the evidence. It now requires an exact count. `_broadcast_request` makes the same check on rank0 before the first collective, because peers size their collectives from the broadcast `ref_sizes` and a mismatch would hang every rank rather than fail. `shutdown` force-stopping the dispatcher left any request still queued holding unconsumed handles; those are now drained and reclaimed. FLUX.2's `forward` docstring still told direct callers that public requests carry "file paths or encoded bytes"; references reach a pipeline as bytes whatever wire form the request declared. The two shared-memory tests now identify blocks by the filename inside each handle rather than diffing `/dev/shm`, which is global to the machine and picks up unrelated processes. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 38 +++++++++- .../visual_gen/models/flux/pipeline_flux2.py | 5 +- .../visual_gen/test_visual_gen_utils.py | 69 ++++++++++++++++--- 3 files changed, 96 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 02a9576dfd37..65bbb14efddb 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -274,7 +274,10 @@ def refs_to_handles(self) -> None: """ if self.params is None: return - handles = [] + # Publish the list before filling it: if a later reference fails to + # reach shared memory, the blocks already taken are still reachable + # for the reclaim path instead of leaking. + self.ref_handles = handles = [] for slot in ("image_reference", "video_reference", "audio_reference"): for index, ref in enumerate(getattr(self.params, slot, None) or []): # bytearray() because the shared storage must be writable. @@ -287,7 +290,8 @@ def refs_to_handles(self) -> None: } ) ref.content = b"" - self.ref_handles = handles or None + if not handles: + self.ref_handles = None def refs_to_bytes(self) -> None: """Restore reference payloads from shared memory, in place (consumer side).""" @@ -318,7 +322,12 @@ def refs_detach(self) -> List[bytes]: def refs_attach(self, payloads: List[bytes]) -> None: """Put the separately broadcast payloads back, in place.""" - for ref, payload in zip(self._refs(), payloads): + refs = list(self._refs()) + if len(refs) != len(payloads): + # zip() would silently leave the tail of either side behind, and + # clearing ref_sizes below would erase the evidence. + raise ValueError(f"expected {len(refs)} reference payloads, got {len(payloads)}.") + for ref, payload in zip(refs, payloads): ref.content = payload self.ref_sizes = None @@ -463,6 +472,13 @@ def _broadcast_request(self, req: Optional[DiffusionRequest]) -> Optional[Diffus object into a tensor first and that copy dominates the collective. """ payloads = req.refs_detach() if self.rank == 0 and req is not None else [] + if self.rank == 0 and len(payloads) != len(getattr(req, "ref_sizes", None) or []): + # Peers derive their collective count from ref_sizes, so a mismatch + # here would hang every rank. Fail on rank0, before the first one. + raise RuntimeError( + f"reference payload/size mismatch: {len(payloads)} payloads, " + f"{len(getattr(req, 'ref_sizes', None) or [])} sizes." + ) obj_list = [req] dist.broadcast_object_list(obj_list, src=0) req = obj_list[0] @@ -1119,6 +1135,21 @@ async def _serve_forever(self): self._cleanup_ipc() + def _reclaim_pending_handles(self) -> None: + """Consume the handles of requests that will never be sent. + + A request abandoned in the queue still holds shared-memory blocks that + nothing downstream will consume, and they stay mapped until this + process exits. + """ + while True: + try: + req = self.pending_requests.get_nowait() + except queue.Empty: + return + if req is not None and req.ref_handles: + req.refs_to_bytes() + def shutdown(self): """Shutdown client and workers.""" logger.info("DiffusionClient: Shutting down") @@ -1129,6 +1160,7 @@ def shutdown(self): logger.warning("DiffusionClient: Force stopping background thread") self.shutdown_event.set() self.background_thread.join(timeout=1.0) + self._reclaim_pending_handles() # Shutdown workers logger.info("DiffusionClient: Stopping workers") 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 b1465590d925..fb1485746fc0 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -450,8 +450,9 @@ def forward( Each prompt's embeddings are repeated and independent noise is sampled, producing N different images per prompt. image: Reference image or shared list of reference images for - image conditioning. Public ``VisualGenParams`` requests use - file paths or encoded bytes; direct calls may also use PIL images. + image conditioning. References arrive here as encoded bytes, + whatever wire form the request declared; direct calls may also + use PIL images. Returns: PipelineOutput with image tensor (B, H, W, C) where diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 88aa5bb9790e..3f50e6e624d2 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -17,9 +17,11 @@ from io import BytesIO from pathlib import Path from typing import Any, Dict, Optional +from unittest import mock import numpy as np import pytest +import torch from fastapi import UploadFile from PIL import Image @@ -1138,24 +1140,69 @@ def test_sizes_let_a_peer_size_its_buffers(self): assert peer.ref_sizes == req.ref_sizes assert peer.ref_sizes[:2] == [1024, 7] + @staticmethod + def _blocks_of(req): + """The shared-memory files this request's handles point at. + + Named per handle rather than scanned out of /dev/shm, which is global + to the machine and picks up unrelated processes. + """ + return [ + Path("/dev/shm") / base64.b64decode(e["handle"]["storage_handle"]).decode().lstrip("/") + for e in req.ref_handles or [] + ] + def test_dropped_request_releases_its_shared_memory(self): - """A handle that is never consumed keeps its shared-memory block mapped - until the process exits, so a request that fails on its way to rank0 - has to consume its own handles on the way out.""" + """An unconsumed handle keeps its block mapped until the process exits, + so a request that fails on its way to rank0 has to consume its own + handles on the way out.""" import gc - import glob - - def blocks(): - return set(glob.glob("/dev/shm/torch_*")) - before = blocks() req = self._request(os.urandom(1024 * 1024)) req.refs_to_handles() - gc.collect() - assert len(blocks() - before) == len(req.ref_handles) + blocks = self._blocks_of(req) + assert blocks and all(b.exists() for b in blocks) # What the sender thread does when the request never reaches rank0. req.refs_to_bytes() del req gc.collect() - assert blocks() - before == set() + assert not any(b.exists() for b in blocks) + + def test_attach_rejects_a_count_mismatch(self): + """Peers size their collectives from ``ref_sizes``, so a payload list + that does not match the reference count is a bug worth raising on + rather than silently leaving references empty.""" + req = self._request(os.urandom(64), os.urandom(64)) + detached = req.refs_detach() + + with pytest.raises(ValueError, match="expected 3 reference payloads, got 2"): + req.refs_attach(detached[:2]) + + def test_partial_handle_failure_stays_reclaimable(self): + """If a later reference cannot reach shared memory, the blocks already + taken must still be reachable, or nothing can free them.""" + import gc + + req = self._request(os.urandom(256 * 1024), os.urandom(256 * 1024)) + real = torch.frombuffer + calls = {"n": 0} + + def flaky(buffer, **kwargs): + calls["n"] += 1 + if calls["n"] == 3: # the third of this request's three references + raise RuntimeError("shared memory exhausted") + return real(buffer, **kwargs) + + with mock.patch.object(torch, "frombuffer", flaky): + with pytest.raises(RuntimeError, match="shared memory exhausted"): + req.refs_to_handles() + + blocks = self._blocks_of(req) + assert len(blocks) == 2, "handles taken before the failure must stay reachable" + assert all(b.exists() for b in blocks) + + req.refs_to_bytes() + del req + gc.collect() + assert not any(b.exists() for b in blocks) From 9d033f049210b1337b04ea797a22b895fcfa38a8 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:09:07 -0700 Subject: [PATCH 29/61] [TRTLLM-15277][chore] Keep the common serving path free of a VisualGen import MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `openai_protocol` named `ContentFormat` and `Role` from `visual_gen.params`, and `visual_gen_utils` imported the reference resolver at module scope. Both are on the path every LLM deployment already walks — the tool parsers import `openai_protocol` — so importing it pulled VisualGen in behind them, undoing the separation #17281 established. Coverage-driven CI triggering reads that import graph, and an edge costs the same there whether it drags three modules or three hundred. The reference wire types are declaration-only, so they move to `tensorrt_llm/media/reference.py`, a leaf that depends on nothing but pydantic. `visual_gen.params` re-exports them, leaving `from tensorrt_llm.visual_gen import MediaRef` and every other public spelling untouched. The resolver, which does reach into VisualGen, is now imported inside the one function that runs when a deprecated `input_reference` arrives. `tests/unittest/llmapi/apps/test_serve_vertical_isolation.py` imports each common serving module in a fresh interpreter and fails if any `visual_gen` module appears in `sys.modules`. Verified in both directions: importing `openai_protocol` pulls three VisualGen modules before this change and none after, and re-adding an edge makes the test name the offending modules. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/media/reference.py | 101 ++++++++++++++++++ tensorrt_llm/serve/openai_protocol.py | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 8 +- tensorrt_llm/visual_gen/params.py | 83 ++------------ .../apps/test_serve_vertical_isolation.py | 92 ++++++++++++++++ 5 files changed, 210 insertions(+), 76 deletions(-) create mode 100644 tensorrt_llm/media/reference.py create mode 100644 tests/unittest/llmapi/apps/test_serve_vertical_isolation.py diff --git a/tensorrt_llm/media/reference.py b/tensorrt_llm/media/reference.py new file mode 100644 index 000000000000..4cb8b4408b89 --- /dev/null +++ b/tensorrt_llm/media/reference.py @@ -0,0 +1,101 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The wire types for a media reference, shared by serve and VisualGen. + +These are declaration-only: the schema a caller fills in, with no resolver, +decoder or engine behind them. They live here rather than under +``visual_gen`` so the common serving protocol can name them without the +request schema of every LLM deployment pulling a vertical in behind it. +Depend on nothing but ``pydantic`` and keep it that way. +""" + +from typing import Any, Optional, Union + +from pydantic import Field, model_validator +from typing_extensions import Literal + +from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status + +Role = Literal["reference", "first_frame", "last_frame"] + +# Wire form of a reference's ``content``. Declared explicitly rather than +# sniffed: a bare string is otherwise ambiguous between a local path and +# base64, and guessing lets a mistyped path silently become base64 (or a +# malformed base64 silently become a filesystem read). +ContentFormat = Literal["path", "url", "base64", "bytes"] + + +@set_api_status("prototype") +class MediaRef(StrictBaseModel): + """A single media reference (image / video / audio). + + Carried by ``image_reference`` / ``video_reference`` / ``audio_reference``; + the field it sits in fixes the modality. ``role`` is required only when the + target model accepts that modality in more than one role (e.g. image first + + last frame); otherwise the pipeline knows the reference's meaning and + ``role`` may be omitted (video/audio are always the single ``reference``). + """ + + content: Union[str, bytes] = Field( + description="The reference payload, in the form declared by ``format``." + ) + format: ContentFormat = Field( + description=( + "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " + "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " + "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." + ) + ) + role: Optional[Role] = Field( + default=None, description="``reference`` | ``first_frame`` | ``last_frame``." + ) + + @model_validator(mode="after") + def _check_content_matches_format(self): + """Reject a ``content`` whose Python type contradicts ``format``. + + ``bytes`` is the only format carrying a binary payload; the other three + name a location or an encoding and are therefore strings. Checking the + pairing here fails at construction — an HTTP 422 or an immediate + ``ValueError`` — instead of deep in the engine's resolve step. + """ + if self.format == "bytes": + if not isinstance(self.content, bytes): + raise ValueError( + f"format='bytes' requires bytes content, got {type(self.content).__name__}." + ) + elif not isinstance(self.content, str): + raise ValueError( + f"format={self.format!r} requires string content, got " + f"{type(self.content).__name__}." + ) + return self + + +def reject_bare_refs(value: Any) -> Any: + """Reject the bare path/bytes shorthand with an actionable message. + + Runs before coercion, so the caller sees what to do instead of a union + mismatch reported against an inner model. A bare string has nowhere to + declare its wire form, and guessing is what ``format`` exists to prevent. + """ + for x in value if isinstance(value, list) else [value]: + if isinstance(x, (str, bytes)): + raise ValueError( + "a reference must declare its wire form; a bare " + f"{type(x).__name__} is no longer accepted. Pass " + 'MediaRef(content=..., format="path"|"url"|"base64"|"bytes").' + ) + return value diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index b60920c5283d..450ce4f3deed 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -58,10 +58,10 @@ from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory +from tensorrt_llm.media.reference import ContentFormat, Role from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy -from tensorrt_llm.visual_gen.params import ContentFormat, Role _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index a98f77f41acc..608741b50290 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -18,7 +18,6 @@ ImageGenerationRequest, VideoGenerationRequest, ) -from tensorrt_llm.visual_gen.media_refs import _resolve_reference if TYPE_CHECKING: from fastapi import UploadFile @@ -406,11 +405,16 @@ def _apply_deprecated_input_reference( logger.warning("'input_reference' is deprecated; use 'image_reference' / 'video_reference'.") if params.image_reference or params.video_reference: return - from tensorrt_llm.visual_gen.params import MediaRef + from tensorrt_llm.media.reference import MediaRef if hasattr(input_reference, "file"): # multipart upload — form implied payload = input_reference.file.read() else: + # Imported here, not at module scope: the resolver reaches into + # VisualGen, and a plain LLM deployment must not pull a vertical in + # behind its request schema. + from tensorrt_llm.visual_gen.media_refs import _resolve_reference + payload = _resolve_reference(input_reference, input_reference_format) kind = sniff_media_kind(payload) if kind == "image": diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index fd1b67d1f5c4..41837edebec6 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -13,83 +13,20 @@ # See the License for the specific language governing permissions and # limitations under the License. import ast -from typing import Any, Dict, List, Literal, Optional, Union +from typing import Any, Dict, List, Optional, Union -from pydantic import Field, field_validator, model_validator +from pydantic import Field, field_validator from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status -Role = Literal["reference", "first_frame", "last_frame"] - -# Wire form of a reference's ``content``. Declared explicitly rather than -# sniffed: a bare string is otherwise ambiguous between a local path and -# base64, and guessing lets a mistyped path silently become base64 (or a -# malformed base64 silently become a filesystem read). -ContentFormat = Literal["path", "url", "base64", "bytes"] - - -@set_api_status("prototype") -class MediaRef(StrictBaseModel): - """A single media reference (image / video / audio). - - Carried by ``image_reference`` / ``video_reference`` / ``audio_reference``; - the field it sits in fixes the modality. ``role`` is required only when the - target model accepts that modality in more than one role (e.g. image first + - last frame); otherwise the pipeline knows the reference's meaning and - ``role`` may be omitted (video/audio are always the single ``reference``). - """ - - content: Union[str, bytes] = Field( - description="The reference payload, in the form declared by ``format``." - ) - format: ContentFormat = Field( - description=( - "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " - "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " - "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." - ) - ) - role: Optional[Role] = Field( - default=None, description="``reference`` | ``first_frame`` | ``last_frame``." - ) - - @model_validator(mode="after") - def _check_content_matches_format(self): - """Reject a ``content`` whose Python type contradicts ``format``. - - ``bytes`` is the only format carrying a binary payload; the other three - name a location or an encoding and are therefore strings. Checking the - pairing here fails at construction — an HTTP 422 or an immediate - ``ValueError`` — instead of deep in the engine's resolve step. - """ - if self.format == "bytes": - if not isinstance(self.content, bytes): - raise ValueError( - f"format='bytes' requires bytes content, got {type(self.content).__name__}." - ) - elif not isinstance(self.content, str): - raise ValueError( - f"format={self.format!r} requires string content, got " - f"{type(self.content).__name__}." - ) - return self - - -def _reject_bare_refs(value: Any) -> Any: - """Reject the bare path/bytes shorthand with an actionable message. - - Runs before coercion, so the caller sees what to do instead of a union - mismatch reported against an inner model. A bare string has nowhere to - declare its wire form, and guessing is what ``format`` exists to prevent. - """ - for x in value if isinstance(value, list) else [value]: - if isinstance(x, (str, bytes)): - raise ValueError( - "a reference must declare its wire form; a bare " - f"{type(x).__name__} is no longer accepted. Pass " - 'MediaRef(content=..., format="path"|"url"|"base64"|"bytes").' - ) - return value +# The reference wire types live in a dependency-neutral leaf so the common +# serving protocol can name them without pulling VisualGen in, but +# ``tensorrt_llm.visual_gen`` stays their public home. The redundant aliases +# mark these as intentional re-exports rather than unused imports. +from tensorrt_llm.media.reference import ContentFormat as ContentFormat +from tensorrt_llm.media.reference import MediaRef as MediaRef +from tensorrt_llm.media.reference import Role as Role +from tensorrt_llm.media.reference import reject_bare_refs as _reject_bare_refs def _normalize_refs(value: Any) -> Optional[list]: diff --git a/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py b/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py new file mode 100644 index 000000000000..b081f78df9ff --- /dev/null +++ b/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py @@ -0,0 +1,92 @@ +# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""The common serving path must not drag a vertical in behind it. + +``openai_protocol`` is imported by components every LLM deployment uses — the +tool parsers, for one — so an import-time edge from it into ``visual_gen`` +makes a DeepSeek test case depend on VisualGen. That matters for +coverage-driven CI triggering, where the import graph decides which stages a +change runs, and the cost is the same whether the edge pulls three modules or +three hundred. + +Shared declaration-only types therefore live in ``tensorrt_llm.media``, and +anything with a VisualGen resolver behind it is imported inside the function +that needs it. +""" + +import subprocess +import sys +import textwrap + +import pytest + +_PROBE = textwrap.dedent( + """ + import importlib, json, sys + importlib.import_module({module!r}) + leaked = sorted( + name for name in sys.modules + if name == "tensorrt_llm.visual_gen" + or name.startswith("tensorrt_llm.visual_gen.") + or name.startswith("tensorrt_llm._torch.visual_gen") + ) + print(json.dumps(leaked)) + """ +) + + +def _visual_gen_modules_pulled_by(module: str) -> list[str]: + """Import *module* in a fresh interpreter and report VisualGen fallout. + + A subprocess because the check is about what a module drags in on its own; + inside pytest the whole suite has already imported half the product. + """ + proc = subprocess.run( + [sys.executable, "-c", _PROBE.format(module=module)], + capture_output=True, + text=True, + timeout=600, + ) + assert proc.returncode == 0, f"probe failed for {module}:\n{proc.stderr[-2000:]}" + return __import__("json").loads(proc.stdout.strip().splitlines()[-1]) + + +@pytest.mark.parametrize( + "module", + [ + "tensorrt_llm.serve.openai_protocol", + "tensorrt_llm.serve.tool_parser.deepseekv3_parser", + ], +) +def test_common_serving_module_does_not_import_visual_gen(module): + leaked = _visual_gen_modules_pulled_by(module) + assert leaked == [], ( + f"{module} pulls VisualGen in at import time: {leaked}. Shared " + "declaration-only types belong in tensorrt_llm.media; import anything " + "with a VisualGen resolver behind it inside the function that uses it." + ) + + +def test_media_reference_types_are_dependency_free(): + """The shared leaf must not acquire a VisualGen import of its own.""" + assert _visual_gen_modules_pulled_by("tensorrt_llm.media.reference") == [] + + +def test_visual_gen_still_exports_the_shared_types(): + """Moving the types must not move the public API they are reached through.""" + from tensorrt_llm.media.reference import MediaRef as LeafMediaRef + from tensorrt_llm.visual_gen import MediaRef + + assert MediaRef is LeafMediaRef From 286e30301abc0e189a0afa49e00c4a7b3afeae49 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:16:16 -0700 Subject: [PATCH 30/61] [TRTLLM-15277][chore] Drop the import-graph guard from this PR The guard answers "how do we stop this regressing?", which is a follow-up the reviewers raised alongside the architecture question, not the dependency edge this PR introduced. Scoping it out keeps the change to the edge itself; the guard belongs with whatever convention that discussion settles on. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../apps/test_serve_vertical_isolation.py | 92 ------------------- 1 file changed, 92 deletions(-) delete mode 100644 tests/unittest/llmapi/apps/test_serve_vertical_isolation.py diff --git a/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py b/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py deleted file mode 100644 index b081f78df9ff..000000000000 --- a/tests/unittest/llmapi/apps/test_serve_vertical_isolation.py +++ /dev/null @@ -1,92 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The common serving path must not drag a vertical in behind it. - -``openai_protocol`` is imported by components every LLM deployment uses — the -tool parsers, for one — so an import-time edge from it into ``visual_gen`` -makes a DeepSeek test case depend on VisualGen. That matters for -coverage-driven CI triggering, where the import graph decides which stages a -change runs, and the cost is the same whether the edge pulls three modules or -three hundred. - -Shared declaration-only types therefore live in ``tensorrt_llm.media``, and -anything with a VisualGen resolver behind it is imported inside the function -that needs it. -""" - -import subprocess -import sys -import textwrap - -import pytest - -_PROBE = textwrap.dedent( - """ - import importlib, json, sys - importlib.import_module({module!r}) - leaked = sorted( - name for name in sys.modules - if name == "tensorrt_llm.visual_gen" - or name.startswith("tensorrt_llm.visual_gen.") - or name.startswith("tensorrt_llm._torch.visual_gen") - ) - print(json.dumps(leaked)) - """ -) - - -def _visual_gen_modules_pulled_by(module: str) -> list[str]: - """Import *module* in a fresh interpreter and report VisualGen fallout. - - A subprocess because the check is about what a module drags in on its own; - inside pytest the whole suite has already imported half the product. - """ - proc = subprocess.run( - [sys.executable, "-c", _PROBE.format(module=module)], - capture_output=True, - text=True, - timeout=600, - ) - assert proc.returncode == 0, f"probe failed for {module}:\n{proc.stderr[-2000:]}" - return __import__("json").loads(proc.stdout.strip().splitlines()[-1]) - - -@pytest.mark.parametrize( - "module", - [ - "tensorrt_llm.serve.openai_protocol", - "tensorrt_llm.serve.tool_parser.deepseekv3_parser", - ], -) -def test_common_serving_module_does_not_import_visual_gen(module): - leaked = _visual_gen_modules_pulled_by(module) - assert leaked == [], ( - f"{module} pulls VisualGen in at import time: {leaked}. Shared " - "declaration-only types belong in tensorrt_llm.media; import anything " - "with a VisualGen resolver behind it inside the function that uses it." - ) - - -def test_media_reference_types_are_dependency_free(): - """The shared leaf must not acquire a VisualGen import of its own.""" - assert _visual_gen_modules_pulled_by("tensorrt_llm.media.reference") == [] - - -def test_visual_gen_still_exports_the_shared_types(): - """Moving the types must not move the public API they are reached through.""" - from tensorrt_llm.media.reference import MediaRef as LeafMediaRef - from tensorrt_llm.visual_gen import MediaRef - - assert MediaRef is LeafMediaRef From baa0d1d6d83187e76cba3ccd0062426a5a838857 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:44:56 -0700 Subject: [PATCH 31/61] [TRTLLM-15277][chore] Name the reference role type, and say what the field is for MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``Role`` was too generic to sit in a shared leaf: ``kv_cache_manager_v2`` already exports a ``Role`` enum for KEY/INDEX_KEY, and ``serve`` imports another from ``openai_harmony`` for ASSISTANT/SYSTEM. ``MediaRole`` pairs with ``MediaRef`` and reads unambiguously next to either. The ``role`` description listed the three literals, which pydantic already publishes from the annotation. It now says what the field decides — which conditioning slot the reference fills — and when it may be omitted. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/pipeline.py | 4 ++-- tensorrt_llm/media/reference.py | 11 ++++++++--- tensorrt_llm/serve/openai_protocol.py | 9 +++++---- tensorrt_llm/visual_gen/params.py | 2 +- .../api_stability/references/trtllm_serve_api.yaml | 2 +- 5 files changed, 17 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index fd14dbd30610..964c9a89a951 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -31,7 +31,7 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel from tensorrt_llm.logger import logger from tensorrt_llm.mapping import Mapping -from tensorrt_llm.visual_gen.params import Role +from tensorrt_llm.visual_gen.params import MediaRole from .cache import CacheDiTAccelerator, TeaCacheAccelerator from .checkpoints import WeightLoader @@ -75,7 +75,7 @@ class ExtraParamSchema(StrictBaseModel): class RoleSpec(StrictBaseModel): """One accepted role for a reference modality, with its count bounds.""" - role: Role = Field(description="Role of the reference input.") + role: MediaRole = Field(description="Role of the reference input.") min: int = Field(default=1, description="Minimum count for this role.") max: Optional[int] = Field( default=1, description="Maximum count for this role (None = unbounded)." diff --git a/tensorrt_llm/media/reference.py b/tensorrt_llm/media/reference.py index 4cb8b4408b89..ed9caeb5e03e 100644 --- a/tensorrt_llm/media/reference.py +++ b/tensorrt_llm/media/reference.py @@ -28,7 +28,7 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status -Role = Literal["reference", "first_frame", "last_frame"] +MediaRole = Literal["reference", "first_frame", "last_frame"] # Wire form of a reference's ``content``. Declared explicitly rather than # sniffed: a bare string is otherwise ambiguous between a local path and @@ -58,8 +58,13 @@ class MediaRef(StrictBaseModel): "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." ) ) - role: Optional[Role] = Field( - default=None, description="``reference`` | ``first_frame`` | ``last_frame``." + role: Optional[MediaRole] = Field( + default=None, + description=( + "Which conditioning slot this reference fills. Required only when the " + "target model accepts this modality in more than one slot; omit it when " + "the model leaves no ambiguity." + ), ) @model_validator(mode="after") diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 450ce4f3deed..cf660a217e45 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -58,7 +58,7 @@ from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory -from tensorrt_llm.media.reference import ContentFormat, Role +from tensorrt_llm.media.reference import ContentFormat, MediaRole from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy @@ -2033,10 +2033,11 @@ class MediaReferenceItem(OpenAIBaseModel): "also accepted). ``bytes`` cannot be carried in JSON — upload the file " "as multipart/form-data instead. Distinct from the top-level " "``format``, which selects the *output* encoding.")) - role: Optional[Role] = Field( + role: Optional[MediaRole] = Field( default=None, - description="Reference role. Required only when the model has more than " - "one required role for the modality; otherwise inferred.", + description="Which conditioning slot this reference fills. Required only " + "when the target model accepts this modality in more than one slot; omit " + "it when the model leaves no ambiguity.", ) @field_validator("format") diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 41837edebec6..419e5bb86969 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -25,7 +25,7 @@ # mark these as intentional re-exports rather than unused imports. from tensorrt_llm.media.reference import ContentFormat as ContentFormat from tensorrt_llm.media.reference import MediaRef as MediaRef -from tensorrt_llm.media.reference import Role as Role +from tensorrt_llm.media.reference import MediaRole as MediaRole from tensorrt_llm.media.reference import reject_bare_refs as _reject_bare_refs diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index b6fd5a84c555..1506a95462d6 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1589,7 +1589,7 @@ models: required: true role: kind: extension - type: Optional[Role] + type: Optional[MediaRole] default: null status: prototype required: false From edf8fe57c220ebe4758592170b302db630982bf2 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 20 Aug 2026 22:54:55 -0700 Subject: [PATCH 32/61] [TRTLLM-15277][chore] Name the reference wire-form type MediaContentFormat MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ``ContentFormat`` collided with ``inputs.content_format.ContentFormat``, an enum that says whether a chat template wants OpenAI dicts or plain strings, and which 30-odd model files already reference as ``ContentFormat.STRING`` / ``.OPENAI`` / ``.PASSTHROUGH``. Ours names the wire form of a reference payload, so ``MediaContentFormat`` it is — same reasoning as ``MediaRole``, and it keeps the shared leaf unambiguous. The LLM-side type is untouched. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/media/reference.py | 4 ++-- tensorrt_llm/serve/openai_protocol.py | 6 +++--- tensorrt_llm/visual_gen/params.py | 2 +- .../unittest/api_stability/references/trtllm_serve_api.yaml | 4 ++-- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/tensorrt_llm/media/reference.py b/tensorrt_llm/media/reference.py index ed9caeb5e03e..0fa513825081 100644 --- a/tensorrt_llm/media/reference.py +++ b/tensorrt_llm/media/reference.py @@ -34,7 +34,7 @@ # sniffed: a bare string is otherwise ambiguous between a local path and # base64, and guessing lets a mistyped path silently become base64 (or a # malformed base64 silently become a filesystem read). -ContentFormat = Literal["path", "url", "base64", "bytes"] +MediaContentFormat = Literal["path", "url", "base64", "bytes"] @set_api_status("prototype") @@ -51,7 +51,7 @@ class MediaRef(StrictBaseModel): content: Union[str, bytes] = Field( description="The reference payload, in the form declared by ``format``." ) - format: ContentFormat = Field( + format: MediaContentFormat = Field( description=( "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index cf660a217e45..b358a534f38c 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -58,7 +58,7 @@ from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory -from tensorrt_llm.media.reference import ContentFormat, MediaRole +from tensorrt_llm.media.reference import MediaContentFormat, MediaRole from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy @@ -2026,7 +2026,7 @@ class MediaReferenceItem(OpenAIBaseModel): content: str = Field( description="The reference payload, in the form declared by ``format``." ) - format: ContentFormat = Field(description=( + format: MediaContentFormat = Field(description=( "Wire form of ``content``: ``path`` (a file readable by the server; a " "``file://`` URI is also accepted), ``url`` (``http(s)``, fetched " "through the SSRF-guarded loader), or ``base64`` (a ``data:`` URI is " @@ -2113,7 +2113,7 @@ class VideoGenerationRequest(OpenAIBaseModel): "ignored whenever a typed ``image_reference`` or ``video_reference`` " "is provided. A string form requires ``input_reference_format``."), ) - input_reference_format: Optional[ContentFormat] = Field( + input_reference_format: Optional[MediaContentFormat] = Field( default=None, description=( "Deprecated, alongside ``input_reference``: the wire form of that " diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 419e5bb86969..0891a05b1550 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -23,7 +23,7 @@ # serving protocol can name them without pulling VisualGen in, but # ``tensorrt_llm.visual_gen`` stays their public home. The redundant aliases # mark these as intentional re-exports rather than unused imports. -from tensorrt_llm.media.reference import ContentFormat as ContentFormat +from tensorrt_llm.media.reference import MediaContentFormat as MediaContentFormat from tensorrt_llm.media.reference import MediaRef as MediaRef from tensorrt_llm.media.reference import MediaRole as MediaRole from tensorrt_llm.media.reference import reject_bare_refs as _reject_bare_refs diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 1506a95462d6..abd2f68b4ee2 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1497,7 +1497,7 @@ models: required: false input_reference_format: kind: extension - type: Optional[ContentFormat] + type: Optional[MediaContentFormat] default: null status: deprecated required: false @@ -1583,7 +1583,7 @@ models: required: true format: kind: extension - type: ContentFormat + type: MediaContentFormat default: null status: prototype required: true From 8f8b8cf34538b7fc763e4962dcaac457aa7547ad Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:22:42 -0700 Subject: [PATCH 33/61] [TRTLLM-15277][fix] Keep the deprecated input_reference working, and fix the perf-sanity payloads MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The new validator required `input_reference_format` beside any string `input_reference`, which rejected the exact shape the deprecated field exists to accept: callers written against the old API send bare base64 and got a 422 before the compatibility path ran. A missing format now resolves to base64 with a deprecation warning. The typed fields still demand an explicit `format` — that asymmetry is deliberate, and a test pins both halves. The three perf-sanity i2v entries were renamed from `input_reference` to `image_reference` without reshaping their values, so they still sent a bare base64 scalar that the typed field rejects; those jobs would have 422'd instead of running. They now send `{content, format: base64}`, and the Wan2.2 I2V one carries `role: first_frame` because that pipeline declares the role with `min=1`. Verified by feeding every perf-sanity `extra_body` through the real request model. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol.py | 21 +++++---- .../visual_gen/ltx2_blackwell.yaml | 4 +- .../visual_gen/wan22_i2v_a14b_blackwell.yaml | 2 +- .../visual_gen/test_visual_gen_params.py | 45 +++++++++++++++++++ 4 files changed, 61 insertions(+), 11 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index b358a534f38c..7396fefebb85 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -58,6 +58,7 @@ from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory +from tensorrt_llm.logger import logger from tensorrt_llm.media.reference import MediaContentFormat, MediaRole from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) @@ -2194,18 +2195,22 @@ def _reject_removed_response_format(cls, value): @model_validator(mode="after") def _check_input_reference_format(self): - """Require the deprecated ``input_reference``'s wire form when it is a string. + """Fill in the deprecated ``input_reference``'s wire form when omitted. - A multipart upload carries its own form, so the sibling is only needed - for the string spelling. + The typed fields require an explicit ``format`` — that is the point of + them. This one exists only so callers written against the old API keep + working, and those callers sent bare base64, so demanding a new sibling + field here would break exactly what the field is for. A multipart + upload carries its own form and needs no sibling either way. """ if isinstance(self.input_reference, str): if self.input_reference_format is None: - raise ValueError( - "'input_reference_format' is required when 'input_reference' is a " - "string; send 'path', 'url' or 'base64' (or upload the file via " - "multipart/form-data)") - if self.input_reference_format == "bytes": + logger.warning( + "'input_reference' without 'input_reference_format' is read as " + "base64; both are deprecated, use 'image_reference' / " + "'video_reference' with an explicit format.") + self.input_reference_format = "base64" + elif self.input_reference_format == "bytes": raise ValueError( "input_reference_format='bytes' cannot be carried in JSON; upload " "the file as multipart/form-data, or send 'base64'") diff --git a/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml b/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml index c85816c826b1..843a14baef02 100644 --- a/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml +++ b/tests/scripts/perf-sanity/visual_gen/ltx2_blackwell.yaml @@ -38,7 +38,7 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: {content: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII=, format: base64} - name: ltx2_2stage_bf16_t2v_cfg2_ulysses4_compile_on model_name: ltx2_bf16 @@ -98,4 +98,4 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: {content: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII=, format: base64} diff --git a/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml b/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml index 4f12c02386ac..4e49506ad25b 100644 --- a/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml +++ b/tests/scripts/perf-sanity/visual_gen/wan22_i2v_a14b_blackwell.yaml @@ -35,4 +35,4 @@ server_configs: max_concurrency: 1 num_prompts: 1 extra_body: - image_reference: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII= + image_reference: {content: iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAQAAAC1HAwCAAAAC0lEQVR42mP8/x8AAwMCAO7Z4h8AAAAASUVORK5CYII=, format: base64, role: first_frame} 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 fda5948de713..66fc23298b4f 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1589,3 +1589,48 @@ def request_warmup_cache_key(req): assert events == ["prepare", "warmup_cache_key", "infer"] executor.pipeline.request_warmup_cache_key.assert_called_once_with(req) executor.pipeline.run_inference.assert_called_once_with(req) + + +class TestDeprecatedInputReferenceStaysCompatible: + """The deprecated field exists so old callers keep working. + + Callers written against the old API sent bare base64 with no sibling + format, so requiring one here would break exactly what the field is for. + The typed fields still demand an explicit format — that distinction is the + point. + """ + + def test_bare_base64_is_read_as_base64(self): + from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + + request = VideoGenerationRequest(prompt="x", input_reference="aGk=") + + assert request.input_reference_format == "base64" + + def test_an_explicit_format_still_wins(self): + from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + + request = VideoGenerationRequest( + prompt="x", input_reference="/tmp/ref.png", input_reference_format="path" + ) + + assert request.input_reference_format == "path" + + def test_bytes_over_json_is_still_rejected(self): + from pydantic import ValidationError + + from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + + with pytest.raises(ValidationError, match="cannot be carried in JSON"): + VideoGenerationRequest( + prompt="x", input_reference="aGk=", input_reference_format="bytes" + ) + + def test_the_typed_field_still_requires_a_format(self): + """Relaxing the deprecated alias must not relax the new API.""" + from pydantic import ValidationError + + from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest + + with pytest.raises(ValidationError): + VideoGenerationRequest(prompt="x", image_reference="aGk=") From 8d1df5c35bb6a3454bdc5e81e0f9e804451df107 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 00:40:39 -0700 Subject: [PATCH 34/61] [TRTLLM-15277][fix] Let Qwen-Image-Edit take a PIL reference again MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Converging the decode sites dropped this loader's PIL branch: every item went through `ImageMediaIO.load_bytes`, which wraps its argument in `BytesIO`, so a caller holding a decoded image got a TypeError. Requests were unaffected — they arrive as bytes — but `forward` is also usable as a library entry point, and the other six migrated pipelines all kept their PIL path. This one was the outlier. Alpha still composites onto white here rather than being dropped, matching the `load_image` behaviour this pipeline had before the migration and differing from FLUX.2 on purpose; a test pins it. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../qwen_image/pipeline_qwen_image_edit.py | 22 ++++++- .../visual_gen/test_qwen_image_pipeline.py | 63 +++++++++++++++++++ 2 files changed, 83 insertions(+), 2 deletions(-) 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 3e2d32d3f6a0..c23be2d70b00 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 @@ -13,13 +13,14 @@ from typing import Any import numpy as np +import PIL.Image import torch import torch.distributed as dist from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline -from tensorrt_llm.inputs.media_io import ImageMediaIO +from tensorrt_llm.inputs.media_io import ImageMediaIO, convert_image_mode from tensorrt_llm.logger import logger from .pipeline_qwen_image import QwenImagePipeline, _calculate_shift @@ -160,11 +161,28 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N @staticmethod def _load_edit_images(image: Any) -> list[Any]: + """Normalize reference inputs to RGB images. + + A request always arrives as encoded bytes; a direct caller may hand + over a PIL image instead, and converting it here costs nothing while + keeping ``forward`` usable as a library entry point. + """ if image is None: raise ValueError("Qwen-Image-Edit requires image_reference.") images = image if isinstance(image, list) else [image] media_io = ImageMediaIO(format="pil") - return [media_io.load_bytes(item) for item in images] + loaded = [] + for index, item in enumerate(images): + if isinstance(item, PIL.Image.Image): + loaded.append(convert_image_mode(item, "RGB")) + elif isinstance(item, bytes): + loaded.append(media_io.load_bytes(item)) + else: + raise ValueError( + "Reference images must be PIL images or encoded bytes; " + f"item {index} has type {type(item).__name__}." + ) + return loaded def _preprocess_edit_images( self, diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py index 4900421d70bd..e3ad3c7c2d04 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py @@ -287,3 +287,66 @@ def test_forward_honors_profile_step_range(tmp_path): str(tmp_path / "visual-gen-trace-rank-0.json") ) cudart.cudaProfilerStop.assert_called_once_with() + + +class TestQwenImageEditReferenceLoading: + """``forward`` doubles as a library entry point, so it takes PIL too. + + A request always arrives as encoded bytes, but a direct caller may already + hold a decoded image and should not have to re-encode it just to get in. + """ + + @staticmethod + def _png(mode="RGB", color=(10, 20, 30)): + import io + + import PIL.Image + + buffer = io.BytesIO() + PIL.Image.new(mode, (8, 8), color).save(buffer, format="PNG") + return buffer.getvalue() + + def test_bytes_and_pil_both_load(self): + import io + + import PIL.Image + + from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import ( + QwenImageEditPlusPipeline, + ) + + data = self._png() + images = QwenImageEditPlusPipeline._load_edit_images( + [data, PIL.Image.open(io.BytesIO(data))] + ) + + assert len(images) == 2 + assert all(image.mode == "RGB" for image in images) + assert all(image.size == (8, 8) for image in images) + + def test_a_path_is_a_type_error(self): + """References reach a pipeline as bytes; a path is not a filesystem read.""" + import pytest + + from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import ( + QwenImageEditPlusPipeline, + ) + + with pytest.raises(ValueError, match="PIL images or encoded bytes"): + QwenImageEditPlusPipeline._load_edit_images(["/tmp/ref.png"]) + + def test_alpha_is_composited_onto_white(self): + """This pipeline went through ``load_image``, which flattens onto white + rather than dropping the channel the way diffusers does.""" + import PIL.Image + + from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import ( + QwenImageEditPlusPipeline, + ) + + transparent = PIL.Image.new("RGBA", (4, 4), (10, 20, 30, 0)) + + (image,) = QwenImageEditPlusPipeline._load_edit_images([transparent]) + + assert image.mode == "RGB" + assert image.getpixel((0, 0)) == (255, 255, 255) From 70ff5ff6a0f5c2e4d88949aa4be02f6ae4a79fb9 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:30:26 -0700 Subject: [PATCH 35/61] [TRTLLM-15277][fix] Classify audio in sniff_media_kind, and stop calling .m4a video MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `sniff_media_kind` only knew image and video, so every audio container fell through to `None`. `None` means "not media" — the value callers reject on — which left an audio reference indistinguishable from a text file, and made `audio_reference` the one modality `_validate_reference_payload` could not check at all. It now returns `"audio"` for WAV, MP3 (tagged or bare frame sync), FLAC, OGG/Opus, M4A and ADTS AAC, and `audio_reference` is validated like the other two. Two things fall out of this. RIFF is no longer read as AVI-or-nothing: the form type at [8:12] separates `AVI ` from `WAVE`, and anything else stays unclassified. And an `.m4a` is an MP4 whose brand says audio-only, so without `_ISOBMFF_AUDIO_BRANDS` it took the video default — that misclassification predates this PR. Verified against files ffmpeg actually produces for all seven audio formats plus MP4 and AVI, with PNG/JPEG and non-media payloads checked for regressions. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/inputs/media_io.py | 37 ++++++++++++++++++- tensorrt_llm/visual_gen/media_refs.py | 8 +++- .../visual_gen/test_visual_gen_utils.py | 31 +++++++++++++++- 3 files changed, 70 insertions(+), 6 deletions(-) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 9bf36b1dda13..719c9106afd0 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -394,6 +394,18 @@ def _get_cv2(): } ) +# ISO-BMFF audio: an `.m4a` is an MP4 whose brand says audio-only. Without +# these it falls through to the video default, which is how an audio file ends +# up classified as video. +_ISOBMFF_AUDIO_BRANDS = frozenset( + { + b"M4A ", + b"M4B ", # iTunes audio / audiobook + b"F4A ", + b"F4B ", # Flash audio / audiobook + } +) + # Work bound, not a format rule. The declared box size is client-controlled, # so without a ceiling the scan below costs O(payload): 1.6 s of interpreter # time for a 64 MB buffer, on the serving event loop. This admits 1020 @@ -478,12 +490,33 @@ def sniff_media_kind(data) -> Optional[str]: # major brand alone is not sufficient to identify still images. if brands & _ISOBMFF_IMAGE_BRANDS: return "image" + if brands & _ISOBMFF_AUDIO_BRANDS: + return "audio" return "video" - if header.startswith(b"RIFF") and header[8:12] == b"AVI ": - return "video" + if header.startswith(b"RIFF"): + # RIFF carries both; the form type at [8:12] is what separates them. + if header[8:12] == b"AVI ": + return "video" + if header[8:12] == b"WAVE": + return "audio" + return None + if header.startswith(b"OggS") or header.startswith(b"fLaC") or header.startswith(b"ID3"): + return "audio" + if _is_mpeg_audio_sync(header): + return "audio" return None +def _is_mpeg_audio_sync(header: bytes) -> bool: + """True for a bare MPEG audio / ADTS AAC frame — an MP3 with no ID3 tag. + + The sync word is eleven set bits, so the second byte carries three bits of + version/layer alongside it; matching the mask rather than a byte list keeps + every MPEG-1/2/2.5 layer and ADTS variant in scope. + """ + return len(header) >= 2 and header[0] == 0xFF and (header[1] & 0xE0) == 0xE0 + + def _select_cv2_stream_buffered_backend() -> Optional[int]: """Return a VideoCapture backend that can read from a Python `BytesIO`. diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index 86041266b134..7c9bc870bee8 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -122,8 +122,12 @@ def _validate_reference_payload(payload: bytes, *, modality: str) -> None: "video_reference is not a recognized media container; supported " "inputs are MP4/AVI video." ) - # audio: no signature sniffing (sniff_media_kind detects only image/video); - # the consuming pipeline validates the audio codec in its worker. + elif modality == "audio": + if sniff_media_kind(payload) != "audio": + raise ValueError( + "audio_reference is not a recognized audio container; supported " + "inputs are WAV/MP3/FLAC/OGG/M4A/AAC." + ) def prepare_reference_slots(params: Any) -> None: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 3f50e6e624d2..1589938faf38 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -622,6 +622,32 @@ def test_json_reference_cannot_declare_bytes(self): class TestMediaBytesProbes: """The in-memory signature probes the serve boundary routes on.""" + def test_sniff_recognizes_audio_containers(self): + """Audio has to classify as itself, not fall through to ``None``. + + ``None`` means "not media", which is what a caller rejects on, so an + audio reference that sniffs to ``None`` is indistinguishable from a + text file. Headers are the real ones ffmpeg emits. + """ + from tensorrt_llm.inputs.media_io import sniff_media_kind + + assert sniff_media_kind(b"RIFF\x2e\x45\x00\x00WAVEfmt ") == "audio" + assert sniff_media_kind(b"fLaC\x00\x00\x00\x22") == "audio" + assert sniff_media_kind(b"OggS\x00\x02\x00\x00\x00\x00\x00\x00") == "audio" + assert sniff_media_kind(b"ID3\x04\x00\x00\x00\x00\x00\x00") == "audio" + # bare MPEG/ADTS frame sync — an MP3 or AAC with no leading tag + assert sniff_media_kind(b"\xff\xf1\x50\x40\x21\x3f\xfc\xde") == "audio" + assert sniff_media_kind(b"\xff\xfb\x90\x00\x00\x00\x00\x00") == "audio" + + def test_sniff_separates_m4a_from_mp4(self): + """An ``.m4a`` is an MP4 whose brand says audio-only; without the brand + check it takes the video default.""" + from tensorrt_llm.inputs.media_io import sniff_media_kind + + assert sniff_media_kind(self._ftyp(b"M4A ", (b"isom", b"mp42"))) == "audio" + assert sniff_media_kind(self._ftyp(b"M4B ")) == "audio" + assert sniff_media_kind(self._ftyp(b"isom", (b"iso2", b"avc1"))) == "video" + def test_sniff_media_kind(self): from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -635,8 +661,9 @@ def test_sniff_media_kind(self): assert sniff_media_kind(TestInputReferenceResolution._avi_bytes()) == "video" assert sniff_media_kind(b"plain text, not media") is None assert sniff_media_kind(b"") is None - # RIFF alone is not AVI (e.g. WAV audio is RIFF too). - assert sniff_media_kind(b"RIFF\x00\x00\x00\x00WAVEfmt ") is None + # RIFF carries both: the form type at [8:12] is what separates them. + assert sniff_media_kind(b"RIFF\x00\x00\x00\x00WAVEfmt ") == "audio" + assert sniff_media_kind(b"RIFF\x00\x00\x00\x00WEBP") is None @staticmethod def _ftyp(major: bytes, compatible: tuple = (), *, size: int = None) -> bytes: From 22ebffa7c4346059e375c04f1f7b7fa63d8bb68c Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:47:50 -0700 Subject: [PATCH 36/61] [TRTLLM-15277][chore] Drop Cosmos3's duplicate video-reference check MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `_validate_video_reference` sniffed the container a second time on the worker. Its docstring still called it a preflight for the `video` extra param, but it reads `video_reference[0].content`, and everything on that path has already been through the coordinator's reference choke point — same sniff, same rejection, one hop earlier and as a synchronous 400 rather than a generation failure. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/models/cosmos3/defaults.py | 11 ----------- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 7 ++++--- 2 files changed, 4 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 0d04c85be07f..9ab33f752848 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -134,17 +134,6 @@ def _validate_output_type(output_type: str) -> None: raise ValueError(f"Cosmos3 output_type must be 'video' or 'image', got {output_type!r}.") -def _validate_video_reference(video) -> None: - """Preflight for the ``video`` extra param: encoded MP4/AVI bytes.""" - if not video: - raise ValueError("Cosmos3 video reference bytes are empty.") - if sniff_media_kind(video) != "video": - raise ValueError( - "Cosmos3 video reference bytes are not a recognized video " - "container (supported: MP4/AVI)." - ) - - # --------------------------------------------------------------------------- # Transfer preflight validators. # diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index fd2953487b0a..abcf55ee0121 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -73,7 +73,6 @@ def tqdm(iterable, **kwargs): _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, resolve_domain_action_config, - _validate_video_reference, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT @@ -759,9 +758,11 @@ def as_given(field_name): return value if field_name in specified else None refs_v = req.params.video_reference + # Container and modality are already checked at the coordinator's + # reference choke point, so the bytes reaching here are known video. video = refs_v[0].content if refs_v else None - if video is not None: - _validate_video_reference(video) + # Container and modality are already checked at the coordinator's + # reference choke point, so the bytes reaching here are known video. is_action = extra_params.get("action_mode") is not None if is_action: # Action resolves its whole recipe in forward() -- the canvas from From 25cdd552db52c98134ab3c963b64f82e0c2932ee Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 01:57:50 -0700 Subject: [PATCH 37/61] [TRTLLM-15277][fix] Bound what a path reference can cost the server MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `format="path"` read whatever it was pointed at, so a remote caller could name `/dev/zero` and the read would never return — a denial of service, not a bad request. `_safe_read_local_file` is the local counterpart of `_safe_request_get`: it requires a regular file within the same 200 MB cap the remote fetch already enforces, and checks the size through `stat` so an oversized file is refused before any bytes come in. `stat` follows symlinks, so a link pointing at a device is refused for what it resolves to rather than what it looks like. This bounds cost, not reach: any regular file the server process can read is still readable, and restricting *which* files a remote caller may name is a deployment-policy question. vLLM answers that one with `--allowed-local-media-path`, which is an allowlist and carries neither of the checks above; sglang-omni has no guard at all and reads local audio unbounded. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/visual_gen/media_refs.py | 43 ++++++++-- .../visual_gen/test_visual_gen_utils.py | 86 +++++++++++++++++++ 2 files changed, 122 insertions(+), 7 deletions(-) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index 7c9bc870bee8..b3a1af7ccb97 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -25,9 +25,11 @@ import base64 from pathlib import Path +from stat import S_ISREG from typing import Any from tensorrt_llm.inputs.media_io import ( + _MAX_RESPONSE_BYTES, _normalize_file_uri, _safe_request_get, is_isobmff_image_bytes, @@ -60,9 +62,39 @@ def _read_reference_payload(reference: str) -> bytes: raise ValueError("reference is not valid base64 data.") from exc -def _local_path(reference: str) -> Path: - """Normalize a ``path`` reference (bare or ``file://``) to a ``Path``.""" - return Path(_normalize_file_uri(reference)) +def _safe_read_local_file(reference: str) -> bytes: + """Read a ``path`` reference, bounding what an unlucky path can cost. + + The counterpart of :func:`_safe_request_get` for the local branch. A + remote caller naming the path is the case worth defending: ``read_bytes`` + on a character device or a FIFO never returns, so an unbounded read is a + denial of service rather than a bad request. Requiring a regular file + within the same size cap the remote fetch uses keeps both branches to one + rule. + + This bounds cost, not reach: any regular file the server process can read + is still readable. Restricting *which* files a remote caller may name is a + deployment-policy question, and belongs with the deployment. + """ + path = Path(_normalize_file_uri(reference)) + try: + stat = path.stat() # follows symlinks, so a link to a device is caught + except OSError as exc: + raise ValueError(f"reference file could not be read: {exc}") from exc + + if not S_ISREG(stat.st_mode): + raise ValueError( + f"reference path is not a regular file: {reference!r}. Character " + "devices, FIFOs and directories cannot be read as media." + ) + if stat.st_size > _MAX_RESPONSE_BYTES: + raise ValueError( + f"reference file is {stat.st_size} bytes, over the {_MAX_RESPONSE_BYTES}-byte limit." + ) + try: + return path.read_bytes() + except OSError as exc: + raise ValueError(f"reference file could not be read: {exc}") from exc def _resolve_reference(content: Any, content_format: str) -> bytes: @@ -90,10 +122,7 @@ def _resolve_reference(content: Any, content_format: str) -> bytes: except Exception as exc: raise ValueError(f"reference URL could not be fetched: {exc}") from exc if content_format == "path": - try: - return _local_path(content).read_bytes() - except OSError as exc: - raise ValueError(f"reference file could not be read: {exc}") from exc + return _safe_read_local_file(content) if content_format == "base64": return _read_reference_payload(content) raise ValueError(f"unsupported reference format: {content_format!r}") diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 1589938faf38..5cd44888d225 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1233,3 +1233,89 @@ def flaky(buffer, **kwargs): del req gc.collect() assert not any(b.exists() for b in blocks) + + +class TestSafeLocalFileRead: + """A ``path`` reference bounds what an unlucky path can cost. + + A remote caller naming the path is the case worth defending: reading a + character device or a FIFO never returns, so an unbounded read is a denial + of service rather than a bad request. + """ + + @staticmethod + def _png(tmp_path): + buffer = BytesIO() + Image.new("RGB", (8, 8)).save(buffer, format="PNG") + target = tmp_path / "ref.png" + target.write_bytes(buffer.getvalue()) + return target + + def test_a_regular_file_reads(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + target = self._png(tmp_path) + + assert _safe_read_local_file(str(target)) == target.read_bytes() + assert _safe_read_local_file(target.as_uri()) == target.read_bytes() + + def test_a_symlink_to_a_regular_file_reads(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + target = self._png(tmp_path) + link = tmp_path / "link.png" + link.symlink_to(target) + + assert _safe_read_local_file(str(link)) == target.read_bytes() + + @pytest.mark.parametrize("device", ["/dev/zero", "/dev/null"]) + def test_a_character_device_is_refused(self, device): + """The unbounded read: `/dev/zero` never reaches EOF.""" + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + with pytest.raises(ValueError, match="not a regular file"): + _safe_read_local_file(device) + + def test_a_symlink_to_a_device_is_refused(self, tmp_path): + """``stat`` follows the link, so the check sees what will be read.""" + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + link = tmp_path / "innocent.png" + link.symlink_to("/dev/zero") + + with pytest.raises(ValueError, match="not a regular file"): + _safe_read_local_file(str(link)) + + def test_a_fifo_is_refused(self, tmp_path): + """Reading a FIFO blocks rather than returning, so size caps cannot help.""" + import os + + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + fifo = tmp_path / "pipe" + os.mkfifo(fifo) + + with pytest.raises(ValueError, match="not a regular file"): + _safe_read_local_file(str(fifo)) + + def test_a_directory_is_refused(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + with pytest.raises(ValueError, match="not a regular file"): + _safe_read_local_file(str(tmp_path)) + + def test_an_oversized_file_is_refused_before_it_is_read(self, tmp_path, monkeypatch): + """The cap is checked against ``stat``, so the bytes never come in.""" + from tensorrt_llm.visual_gen import media_refs + + target = self._png(tmp_path) + monkeypatch.setattr(media_refs, "_MAX_RESPONSE_BYTES", 8) + + with pytest.raises(ValueError, match="over the 8-byte limit"): + media_refs._safe_read_local_file(str(target)) + + def test_a_missing_file_is_a_client_error(self, tmp_path): + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + + with pytest.raises(ValueError, match="could not be read"): + _safe_read_local_file(str(tmp_path / "nope.png")) From 2ef694bedcfa2ee89ecd25f884f7cc41e8d0c733 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:08:41 -0700 Subject: [PATCH 38/61] [TRTLLM-15277][chore] Drop the size cap on path references Reusing the remote fetch's 200 MB ceiling put a threshold where there is no line to draw. `path` exists for the local Python API, where naming a large file of one's own is the normal case, and a legitimate V2V reference can exceed the limit that was borrowed from a setting neither vLLM nor SGLang has an equivalent of. The regular-file check stays and is the one that matters: a character device never reaches EOF and a FIFO blocks, so those are unbounded rather than merely large. A regular file is finite, which is the property worth requiring. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/visual_gen/media_refs.py | 30 ++++++++----------- .../visual_gen/test_visual_gen_utils.py | 10 ------- 2 files changed, 13 insertions(+), 27 deletions(-) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py index b3a1af7ccb97..1d93fbb8e990 100644 --- a/tensorrt_llm/visual_gen/media_refs.py +++ b/tensorrt_llm/visual_gen/media_refs.py @@ -29,7 +29,6 @@ from typing import Any from tensorrt_llm.inputs.media_io import ( - _MAX_RESPONSE_BYTES, _normalize_file_uri, _safe_request_get, is_isobmff_image_bytes, @@ -63,18 +62,19 @@ def _read_reference_payload(reference: str) -> bytes: def _safe_read_local_file(reference: str) -> bytes: - """Read a ``path`` reference, bounding what an unlucky path can cost. - - The counterpart of :func:`_safe_request_get` for the local branch. A - remote caller naming the path is the case worth defending: ``read_bytes`` - on a character device or a FIFO never returns, so an unbounded read is a - denial of service rather than a bad request. Requiring a regular file - within the same size cap the remote fetch uses keeps both branches to one - rule. - - This bounds cost, not reach: any regular file the server process can read - is still readable. Restricting *which* files a remote caller may name is a - deployment-policy question, and belongs with the deployment. + """Read a ``path`` reference, refusing anything that has no end. + + The counterpart of :func:`_safe_request_get` for the local branch, and it + guards the one case that is unbounded rather than merely large: reading a + character device never reaches EOF and reading a FIFO blocks, so either + turns a request into a denial of service. A regular file is finite, which + is the property required here. + + Size is deliberately not capped. ``path`` exists for the local Python API, + where naming a large file of one's own is the normal case, and no + threshold separates that from an abusive one. This bounds the shape of + what may be read, then, not its size or its reach: any regular file the + server process can read is still readable. """ path = Path(_normalize_file_uri(reference)) try: @@ -87,10 +87,6 @@ def _safe_read_local_file(reference: str) -> bytes: f"reference path is not a regular file: {reference!r}. Character " "devices, FIFOs and directories cannot be read as media." ) - if stat.st_size > _MAX_RESPONSE_BYTES: - raise ValueError( - f"reference file is {stat.st_size} bytes, over the {_MAX_RESPONSE_BYTES}-byte limit." - ) try: return path.read_bytes() except OSError as exc: diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 5cd44888d225..5d3a6c871c04 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1304,16 +1304,6 @@ def test_a_directory_is_refused(self, tmp_path): with pytest.raises(ValueError, match="not a regular file"): _safe_read_local_file(str(tmp_path)) - def test_an_oversized_file_is_refused_before_it_is_read(self, tmp_path, monkeypatch): - """The cap is checked against ``stat``, so the bytes never come in.""" - from tensorrt_llm.visual_gen import media_refs - - target = self._png(tmp_path) - monkeypatch.setattr(media_refs, "_MAX_RESPONSE_BYTES", 8) - - with pytest.raises(ValueError, match="over the 8-byte limit"): - media_refs._safe_read_local_file(str(target)) - def test_a_missing_file_is_a_client_error(self, tmp_path): from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file From 09c26fdc4db6489313c584e3e9eced433f9a9a71 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:20:55 -0700 Subject: [PATCH 39/61] [TRTLLM-15277][feat] Let a deployment turn off path references over HTTP `format="path"` asks the server to read its own disk, which is what a co-located client wants and what a remote one has no business doing. Nothing in the code distinguishes those two deployments, so the gate is theirs to set: `TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1` refuses `path` at the HTTP boundary with a 400 naming the alternatives, and it is enabled by default so a working co-located setup keeps working. Shaped after `TRTLLM_DISABLE_RESPONSE_FORMAT_PATH`, which gates the output side of the same concern, down to warning on a value that is neither `0` nor `1` rather than silently reading a typo as "off". The local Python API is unaffected: the gate lives at the serve boundary, not in the shared resolver, so a script naming its own file still works. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 34 ++++++++++++ .../visual_gen/test_visual_gen_utils.py | 54 +++++++++++++++++++ 3 files changed, 89 insertions(+), 1 deletion(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 7b065ac2b739..eef8af0a2624 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -125,7 +125,7 @@ Conditioning references are supplied through the typed, per-modality fields `ima | `format` | Content | Notes | |---|---|---| -| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must exist; it is read once on the coordinator and is never modified or deleted. | +| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must be a regular file and must exist; it is read once on the coordinator and is never modified or deleted. Over HTTP this reads a file on the *server*, so it is only meaningful for a co-located client and can be turned off with `TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1`; the Python API is unaffected. | | `url` | An `http(s)` URL | Fetched on the coordinator through the SSRF-guarded loader. | | `base64` | Base64 text | A `data:` URI is also accepted. | | `bytes` | Raw `bytes` | Python API only. Rejected over JSON (HTTP 422) — send `base64` or upload the file via multipart. | diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 608741b50290..206e0c69ba19 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -109,6 +109,26 @@ def _merge_extra_params( params.extra_params = None +def _reference_path_is_disabled() -> bool: + """Whether ``format='path'`` is turned off for HTTP requests. + + A ``path`` reference asks the server to read its own disk, which is what a + co-located client wants and what a remote one has no business doing. Which + of the two a deployment has is not something the code can know, so it is + enabled by default and disabled with + ``TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1``. The local Python API is + unaffected either way: this gate is the HTTP boundary's. + """ + raw = os.environ.get("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "0") + if raw not in ("0", "1"): + logger.warning( + "Unrecognized value for TRTLLM_DISABLE_REFERENCE_FORMAT_PATH: " + f"{raw!r}. Expected '0' or '1'. Treating as '0' " + "(reference format='path' enabled)." + ) + return raw == "1" + + def _reference_transport(ref: Any) -> tuple[Any, str, Optional[str]]: """Extract ``(content, format, role)`` from one raw HTTP reference. @@ -123,6 +143,13 @@ def _reference_transport(ref: Any) -> tuple[Any, str, Optional[str]]: content = getattr(ref, "content", None) if not isinstance(content, str): raise ValueError("reference item must carry a 'content' string.") + if ref.format == "path" and _reference_path_is_disabled(): + raise ValueError( + "reference format='path' is disabled on this server " + "(TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1); it reads server-side " + "files and is only meaningful for co-located clients. Send the " + "file as base64 or upload it via multipart/form-data." + ) return content, ref.format, getattr(ref, "role", None) @@ -410,6 +437,13 @@ def _apply_deprecated_input_reference( if hasattr(input_reference, "file"): # multipart upload — form implied payload = input_reference.file.read() else: + if input_reference_format == "path" and _reference_path_is_disabled(): + raise ValueError( + "reference format='path' is disabled on this server " + "(TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1); it reads server-side " + "files and is only meaningful for co-located clients. Send the " + "file as base64 or upload it via multipart/form-data." + ) # Imported here, not at module scope: the resolver reaches into # VisualGen, and a plain LLM deployment must not pull a vertical in # behind its request schema. diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 5d3a6c871c04..abf4f725e72d 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1309,3 +1309,57 @@ def test_a_missing_file_is_a_client_error(self, tmp_path): with pytest.raises(ValueError, match="could not be read"): _safe_read_local_file(str(tmp_path / "nope.png")) + + +class TestReferencePathCanBeDisabled: + """``format='path'`` reads server-side files, so a deployment can refuse it. + + Enabled by default: a co-located client naming a shared path is a real + setup, and the code cannot tell that deployment from an untrusted one. + """ + + @staticmethod + def _request(fmt="path", content="/tmp/ref.png"): + return VideoGenerationRequest( + prompt="x", image_reference={"content": content, "format": fmt} + ) + + def test_path_is_accepted_by_default(self, monkeypatch): + monkeypatch.delenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", raising=False) + params = parse_visual_gen_params(self._request(), _StubVisualGen()) + + assert params.image_reference[0].format == "path" + + def test_path_is_refused_when_disabled(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + + with pytest.raises(ValueError, match="is disabled on this server"): + parse_visual_gen_params(self._request(), _StubVisualGen()) + + def test_disabling_path_leaves_the_other_formats_alone(self, monkeypatch): + """The gate is about reading server-side files, not about references.""" + monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + + params = parse_visual_gen_params( + self._request(fmt="base64", content="aGk="), _StubVisualGen() + ) + + assert params.image_reference[0].format == "base64" + + def test_an_unrecognized_value_warns_and_stays_enabled(self, monkeypatch, caplog): + """Treating a typo as "on" would silently break working deployments; + treating it as "off" silently is worse, hence the warning.""" + monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "true") + + params = parse_visual_gen_params(self._request(), _StubVisualGen()) + + assert params.image_reference[0].format == "path" + + def test_the_deprecated_field_is_gated_too(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + request = VideoGenerationRequest( + prompt="x", input_reference="/tmp/ref.png", input_reference_format="path" + ) + + with pytest.raises(ValueError, match="is disabled on this server"): + parse_visual_gen_params(request, _StubVisualGen()) From 34f348c690241052a9ccaa547b9b31d86b0cb72d Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 02:28:29 -0700 Subject: [PATCH 40/61] [TRTLLM-15277][chore] Name the path gate after what it guards The gate turns off reading a server-side file for an HTTP request, so name it for that rather than for the `format` enum value that happens to reach it today: `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH`. This matches the vocabulary vLLM already uses for the same concern (`--allowed-local-media-path`) and leaves room for a second format to reach the same guard without the name contradicting it. Also make the unrecognized-value test verify the warning it is named for. `Logger` sets `propagate = False`, so `caplog`, which collects from the root logger, never received the record and the assertion was never made; the fixture was captured but unused. Collecting through `logger.warning` instead makes the test fail when the warning is silenced, which was confirmed by mutation. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 30 ++++++++--------- .../visual_gen/test_visual_gen_utils.py | 33 +++++++++++-------- 3 files changed, 35 insertions(+), 30 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index eef8af0a2624..17a5318edf83 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -125,7 +125,7 @@ Conditioning references are supplied through the typed, per-modality fields `ima | `format` | Content | Notes | |---|---|---| -| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must be a regular file and must exist; it is read once on the coordinator and is never modified or deleted. Over HTTP this reads a file on the *server*, so it is only meaningful for a co-located client and can be turned off with `TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1`; the Python API is unaffected. | +| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must be a regular file and must exist; it is read once on the coordinator and is never modified or deleted. Over HTTP this reads a file on the *server*, so it is only meaningful for a co-located client and can be turned off with `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1`; the Python API is unaffected. | | `url` | An `http(s)` URL | Fetched on the coordinator through the SSRF-guarded loader. | | `base64` | Base64 text | A `data:` URI is also accepted. | | `bytes` | Raw `bytes` | Python API only. Rejected over JSON (HTTP 422) — send `base64` or upload the file via multipart. | diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 206e0c69ba19..a9d4594bb819 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -109,22 +109,22 @@ def _merge_extra_params( params.extra_params = None -def _reference_path_is_disabled() -> bool: +def _local_media_path_is_disallowed() -> bool: """Whether ``format='path'`` is turned off for HTTP requests. A ``path`` reference asks the server to read its own disk, which is what a co-located client wants and what a remote one has no business doing. Which of the two a deployment has is not something the code can know, so it is - enabled by default and disabled with - ``TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1``. The local Python API is - unaffected either way: this gate is the HTTP boundary's. + allowed by default and turned off with + ``TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1``. The local Python API is unaffected + either way: this gate is the HTTP boundary's. """ - raw = os.environ.get("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "0") + raw = os.environ.get("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "0") if raw not in ("0", "1"): logger.warning( - "Unrecognized value for TRTLLM_DISABLE_REFERENCE_FORMAT_PATH: " + "Unrecognized value for TRTLLM_DISALLOW_LOCAL_MEDIA_PATH: " f"{raw!r}. Expected '0' or '1'. Treating as '0' " - "(reference format='path' enabled)." + "(reference format='path' allowed)." ) return raw == "1" @@ -143,12 +143,12 @@ def _reference_transport(ref: Any) -> tuple[Any, str, Optional[str]]: content = getattr(ref, "content", None) if not isinstance(content, str): raise ValueError("reference item must carry a 'content' string.") - if ref.format == "path" and _reference_path_is_disabled(): + if ref.format == "path" and _local_media_path_is_disallowed(): raise ValueError( - "reference format='path' is disabled on this server " - "(TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1); it reads server-side " - "files and is only meaningful for co-located clients. Send the " - "file as base64 or upload it via multipart/form-data." + "reference format='path' is disallowed on this server " + "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it reads server-side files " + "and is only meaningful for co-located clients. Send the file as " + "base64 or upload it via multipart/form-data." ) return content, ref.format, getattr(ref, "role", None) @@ -437,10 +437,10 @@ def _apply_deprecated_input_reference( if hasattr(input_reference, "file"): # multipart upload — form implied payload = input_reference.file.read() else: - if input_reference_format == "path" and _reference_path_is_disabled(): + if input_reference_format == "path" and _local_media_path_is_disallowed(): raise ValueError( - "reference format='path' is disabled on this server " - "(TRTLLM_DISABLE_REFERENCE_FORMAT_PATH=1); it reads server-side " + "reference format='path' is disallowed on this server " + "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it reads server-side " "files and is only meaningful for co-located clients. Send the " "file as base64 or upload it via multipart/form-data." ) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index abf4f725e72d..b1d3e89f0d21 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -25,6 +25,7 @@ from fastapi import UploadFile from PIL import Image +from tensorrt_llm.serve import visual_gen_utils from tensorrt_llm.serve.openai_protocol import ImageGenerationRequest, VideoGenerationRequest from tensorrt_llm.serve.visual_gen_utils import ( _merge_extra_params, @@ -1311,10 +1312,10 @@ def test_a_missing_file_is_a_client_error(self, tmp_path): _safe_read_local_file(str(tmp_path / "nope.png")) -class TestReferencePathCanBeDisabled: +class TestLocalMediaPathCanBeDisallowed: """``format='path'`` reads server-side files, so a deployment can refuse it. - Enabled by default: a co-located client naming a shared path is a real + Allowed by default: a co-located client naming a shared path is a real setup, and the code cannot tell that deployment from an untrusted one. """ @@ -1325,20 +1326,20 @@ def _request(fmt="path", content="/tmp/ref.png"): ) def test_path_is_accepted_by_default(self, monkeypatch): - monkeypatch.delenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", raising=False) + monkeypatch.delenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", raising=False) params = parse_visual_gen_params(self._request(), _StubVisualGen()) assert params.image_reference[0].format == "path" - def test_path_is_refused_when_disabled(self, monkeypatch): - monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + def test_path_is_refused_when_disallowed(self, monkeypatch): + monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") - with pytest.raises(ValueError, match="is disabled on this server"): + with pytest.raises(ValueError, match="is disallowed on this server"): parse_visual_gen_params(self._request(), _StubVisualGen()) - def test_disabling_path_leaves_the_other_formats_alone(self, monkeypatch): + def test_disallowing_path_leaves_the_other_formats_alone(self, monkeypatch): """The gate is about reading server-side files, not about references.""" - monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") params = parse_visual_gen_params( self._request(fmt="base64", content="aGk="), _StubVisualGen() @@ -1346,20 +1347,24 @@ def test_disabling_path_leaves_the_other_formats_alone(self, monkeypatch): assert params.image_reference[0].format == "base64" - def test_an_unrecognized_value_warns_and_stays_enabled(self, monkeypatch, caplog): - """Treating a typo as "on" would silently break working deployments; - treating it as "off" silently is worse, hence the warning.""" - monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "true") + def test_an_unrecognized_value_warns_and_stays_allowed(self, monkeypatch): + """Silently reading a typo as "1" would break working deployments, and + silently reading it as "0" would leave one that believes it is locked + down wide open. Neither is safe to do quietly.""" + monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "true") + warnings: list[str] = [] + monkeypatch.setattr(visual_gen_utils.logger, "warning", warnings.append) params = parse_visual_gen_params(self._request(), _StubVisualGen()) assert params.image_reference[0].format == "path" + assert any("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH" in w for w in warnings) def test_the_deprecated_field_is_gated_too(self, monkeypatch): - monkeypatch.setenv("TRTLLM_DISABLE_REFERENCE_FORMAT_PATH", "1") + monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") request = VideoGenerationRequest( prompt="x", input_reference="/tmp/ref.png", input_reference_format="path" ) - with pytest.raises(ValueError, match="is disabled on this server"): + with pytest.raises(ValueError, match="is disallowed on this server"): parse_visual_gen_params(request, _StubVisualGen()) From 5a8a5e7fbc744c4d3336ae6dc6c55d8a8d52f138 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:06:51 -0700 Subject: [PATCH 41/61] [TRTLLM-15277][fix] Carry image-edit inputs as bytes without losing their checks MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Upstream shipped a real `/v1/images/edits` while this branch was open, and the branch was written against the 501 stub that endpoint used to be. Rebasing put the two designs on top of each other; this makes the endpoint whole again. `parse_visual_gen_params` no longer writes conditioning inputs to media storage. It decodes each entry — the same read a multipart upload already gets — and hands the engine a `MediaRef(format="bytes")`, so image edit reaches the pipeline the same way video references do. Every check the materialize path performed survives at the boundary: paths and URLs are refused, per-image and total byte limits apply, and the input must still be a PNG or JPEG. Only the disk write is gone, which takes `cleanup_materialized_conditioning_inputs` and the four helpers behind it with it. `HunyuanVideo15Pipeline.infer` guarded text-to-video with `params.image`, a field this branch removes; the guard now reads `image_reference`. That pipeline landed upstream after the migration commits, so it was never converted. The mock generator declared a first-frame slot for every test, which is a video shape; image edit takes joint conditioning images. `ref_slot_specs` is now per-test and the edit tests mirror Qwen-Image-Edit's own declaration. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../pipeline_hunyuan_video1_5.py | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 169 ++++++------------ .../visual_gen/test_trtllm_serve_endpoints.py | 24 ++- 3 files changed, 77 insertions(+), 118 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py index 3c1b806391e2..015dc4452f50 100644 --- a/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py +++ b/tensorrt_llm/_torch/visual_gen/models/hunyuan_video1_5/pipeline_hunyuan_video1_5.py @@ -548,7 +548,7 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N def infer(self, req): """Run inference from a DiffusionRequest (serve / high-level API path).""" - if req.params.image is not None: + if req.params.image_reference: raise ValueError( "HunyuanVideo 1.5 currently supports text-to-video only; " "image conditioning (I2V) is not supported." diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index a9d4594bb819..6673161d8441 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -175,28 +175,81 @@ def _build_reference_list(value: Any) -> Optional[list]: return refs +def _read_image_edit_upload(value: Any) -> bytes: + """Read a multipart image-edit upload, capping it while it streams.""" + 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 b"".join(chunks) + + +def _decode_image_edit_string(value: str) -> bytes: + """Decode one base64 image-edit input, refusing paths and URLs. + + OpenAI's ``image`` field is base64 or an upload. A path or URL there would + ask the server to fetch on the client's behalf, which this endpoint has + never offered. + """ + decoded = _decode_base64_media(value) + if decoded is not None: + return decoded + if urlparse(value).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.") + + def _build_image_edit_reference_list(value: Any) -> Optional[list]: """Build references from an image-edit request's OpenAI-shaped ``image``. That field follows OpenAI's schema, which has no place to declare a wire - form: an entry is a bare base64 string or a multipart upload, so the format - is implied by the transport instead of read off the item. + form, so each entry is decoded here — the same read a multipart upload + already gets — and carried on as ``bytes``. The size and PNG/JPEG checks + stay at the boundary because only this endpoint imposes them. """ if value is None: return None from tensorrt_llm.visual_gen.params import MediaRef refs = [] + total_bytes = 0 for item in value if isinstance(value, list) else [value]: if hasattr(item, "file"): # multipart UploadFile - refs.append(MediaRef(content=item.file.read(), format="bytes")) + payload = _read_image_edit_upload(item) elif isinstance(item, str): - refs.append(MediaRef(content=item, format="base64")) + payload = _decode_image_edit_string(item) else: raise ValueError( - "image edit inputs must be base64-encoded images or uploaded files, " + "Image edit inputs must be base64-encoded images or uploaded files, " f"got {type(item).__name__}." ) + if len(payload) > IMAGE_EDIT_MAX_IMAGE_BYTES: + raise ValueError( + "Image edit input exceeds the per-image byte limit " + f"({len(payload)} > {IMAGE_EDIT_MAX_IMAGE_BYTES})." + ) + _validate_png_jpeg_image(payload) + total_bytes += len(payload) + 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})." + ) + refs.append(MediaRef(content=payload, format="bytes")) return refs @@ -250,19 +303,6 @@ def _decode_base64_media(value: str) -> Optional[bytes]: 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: @@ -274,58 +314,6 @@ def _validate_png_jpeg_image(value: bytes) -> None: 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, @@ -372,47 +360,6 @@ def _validate_image_edit_request_limits( 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 _apply_deprecated_input_reference( input_reference: str | UploadFile | None, params: VisualGenParams, 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 44fe4e7e5803..073b3d5fd5e7 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -189,6 +189,7 @@ def __init__( extra_param_specs: Optional[dict] = None, model: str = "test-model", supports_image_edit: bool = False, + ref_slot_specs: Optional[dict] = None, ): from types import SimpleNamespace @@ -236,7 +237,9 @@ def __init__( extra_param_specs=extra_param_specs or {"stg_scale": ExtraParamSchema(type="float", default=1.0)}, supports_image_edit=supports_image_edit, - ref_slot_specs={ + ref_slot_specs=ref_slot_specs + if ref_slot_specs is not None + else { "image_reference": RefSlotSpec( modality="image", roles=[RoleSpec(role="first_frame", min=0, max=1)] ), @@ -994,12 +997,21 @@ def _client( should_fail: bool = False, supports_image_edit: bool = True, ): + from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec + 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, + # Edit pipelines take joint conditioning images, not a first frame; + # mirrors Qwen-Image-Edit's own slot declaration. + ref_slot_specs={ + "image_reference": RefSlotSpec( + modality="image", roles=[RoleSpec(role="reference", min=1, max=None)] + ), + }, ) monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) return _create_server(gen, model_name=model), gen @@ -1057,8 +1069,8 @@ def test_image_edit_accepts_json_base64_image(self, tmp_path, monkeypatch): ) 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.image_reference[0].format == "bytes" + assert isinstance(gen.last_params.image_reference[0].content, bytes) assert gen.last_params.num_images_per_prompt == 2 body = resp.json() assert body["output_format"] == "webp" @@ -1107,8 +1119,8 @@ def test_image_edit_default_url_returns_fetchable_output(self, tmp_path, monkeyp body = resp.json() 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) + assert gen.last_params.image_reference[0].format == "bytes" + assert isinstance(gen.last_params.image_reference[0].content, bytes) path = url.split("//", 1)[-1].split("/", 1)[1] content = client.get("/" + path) @@ -1329,7 +1341,7 @@ def test_image_edit_allows_max_input_images_without_output_fanout(self, tmp_path ) assert resp.status_code == 200 - assert len(gen.last_params.image) == 16 + assert len(gen.last_params.image_reference) == 16 assert len(resp.json()["data"]) == 1 assert list(tmp_path.iterdir()) == [] From 015a1471f95a7a4040d9b868b1533801f084454b Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 21 Aug 2026 05:27:27 -0700 Subject: [PATCH 42/61] [TRTLLM-15277][fix] Hand back the shared memory of a request that never ships `refs_to_handles` publishes its handle list before filling it so that a reference failing partway leaves the blocks it already took reachable rather than orphaned. Nothing used that: the sender thread reclaims what it fails to send and shutdown reclaims what is still pending, but a failure between minting the handles and handing the request to the executor fell through both, and an unconsumed handle keeps its block mapped until the process exits. Injecting a failure on the second of two references leaks one block; with the call site reclaiming, none survive. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/visual_gen/visual_gen.py | 11 +++-- .../visual_gen/test_visual_gen_utils.py | 42 +++++++++++++++++++ 2 files changed, 50 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index e4c64a749a5e..7fb4d43bbc5f 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -452,9 +452,14 @@ def generate_async( # Hand the reference payloads to rank0 through shared memory instead of # through the request pickle, which copies every reference byte to # cross one process boundary. - request.refs_to_handles() - - self.executor.enqueue_requests([request]) + try: + request.refs_to_handles() + self.executor.enqueue_requests([request]) + except Exception: + # The request never reached rank0, so nothing downstream will + # consume the blocks it already took. + request.refs_to_bytes() + raise return VisualGenResult(req_id, self.executor, batch_size=batch_size) @staticmethod diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index b1d3e89f0d21..23506844713b 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1235,6 +1235,48 @@ def flaky(buffer, **kwargs): gc.collect() assert not any(b.exists() for b in blocks) + def test_a_request_that_never_ships_hands_its_blocks_back(self): + """The blocks staying reachable is only half of it; the call site is + what has to hand them back when the request never reaches rank0.""" + import gc + import itertools + from types import SimpleNamespace + + from tensorrt_llm.visual_gen import VisualGen, VisualGenParams + from tensorrt_llm.visual_gen.params import MediaRef + + taken = {} + blocks_of = self._blocks_of + + class _DeadExecutor: + default_generation_params = {} + extra_param_specs = {} + ref_slot_specs = None + + def enqueue_requests(self, requests): + # Read the handles while they still exist, then fail the way a + # dead worker queue would. + taken["blocks"] = blocks_of(requests[0]) + raise RuntimeError("worker queue is gone") + + caller = SimpleNamespace( + _req_counter=itertools.count(), + default_params=VisualGenParams(), + executor=_DeadExecutor(), + ) + buf = BytesIO() + # A real PNG: the engine choke point content-checks before the handles + # are minted, so random bytes never get far enough to take a block. + Image.new("RGB", (256, 256)).save(buf, format="PNG") + params = VisualGenParams(image_reference=[MediaRef(content=buf.getvalue(), format="bytes")]) + + with pytest.raises(RuntimeError, match="worker queue is gone"): + VisualGen.generate_async(caller, "x", params) + + assert taken["blocks"], "the request must have taken a block to begin with" + gc.collect() + assert not any(b.exists() for b in taken["blocks"]) + class TestSafeLocalFileRead: """A ``path`` reference bounds what an unlucky path can cost. From 33872e00f4369aeff2405ead715c9d77a8dd927a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Sun, 23 Aug 2026 21:36:31 -0700 Subject: [PATCH 43/61] [TRTLLM-15277][fix] Release reference blocks the workers never take Consuming a handle is what frees its shared-memory block, so a handle nobody consumes keeps its block mapped until the process exits. Two paths stopped short of consuming and nothing else picked them up. The restore loop is also the reclaim path, but it stopped at the first handle it could not rebuild, stranding every block behind it. Each handle is now taken independently and the failures are reported together. Blocks of a request already handed to the workers were released by nothing if the workers then died. The client tracks what it has sent, forgets a request as soon as a response for it arrives, and releases the rest once the worker processes are gone or at shutdown. Release unlinks the block by name rather than rebuilding it, which makes it idempotent and safe to reach from more than one place. It is keyed on the workers being gone rather than on the caller losing interest: rank0 keeps going after a cancelled wait and still consumes the handle, and unlinking a block it has not opened yet would turn its rebuild into a failure. CUDA handles carry no unlinkable name, so this covers CPU blocks only. Releasing one early needs a pool the sender returns slots to instead of a handle minted per request, which belongs in the shared-tensor layer the LLM path sits on too. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 98 ++++++++++++++++++- .../test_executor_shared_tensor_ipc.py | 83 ++++++++++++++++ tests/unittest/visual_gen/test_output.py | 1 + 3 files changed, 177 insertions(+), 5 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 65bbb14efddb..461ded41f4c9 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -1,4 +1,5 @@ import asyncio +import base64 import os import queue import socket @@ -15,7 +16,10 @@ import torch.multiprocessing as mp import zmq -from tensorrt_llm._torch.shared_tensor import SharedTensorContainer +from tensorrt_llm._torch.shared_tensor import ( + SharedTensorContainer, + _SharedTensorRebuildMethodRegistry, +) from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.executor.ipc import ZeroMqQueue @@ -26,6 +30,39 @@ if TYPE_CHECKING: from tensorrt_llm.visual_gen.params import VisualGenParams + +def _release_cpu_ref_blocks(handles: Optional[List[Dict]]) -> int: + """Unlink the shared-memory blocks behind reference handles nobody will take. + + Only call this once the consumer is known never to open them: unlinking a + block the consumer has not yet opened turns its rebuild into a failure. A + dead worker is that signal; a request merely abandoned by its caller is not, + because rank0 keeps going and still consumes the handle. + + Unlinking rather than rebuilding keeps it idempotent -- a block the consumer + already released is simply absent. CUDA handles carry no unlinkable name and + are left alone; the caching allocator reclaims them only once the consumer + process is gone. + """ + released = 0 + for entry in handles or []: + handle = entry.get("handle") or {} + if handle.get("method_key") != _SharedTensorRebuildMethodRegistry.REBUILD_CPU: + continue + raw = handle.get("storage_handle") + if not raw: + continue + try: + name = base64.b64decode(raw).decode().lstrip("/") + os.unlink(os.path.join("/dev/shm", name)) + released += 1 + except FileNotFoundError: + pass + except OSError as exc: + logger.warning(f"DiffusionClient: could not release shared block {raw!r}: {exc}") + return released + + # Timeouts (seconds) for the client-side coordinator. POLL_TIMEOUT = 0.01 AWAIT_TIMEOUT = 0.05 @@ -294,12 +331,25 @@ def refs_to_handles(self) -> None: self.ref_handles = None def refs_to_bytes(self) -> None: - """Restore reference payloads from shared memory, in place (consumer side).""" + """Restore reference payloads from shared memory, in place (consumer side). + + Each handle is taken independently: one that cannot be rebuilt must not + strand the blocks behind it, since this is also the reclaim path and a + handle nobody consumes keeps its block mapped until the process exits. + """ + failures = [] for entry in self.ref_handles or []: - ref = getattr(self.params, entry["slot"])[entry["index"]] - container = SharedTensorContainer.from_dict(entry["handle"]) - ref.content = container.get_local_view().numpy().tobytes() + try: + ref = getattr(self.params, entry["slot"])[entry["index"]] + container = SharedTensorContainer.from_dict(entry["handle"]) + ref.content = container.get_local_view().numpy().tobytes() + except Exception as exc: + failures.append(f"{entry['slot']}[{entry['index']}]: {exc}") self.ref_handles = None + if failures: + raise RuntimeError( + "failed to restore reference payloads from shared memory: " + "; ".join(failures) + ) def _refs(self): for slot in ("image_reference", "video_reference", "audio_reference"): @@ -807,6 +857,10 @@ def __init__( # full PipelineOutput tensor does not pin in completed_responses for # the process lifetime. self._abandoned_request_ids: Set[int] = set() + # Handles of requests already handed to the workers. rank0 frees a + # block by consuming it, so an entry here means "sent, not yet known + # consumed" -- what has to be released if the workers die first. + self._sent_ref_handles: Dict[int, List[Dict]] = {} # Iteration-stats tracker — populated on lifecycle events (enqueue, # request started, response received) and drained by @@ -1003,6 +1057,8 @@ def _process_requests(self): logger.info(f"DiffusionClient: Sending request {req.request_id}") self.requests_ipc.put(req) + if req.ref_handles: + self._sent_ref_handles[req.request_id] = req.ref_handles # Once the request has been handed to the workers it becomes the # in-flight ("active") request from the client's perspective. self._iter_stats.record_request_started(req.request_id, self.pending_requests.qsize()) @@ -1036,6 +1092,22 @@ def _process_responses(self): except Exception as e: logger.error(f"DiffusionClient: Error processing response: {e}") + def _forget_sent_handles(self, request_id: int) -> None: + """A response means rank0 got that far, so its blocks are its own now.""" + self._sent_ref_handles.pop(request_id, None) + + def _release_sent_handles(self, reason: str) -> None: + """Release the blocks of every request the workers can no longer take.""" + if not self._sent_ref_handles: + return + outstanding, self._sent_ref_handles = self._sent_ref_handles, {} + released = sum(_release_cpu_ref_blocks(h) for h in outstanding.values()) + if released: + logger.warning( + f"DiffusionClient: released {released} shared block(s) from " + f"{len(outstanding)} unfinished request(s) after {reason}" + ) + async def _store_response(self, response: DiffusionResponse): """Store response in the completed_responses dict (async helper). @@ -1043,6 +1115,7 @@ async def _store_response(self, response: DiffusionResponse): late-arriving responses for timed-out requests do not leak into ``completed_responses`` for the process lifetime. """ + self._forget_sent_handles(response.request_id) async with self.lock: if response.request_id in self._abandoned_request_ids: self._abandoned_request_ids.discard(response.request_id) @@ -1131,10 +1204,23 @@ async def _serve_forever(self): while not self.shutdown_event.is_set(): self._process_requests() self._process_responses() + if self._sent_ref_handles and self._workers_gone(): + # A dead worker cannot open a block, so the ones it never got + # to are safe to release -- and nothing else will. + self._release_sent_handles("worker exit") await asyncio.sleep(0.001) # Yield control to allow other coroutines to run self._cleanup_ipc() + def _workers_gone(self) -> bool: + """Whether the processes (or the external-launch thread) that consume + requests have exited.""" + if self.worker_processes: + return any(not p.is_alive() for p in self.worker_processes) + if self._ext_worker_thread is not None: + return not self._ext_worker_thread.is_alive() + return False + def _reclaim_pending_handles(self) -> None: """Consume the handles of requests that will never be sent. @@ -1162,6 +1248,8 @@ def shutdown(self): self.background_thread.join(timeout=1.0) self._reclaim_pending_handles() + self._reclaim_pending_handles() + # Shutdown workers logger.info("DiffusionClient: Stopping workers") for p in self.worker_processes: diff --git a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py index 35f84a1c7556..d29a249b746a 100644 --- a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py +++ b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py @@ -8,18 +8,25 @@ main process, mirroring the real worker/client split. """ +import base64 import multiprocessing as mp import os import unittest +from pathlib import Path from unittest import mock import torch import zmq +from tensorrt_llm._torch.shared_tensor import ( + SharedTensorContainer, + _SharedTensorRebuildMethodRegistry, +) from tensorrt_llm._torch.visual_gen.executor import ( DiffusionExecutor, DiffusionRequest, DiffusionResponse, + _release_cpu_ref_blocks, find_free_port, run_diffusion_worker, ) @@ -258,6 +265,82 @@ def test_spawn_worker_response_rebuilds_in_client(self): q.close() +class TestUnconsumedReferenceBlocks(unittest.TestCase): + """A handle nobody consumes keeps its block mapped until the process exits. + + Consuming is the normal release, so these cover the two ways a request can + stop short of it: a rebuild that fails partway, and workers that die before + reading what was already sent to them. Release is keyed on the workers being + gone rather than on the caller losing interest -- unlinking a block rank0 has + not opened yet would turn its rebuild into a failure. + """ + + @staticmethod + def _request(*payloads): + from tensorrt_llm.visual_gen import MediaRef, VisualGenParams + + params = VisualGenParams( + image_reference=[MediaRef(content=p, format="bytes") for p in payloads] + ) + return DiffusionRequest(request_id=7, prompt=["x"], params=params) + + @staticmethod + def _blocks_of(handles): + return [ + Path("/dev/shm") / base64.b64decode(e["handle"]["storage_handle"]).decode().lstrip("/") + for e in handles or [] + ] + + def test_one_bad_handle_does_not_strand_the_rest(self): + import gc + + req = self._request(os.urandom(4096), os.urandom(4096), os.urandom(4096)) + req.refs_to_handles() + blocks = self._blocks_of(req.ref_handles) + self.assertEqual(len(blocks), 3) + + real = SharedTensorContainer.from_dict + calls = {"n": 0} + + def flaky(handle): + calls["n"] += 1 + if calls["n"] == 2: + raise RuntimeError("handle is unusable") + return real(handle) + + with mock.patch.object(SharedTensorContainer, "from_dict", staticmethod(flaky)): + with self.assertRaises(RuntimeError): + req.refs_to_bytes() + + self.assertEqual(calls["n"], 3, "a failure must not stop the loop at the bad entry") + del req + gc.collect() + self.assertLessEqual(sum(b.exists() for b in blocks), 1) + + def test_blocks_of_a_dead_workers_request_are_released(self): + req = self._request(os.urandom(64 * 1024), os.urandom(64 * 1024)) + req.refs_to_handles() + blocks = self._blocks_of(req.ref_handles) + self.assertTrue(all(b.exists() for b in blocks)) + + self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 2) + self.assertFalse(any(b.exists() for b in blocks)) + + def test_releasing_twice_is_harmless(self): + """Reclaim can fire from more than one place; it must not care.""" + req = self._request(os.urandom(4096)) + req.refs_to_handles() + + self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 1) + self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 0) + + def test_a_cuda_handle_is_left_alone(self): + """CUDA blocks carry no unlinkable name; skipping them is the contract.""" + cuda_key = _SharedTensorRebuildMethodRegistry.REBUILD_CUDA + cuda_shaped = [{"slot": "image_reference", "index": 0, "handle": {"method_key": cuda_key}}] + self.assertEqual(_release_cpu_ref_blocks(cuda_shaped), 0) + + if __name__ == "__main__": mp.set_start_method("spawn", force=True) unittest.main() diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index bfbaafdc4380..411f55f7b49b 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -669,6 +669,7 @@ def _make_minimal_client_state(): client = DiffusionRemoteClient.__new__(DiffusionRemoteClient) client.completed_responses = {} client._abandoned_request_ids = set() + client._sent_ref_handles = {} client._iter_stats = _IterationStatsTracker() client.pending_requests = queue.Queue() From 2deea288a66f76f8de6ed5bad6f5dbd049d72a66 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 24 Aug 2026 00:17:45 -0700 Subject: [PATCH 44/61] [TRTLLM-15277][doc] Align reference-input docs with each page's scope The reference-input section explained its own rationale where the rest of the page states behavior, and was the only section carrying curl examples -- visual-generation.md documents the Python side and points at examples/visual_gen/serve/ for request examples. Keep the Python API and the format table there, move the serve-only material (multipart, the top-level output format, input_reference) to the serve README, and give that README the role example it was missing. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 48 +++++-------------------- examples/visual_gen/serve/README.md | 19 +++++++--- 2 files changed, 24 insertions(+), 43 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 17a5318edf83..9435ce97f616 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -121,47 +121,29 @@ The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queue ### Reference Inputs -Conditioning references are supplied through the typed, per-modality fields `image_reference`, `video_reference`, and `audio_reference`. These fields share the **same names and shapes** across the Python API (`VisualGenParams`) and the serve request (`VideoGenerationRequest`), and each accepts a single reference or a list. A reference always declares the wire form of its content — `MediaRef(content=..., format=...)` in Python, `{"content": ..., "format": ...}` in JSON. `format` is **required** and nothing is guessed: a bare string or bare bytes is rejected, so a mistyped path can never be silently read as base64. Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`; a request is validated against that declaration before generation begins, so a missing required reference, an excess reference, or an unsupported role is rejected at the boundary. Whatever form a reference is declared in, it is resolved to raw bytes on the coordinator before the request is broadcast, so a worker never needs a filesystem shared with the client; `http(s)` URLs are fetched through the same SSRF-guarded loader as the LLM multimodal path (private-address block, redirect re-validation, timeout, and size cap). +Conditioning references are supplied through the typed fields `image_reference`, `video_reference`, and `audio_reference`. Each field takes a single reference or a list. A reference is `MediaRef(content=..., format=...)`, and `format` is required. | `format` | Content | Notes | |---|---|---| -| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. The file must be a regular file and must exist; it is read once on the coordinator and is never modified or deleted. Over HTTP this reads a file on the *server*, so it is only meaningful for a co-located client and can be turned off with `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1`; the Python API is unaffected. | +| `path` | A local file readable by the coordinator process | Bare path or `file://` URI. | | `url` | An `http(s)` URL | Fetched on the coordinator through the SSRF-guarded loader. | | `base64` | Base64 text | A `data:` URI is also accepted. | -| `bytes` | Raw `bytes` | Python API only. Rejected over JSON (HTTP 422) — send `base64` or upload the file via multipart. | +| `bytes` | Raw `bytes` | Python API only. | -The `file://` and `data:` prefixes are still accepted, but are no longer needed to disambiguate: `format` already states which form the content is in. A multipart file upload carries no `format` — the server implies `bytes` from the transport. A reference item's `format` is distinct from the request's top-level `format` field, which selects the *output* encoding (`mp4`, `png`, …). +Every pipeline declares the reference slots and roles it accepts through `ref_slot_specs`, and a request is validated against that declaration before generation begins. References are resolved to raw bytes on the coordinator, so a worker never needs a filesystem shared with the client. -Most models take a single reference whose role is unambiguous, so no `role` is specified: +Most models take a single reference whose role is unambiguous: ```python from tensorrt_llm import VisualGen, MediaRef -# The image conditions the generated video's first frame. vg = VisualGen(model="Wan-AI/Wan2.2-TI2V-5B-Diffusers") params = vg.default_params params.image_reference = MediaRef(content="start.png", format="path") output = vg.generate(inputs="the scene comes alive with gentle motion", params=params) - -# Cosmos conditions generation on a reference video. -vg = VisualGen(model="nvidia/Cosmos3-Super") -params = vg.default_params -params.video_reference = MediaRef(content="clip.mp4", format="path") ``` -The equivalent serve request uploads the file, or sends a `{content, format}` object as the field value in a JSON body: - -```bash -# multipart file upload (raw bytes, no base64) -curl http://localhost:8000/v1/videos -F "prompt=the scene comes alive" -F "image_reference=@start.png" -curl http://localhost:8000/v1/videos -F "prompt=continue the scene" -F "video_reference=@clip.mp4" - -# JSON body: the content plus the format it is in (path, url, or base64) -curl http://localhost:8000/v1/videos -H 'content-type: application/json' \ - -d '{"prompt": "the scene comes alive", "image_reference": {"content": "https://example.com/start.png", "format": "url"}}' -``` - -When a model accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame — the `role` is required to disambiguate: +Models that accept the same modality in more than one role need `role`. Wan 2.1 I2V takes a first frame and an optional last frame: ```python from tensorrt_llm import VisualGen, MediaRef @@ -170,25 +152,13 @@ vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") params = vg.default_params params.image_reference = [ MediaRef(content="start.png", format="path", role="first_frame"), - MediaRef(content="end.png", format="path", role="last_frame"), # optional + MediaRef(content="end.png", format="path", role="last_frame"), ] ``` -A JSON serve request carries the role and lists; a multipart upload is limited to a single file with no role: - -```bash -curl http://localhost:8000/v1/videos -H 'content-type: application/json' -d '{ - "prompt": "the subject comes alive", - "image_reference": [ - {"content": "", "format": "base64", "role": "first_frame"}, - {"content": "", "format": "base64", "role": "last_frame"} - ] -}' -``` - -FLUX.2 and Qwen-Image-Edit accept multiple reference images as a list on the same `image_reference` field through the Python API. +FLUX.2 and Qwen-Image-Edit accept a list of reference images on `image_reference`. -A single `input_reference` field (deprecated) is still accepted on the serve video endpoints for backward compatibility; it is routed by content signature to image-to-video or video-to-video, and is ignored when a typed `image_reference` / `video_reference` is also provided. Being a bare value, it declares its wire form through the sibling `input_reference_format` field (`path` / `url` / `base64`), which is required when `input_reference` is a string and implied for a multipart upload. Prefer the typed fields. +The same fields carry references over `trtllm-serve`; see [`examples/visual_gen/serve/`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/visual_gen/serve) for request examples. ## Optimizations diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index 5e605f800568..e789fb746627 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -286,14 +286,26 @@ You can customize these by: - `frame_rate` (canonical) or `fps` (alias): frames per second - `num_frames`: when set, wins over the `seconds * frame_rate` derivation - `seed`, `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls -- `image_reference`: Reference image(s) for I2V/TI2V. `video_reference`: reference video(s) for V2V. `audio_reference`: reference audio(s). In JSON each accepts a `{content, format, role}` object or a list of them; `format` is required and declares how to read `content` — `"path"` (a file readable by the server; a `file://` URI is also accepted), `"url"` (`http(s)`), or `"base64"` (a `data:` URI is also accepted). Nothing is guessed, so a bare string is rejected, and `"bytes"` is rejected over JSON — upload the file instead. A multipart file upload needs no `format`: the transport implies it. +- `image_reference`, `video_reference`, `audio_reference`: reference image(s) for I2V/TI2V, video(s) for V2V, audio(s). In JSON each takes a `{content, format, role}` object or a list of them; a multipart file upload needs no `format`. + - `format` declares how to read `content`: `"path"` (a file readable by the server, or a `file://` URI), `"url"` (`http(s)`), or `"base64"` (or a `data:` URI). It is required in JSON, where `"bytes"` is rejected — upload the file instead. ```json {"image_reference": {"content": "iVBORw0KGgoAAAANSUhEUg...", "format": "base64"}} ``` + - `"path"` reads a file on the *server*, so it is only meaningful for a co-located client; set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject it. `"url"` is fetched through the SSRF-guarded loader (private-address block, redirect re-validation, timeout, size cap). - `format` here is the *input* wire form; the top-level `format` selects the *output* encoding. + - `role` disambiguates a model that accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame. Roles and lists need a JSON body; a multipart upload is a single file with no role. + + ```json + {"image_reference": [ + {"content": "", "format": "base64", "role": "first_frame"}, + {"content": "", "format": "base64", "role": "last_frame"} + ]} + ``` + - **Supported formats**: PNG and JPEG images; MP4 and AVI video, with H.264 the tested codec and others best-effort. HEIF/AVIF are not supported. +- `input_reference` (deprecated): a single image or video reference, routed by content signature to I2V or V2V. Declares its wire form through the sibling `input_reference_format` field, and is ignored when `image_reference` / `video_reference` is also given. Prefer the typed fields. - `extra_params`: model-specific overflow (see below) - `response_format`: `"file"` (default; `FileResponse` byte download) or `"path"` (server-side output path JSON, for co-located clients) - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). @@ -399,9 +411,8 @@ curl -X POST "http://localhost:8000/v1/videos" \ ### Video-to-Video (Multipart with File Upload, Cosmos3) ```bash -# The modality is declared by the field name: image_reference -> I2V, -# video_reference -> V2V. V2V conditioning knobs ride in extra_params -# (values below are the defaults). +# Modality comes from the field name: image_reference -> I2V, video_reference -> V2V. +# V2V conditioning knobs ride in extra_params (values below are the defaults). curl -X POST "http://localhost:8000/v1/videos" \ -F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \ -F "video_reference=@./media/reference.mp4" \ From 8f7699bda3638f6a0f1ffc5a5a81e35189264cb8 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 24 Aug 2026 01:34:43 -0700 Subject: [PATCH 45/61] [TRTLLM-15277][perf] Drop the payload copy that only fed frombuffer ``bytearray(ref.content)`` existed to satisfy torch.frombuffer, which refuses a read-only buffer -- and ``bytes`` is always read-only. The copy carried no information: from_tensor() immediately copies again, into shared memory, so every reference was duplicated twice on its way to rank0. Hand frombuffer the bytes directly. It only reads them, and so does the broadcast on the src rank. Measured on a 256MB payload: 161.5ms -> 29.3ms; a 2MB image reference goes 0.76ms -> 0.48ms. torch warns once per process that the buffer is not writable, then suppresses itself. Nothing writes through these tensors. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 461ded41f4c9..a9deb4c9e065 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -317,8 +317,9 @@ def refs_to_handles(self) -> None: self.ref_handles = handles = [] for slot in ("image_reference", "video_reference", "audio_reference"): for index, ref in enumerate(getattr(self.params, slot, None) or []): - # bytearray() because the shared storage must be writable. - buffer = torch.frombuffer(bytearray(ref.content), dtype=torch.uint8) + # A read-only view: the one copy happens inside from_tensor(), + # which is what moves the payload into shared memory. + buffer = torch.frombuffer(ref.content, dtype=torch.uint8) handles.append( { "slot": slot, @@ -536,8 +537,8 @@ def _broadcast_request(self, req: Optional[DiffusionRequest]) -> Optional[Diffus return None if self.rank == 0: - # bytearray() because a tensor over an immutable buffer is read-only. - buffers = [torch.frombuffer(bytearray(p), dtype=torch.uint8) for p in payloads] + # Read-only views: the src rank only reads its buffer. + buffers = [torch.frombuffer(p, dtype=torch.uint8) for p in payloads] else: buffers = [torch.empty(n, dtype=torch.uint8) for n in req.ref_sizes or []] for buffer in buffers: From 00959e32cd0c7a5b69b782c2bcace9be5a5bab8c Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 24 Aug 2026 03:03:18 -0700 Subject: [PATCH 46/61] [TRTLLM-15277][chore] Name the shared-memory hop for what it does refs_to_handles/refs_to_bytes named the representation on each side of the coordinator->rank0 hop, which left the pair reading as a format conversion. Their own docstrings already said "into shared memory" and "from shared memory"; the names now say the same thing. The second hop keeps detach/attach: it lifts payloads out of the object for the collective rather than moving them anywhere. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 12 ++++----- tensorrt_llm/visual_gen/visual_gen.py | 4 +-- .../visual_gen/test_visual_gen_utils.py | 26 +++++++++---------- .../test_executor_shared_tensor_ipc.py | 8 +++--- 4 files changed, 25 insertions(+), 25 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index a9deb4c9e065..cf627bc33028 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -296,13 +296,13 @@ class DiffusionRequest: params: Optional["VisualGenParams"] = None prepared_inputs: Dict[str, Any] = field(default_factory=dict, repr=False) # Set only between the two ends of the coordinator -> rank0 hop; see - # ``refs_to_handles``. + # ``refs_to_shm``. ref_handles: Optional[List[Dict[str, Any]]] = field(default=None, repr=False) # Set only while the request is in flight on the rank0 -> N-rank hop; see # ``refs_detach``. ref_sizes: Optional[List[int]] = field(default=None, repr=False) - def refs_to_handles(self) -> None: + def refs_to_shm(self) -> None: """Move reference payloads into shared memory, in place (producer side). Only the coordinator -> rank0 hop travels as handles: rank0 restores the @@ -331,7 +331,7 @@ def refs_to_handles(self) -> None: if not handles: self.ref_handles = None - def refs_to_bytes(self) -> None: + def refs_from_shm(self) -> None: """Restore reference payloads from shared memory, in place (consumer side). Each handle is taken independently: one that cannot be rebuilt must not @@ -556,7 +556,7 @@ def serve_forever(self): if self.rank == 0: req = self.requests_ipc.get() if req is not None: - req.refs_to_bytes() + req.refs_from_shm() logger.info(f"Worker {self.device_id}: Request available") # Broadcast to all ranks. ``req.params.seed`` is already a @@ -1072,7 +1072,7 @@ def _process_requests(self): # The request never reached rank0, so nothing downstream will # consume its handles, and an unconsumed handle keeps its # shared-memory block mapped until this process exits. - req.refs_to_bytes() + req.refs_from_shm() def _process_responses(self): """Poll and process responses.""" @@ -1235,7 +1235,7 @@ def _reclaim_pending_handles(self) -> None: except queue.Empty: return if req is not None and req.ref_handles: - req.refs_to_bytes() + req.refs_from_shm() def shutdown(self): """Shutdown client and workers.""" diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 7fb4d43bbc5f..809e46d62f36 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -453,12 +453,12 @@ def generate_async( # through the request pickle, which copies every reference byte to # cross one process boundary. try: - request.refs_to_handles() + request.refs_to_shm() self.executor.enqueue_requests([request]) except Exception: # The request never reached rank0, so nothing downstream will # consume the blocks it already took. - request.refs_to_bytes() + request.refs_from_shm() raise return VisualGenResult(req_id, self.executor, batch_size=batch_size) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 23506844713b..bdbf3df6da18 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1065,8 +1065,8 @@ def test_round_trip_is_byte_identical(self): payloads = (b"\x89PNG\r\n\x1a\n" + os.urandom(4096), os.urandom(1024)) req = self._request(*payloads) - req.refs_to_handles() - req.refs_to_bytes() + req.refs_to_shm() + req.refs_from_shm() assert tuple(r.content for r in req.params.image_reference) == payloads assert all(r.format == "bytes" for r in req.params.image_reference) @@ -1078,7 +1078,7 @@ def test_payload_leaves_the_request_pickle(self): req = self._request(payload) before = len(pickle.dumps(req)) - req.refs_to_handles() + req.refs_to_shm() after = len(pickle.dumps(req)) assert after < before - len(payload) // 2 @@ -1089,10 +1089,10 @@ def test_survives_a_real_pickle_round_trip(self): and rebuilt, which is what the IPC queue does to it.""" payload = os.urandom(8192) req = self._request(payload) - req.refs_to_handles() + req.refs_to_shm() received = pickle.loads(pickle.dumps(req)) - received.refs_to_bytes() + received.refs_from_shm() assert received.params.image_reference[0].content == payload @@ -1102,7 +1102,7 @@ def test_no_references_costs_nothing(self): from tensorrt_llm.visual_gen import VisualGenParams req = DiffusionRequest(request_id=1, prompt=["x"], params=VisualGenParams()) - req.refs_to_handles() + req.refs_to_shm() assert req.ref_handles is None def test_restore_is_idempotent(self): @@ -1110,9 +1110,9 @@ def test_restore_is_idempotent(self): handle that has already been resolved.""" payload = os.urandom(2048) req = self._request(payload) - req.refs_to_handles() - req.refs_to_bytes() - req.refs_to_bytes() + req.refs_to_shm() + req.refs_from_shm() + req.refs_from_shm() assert req.params.image_reference[0].content == payload @@ -1187,12 +1187,12 @@ def test_dropped_request_releases_its_shared_memory(self): import gc req = self._request(os.urandom(1024 * 1024)) - req.refs_to_handles() + req.refs_to_shm() blocks = self._blocks_of(req) assert blocks and all(b.exists() for b in blocks) # What the sender thread does when the request never reaches rank0. - req.refs_to_bytes() + req.refs_from_shm() del req gc.collect() assert not any(b.exists() for b in blocks) @@ -1224,13 +1224,13 @@ def flaky(buffer, **kwargs): with mock.patch.object(torch, "frombuffer", flaky): with pytest.raises(RuntimeError, match="shared memory exhausted"): - req.refs_to_handles() + req.refs_to_shm() blocks = self._blocks_of(req) assert len(blocks) == 2, "handles taken before the failure must stay reachable" assert all(b.exists() for b in blocks) - req.refs_to_bytes() + req.refs_from_shm() del req gc.collect() assert not any(b.exists() for b in blocks) diff --git a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py index d29a249b746a..6efa8ed0c473 100644 --- a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py +++ b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py @@ -295,7 +295,7 @@ def test_one_bad_handle_does_not_strand_the_rest(self): import gc req = self._request(os.urandom(4096), os.urandom(4096), os.urandom(4096)) - req.refs_to_handles() + req.refs_to_shm() blocks = self._blocks_of(req.ref_handles) self.assertEqual(len(blocks), 3) @@ -310,7 +310,7 @@ def flaky(handle): with mock.patch.object(SharedTensorContainer, "from_dict", staticmethod(flaky)): with self.assertRaises(RuntimeError): - req.refs_to_bytes() + req.refs_from_shm() self.assertEqual(calls["n"], 3, "a failure must not stop the loop at the bad entry") del req @@ -319,7 +319,7 @@ def flaky(handle): def test_blocks_of_a_dead_workers_request_are_released(self): req = self._request(os.urandom(64 * 1024), os.urandom(64 * 1024)) - req.refs_to_handles() + req.refs_to_shm() blocks = self._blocks_of(req.ref_handles) self.assertTrue(all(b.exists() for b in blocks)) @@ -329,7 +329,7 @@ def test_blocks_of_a_dead_workers_request_are_released(self): def test_releasing_twice_is_harmless(self): """Reclaim can fire from more than one place; it must not care.""" req = self._request(os.urandom(4096)) - req.refs_to_handles() + req.refs_to_shm() self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 1) self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 0) From 1aa72f51afe7b70f2e59633f5e406c4e662e016a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 24 Aug 2026 04:44:36 -0700 Subject: [PATCH 47/61] [TRTLLM-15277][chore] Fold the broadcast split into the hop that needs it refs_detach/refs_attach and the _refs generator behind them existed for a single caller. Inlining them into _broadcast_request puts the whole hop in one place: the payloads leave the object, the collective runs, and they go back, in reading order. Taking them out is now one pass instead of two -- collecting a payload and clearing its reference happen together. Drop the payload/size check that followed the take-out. It compared a length against one derived from the same list two lines earlier. It was reachable only because the take-out skipped its bookkeeping when params was None, leaving a stale ref_sizes behind; that path now records zero sizes, so the request is consistent on every path and the check has nothing left to catch. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 64 ++++++--------- .../visual_gen/test_visual_gen_utils.py | 77 +++++++++++++------ 2 files changed, 76 insertions(+), 65 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index cf627bc33028..5675c331ca9c 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -299,7 +299,7 @@ class DiffusionRequest: # ``refs_to_shm``. ref_handles: Optional[List[Dict[str, Any]]] = field(default=None, repr=False) # Set only while the request is in flight on the rank0 -> N-rank hop; see - # ``refs_detach``. + # ``DiffusionExecutor._broadcast_request``. ref_sizes: Optional[List[int]] = field(default=None, repr=False) def refs_to_shm(self) -> None: @@ -352,36 +352,6 @@ def refs_from_shm(self) -> None: "failed to restore reference payloads from shared memory: " + "; ".join(failures) ) - def _refs(self): - for slot in ("image_reference", "video_reference", "audio_reference"): - yield from getattr(self.params, slot, None) or [] - - def refs_detach(self) -> List[bytes]: - """Take the reference payloads out of the request and record their sizes. - - Broadcasting them inside the request object would make - ``broadcast_object_list`` serialize every reference byte into a tensor - first, which costs more than the collective that follows. - """ - if self.params is None: - return [] - payloads = [ref.content for ref in self._refs()] - for ref in self._refs(): - ref.content = b"" - self.ref_sizes = [len(p) for p in payloads] - return payloads - - def refs_attach(self, payloads: List[bytes]) -> None: - """Put the separately broadcast payloads back, in place.""" - refs = list(self._refs()) - if len(refs) != len(payloads): - # zip() would silently leave the tail of either side behind, and - # clearing ref_sizes below would erase the evidence. - raise ValueError(f"expected {len(refs)} reference payloads, got {len(payloads)}.") - for ref, payload in zip(refs, payloads): - ref.content = payload - self.ref_sizes = None - @dataclass class DiffusionResponse: @@ -522,14 +492,16 @@ def _broadcast_request(self, req: Optional[DiffusionRequest]) -> Optional[Diffus rather than inside it, because ``broadcast_object_list`` pickles the object into a tensor first and that copy dominates the collective. """ - payloads = req.refs_detach() if self.rank == 0 and req is not None else [] - if self.rank == 0 and len(payloads) != len(getattr(req, "ref_sizes", None) or []): - # Peers derive their collective count from ref_sizes, so a mismatch - # here would hang every rank. Fail on rank0, before the first one. - raise RuntimeError( - f"reference payload/size mismatch: {len(payloads)} payloads, " - f"{len(getattr(req, 'ref_sizes', None) or [])} sizes." - ) + payloads = [] + if self.rank == 0 and req is not None: + # Take the payloads out before the object is pickled, and leave their + # sizes behind so the peers can size their receive buffers. + for slot in ("image_reference", "video_reference", "audio_reference"): + for ref in getattr(req.params, slot, None) or []: + payloads.append(ref.content) + ref.content = b"" + req.ref_sizes = [len(p) for p in payloads] + obj_list = [req] dist.broadcast_object_list(obj_list, src=0) req = obj_list[0] @@ -546,7 +518,19 @@ def _broadcast_request(self, req: Optional[DiffusionRequest]) -> Optional[Diffus if self.rank != 0: payloads = [b.numpy().tobytes() for b in buffers] - req.refs_attach(payloads) + + refs = [ + ref + for slot in ("image_reference", "video_reference", "audio_reference") + for ref in getattr(req.params, slot, None) or [] + ] + if len(refs) != len(payloads): + # zip() would silently leave the tail of either side behind, and + # clearing ref_sizes below would erase the evidence. + raise ValueError(f"expected {len(refs)} reference payloads, got {len(payloads)}.") + for ref, payload in zip(refs, payloads): + ref.content = payload + req.ref_sizes = None return req def serve_forever(self): diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index bdbf3df6da18..b4e06d77eba0 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -1135,38 +1135,75 @@ def _request(*payloads: bytes): ) return DiffusionRequest(request_id=1, prompt=["x"], params=params) - def test_detach_empties_the_object_and_records_sizes(self): + @staticmethod + def _drive_broadcast(req, rank=0, peer_sizes=None): + """Run ``_broadcast_request`` with the collectives stubbed out. + + Returns ``(result, pickled_object)`` so a test can inspect exactly what + ``broadcast_object_list`` was handed. + """ + from unittest import mock + + from tensorrt_llm._torch.visual_gen import executor as ex + + captured = {} + + def fake_obj_list(obj_list, src=0): + captured["obj"] = pickle.dumps(obj_list[0]) + if rank != 0: + peer = pickle.loads(captured["obj"]) + peer.ref_sizes = peer_sizes if peer_sizes is not None else peer.ref_sizes + obj_list[0] = peer + + self_ = mock.Mock(rank=rank) + with ( + mock.patch.object(ex.dist, "broadcast_object_list", side_effect=fake_obj_list), + mock.patch.object(ex.dist, "broadcast", lambda *a, **k: None), + ): + out = ex.DiffusionExecutor._broadcast_request(self_, req) + return out, captured.get("obj") + + def test_payload_never_reaches_the_object_pickle(self): + """The whole point of the split: the bytes must not be in what + ``broadcast_object_list`` serializes.""" payloads = (os.urandom(4096), os.urandom(64)) req = self._request(*payloads) - detached = req.refs_detach() + out, pickled = self._drive_broadcast(req) - assert detached[:2] == list(payloads) - assert req.ref_sizes == [len(p) for p in detached] - assert all(r.content == b"" for r in req.params.image_reference) - assert payloads[0] not in pickle.dumps(req) + assert payloads[0] not in pickled + assert payloads[1] not in pickled + assert out is req - def test_attach_restores_every_slot_in_order(self): + def test_payloads_come_back_to_their_own_slots(self): payloads = (os.urandom(2048), os.urandom(128)) req = self._request(*payloads) - detached = req.refs_detach() - req.refs_attach(detached) + out, _ = self._drive_broadcast(req) - assert tuple(r.content for r in req.params.image_reference) == payloads - assert req.params.video_reference[0].content == b"\x00\x00\x00\x18ftypmp42" - assert req.ref_sizes is None + assert tuple(r.content for r in out.params.image_reference) == payloads + assert out.params.video_reference[0].content == b"\x00\x00\x00\x18ftypmp42" + assert out.ref_sizes is None def test_sizes_let_a_peer_size_its_buffers(self): """Non-source ranks allocate from ``ref_sizes`` alone, so it has to survive the object hop and match the payloads exactly.""" payloads = (os.urandom(1024), os.urandom(7)) req = self._request(*payloads) - req.refs_detach() - peer = pickle.loads(pickle.dumps(req)) - assert peer.ref_sizes == req.ref_sizes + _, pickled = self._drive_broadcast(req) + + peer = pickle.loads(pickled) assert peer.ref_sizes[:2] == [1024, 7] + assert all(r.content == b"" for r in peer.params.image_reference) + + def test_a_size_count_mismatch_is_refused(self): + """A peer whose ref_sizes disagrees with its slots would pair payloads + to the wrong references; zip() would do it silently.""" + req = self._request(os.urandom(64), os.urandom(32)) + + with pytest.raises(ValueError, match="reference payloads"): + self._drive_broadcast(req, rank=1, peer_sizes=[64]) @staticmethod def _blocks_of(req): @@ -1197,16 +1234,6 @@ def test_dropped_request_releases_its_shared_memory(self): gc.collect() assert not any(b.exists() for b in blocks) - def test_attach_rejects_a_count_mismatch(self): - """Peers size their collectives from ``ref_sizes``, so a payload list - that does not match the reference count is a bug worth raising on - rather than silently leaving references empty.""" - req = self._request(os.urandom(64), os.urandom(64)) - detached = req.refs_detach() - - with pytest.raises(ValueError, match="expected 3 reference payloads, got 2"): - req.refs_attach(detached[:2]) - def test_partial_handle_failure_stays_reclaimable(self): """If a later reference cannot reach shared memory, the blocks already taken must still be reachable, or nothing can free them.""" From fa337fd0421d0630e7a1110ed64f3bf6c95f864c Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 02:21:27 -0700 Subject: [PATCH 48/61] [TRTLLM-15277][chore] Leave the shared blocks to the process that owns them Measuring the block lifetime removed the reason for all of this. When the producer process dies, torch's shm manager unlinks everything it registered -- verified under SIGTERM and SIGKILL -- so nothing survives a crash. What could accumulate inside a live coordinator turns out not to: a malformed payload never mints a block, because the reference choke point rejects it first, and a dead worker does not make the sender fail, it makes it block, since ZeroMqQueue is a zmq.PAIR socket and PAIR blocks in send() with no peer. That leaves shared-memory exhaustion mid-mint, in a deployment already too small to run this feature, and a shutdown race on a closed socket, which the process exit cleans up anyway. Machinery that only fires in states that clean themselves is worse than none: it reads as if the case were handled. Consuming a handle is now the one thing that frees a block, and rank0 is the one caller. The test that covers it says so. Two findings for follow-up, both outside this code: a dead worker silently deadlocks the sender thread with no error or timeout, and SharedTensorContainer has no producer-side release -- freeing a handle from outside the library means reaching into method_key, storage_handle and the base64 filename, and a CUDA handle has no unlinkable name at all. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 106 +----------------- tensorrt_llm/visual_gen/visual_gen.py | 10 +- .../visual_gen/test_visual_gen_utils.py | 79 +------------ .../test_executor_shared_tensor_ipc.py | 83 -------------- tests/unittest/visual_gen/test_output.py | 1 - 5 files changed, 8 insertions(+), 271 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 5675c331ca9c..5974a1d9aa6c 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -1,5 +1,4 @@ import asyncio -import base64 import os import queue import socket @@ -16,10 +15,7 @@ import torch.multiprocessing as mp import zmq -from tensorrt_llm._torch.shared_tensor import ( - SharedTensorContainer, - _SharedTensorRebuildMethodRegistry, -) +from tensorrt_llm._torch.shared_tensor import SharedTensorContainer from tensorrt_llm._torch.visual_gen.output import PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.executor.ipc import ZeroMqQueue @@ -31,38 +27,6 @@ from tensorrt_llm.visual_gen.params import VisualGenParams -def _release_cpu_ref_blocks(handles: Optional[List[Dict]]) -> int: - """Unlink the shared-memory blocks behind reference handles nobody will take. - - Only call this once the consumer is known never to open them: unlinking a - block the consumer has not yet opened turns its rebuild into a failure. A - dead worker is that signal; a request merely abandoned by its caller is not, - because rank0 keeps going and still consumes the handle. - - Unlinking rather than rebuilding keeps it idempotent -- a block the consumer - already released is simply absent. CUDA handles carry no unlinkable name and - are left alone; the caching allocator reclaims them only once the consumer - process is gone. - """ - released = 0 - for entry in handles or []: - handle = entry.get("handle") or {} - if handle.get("method_key") != _SharedTensorRebuildMethodRegistry.REBUILD_CPU: - continue - raw = handle.get("storage_handle") - if not raw: - continue - try: - name = base64.b64decode(raw).decode().lstrip("/") - os.unlink(os.path.join("/dev/shm", name)) - released += 1 - except FileNotFoundError: - pass - except OSError as exc: - logger.warning(f"DiffusionClient: could not release shared block {raw!r}: {exc}") - return released - - # Timeouts (seconds) for the client-side coordinator. POLL_TIMEOUT = 0.01 AWAIT_TIMEOUT = 0.05 @@ -311,9 +275,6 @@ def refs_to_shm(self) -> None: """ if self.params is None: return - # Publish the list before filling it: if a later reference fails to - # reach shared memory, the blocks already taken are still reachable - # for the reclaim path instead of leaking. self.ref_handles = handles = [] for slot in ("image_reference", "video_reference", "audio_reference"): for index, ref in enumerate(getattr(self.params, slot, None) or []): @@ -334,9 +295,8 @@ def refs_to_shm(self) -> None: def refs_from_shm(self) -> None: """Restore reference payloads from shared memory, in place (consumer side). - Each handle is taken independently: one that cannot be rebuilt must not - strand the blocks behind it, since this is also the reclaim path and a - handle nobody consumes keeps its block mapped until the process exits. + Each handle is taken independently so one that cannot be rebuilt does + not strand the blocks behind it. """ failures = [] for entry in self.ref_handles or []: @@ -842,11 +802,6 @@ def __init__( # full PipelineOutput tensor does not pin in completed_responses for # the process lifetime. self._abandoned_request_ids: Set[int] = set() - # Handles of requests already handed to the workers. rank0 frees a - # block by consuming it, so an entry here means "sent, not yet known - # consumed" -- what has to be released if the workers die first. - self._sent_ref_handles: Dict[int, List[Dict]] = {} - # Iteration-stats tracker — populated on lifecycle events (enqueue, # request started, response received) and drained by # ``get_iteration_stats`` for the /metrics HTTP endpoint. Mirrors @@ -1042,8 +997,6 @@ def _process_requests(self): logger.info(f"DiffusionClient: Sending request {req.request_id}") self.requests_ipc.put(req) - if req.ref_handles: - self._sent_ref_handles[req.request_id] = req.ref_handles # Once the request has been handed to the workers it becomes the # in-flight ("active") request from the client's perspective. self._iter_stats.record_request_started(req.request_id, self.pending_requests.qsize()) @@ -1052,11 +1005,6 @@ def _process_requests(self): except Exception as e: logger.error(f"DiffusionClient: Error sending request: {e}") logger.error(traceback.format_exc()) - if req is not None and req.ref_handles: - # The request never reached rank0, so nothing downstream will - # consume its handles, and an unconsumed handle keeps its - # shared-memory block mapped until this process exits. - req.refs_from_shm() def _process_responses(self): """Poll and process responses.""" @@ -1077,22 +1025,6 @@ def _process_responses(self): except Exception as e: logger.error(f"DiffusionClient: Error processing response: {e}") - def _forget_sent_handles(self, request_id: int) -> None: - """A response means rank0 got that far, so its blocks are its own now.""" - self._sent_ref_handles.pop(request_id, None) - - def _release_sent_handles(self, reason: str) -> None: - """Release the blocks of every request the workers can no longer take.""" - if not self._sent_ref_handles: - return - outstanding, self._sent_ref_handles = self._sent_ref_handles, {} - released = sum(_release_cpu_ref_blocks(h) for h in outstanding.values()) - if released: - logger.warning( - f"DiffusionClient: released {released} shared block(s) from " - f"{len(outstanding)} unfinished request(s) after {reason}" - ) - async def _store_response(self, response: DiffusionResponse): """Store response in the completed_responses dict (async helper). @@ -1100,7 +1032,6 @@ async def _store_response(self, response: DiffusionResponse): late-arriving responses for timed-out requests do not leak into ``completed_responses`` for the process lifetime. """ - self._forget_sent_handles(response.request_id) async with self.lock: if response.request_id in self._abandoned_request_ids: self._abandoned_request_ids.discard(response.request_id) @@ -1189,38 +1120,10 @@ async def _serve_forever(self): while not self.shutdown_event.is_set(): self._process_requests() self._process_responses() - if self._sent_ref_handles and self._workers_gone(): - # A dead worker cannot open a block, so the ones it never got - # to are safe to release -- and nothing else will. - self._release_sent_handles("worker exit") await asyncio.sleep(0.001) # Yield control to allow other coroutines to run self._cleanup_ipc() - def _workers_gone(self) -> bool: - """Whether the processes (or the external-launch thread) that consume - requests have exited.""" - if self.worker_processes: - return any(not p.is_alive() for p in self.worker_processes) - if self._ext_worker_thread is not None: - return not self._ext_worker_thread.is_alive() - return False - - def _reclaim_pending_handles(self) -> None: - """Consume the handles of requests that will never be sent. - - A request abandoned in the queue still holds shared-memory blocks that - nothing downstream will consume, and they stay mapped until this - process exits. - """ - while True: - try: - req = self.pending_requests.get_nowait() - except queue.Empty: - return - if req is not None and req.ref_handles: - req.refs_from_shm() - def shutdown(self): """Shutdown client and workers.""" logger.info("DiffusionClient: Shutting down") @@ -1231,9 +1134,6 @@ def shutdown(self): logger.warning("DiffusionClient: Force stopping background thread") self.shutdown_event.set() self.background_thread.join(timeout=1.0) - self._reclaim_pending_handles() - - self._reclaim_pending_handles() # Shutdown workers logger.info("DiffusionClient: Stopping workers") diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 809e46d62f36..ad595a44fab8 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -452,14 +452,8 @@ def generate_async( # Hand the reference payloads to rank0 through shared memory instead of # through the request pickle, which copies every reference byte to # cross one process boundary. - try: - request.refs_to_shm() - self.executor.enqueue_requests([request]) - except Exception: - # The request never reached rank0, so nothing downstream will - # consume the blocks it already took. - request.refs_from_shm() - raise + request.refs_to_shm() + self.executor.enqueue_requests([request]) return VisualGenResult(req_id, self.executor, batch_size=batch_size) @staticmethod diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index b4e06d77eba0..870933842e49 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -17,11 +17,9 @@ from io import BytesIO from pathlib import Path from typing import Any, Dict, Optional -from unittest import mock import numpy as np import pytest -import torch from fastapi import UploadFile from PIL import Image @@ -1217,10 +1215,9 @@ def _blocks_of(req): for e in req.ref_handles or [] ] - def test_dropped_request_releases_its_shared_memory(self): - """An unconsumed handle keeps its block mapped until the process exits, - so a request that fails on its way to rank0 has to consume its own - handles on the way out.""" + def test_consuming_a_handle_releases_its_block(self): + """rank0 frees a block by consuming its handle; nothing else does, so an + unconsumed one stays resident for the life of the process.""" import gc req = self._request(os.urandom(1024 * 1024)) @@ -1234,76 +1231,6 @@ def test_dropped_request_releases_its_shared_memory(self): gc.collect() assert not any(b.exists() for b in blocks) - def test_partial_handle_failure_stays_reclaimable(self): - """If a later reference cannot reach shared memory, the blocks already - taken must still be reachable, or nothing can free them.""" - import gc - - req = self._request(os.urandom(256 * 1024), os.urandom(256 * 1024)) - real = torch.frombuffer - calls = {"n": 0} - - def flaky(buffer, **kwargs): - calls["n"] += 1 - if calls["n"] == 3: # the third of this request's three references - raise RuntimeError("shared memory exhausted") - return real(buffer, **kwargs) - - with mock.patch.object(torch, "frombuffer", flaky): - with pytest.raises(RuntimeError, match="shared memory exhausted"): - req.refs_to_shm() - - blocks = self._blocks_of(req) - assert len(blocks) == 2, "handles taken before the failure must stay reachable" - assert all(b.exists() for b in blocks) - - req.refs_from_shm() - del req - gc.collect() - assert not any(b.exists() for b in blocks) - - def test_a_request_that_never_ships_hands_its_blocks_back(self): - """The blocks staying reachable is only half of it; the call site is - what has to hand them back when the request never reaches rank0.""" - import gc - import itertools - from types import SimpleNamespace - - from tensorrt_llm.visual_gen import VisualGen, VisualGenParams - from tensorrt_llm.visual_gen.params import MediaRef - - taken = {} - blocks_of = self._blocks_of - - class _DeadExecutor: - default_generation_params = {} - extra_param_specs = {} - ref_slot_specs = None - - def enqueue_requests(self, requests): - # Read the handles while they still exist, then fail the way a - # dead worker queue would. - taken["blocks"] = blocks_of(requests[0]) - raise RuntimeError("worker queue is gone") - - caller = SimpleNamespace( - _req_counter=itertools.count(), - default_params=VisualGenParams(), - executor=_DeadExecutor(), - ) - buf = BytesIO() - # A real PNG: the engine choke point content-checks before the handles - # are minted, so random bytes never get far enough to take a block. - Image.new("RGB", (256, 256)).save(buf, format="PNG") - params = VisualGenParams(image_reference=[MediaRef(content=buf.getvalue(), format="bytes")]) - - with pytest.raises(RuntimeError, match="worker queue is gone"): - VisualGen.generate_async(caller, "x", params) - - assert taken["blocks"], "the request must have taken a block to begin with" - gc.collect() - assert not any(b.exists() for b in taken["blocks"]) - class TestSafeLocalFileRead: """A ``path`` reference bounds what an unlucky path can cost. diff --git a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py index 6efa8ed0c473..35f84a1c7556 100644 --- a/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py +++ b/tests/unittest/visual_gen/test_executor_shared_tensor_ipc.py @@ -8,25 +8,18 @@ main process, mirroring the real worker/client split. """ -import base64 import multiprocessing as mp import os import unittest -from pathlib import Path from unittest import mock import torch import zmq -from tensorrt_llm._torch.shared_tensor import ( - SharedTensorContainer, - _SharedTensorRebuildMethodRegistry, -) from tensorrt_llm._torch.visual_gen.executor import ( DiffusionExecutor, DiffusionRequest, DiffusionResponse, - _release_cpu_ref_blocks, find_free_port, run_diffusion_worker, ) @@ -265,82 +258,6 @@ def test_spawn_worker_response_rebuilds_in_client(self): q.close() -class TestUnconsumedReferenceBlocks(unittest.TestCase): - """A handle nobody consumes keeps its block mapped until the process exits. - - Consuming is the normal release, so these cover the two ways a request can - stop short of it: a rebuild that fails partway, and workers that die before - reading what was already sent to them. Release is keyed on the workers being - gone rather than on the caller losing interest -- unlinking a block rank0 has - not opened yet would turn its rebuild into a failure. - """ - - @staticmethod - def _request(*payloads): - from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - - params = VisualGenParams( - image_reference=[MediaRef(content=p, format="bytes") for p in payloads] - ) - return DiffusionRequest(request_id=7, prompt=["x"], params=params) - - @staticmethod - def _blocks_of(handles): - return [ - Path("/dev/shm") / base64.b64decode(e["handle"]["storage_handle"]).decode().lstrip("/") - for e in handles or [] - ] - - def test_one_bad_handle_does_not_strand_the_rest(self): - import gc - - req = self._request(os.urandom(4096), os.urandom(4096), os.urandom(4096)) - req.refs_to_shm() - blocks = self._blocks_of(req.ref_handles) - self.assertEqual(len(blocks), 3) - - real = SharedTensorContainer.from_dict - calls = {"n": 0} - - def flaky(handle): - calls["n"] += 1 - if calls["n"] == 2: - raise RuntimeError("handle is unusable") - return real(handle) - - with mock.patch.object(SharedTensorContainer, "from_dict", staticmethod(flaky)): - with self.assertRaises(RuntimeError): - req.refs_from_shm() - - self.assertEqual(calls["n"], 3, "a failure must not stop the loop at the bad entry") - del req - gc.collect() - self.assertLessEqual(sum(b.exists() for b in blocks), 1) - - def test_blocks_of_a_dead_workers_request_are_released(self): - req = self._request(os.urandom(64 * 1024), os.urandom(64 * 1024)) - req.refs_to_shm() - blocks = self._blocks_of(req.ref_handles) - self.assertTrue(all(b.exists() for b in blocks)) - - self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 2) - self.assertFalse(any(b.exists() for b in blocks)) - - def test_releasing_twice_is_harmless(self): - """Reclaim can fire from more than one place; it must not care.""" - req = self._request(os.urandom(4096)) - req.refs_to_shm() - - self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 1) - self.assertEqual(_release_cpu_ref_blocks(req.ref_handles), 0) - - def test_a_cuda_handle_is_left_alone(self): - """CUDA blocks carry no unlinkable name; skipping them is the contract.""" - cuda_key = _SharedTensorRebuildMethodRegistry.REBUILD_CUDA - cuda_shaped = [{"slot": "image_reference", "index": 0, "handle": {"method_key": cuda_key}}] - self.assertEqual(_release_cpu_ref_blocks(cuda_shaped), 0) - - if __name__ == "__main__": mp.set_start_method("spawn", force=True) unittest.main() diff --git a/tests/unittest/visual_gen/test_output.py b/tests/unittest/visual_gen/test_output.py index 411f55f7b49b..bfbaafdc4380 100644 --- a/tests/unittest/visual_gen/test_output.py +++ b/tests/unittest/visual_gen/test_output.py @@ -669,7 +669,6 @@ def _make_minimal_client_state(): client = DiffusionRemoteClient.__new__(DiffusionRemoteClient) client.completed_responses = {} client._abandoned_request_ids = set() - client._sent_ref_handles = {} client._iter_stats = _IterationStatsTracker() client.pending_requests = queue.Queue() From d10ef78e8cd9fc979c3f532f9657d878a8621055 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 21:27:36 -0700 Subject: [PATCH 49/61] [TRTLLM-15277][fix] Decode reference images the way they were decoded before Routing image decoding through ImageMediaIO changed what an RGBA reference becomes. Its default composites the alpha onto white, while the PIL.Image.open(...).convert("RGB") it replaced discards the channel and keeps the stored RGB. Qwen-Image-Edit had exactly that substitution, so a transparent PNG came out white-filled with nothing to signal it. A fully transparent red pixel reads (255, 0, 0) before and (255, 255, 255) after. The other call sites passed drop_alpha=True to opt back into the original behaviour, which is the tell: the parameter existed only so the shared helper could reproduce what plain PIL already does, and it had to be threaded through convert_image_mode, _load_and_convert_image and four ImageMediaIO methods to get there. On this path the helper adds no format check, no size bound and no EXIF handling, so it bought nothing and cost a default that was wrong for us. Decode with PIL at the nine call sites and hand media_io back its original signatures. The audio sniffing stays: audio_reference needs it to validate a payload's modality at the coordinator. media/decoding.py returns to upstream verbatim. Its FrameSelector, WindowSelector and NvdecVideoMediaIO had no production caller -- only tests, one of which defined its own selector subclass to exercise the extension point -- and the function they were factored out of already took bytes, so the reference work never needed them. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../models/cosmos3/pipeline_cosmos3.py | 4 +- .../visual_gen/models/flux/pipeline_flux2.py | 9 +-- .../visual_gen/models/ltx2/pipeline_ltx2.py | 4 +- .../qwen_image/pipeline_qwen_image_edit.py | 7 +- .../pipeline_qwen_image_layered.py | 5 +- .../visual_gen/models/wan/pipeline_wan.py | 4 +- .../visual_gen/models/wan/pipeline_wan_i2v.py | 13 ++-- tensorrt_llm/inputs/media_io.py | 43 +++-------- .../_torch/visual_gen/test_media_decode.py | 73 ------------------- 9 files changed, 29 insertions(+), 133 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index abcf55ee0121..f36af8b3f97c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -18,6 +18,7 @@ import math import os import time +from io import BytesIO from typing import Any, Iterable, List, Optional, Union import PIL.Image @@ -45,7 +46,6 @@ def tqdm(iterable, **kwargs): synchronize_media_prepare_status, ) from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from tensorrt_llm.media.decoding import decode_video_reference_window, video_stream_info @@ -246,7 +246,7 @@ def _load_reference_image(data: bytes): upload would be reported as a server fault. """ try: - return ImageMediaIO(format="pil").load_bytes(data) + return PIL.Image.open(BytesIO(data)).convert("RGB") except OSError as exc: raise ValueError( f"Image reference could not be decoded; it may be truncated, " 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 fb1485746fc0..3f72728d72c4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -22,6 +22,7 @@ import os import time from contextlib import contextmanager +from io import BytesIO from typing import Any, Iterator, List, Optional, Tuple, Union import numpy as np @@ -45,7 +46,6 @@ from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline -from tensorrt_llm.inputs.media_io import ImageMediaIO, convert_image_mode from tensorrt_llm.logger import logger from .transformer_flux2 import Flux2Transformer2DModel @@ -763,16 +763,13 @@ def _load_reference_images( if not inputs: raise ValueError("`image` must contain at least one reference image.") - # drop_alpha keeps diffusers' semantics: convert("RGB") drops the alpha - # channel rather than compositing it onto white. - media_io = ImageMediaIO(format="pil", drop_alpha=True) images = [] for index, item in enumerate(inputs): try: if isinstance(item, PIL.Image.Image): - images.append(convert_image_mode(item, "RGB", drop_alpha=True)) + images.append(item.convert("RGB")) elif isinstance(item, bytes): - images.append(media_io.load_bytes(item)) + images.append(PIL.Image.open(BytesIO(item)).convert("RGB")) else: raise ValueError( "Reference images must be PIL images or encoded bytes; " 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 b2a25cc0a432..53f996322059 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -7,6 +7,7 @@ import json import os import time +from io import BytesIO from pathlib import Path from typing import Any, Dict, List, Optional, Set, Tuple, Union @@ -28,7 +29,6 @@ ) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor -from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .ltx2_core.audio_vae import AudioDecoderConfigurator, VocoderConfigurator, decode_audio @@ -1237,7 +1237,7 @@ def _load_and_preprocess_image( if isinstance(image, bytes): from PIL import Image - pil_img = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) + pil_img = Image.open(BytesIO(image)).convert("RGB") pil_img = pil_img.resize((width, height), Image.LANCZOS) import numpy as np 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 c23be2d70b00..7cf64dabb121 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 @@ -10,6 +10,7 @@ import math import time +from io import BytesIO from typing import Any import numpy as np @@ -20,7 +21,6 @@ from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import register_pipeline -from tensorrt_llm.inputs.media_io import ImageMediaIO, convert_image_mode from tensorrt_llm.logger import logger from .pipeline_qwen_image import QwenImagePipeline, _calculate_shift @@ -170,13 +170,12 @@ def _load_edit_images(image: Any) -> list[Any]: if image is None: raise ValueError("Qwen-Image-Edit requires image_reference.") images = image if isinstance(image, list) else [image] - media_io = ImageMediaIO(format="pil") loaded = [] for index, item in enumerate(images): if isinstance(item, PIL.Image.Image): - loaded.append(convert_image_mode(item, "RGB")) + loaded.append(item.convert("RGB")) elif isinstance(item, bytes): - loaded.append(media_io.load_bytes(item)) + loaded.append(PIL.Image.open(BytesIO(item)).convert("RGB")) else: raise ValueError( "Reference images must be PIL images or encoded bytes; " 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 35dd70734b52..547e4f960811 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 @@ -16,9 +16,11 @@ import math import time +from io import BytesIO from typing import List, Optional, Tuple, Union import numpy as np +import PIL.Image import torch from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput @@ -29,7 +31,6 @@ RoleSpec, ) from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline -from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .transformer_qwen_image_layered import QwenImageLayeredTransformer2DModel @@ -365,7 +366,7 @@ def _load_image_input(image): return [QwenImageLayeredPipeline._load_image_input(item) for item in image] if isinstance(image, bytes): # Layer decomposition needs the alpha channel, not a flattened RGB. - return ImageMediaIO(format="pil", mode="RGBA").load_bytes(image) + return PIL.Image.open(BytesIO(image)).convert("RGBA") if hasattr(image, "convert") and getattr(image, "mode", None) != "RGBA": return image.convert("RGBA") return image 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 f5ce9dee270d..c2d7a8a6846e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -15,6 +15,7 @@ import os import time +from io import BytesIO from typing import List, Optional, Union import diffusers @@ -43,7 +44,6 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm._utils import nvtx_range -from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger from .transformer_wan import WanTransformer3DModel @@ -816,7 +816,7 @@ def _prepare_latents_wan22_5B_i2v( # Load and preprocess image if isinstance(image, bytes): - image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) + image = PIL.Image.open(BytesIO(image)).convert("RGB") image = ( self.video_processor.preprocess(image, height=height, width=width) .to(self.device, dtype=self.vae.dtype) 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 b3e9460590e9..ffed1732355f 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 @@ -16,6 +16,7 @@ import json import os import time +from io import BytesIO from typing import List, Optional, Tuple, Union import diffusers @@ -39,7 +40,6 @@ from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline, RefSlotSpec, RoleSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor -from tensorrt_llm.inputs.media_io import ImageMediaIO from tensorrt_llm.logger import logger # Supported Wan I2V 14B models: @@ -728,13 +728,10 @@ def _encode_image( last_image: Optional[Union[PIL.Image.Image, torch.Tensor, bytes]] = None, ) -> torch.Tensor: """Encode image(s) using CLIP image encoder (Wan 2.1 I2V only).""" - # drop_alpha keeps diffusers' semantics: convert("RGB") drops the alpha - # channel rather than compositing it onto white. - media_io = ImageMediaIO(format="pil", drop_alpha=True) if isinstance(image, bytes): - image = media_io.load_bytes(image) + image = PIL.Image.open(BytesIO(image)).convert("RGB") if isinstance(last_image, bytes): - last_image = media_io.load_bytes(last_image) + last_image = PIL.Image.open(BytesIO(last_image)).convert("RGB") images_to_encode = [image] if last_image is None else [image, last_image] @@ -766,14 +763,14 @@ def _prepare_latents( # Load and preprocess image(s) if isinstance(image, bytes): - image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(image) + image = PIL.Image.open(BytesIO(image)).convert("RGB") image = self.video_processor.preprocess(image, height=height, width=width).to( self.device, dtype=torch.float32 ) if last_image is not None: if isinstance(last_image, bytes): - last_image = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(last_image) + last_image = PIL.Image.open(BytesIO(last_image)).convert("RGB") last_image = self.video_processor.preprocess(last_image, height=height, width=width).to( self.device, dtype=torch.float32 ) diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 719c9106afd0..3ce027cd2e54 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -60,19 +60,11 @@ def rgba_to_rgb( return converted -def convert_image_mode(image: Image.Image, to_mode: str, drop_alpha: bool = False) -> Image.Image: - """Convert image to specified mode with proper handling of RGBA to RGB conversion. - - ``drop_alpha`` selects how an RGBA source sheds its alpha channel on the way - to RGB: ``True`` discards the channel and keeps the stored RGB, which is what - PIL's own ``convert("RGB")`` and diffusers' ``load_image`` do; ``False`` - composites onto a white background, which changes every pixel with - ``alpha < 255``. It only applies to that one direction, and defaults to - compositing so existing callers are unaffected. - """ +def convert_image_mode(image: Image.Image, to_mode: str) -> Image.Image: + """Convert image to specified mode with proper handling of RGBA to RGB conversion.""" if image.mode == to_mode: return image - elif image.mode == "RGBA" and to_mode == "RGB" and not drop_alpha: + elif image.mode == "RGBA" and to_mode == "RGB": return rgba_to_rgb(image) else: return image.convert(to_mode) @@ -250,10 +242,10 @@ async def _fetch(fetch_session: aiohttp.ClientSession) -> bytes: return await _fetch(owned_session) -def _load_and_convert_image(image, mode: str = "RGB", drop_alpha: bool = False): +def _load_and_convert_image(image): image = Image.open(image) image.load() - return convert_image_mode(image, mode, drop_alpha) + return convert_image_mode(image, "RGB") def _audio_frame_to_array(frame, mono: bool) -> np.ndarray: @@ -866,22 +858,11 @@ async def _run_in_executor(fn, *args, **kwargs): class ImageMediaIO(BaseMediaIO[Union[Image.Image, torch.Tensor, np.ndarray]]): """I/O for the image modality.""" - def __init__( - self, - format: str = "pt", - device: str = "cpu", - mode: str = "RGB", - drop_alpha: bool = False, - ) -> None: + def __init__(self, format: str = "pt", device: str = "cpu") -> None: if format not in _SUPPORTED_IMAGE_FORMATS: raise ValueError(f"format must be one of {_SUPPORTED_IMAGE_FORMATS}, got {format!r}") self._format = format self._device = device - # Target PIL mode, plus how RGBA sheds its alpha en route to RGB. A - # consumer that composites layers asks for mode="RGBA"; one that must - # match diffusers' preprocessing asks for drop_alpha=True. - self._mode = mode - self._drop_alpha = drop_alpha def _postprocess(self, image: Image.Image) -> Union[Image.Image, torch.Tensor, np.ndarray]: if self._format == "pt": @@ -894,21 +875,15 @@ def _postprocess(self, image: Image.Image) -> Union[Image.Image, torch.Tensor, n return image def load_bytes(self, data: bytes) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess( - _load_and_convert_image(BytesIO(data), self._mode, self._drop_alpha) - ) + return self._postprocess(_load_and_convert_image(BytesIO(data))) def load_base64( self, media_type: str, data: str ) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess( - _load_and_convert_image(BytesIO(base64.b64decode(data)), self._mode, self._drop_alpha) - ) + return self._postprocess(_load_and_convert_image(BytesIO(base64.b64decode(data)))) def load_file(self, url: str) -> Union[Image.Image, torch.Tensor, np.ndarray]: - return self._postprocess( - _load_and_convert_image(Path(_normalize_file_uri(url)), self._mode, self._drop_alpha) - ) + return self._postprocess(_load_and_convert_image(Path(_normalize_file_uri(url)))) class AudioMediaIO(BaseMediaIO[Tuple[np.ndarray, int]]): diff --git a/tests/unittest/_torch/visual_gen/test_media_decode.py b/tests/unittest/_torch/visual_gen/test_media_decode.py index f157731c65a8..0c88f47ff86b 100644 --- a/tests/unittest/_torch/visual_gen/test_media_decode.py +++ b/tests/unittest/_torch/visual_gen/test_media_decode.py @@ -27,11 +27,7 @@ from tensorrt_llm._torch.visual_gen.utils import synchronize_media_prepare_status from tensorrt_llm.media.decoding import ( - FrameSelector, - NvdecVideoMediaIO, - WindowSelector, _lanczos_taps, - _nvdec_decode, decode_video_reference_window, resize_center_crop_uint8, resize_fit_pad_uint8, @@ -383,72 +379,3 @@ def test_resize_perf_representative(self): torch.cuda.synchronize() per_frame = (time.perf_counter() - start) / 10 assert per_frame < 0.25, f"resize took {per_frame * 1e3:.1f} ms/frame" - - -class TestSelectorAndMediaIO: - """The decode mechanism, its explicit frame selector, and the MediaIO leaf.""" - - _DEVICE = torch.device("cuda:0") - - def test_window_selector_rejects_mixed_and_reversed_ranges(self): - """The range is validated when the selector is built, not at decode.""" - with pytest.raises(ValueError, match="both count from"): - WindowSelector(0, -1) - with pytest.raises(ValueError, match="must not exceed"): - WindowSelector(3, 1) - - def test_unimplemented_selector_is_rejected(self): - """Selection has no default: an unknown strategy must not silently - fall back to a window.""" - - class FpsSelector(FrameSelector): - pass - - with pytest.raises(NotImplementedError, match="FpsSelector"): - _nvdec_decode(_MP4.read_bytes(), selector=FpsSelector(), device=self._DEVICE) - - @pytest.mark.parametrize("fixture", [_MP4, _AVI], ids=["mp4", "avi"]) - def test_media_io_matches_the_window_function(self, fixture): - """The MediaIO leaf and the legacy free function are the same decode.""" - data = fixture.read_bytes() - reference = decode_video_reference_window( - data, first_frame=0, last_frame=4, target_h=64, target_w=64, device=self._DEVICE - ) - got = NvdecVideoMediaIO( - selector=WindowSelector(0, 4), device=self._DEVICE, target_hw=(64, 64) - ).load_bytes(data) - assert torch.equal(got, reference) - - def test_media_io_load_file_matches_load_bytes(self, tmp_path): - io = NvdecVideoMediaIO( - selector=WindowSelector(0, 4), device=self._DEVICE, target_hw=(64, 64) - ) - assert torch.equal(io.load_file(str(_MP4)), io.load_bytes(_MP4.read_bytes())) - - def test_no_target_keeps_the_source_resolution(self): - """Resize only happens when a target is given. - - The fixture is natively 64x64, so the target has to be a different - size for the two paths to be distinguishable at all. - """ - data = _MP4.read_bytes() - native = _nvdec_decode(data, selector=WindowSelector(0, 0), device=self._DEVICE) - resized = _nvdec_decode( - data, selector=WindowSelector(0, 0), device=self._DEVICE, target_hw=(32, 32) - ) - assert native.shape == (1, 64, 64, 3) - assert resized.shape == (1, 32, 32, 3) - assert native.dtype == torch.uint8 - - def test_range_past_the_clip_returns_empty_not_error(self): - """A window beyond the clip yields what exists — here, nothing. - - The target's spatial dims are kept so the caller can still pad. - """ - window = _nvdec_decode( - _MP4.read_bytes(), - selector=WindowSelector(100, 104), - device=self._DEVICE, - target_hw=(64, 64), - ) - assert window.shape == (0, 64, 64, 3) From 5183bbd55f29a6e9949cd00957869c5930e7cf6e Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:18:27 -0700 Subject: [PATCH 50/61] [TRTLLM-15277][chore] Say each thing once, and in the right place The reference work left comments that repeat each other across files, restate the literal on the next line, or explain a sibling's design. Cut them to the part the code cannot say. The role and arity rules were a nine-line block inside validate_visual_gen_params while its docstring, which lists every condition the function raises on, did not mention references at all. Move the contract to the docstring and leave one line in the body. Cosmos3 keeps compositing an RGBA reference onto white. It reached load_image with a path, which flattens; decoding the bytes with convert("RGB") instead would have dropped the channel and changed the image. The alpha test now asserts what the code replaced -- the channel is dropped, not flattened -- and covers the bytes branch as well as the PIL one, which is the branch that had no test when it changed. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/_torch/visual_gen/executor.py | 8 ++--- .../models/cosmos3/pipeline_cosmos3.py | 8 +++-- .../visual_gen/models/flux/pipeline_flux2.py | 4 +-- .../visual_gen/models/ltx2/pipeline_ltx2.py | 1 - .../visual_gen/models/wan/pipeline_wan.py | 6 ++-- .../visual_gen/models/wan/pipeline_wan_i2v.py | 2 +- tensorrt_llm/serve/openai_video_routes.py | 24 +++++--------- tensorrt_llm/serve/visual_gen_utils.py | 4 +-- tensorrt_llm/visual_gen/params.py | 33 +++++++------------ .../visual_gen/test_qwen_image_pipeline.py | 18 ++++++---- 10 files changed, 44 insertions(+), 64 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 5974a1d9aa6c..04efe4b0efbd 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -503,12 +503,8 @@ def serve_forever(self): req.refs_from_shm() logger.info(f"Worker {self.device_id}: Request available") - # Broadcast to all ranks. ``req.params.seed`` is already a - # concrete int — resolved once on the coordinator process at - # :meth:`VisualGen.generate_async` entry — so the broadcast - # propagates the same value to every rank. Single-rank runs skip - # it: there is no peer, and the object broadcast would still - # serialize the whole request to a tensor before discovering that. + # Skipped at world_size 1: with no peer, the object broadcast would + # still serialize the whole request to a tensor before finding out. if self.world_size > 1: req = self._broadcast_request(req) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index f36af8b3f97c..932e00de4981 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -46,6 +46,7 @@ def tqdm(iterable, **kwargs): synchronize_media_prepare_status, ) from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.inputs.media_io import convert_image_mode from tensorrt_llm.logger import logger from tensorrt_llm.media.decoding import decode_video_reference_window, video_stream_info @@ -246,7 +247,9 @@ def _load_reference_image(data: bytes): upload would be reported as a server fault. """ try: - return PIL.Image.open(BytesIO(data)).convert("RGB") + # convert_image_mode, not convert("RGB"): it composites RGBA onto white, + # which is what this pipeline's references went through before. + return convert_image_mode(PIL.Image.open(BytesIO(data)), "RGB") except OSError as exc: raise ValueError( f"Image reference could not be decoded; it may be truncated, " @@ -605,8 +608,7 @@ def extra_param_specs(self): @property def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { - # image (I2V) and video (V2V) are both optional; Cosmos3 also runs - # T2V with neither. + # Both optional: Cosmos3 also runs T2V with neither. "image_reference": RefSlotSpec( modality="image", roles=[RoleSpec(role="first_frame", min=0, max=1)], 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 3f72728d72c4..724695babc08 100644 --- a/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py +++ b/tensorrt_llm/_torch/visual_gen/models/flux/pipeline_flux2.py @@ -372,8 +372,8 @@ def default_generation_params(self): @property def ref_slot_specs(self) -> dict[str, RefSlotSpec]: - # Optional reference image(s): "reference" role, count 0..N (multi-subject); - # FLUX.2 also runs plain text-to-image with none. + # Unbounded: one prompt can reference several subjects. FLUX.2 also + # runs plain text-to-image with no reference at all. return { "image_reference": RefSlotSpec( modality="image", 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 53f996322059..de6411754878 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/pipeline_ltx2.py @@ -1385,7 +1385,6 @@ def ref_slot_specs(self) -> dict[str, RefSlotSpec]: return { "image_reference": RefSlotSpec( modality="image", - # Optional first-frame conditioning (min=0); LTX-2 also runs T2V. roles=[RoleSpec(role="first_frame", min=0, max=1)], ), } 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 c2d7a8a6846e..fbf9684736f1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/pipeline_wan.py @@ -427,12 +427,10 @@ def extra_param_specs(self): @property def ref_slot_specs(self) -> dict[str, RefSlotSpec]: - # Only Wan 2.2 TI2V-5B conditions on a first frame; the T2V variants - # accept no reference (forward() rejects an image otherwise), so they - # declare no slot and unsupported image requests fail at preflight. + # Only TI2V-5B takes a reference; the T2V variants declare no slot, so + # an image request fails at preflight instead of inside forward(). if not self.is_wan22_5b: return {} - # Optional single conditioning image (first frame); T2V when absent. return { "image_reference": RefSlotSpec( modality="image", 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 ffed1732355f..332fb7dc530c 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 @@ -407,7 +407,7 @@ def extra_param_specs(self): @property def ref_slot_specs(self) -> dict[str, RefSlotSpec]: - # I2V first frame (required) + optional last frame for interpolation. + # The last frame is what interpolation conditions on. return { "image_reference": RefSlotSpec( modality="image", diff --git a/tensorrt_llm/serve/openai_video_routes.py b/tensorrt_llm/serve/openai_video_routes.py index 2d61ce2a7ae1..ba61e83a2d57 100644 --- a/tensorrt_llm/serve/openai_video_routes.py +++ b/tensorrt_llm/serve/openai_video_routes.py @@ -148,15 +148,11 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time - # Names this request's output files (``{video_id}_{i}``) and the b64 - # response id; references are keyed and reclaimed by the engine instead. + # Prefixes this request's output files (``{video_id}_{i}``). video_id = f"video_{uuid.uuid4().hex}" try: - # Client-side ValueErrors from content-type parsing, request - # translation, encoder-format preflight, parameter validation, - # and the engine call return 400. Serialization / encoder failures - # further down (server-side) fall through to the outer - # ``except Exception`` → 500. + # ValueError here is the client's fault and returns 400; anything + # further down falls through to the outer handler as a 500. try: # Parse request based on content-type request = await self._parse_video_generation_request(raw_request) @@ -172,9 +168,8 @@ async def openai_video_generation_sync(self, raw_request: Request) -> Response: f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) sync_video_start = time.perf_counter() - # Offload the blocking resolve/enqueue off the event loop but - # await it, so bad media / bad params still surface as 400 - # here; then await generation on the executor's loop. + # Awaited, not fire-and-forget: bad media / bad params must + # surface as a 400 rather than a failed job. handle = await asyncio.to_thread( self.generator.generate_async, request.prompt, params ) @@ -396,8 +391,7 @@ async def openai_video_generation_async( - Multipart: Send form fields + optional image_reference / video_reference file """ request_received = raw_request.state.server_arrival_time - # Names this request's output files and VIDEO_STORE entry; references - # are keyed and reclaimed by the engine when the task awaits the handle. + # Prefixes this request's output files and keys its VIDEO_STORE entry. video_id = f"video_{uuid.uuid4().hex}" try: # Parse request based on content-type @@ -428,10 +422,8 @@ async def openai_video_generation_async( f"Generating video: {video_id} with params: {params} and prompt: {request.prompt}" ) - # Resolve references, validate params, and enqueue in the - # foreground (offloaded but awaited) so bad media / unknown - # extra_params surface as 400 here, before the 202 — not as a queued - # job that later fails. + # Awaited before the 202, so bad media / unknown extra_params are a + # 400 rather than a queued job that later fails. handle = await asyncio.to_thread(self.generator.generate_async, request.prompt, params) # Persist the queued job before scheduling the background task so diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 6673161d8441..43f3f91d3068 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -391,9 +391,7 @@ def _apply_deprecated_input_reference( "files and is only meaningful for co-located clients. Send the " "file as base64 or upload it via multipart/form-data." ) - # Imported here, not at module scope: the resolver reaches into - # VisualGen, and a plain LLM deployment must not pull a vertical in - # behind its request schema. + # Local import, for the reason given at the first one. from tensorrt_llm.visual_gen.media_refs import _resolve_reference payload = _resolve_reference(input_reference, input_reference_format) diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 0891a05b1550..73a6c59980b6 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -19,10 +19,8 @@ from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status -# The reference wire types live in a dependency-neutral leaf so the common -# serving protocol can name them without pulling VisualGen in, but -# ``tensorrt_llm.visual_gen`` stays their public home. The redundant aliases -# mark these as intentional re-exports rather than unused imports. +# Defined in a dependency-neutral leaf so the serving protocol can name them +# without importing VisualGen; re-exported here as their public home. from tensorrt_llm.media.reference import MediaContentFormat as MediaContentFormat from tensorrt_llm.media.reference import MediaRef as MediaRef from tensorrt_llm.media.reference import MediaRole as MediaRole @@ -194,6 +192,10 @@ def validate_visual_gen_params( field set can still validate ``extra_params``. - Type mismatches for ``extra_params`` values. - Out-of-range ``extra_params`` values. + - References in a slot the pipeline does not declare, in an unsupported + role, or in counts outside the role's ``min``/``max``. Skipped when + ``ref_slot_specs`` is ``None``; an empty mapping declares no slots and + so rejects every reference. """ messages: List[str] = [] specs = extra_param_specs @@ -267,31 +269,19 @@ def validate_visual_gen_params( f"extra_params['{key}'] value {value} is out of range [{lo}, {hi}]" ) - # --- reference role / arity checks (duck-typed RefSlotSpec) --- - # ``ref_slot_specs`` maps a reference field name to a spec exposing - # ``.roles`` (a list of role specs with ``.role`` / ``.min`` / ``.max``). - # role must be explicit only when the assignment is ambiguous (a multi-role - # slot with more than one required role); a single-role slot or a single - # required role is inferred. Reference fields are - # already normalized to ``list[*Ref]`` by the field validators. An empty - # (but non-None) mapping means the pipeline declares no slots, so any - # reference the client sent is rejected; only ``None`` skips validation. + # The field validators have already normalized each slot to ``list[MediaRef]``. if ref_slot_specs is not None: for field in ("image_reference", "video_reference", "audio_reference"): refs = getattr(params, field, None) or [] spec = ref_slot_specs.get(field) if spec is None: - # An undeclared slot is only an error if the client actually - # sent one; an absent undeclared slot is fine. if refs: messages.append(f"'{field}' is not accepted by the loaded pipeline.") continue role_specs = list(spec.roles) allowed = {rs.role for rs in role_specs} - # A role-less ref is inferred when unambiguous: a single-role slot, - # or a multi-role slot with exactly one required role (min >= 1) — - # e.g. i2v's first_frame — matching the pipeline's own default. Only - # a genuinely ambiguous slot (multiple required roles) demands one. + # A missing role is inferred while it is unambiguous, so only a slot + # with more than one required role forces the caller to name one. required_roles = [rs.role for rs in role_specs if rs.min >= 1] counts: Dict[str, int] = {} for r in refs: @@ -313,9 +303,8 @@ def validate_visual_gen_params( ) continue counts[role] = counts.get(role, 0) + 1 - # Arity runs even for an absent slot: a role with ``min >= 1`` is a - # required reference, enforced here as a clean 400 instead of a deep - # worker crash. ``min == 0`` leaves the slot optional. + # Runs for an absent slot too, which is what catches a missing + # required reference before the worker sees the request. for rs in role_specs: n = counts.get(rs.role, 0) if n < rs.min or (rs.max is not None and n > rs.max): diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py index e3ad3c7c2d04..2de382582c53 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py @@ -335,9 +335,13 @@ def test_a_path_is_a_type_error(self): with pytest.raises(ValueError, match="PIL images or encoded bytes"): QwenImageEditPlusPipeline._load_edit_images(["/tmp/ref.png"]) - def test_alpha_is_composited_onto_white(self): - """This pipeline went through ``load_image``, which flattens onto white - rather than dropping the channel the way diffusers does.""" + def test_alpha_is_dropped_not_composited(self): + """A transparent pixel keeps its stored RGB. + + ``convert("RGB")`` discards the channel; compositing onto white instead + rewrites every pixel with ``alpha < 255``, silently, so a reference with + transparency would reach the model as a different image. + """ import PIL.Image from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import ( @@ -345,8 +349,10 @@ def test_alpha_is_composited_onto_white(self): ) transparent = PIL.Image.new("RGBA", (4, 4), (10, 20, 30, 0)) + encoded = self._png(mode="RGBA", color=(10, 20, 30, 0)) - (image,) = QwenImageEditPlusPipeline._load_edit_images([transparent]) + from_pil, from_bytes = QwenImageEditPlusPipeline._load_edit_images([transparent, encoded]) - assert image.mode == "RGB" - assert image.getpixel((0, 0)) == (255, 255, 255) + assert from_pil.mode == "RGB" and from_bytes.mode == "RGB" + assert from_pil.getpixel((0, 0)) == (10, 20, 30) + assert from_bytes.getpixel((0, 0)) == (10, 20, 30) From ac3f7bfbdabe1f45aae232f2b8204ae491f50a6a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 22:30:46 -0700 Subject: [PATCH 51/61] [TRTLLM-15277][fix] Let one switch answer the local-path question main added TRTLLM_DISALLOW_LOCAL_MEDIA_PATH for response_format='path', which discloses where the server wrote a file. This branch added the same name for format='path' on a reference, which has the server read one. Two readers of the same variable, each with its own default handling and its own warning, and neither mentioning the other. They are the same question -- whether a client may name a path on the server's filesystem -- so they now share one predicate. An operator who turns it off gets both directions, which is what the name promises. Reading the variable twice was not only duplication: each copy validated the value and warned on its own, so the two could disagree about anything but "0" and "1". Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 2 +- examples/visual_gen/serve/README.md | 4 ++-- tensorrt_llm/serve/openai_server.py | 17 +++++------------ tensorrt_llm/serve/visual_gen_utils.py | 25 +++++++++++++------------ 4 files changed, 21 insertions(+), 27 deletions(-) diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 9435ce97f616..22e731b88423 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -117,7 +117,7 @@ When served via `trtllm-serve`, the following OpenAI-compatible endpoints are av The asynchronous `/v1/videos` job advances through `GET /v1/videos/{id}`: `queued` → `generating` (model inference) → `postprocessing` (encode the media and/or write the output file) → `completed`. The `generating` → `postprocessing` transition marks the end of inference; the video is downloadable via `/content` once `completed`. -`response_format="path"` returns the generated file's server-side path (under `TRTLLM_MEDIA_STORAGE_PATH`) for co-located clients, enabled by default. Set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject such requests with HTTP 400. See the [serve examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/visual_gen/serve) for the full `response_format` reference. +`response_format="path"` returns the generated file's server-side path (under `TRTLLM_MEDIA_STORAGE_PATH`) for co-located clients, enabled by default. Set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject such requests with HTTP 400; the same switch also rejects a reference sent with `format="path"`, since both ask the server to trust a local filesystem path. See the [serve examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/visual_gen/serve) for the full `response_format` reference. ### Reference Inputs diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index e789fb746627..aad4bfd5ce22 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -293,7 +293,7 @@ You can customize these by: {"image_reference": {"content": "iVBORw0KGgoAAAANSUhEUg...", "format": "base64"}} ``` - - `"path"` reads a file on the *server*, so it is only meaningful for a co-located client; set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject it. `"url"` is fetched through the SSRF-guarded loader (private-address block, redirect re-validation, timeout, size cap). + - `"path"` reads a file on the *server*, so it is only meaningful for a co-located client; set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject it (the same switch also disables `response_format="path"`). `"url"` is fetched through the SSRF-guarded loader (private-address block, redirect re-validation, timeout, size cap). - `format` here is the *input* wire form; the top-level `format` selects the *output* encoding. - `role` disambiguates a model that accepts the same modality in more than one role — Wan 2.1 I2V takes a first frame and an optional last frame. Roles and lists need a JSON body; a multipart upload is a single file with no role. @@ -310,7 +310,7 @@ You can customize these by: - `response_format`: `"file"` (default; `FileResponse` byte download) or `"path"` (server-side output path JSON, for co-located clients) - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). -> **`response_format="path"`** (image and video) returns absolute server-side file paths under the server's media-storage directory (`TRTLLM_MEDIA_STORAGE_PATH`), for clients co-located with the server (shared filesystem). Enabled by default; set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject `path` requests with HTTP 400. +> **`response_format="path"`** (image and video) returns absolute server-side file paths under the server's media-storage directory (`TRTLLM_MEDIA_STORAGE_PATH`), for clients co-located with the server (shared filesystem). Enabled by default; set `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1` to reject `path` requests with HTTP 400. One switch covers both directions: it also rejects a reference sent with `format="path"`. #### Tensor-format consumer contract diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 7fc7f392b153..88bbafed5fc7 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -107,7 +107,8 @@ from tensorrt_llm.serve.tool_parser.tool_parser_factory import ToolParserFactory from tensorrt_llm.serve.visual_gen_metrics import ( build_visual_gen_server_timings, 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 (local_media_path_is_disallowed, + parse_visual_gen_params) from tensorrt_llm.usage import TerminalOutcome, record_termination_observation from tensorrt_llm.version import __version__ as VERSION @@ -3226,20 +3227,12 @@ def _reject_disabled_path( self, response_format: Optional[str]) -> Optional[Response]: """Return a 400 when ``response_format='path'`` but it is disabled. - ``path`` discloses absolute server-side filesystem paths, so it can be - turned off via ``TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1`` on shared / - untrusted deployments (enabled by default). Returns ``None`` when - allowed. + Shares the switch with ``format='path'`` on the request side; see + :func:`local_media_path_is_disallowed`. Returns ``None`` when allowed. """ if response_format != "path": return None - raw = os.environ.get("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "0") - if raw not in ("0", "1"): - logger.warning( - "Unrecognized value for TRTLLM_DISALLOW_LOCAL_MEDIA_PATH: " - f"{raw!r}. Expected '0' or '1'. Treating as '0' " - "(response_format='path' enabled).") - if raw == "1": + if local_media_path_is_disallowed(): return self.create_error_response( "response_format='path' is disabled on this server " "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it returns " diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 43f3f91d3068..f73d0391cd0a 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -109,22 +109,23 @@ def _merge_extra_params( params.extra_params = None -def _local_media_path_is_disallowed() -> bool: - """Whether ``format='path'`` is turned off for HTTP requests. - - A ``path`` reference asks the server to read its own disk, which is what a - co-located client wants and what a remote one has no business doing. Which - of the two a deployment has is not something the code can know, so it is - allowed by default and turned off with - ``TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1``. The local Python API is unaffected - either way: this gate is the HTTP boundary's. +def local_media_path_is_disallowed() -> bool: + """Whether server-side filesystem paths are refused at the HTTP boundary. + + One switch for both directions, because both are the same trust question: + a ``path`` reference has the server read its own disk, and + ``response_format='path'`` has it disclose where it wrote. Either is what a + co-located client wants and what a remote one has no business doing, and + the code cannot tell the two deployments apart, so both are allowed by + default and turned off together with ``TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1``. + The local Python API is unaffected: this gate is the HTTP boundary's. """ raw = os.environ.get("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "0") if raw not in ("0", "1"): logger.warning( "Unrecognized value for TRTLLM_DISALLOW_LOCAL_MEDIA_PATH: " f"{raw!r}. Expected '0' or '1'. Treating as '0' " - "(reference format='path' allowed)." + "(server-side paths allowed)." ) return raw == "1" @@ -143,7 +144,7 @@ def _reference_transport(ref: Any) -> tuple[Any, str, Optional[str]]: content = getattr(ref, "content", None) if not isinstance(content, str): raise ValueError("reference item must carry a 'content' string.") - if ref.format == "path" and _local_media_path_is_disallowed(): + if ref.format == "path" and local_media_path_is_disallowed(): raise ValueError( "reference format='path' is disallowed on this server " "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it reads server-side files " @@ -384,7 +385,7 @@ def _apply_deprecated_input_reference( if hasattr(input_reference, "file"): # multipart upload — form implied payload = input_reference.file.read() else: - if input_reference_format == "path" and _local_media_path_is_disallowed(): + if input_reference_format == "path" and local_media_path_is_disallowed(): raise ValueError( "reference format='path' is disallowed on this server " "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it reads server-side " From e04226625c34fa87dba172e61329439c3a551d6f Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 23:14:54 -0700 Subject: [PATCH 52/61] [TRTLLM-15277][fix] Decode the reference before accepting it PIL.Image.open reads the header; nothing decodes until something asks for pixels. The load_image path this replaced called image.load() to force that, so a truncated upload raised OSError and became a client-side ValueError here. convert_image_mode returns an already-RGB image untouched, so with the load() gone the truncated bytes passed the acceptance check and failed later, as a server fault. The other call sites spell the conversion as .convert("RGB"), which loads first, so this was the one place the check went missing. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 932e00de4981..f316e01aa7f2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -247,9 +247,13 @@ def _load_reference_image(data: bytes): upload would be reported as a server fault. """ try: - # convert_image_mode, not convert("RGB"): it composites RGBA onto white, - # which is what this pipeline's references went through before. - return convert_image_mode(PIL.Image.open(BytesIO(data)), "RGB") + image = PIL.Image.open(BytesIO(data)) + # open() reads the header only. Truncated pixel data raises here and + # nowhere else: an already-RGB image never reaches a decoding convert. + image.load() + # convert_image_mode composites RGBA onto white, which is what this + # pipeline's references went through before. + return convert_image_mode(image, "RGB") except OSError as exc: raise ValueError( f"Image reference could not be decoded; it may be truncated, " From 6127043e392d3dbfb333f2e6105419f8164cc219 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Tue, 25 Aug 2026 23:57:17 -0700 Subject: [PATCH 53/61] [TRTLLM-15277][fix] Send the async example's reference as base64 client.videos.create(image_reference=...) raises TypeError before anything is sent: the OpenAI SDK declares one file parameter for videos, input_reference, and this example was never run. extra_body does not carry a file either. The SDK only emits a multipart file part for parameters it declares as FileTypes, so a file object placed there reaches the wire as a form field holding the object's repr. A nested dict is no better -- it is flattened into image_reference[content] and image_reference[format], which match no field the server knows. Encode the file and pass the reference as a JSON string instead. That is the spelling the multipart parser already accepts for a reference sent as a text part, it needs no file parameter, and it does not fall back to input_reference, which this PR deprecates. sync_video_gen.py posts multipart with requests directly, so its image_reference is a real file part and needs no change. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- examples/visual_gen/serve/async_video_gen.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/examples/visual_gen/serve/async_video_gen.py b/examples/visual_gen/serve/async_video_gen.py index 9dfd50925994..4e9aaf2b2312 100755 --- a/examples/visual_gen/serve/async_video_gen.py +++ b/examples/visual_gen/serve/async_video_gen.py @@ -17,6 +17,8 @@ """ import argparse +import base64 +import json import sys import time from pathlib import Path @@ -81,17 +83,19 @@ def test_async_video_generation( }, } - # Add input reference if provided (TI2V mode). Keep the create call - # inside the file's context so the handle closes once the request is sent. + # Add the conditioning image if provided (TI2V mode). The OpenAI SDK + # only knows one file parameter, `input_reference`, so the typed field + # travels base64-encoded in extra_body. It has to be a JSON string: a + # nested dict would be flattened into `image_reference[content]`. if image_reference: if not Path(image_reference).exists(): print(f"\n❌ Error: Input reference image not found: {image_reference}") return False - with open(image_reference, "rb") as ref_file: - create_params["image_reference"] = ref_file - job = client.videos.create(**create_params) - else: - job = client.videos.create(**create_params) + encoded = base64.b64encode(Path(image_reference).read_bytes()).decode() + create_params["extra_body"]["image_reference"] = json.dumps( + {"content": encoded, "format": "base64"} + ) + job = client.videos.create(**create_params) print("Video generation started: \n", job.model_dump_json(indent=2)) From 7ca16ce261e366b189d92be9c565ca22071fcb3a Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 26 Aug 2026 00:05:26 -0700 Subject: [PATCH 54/61] [TRTLLM-15277][chore] Drop reference tests that another test already covers Six cases, each with a named survivor. The choke point's format matrix and its missing-path case repeat TestResolveReference, which covers six spellings against three and reaches the same error. What only the choke point can show -- that resolving rewrites format to "bytes" -- moves to the test that already asserts nothing is written to disk. test_restore_is_idempotent asserted that a second refs_from_shm is harmless. That mattered when a reclaim path could call it after the consumer; there is one caller now, once per request, so nothing relies on it. A character device, a FIFO and a directory all exercise one "not a regular file" branch, and the two multipart video cases differ only in the container, so each set becomes one parametrized case. The symlink cases stay: they assert that stat follows the link, which is a different mechanism. Two comments in those tests still said the payload was persisted to a file. It is not. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../visual_gen/test_visual_gen_utils.py | 108 ++++-------------- 1 file changed, 25 insertions(+), 83 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 870933842e49..1a946258e2f8 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -333,29 +333,18 @@ def _avi_bytes() -> bytes: TestInputReferenceResolution._TEST_DATA / "cosmos3_v2v_ref_9f_bframes.avi" ).read_bytes() - def test_multipart_avi_video_reference_resolves_to_bytes(self, tmp_path): - # The AVI container survives the boundary and is persisted as untouched - # encoded bytes for the worker to demux. + @pytest.mark.parametrize("container", ["avi", "mp4"]) + def test_multipart_video_reference_resolves_to_bytes(self, container): + """The boundary never decodes video: the encoded bytes cross it + untouched, and the worker demuxes the conditioning window itself.""" generator = _StubVisualGen() - payload = self._avi_bytes() - upload = UploadFile(file=BytesIO(payload), filename="clip.avi") + payload = self._avi_bytes() if container == "avi" else self._mp4_bytes() + upload = UploadFile(file=BytesIO(payload), filename=f"clip.{container}") request = VideoGenerationRequest(prompt="x", video_reference=upload) params = _parse_and_prepare(request, generator) assert params.image_reference is None assert params.video_reference[0].content == payload - def test_multipart_mp4_video_reference_resolves_to_bytes(self, tmp_path): - generator = _StubVisualGen() - payload = self._mp4_bytes() - upload = UploadFile(file=BytesIO(payload), filename="clip.mp4") - request = VideoGenerationRequest(prompt="x", video_reference=upload) - params = _parse_and_prepare(request, generator) - # Encoded payload is persisted byte-identical — the boundary never - # decodes video; the worker demuxes/NVDEC-decodes the conditioning - # window from the stored file. - assert params.image_reference is None - assert params.video_reference[0].content == payload - def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): # The deprecated single input_reference is sniff-routed to the typed slot. generator = _StubVisualGen() @@ -893,31 +882,10 @@ def test_malformed_base64_is_a_client_error(self): class TestPrepareReferenceSlots: """The engine choke point: every declared form resolves to raw bytes.""" - def test_every_format_resolves_to_the_same_bytes(self, tmp_path): - """path / base64 / bytes all name the same payload, so all three must - land on byte-identical content with format rewritten to "bytes".""" - from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots - - buf = BytesIO() - Image.new("RGB", (4, 4), (7, 8, 9)).save(buf, format="PNG") - png = buf.getvalue() - src = tmp_path / "ref.png" - src.write_bytes(png) - - for ref in ( - MediaRef(content=str(src), format="path"), - MediaRef(content=base64.b64encode(png).decode(), format="base64"), - MediaRef(content=png, format="bytes"), - ): - params = VisualGenParams(image_reference=ref) - prepare_reference_slots(params) - assert params.image_reference[0].content == png - assert params.image_reference[0].format == "bytes" - - def test_nothing_is_resolves_to_bytes(self, tmp_path, monkeypatch): + def test_resolving_writes_nothing_to_disk(self, tmp_path, monkeypatch): """References never touch the filesystem, so a worker needs no shared - filesystem to read what the coordinator resolved.""" + filesystem to read what the coordinator resolved. The rewritten + ``format`` is what stops a worker being handed a stale spelling.""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots @@ -931,6 +899,7 @@ def test_nothing_is_resolves_to_bytes(self, tmp_path, monkeypatch): ) prepare_reference_slots(params) assert list(tmp_path.iterdir()) == [] + assert params.image_reference[0].format == "bytes" def test_wrong_modality_is_rejected(self, tmp_path): """Content is validated against the slot's modality before dispatch.""" @@ -946,16 +915,6 @@ def test_wrong_modality_is_rejected(self, tmp_path): with pytest.raises(ValueError, match="video_reference is not a recognized"): prepare_reference_slots(params) - def test_missing_path_is_a_client_error(self, tmp_path): - from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots - - params = VisualGenParams( - image_reference=MediaRef(content=str(tmp_path / "nope.png"), format="path") - ) - with pytest.raises(ValueError, match="could not be read"): - prepare_reference_slots(params) - class TestResolveReference: """``_resolve_reference`` dispatches on the declared format, never on the value.""" @@ -1103,16 +1062,6 @@ def test_no_references_costs_nothing(self): req.refs_to_shm() assert req.ref_handles is None - def test_restore_is_idempotent(self): - """rank0 restores unconditionally; a second call must not re-consume a - handle that has already been resolved.""" - payload = os.urandom(2048) - req = self._request(payload) - req.refs_to_shm() - req.refs_from_shm() - req.refs_from_shm() - assert req.params.image_reference[0].content == payload - class TestReferenceBroadcastSplit: """Reference payloads leave the object before the rank0 -> N-rank hop. @@ -1265,13 +1214,24 @@ def test_a_symlink_to_a_regular_file_reads(self, tmp_path): assert _safe_read_local_file(str(link)) == target.read_bytes() - @pytest.mark.parametrize("device", ["/dev/zero", "/dev/null"]) - def test_a_character_device_is_refused(self, device): - """The unbounded read: `/dev/zero` never reaches EOF.""" + @pytest.mark.parametrize("kind", ["chardev", "fifo", "directory"]) + def test_a_non_regular_file_is_refused(self, tmp_path, kind): + """Only a regular file has a size the read can trust: a character + device never reaches EOF and a FIFO blocks instead of returning.""" + import os + from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + if kind == "chardev": + target = "/dev/zero" + elif kind == "fifo": + target = str(tmp_path / "pipe") + os.mkfifo(target) + else: + target = str(tmp_path) + with pytest.raises(ValueError, match="not a regular file"): - _safe_read_local_file(device) + _safe_read_local_file(target) def test_a_symlink_to_a_device_is_refused(self, tmp_path): """``stat`` follows the link, so the check sees what will be read.""" @@ -1283,24 +1243,6 @@ def test_a_symlink_to_a_device_is_refused(self, tmp_path): with pytest.raises(ValueError, match="not a regular file"): _safe_read_local_file(str(link)) - def test_a_fifo_is_refused(self, tmp_path): - """Reading a FIFO blocks rather than returning, so size caps cannot help.""" - import os - - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file - - fifo = tmp_path / "pipe" - os.mkfifo(fifo) - - with pytest.raises(ValueError, match="not a regular file"): - _safe_read_local_file(str(fifo)) - - def test_a_directory_is_refused(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file - - with pytest.raises(ValueError, match="not a regular file"): - _safe_read_local_file(str(tmp_path)) - def test_a_missing_file_is_a_client_error(self, tmp_path): from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file From 669facacd09779c7b1099171bf89102ddfb836b3 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 26 Aug 2026 11:11:50 -0700 Subject: [PATCH 55/61] [TRTLLM-15277][chore] Home the media reference types and their resolution MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback: `tensorrt_llm/media/` hosts modality utilities that are not VisualGen-specific, and the top-level API stays small. `MediaRef`, `MediaRole`, `MediaContentFormat` and the bare-reference rejector move into `visual_gen/params.py`, next to the three `*_reference` fields that carry them, and `MediaRef` leaves the top-level `__all__`. Reference resolution follows them: it runs in `generate_async` before the request exists, alongside `validate_visual_gen_params`, and answers the same question — is this request acceptable — while decoding stays in the pipelines. `_safe_read_local_file` instead joins `_normalize_file_uri` and `_safe_request_get` in `inputs/media_io.py`; reading a local file safely is not a reference concern, and its wording no longer says otherwise. `tensorrt_llm/media/` and `tensorrt_llm/__init__.py` are back to no diff against main, and the branch adds no new file under `visual_gen/`. Tests: drop three cases another test already covers (a slot accepted without `ref_slot_specs`, a list of references, a valid content/format pairing) and add the one guard that was missing — a role the slot never declared must be refused, verified by reverting the check and watching the case fail. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- docs/source/models/visual-generation.md | 6 +- examples/visual_gen/models/cosmos3/cosmos3.py | 3 +- examples/visual_gen/models/flux2.py | 3 +- examples/visual_gen/models/qwen_image_edit.py | 3 +- .../visual_gen/models/qwen_image_layered.py | 3 +- examples/visual_gen/models/wan_i2v.py | 3 +- tensorrt_llm/__init__.py | 8 +- tensorrt_llm/inputs/media_io.py | 34 ++ tensorrt_llm/media/reference.py | 106 ------ tensorrt_llm/serve/openai_protocol.py | 2 +- tensorrt_llm/serve/visual_gen_utils.py | 4 +- tensorrt_llm/visual_gen/media_refs.py | 184 --------- tensorrt_llm/visual_gen/params.py | 208 +++++++++- tensorrt_llm/visual_gen/visual_gen.py | 7 +- .../test_flux2_image_conditioning.py | 7 - .../visual_gen/test_qwen_image_pipeline.py | 11 - .../visual_gen/test_trtllm_serve_endpoints.py | 3 +- .../visual_gen/test_visual_gen_params.py | 48 +-- .../visual_gen/test_visual_gen_utils.py | 354 ++---------------- tests/unittest/llmapi/apps/test_media_io.py | 88 +---- 20 files changed, 306 insertions(+), 779 deletions(-) delete mode 100644 tensorrt_llm/media/reference.py delete mode 100644 tensorrt_llm/visual_gen/media_refs.py diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 22e731b88423..3e80c120b2ef 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -135,7 +135,8 @@ Every pipeline declares the reference slots and roles it accepts through `ref_sl Most models take a single reference whose role is unambiguous: ```python -from tensorrt_llm import VisualGen, MediaRef +from tensorrt_llm import VisualGen +from tensorrt_llm.visual_gen import MediaRef vg = VisualGen(model="Wan-AI/Wan2.2-TI2V-5B-Diffusers") params = vg.default_params @@ -146,7 +147,8 @@ output = vg.generate(inputs="the scene comes alive with gentle motion", params=p Models that accept the same modality in more than one role need `role`. Wan 2.1 I2V takes a first frame and an optional last frame: ```python -from tensorrt_llm import VisualGen, MediaRef +from tensorrt_llm import VisualGen +from tensorrt_llm.visual_gen import MediaRef vg = VisualGen(model="Wan-AI/Wan2.1-I2V-14B-480P-Diffusers") params = vg.default_params diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 402239944b23..af390395dc89 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -27,8 +27,9 @@ from pathlib import Path from typing import Any, Dict, Optional -from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs +from tensorrt_llm import VisualGen, VisualGenArgs from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS +from tensorrt_llm.visual_gen import MediaRef _SCRIPT_DIR = Path(__file__).resolve().parent _ACTION_MODES = ("policy", "forward_dynamics", "inverse_dynamics") diff --git a/examples/visual_gen/models/flux2.py b/examples/visual_gen/models/flux2.py index adf50af3bf8f..652f350cef0a 100644 --- a/examples/visual_gen/models/flux2.py +++ b/examples/visual_gen/models/flux2.py @@ -25,7 +25,8 @@ import argparse from pathlib import Path -from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm.visual_gen import MediaRef def _output_paths(output_path: str, num_images: int) -> str | list[str]: diff --git a/examples/visual_gen/models/qwen_image_edit.py b/examples/visual_gen/models/qwen_image_edit.py index d48904f06b6d..c0b389ca2f45 100644 --- a/examples/visual_gen/models/qwen_image_edit.py +++ b/examples/visual_gen/models/qwen_image_edit.py @@ -23,7 +23,8 @@ import argparse -from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm.visual_gen import MediaRef def parse_args() -> argparse.Namespace: diff --git a/examples/visual_gen/models/qwen_image_layered.py b/examples/visual_gen/models/qwen_image_layered.py index 46115c7c15e1..c1485c365ad3 100644 --- a/examples/visual_gen/models/qwen_image_layered.py +++ b/examples/visual_gen/models/qwen_image_layered.py @@ -23,7 +23,8 @@ import argparse from pathlib import Path -from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm.visual_gen import MediaRef def parse_args() -> argparse.Namespace: diff --git a/examples/visual_gen/models/wan_i2v.py b/examples/visual_gen/models/wan_i2v.py index 2438680157ca..ae7101e0ad2e 100644 --- a/examples/visual_gen/models/wan_i2v.py +++ b/examples/visual_gen/models/wan_i2v.py @@ -23,7 +23,8 @@ import argparse import os -from tensorrt_llm import MediaRef, VisualGen, VisualGenArgs +from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm.visual_gen import MediaRef _DEFAULT_IMAGE = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "cat_piano.png") diff --git a/tensorrt_llm/__init__.py b/tensorrt_llm/__init__.py index e878bf7ad36b..516e31ad2c52 100644 --- a/tensorrt_llm/__init__.py +++ b/tensorrt_llm/__init__.py @@ -62,9 +62,9 @@ from .mapping import Mapping from .models.automodel import AutoConfig, AutoModelForCausalLM from .sampling_params import SamplingParams - from .visual_gen import (ExtraParamSchema, MediaRef, VisualGen, - VisualGenArgs, VisualGenMetrics, VisualGenOutput, - VisualGenParams, VisualGenResult) + from .visual_gen import (ExtraParamSchema, VisualGen, VisualGenArgs, + VisualGenMetrics, VisualGenOutput, VisualGenParams, + VisualGenResult) # Public name -> (source module, attribute); attribute None = the module itself. _LAZY_ATTRS = { @@ -105,7 +105,6 @@ 'VisualGenMetrics': ('tensorrt_llm.visual_gen', 'VisualGenMetrics'), 'VisualGenOutput': ('tensorrt_llm.visual_gen', 'VisualGenOutput'), 'VisualGenParams': ('tensorrt_llm.visual_gen', 'VisualGenParams'), - 'MediaRef': ('tensorrt_llm.visual_gen', 'MediaRef'), 'VisualGenResult': ('tensorrt_llm.visual_gen', 'VisualGenResult'), } @@ -174,7 +173,6 @@ def __dir__(): 'math_utils', 'VisualGen', 'VisualGenParams', - 'MediaRef', '__version__', ] diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 3ce027cd2e54..bdfa472e9afc 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -14,6 +14,7 @@ from concurrent.futures import Executor from io import BytesIO from pathlib import Path +from stat import S_ISREG from types import MappingProxyType from typing import ( Any, @@ -749,6 +750,39 @@ def _normalize_file_uri(uri: str) -> str: return uri +def _safe_read_local_file(location: str) -> bytes: + """Read a local file, refusing anything that has no end. + + Takes a bare path or a ``file://`` URI. + + The counterpart of :func:`_safe_request_get` for the local branch, and it + guards the one case that is unbounded rather than merely large: reading a + character device never reaches EOF and reading a FIFO blocks, so either + turns a caller into a denial of service. A regular file is finite, which is + the property required here. + + Size is deliberately not capped: naming a large file of one's own is the + normal case for a local caller, and no threshold separates that from an + abusive one. This bounds the shape of what may be read, then, not its size + or its reach — any regular file the process can read is still readable. + """ + path = Path(_normalize_file_uri(location)) + try: + stat = path.stat() # follows symlinks, so a link to a device is caught + except OSError as exc: + raise ValueError(f"file could not be read: {exc}") from exc + + if not S_ISREG(stat.st_mode): + raise ValueError( + f"path is not a regular file: {location!r}. Character devices, " + "FIFOs and directories cannot be read as media." + ) + try: + return path.read_bytes() + except OSError as exc: + raise ValueError(f"file could not be read: {exc}") from exc + + _MediaT = TypeVar("_MediaT") diff --git a/tensorrt_llm/media/reference.py b/tensorrt_llm/media/reference.py deleted file mode 100644 index 0fa513825081..000000000000 --- a/tensorrt_llm/media/reference.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""The wire types for a media reference, shared by serve and VisualGen. - -These are declaration-only: the schema a caller fills in, with no resolver, -decoder or engine behind them. They live here rather than under -``visual_gen`` so the common serving protocol can name them without the -request schema of every LLM deployment pulling a vertical in behind it. -Depend on nothing but ``pydantic`` and keep it that way. -""" - -from typing import Any, Optional, Union - -from pydantic import Field, model_validator -from typing_extensions import Literal - -from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status - -MediaRole = Literal["reference", "first_frame", "last_frame"] - -# Wire form of a reference's ``content``. Declared explicitly rather than -# sniffed: a bare string is otherwise ambiguous between a local path and -# base64, and guessing lets a mistyped path silently become base64 (or a -# malformed base64 silently become a filesystem read). -MediaContentFormat = Literal["path", "url", "base64", "bytes"] - - -@set_api_status("prototype") -class MediaRef(StrictBaseModel): - """A single media reference (image / video / audio). - - Carried by ``image_reference`` / ``video_reference`` / ``audio_reference``; - the field it sits in fixes the modality. ``role`` is required only when the - target model accepts that modality in more than one role (e.g. image first + - last frame); otherwise the pipeline knows the reference's meaning and - ``role`` may be omitted (video/audio are always the single ``reference``). - """ - - content: Union[str, bytes] = Field( - description="The reference payload, in the form declared by ``format``." - ) - format: MediaContentFormat = Field( - description=( - "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " - "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " - "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." - ) - ) - role: Optional[MediaRole] = Field( - default=None, - description=( - "Which conditioning slot this reference fills. Required only when the " - "target model accepts this modality in more than one slot; omit it when " - "the model leaves no ambiguity." - ), - ) - - @model_validator(mode="after") - def _check_content_matches_format(self): - """Reject a ``content`` whose Python type contradicts ``format``. - - ``bytes`` is the only format carrying a binary payload; the other three - name a location or an encoding and are therefore strings. Checking the - pairing here fails at construction — an HTTP 422 or an immediate - ``ValueError`` — instead of deep in the engine's resolve step. - """ - if self.format == "bytes": - if not isinstance(self.content, bytes): - raise ValueError( - f"format='bytes' requires bytes content, got {type(self.content).__name__}." - ) - elif not isinstance(self.content, str): - raise ValueError( - f"format={self.format!r} requires string content, got " - f"{type(self.content).__name__}." - ) - return self - - -def reject_bare_refs(value: Any) -> Any: - """Reject the bare path/bytes shorthand with an actionable message. - - Runs before coercion, so the caller sees what to do instead of a union - mismatch reported against an inner model. A bare string has nowhere to - declare its wire form, and guessing is what ``format`` exists to prevent. - """ - for x in value if isinstance(value, list) else [value]: - if isinstance(x, (str, bytes)): - raise ValueError( - "a reference must declare its wire form; a bare " - f"{type(x).__name__} is no longer accepted. Pass " - 'MediaRef(content=..., format="path"|"url"|"base64"|"bytes").' - ) - return value diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 7396fefebb85..acedd2bf64e4 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -59,10 +59,10 @@ SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory from tensorrt_llm.logger import logger -from tensorrt_llm.media.reference import MediaContentFormat, MediaRole from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy +from tensorrt_llm.visual_gen.params import MediaContentFormat, MediaRole _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index f73d0391cd0a..e7fb9f81c285 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -380,7 +380,7 @@ def _apply_deprecated_input_reference( logger.warning("'input_reference' is deprecated; use 'image_reference' / 'video_reference'.") if params.image_reference or params.video_reference: return - from tensorrt_llm.media.reference import MediaRef + from tensorrt_llm.visual_gen.params import MediaRef if hasattr(input_reference, "file"): # multipart upload — form implied payload = input_reference.file.read() @@ -393,7 +393,7 @@ def _apply_deprecated_input_reference( "file as base64 or upload it via multipart/form-data." ) # Local import, for the reason given at the first one. - from tensorrt_llm.visual_gen.media_refs import _resolve_reference + from tensorrt_llm.visual_gen.params import _resolve_reference payload = _resolve_reference(input_reference, input_reference_format) kind = sniff_media_kind(payload) diff --git a/tensorrt_llm/visual_gen/media_refs.py b/tensorrt_llm/visual_gen/media_refs.py deleted file mode 100644 index 1d93fbb8e990..000000000000 --- a/tensorrt_llm/visual_gen/media_refs.py +++ /dev/null @@ -1,184 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Reference-media resolution, shared by serve and engine. - -Every declared wire form resolves to raw bytes, the canonical form carried all -the way to the pipeline. Used by both the serve boundary -(``tensorrt_llm/serve``) and the engine frontend (``VisualGen.generate_async``), -so this lives here rather than under ``serve`` to avoid an engine -> serve -import. -""" - -from __future__ import annotations - -import base64 -from pathlib import Path -from stat import S_ISREG -from typing import Any - -from tensorrt_llm.inputs.media_io import ( - _normalize_file_uri, - _safe_request_get, - is_isobmff_image_bytes, - sniff_media_kind, -) - - -def _read_reference_payload(reference: str) -> bytes: - """Decode one base64 (optionally ``data:`` URI) reference string to bytes. - - Payload size is deliberately not checked here: encoded size is not part - of the request-validity contract, and body limits belong to the - proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so - malformed encodings — not sizes — are rejected. - """ - data = reference - if data.startswith("data:"): - comma = data.find(",") - if comma == -1: - raise ValueError("reference data: URI is malformed (missing comma).") - # Match the LLM loader: only base64 payloads are supported, and saying so - # beats letting a percent-encoded body fail as "not valid base64". - if "base64" not in data[:comma].split(";")[1:]: - raise ValueError("only base64 data: URIs are supported for references.") - data = data[comma + 1 :] - try: - return base64.b64decode(data, validate=True) - except ValueError as exc: - # binascii.Error subclasses ValueError. - raise ValueError("reference is not valid base64 data.") from exc - - -def _safe_read_local_file(reference: str) -> bytes: - """Read a ``path`` reference, refusing anything that has no end. - - The counterpart of :func:`_safe_request_get` for the local branch, and it - guards the one case that is unbounded rather than merely large: reading a - character device never reaches EOF and reading a FIFO blocks, so either - turns a request into a denial of service. A regular file is finite, which - is the property required here. - - Size is deliberately not capped. ``path`` exists for the local Python API, - where naming a large file of one's own is the normal case, and no - threshold separates that from an abusive one. This bounds the shape of - what may be read, then, not its size or its reach: any regular file the - server process can read is still readable. - """ - path = Path(_normalize_file_uri(reference)) - try: - stat = path.stat() # follows symlinks, so a link to a device is caught - except OSError as exc: - raise ValueError(f"reference file could not be read: {exc}") from exc - - if not S_ISREG(stat.st_mode): - raise ValueError( - f"reference path is not a regular file: {reference!r}. Character " - "devices, FIFOs and directories cannot be read as media." - ) - try: - return path.read_bytes() - except OSError as exc: - raise ValueError(f"reference file could not be read: {exc}") from exc - - -def _resolve_reference(content: Any, content_format: str) -> bytes: - """Resolve one reference to raw bytes using its declared wire form. - - Dispatch is on the caller-declared ``format``, never on the shape of the - value: a bare string is otherwise ambiguous between a local path and - base64, and guessing lets a mistyped path become base64 (or a malformed - base64 become a filesystem read). Fetch/read/decode failures become - ``ValueError`` so a bad reference is a client 400, not a server 500. - """ - if content_format == "bytes": - if not isinstance(content, bytes): - raise ValueError( - f"format='bytes' requires bytes content, got {type(content).__name__}." - ) - return content - if not isinstance(content, str): - raise ValueError( - f"format={content_format!r} requires string content, got {type(content).__name__}." - ) - if content_format == "url": - try: - return _safe_request_get(content).content - except Exception as exc: - raise ValueError(f"reference URL could not be fetched: {exc}") from exc - if content_format == "path": - return _safe_read_local_file(content) - if content_format == "base64": - return _read_reference_payload(content) - raise ValueError(f"unsupported reference format: {content_format!r}") - - -def _validate_reference_payload(payload: bytes, *, modality: str) -> None: - """Reject a payload whose container does not match the declared modality. - - HEIF/AVIF images are rejected on signature alone (Pillow support depends - on optional plugins the worker need not share). Video acceptance beyond the - container signature happens in the worker's NVDEC demux. - """ - if modality == "image": - if sniff_media_kind(payload) != "image": - raise ValueError( - "image_reference is not a recognized image; supported inputs are PNG/JPEG." - ) - if is_isobmff_image_bytes(payload): - raise ValueError( - "image_reference is a HEIF/AVIF image, which is not a supported " - "reference format; convert it to PNG or JPEG." - ) - elif modality == "video": - if sniff_media_kind(payload) != "video": - raise ValueError( - "video_reference is not a recognized media container; supported " - "inputs are MP4/AVI video." - ) - elif modality == "audio": - if sniff_media_kind(payload) != "audio": - raise ValueError( - "audio_reference is not a recognized audio container; supported " - "inputs are WAV/MP3/FLAC/OGG/M4A/AAC." - ) - - -def prepare_reference_slots(params: Any) -> None: - """Resolve every reference to raw bytes, in place. - - The single reference choke point, used by the engine (``generate_async``) - so serve and the standalone Python API share one path. Dispatch is on each - reference's declared ``format``, never on the shape of its content: the - declared form is resolved to bytes, content-validated against the slot's - modality, and written back with ``format`` set to ``"bytes"``. - - ``format`` is rewritten alongside ``content`` because the mutated params - object is what gets broadcast to the workers; a stale format would tell a - worker it is holding base64 when it is holding raw bytes. - - Bytes are the canonical form all the way to the pipeline, so a reference - never touches the filesystem: there is nothing to clean up afterwards, and - a worker needs no shared filesystem to read what the coordinator resolved. - - Runs before the coordinator broadcasts the request, so a bad reference - raises ``ValueError`` synchronously and serve keeps its immediate 400. - """ - for slot in ("image_reference", "video_reference", "audio_reference"): - modality = slot.split("_", 1)[0] - for ref in getattr(params, slot, None) or []: - data = _resolve_reference(ref.content, ref.format) - _validate_reference_payload(data, modality=modality) - ref.content = data - ref.format = "bytes" diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 73a6c59980b6..2bf0a05d63b3 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -13,18 +13,96 @@ # See the License for the specific language governing permissions and # limitations under the License. import ast +import base64 from typing import Any, Dict, List, Optional, Union -from pydantic import Field, field_validator +from pydantic import Field, field_validator, model_validator +from typing_extensions import Literal +from tensorrt_llm.inputs.media_io import ( + _safe_read_local_file, + _safe_request_get, + is_isobmff_image_bytes, + sniff_media_kind, +) from tensorrt_llm.llmapi.utils import StrictBaseModel, set_api_status -# Defined in a dependency-neutral leaf so the serving protocol can name them -# without importing VisualGen; re-exported here as their public home. -from tensorrt_llm.media.reference import MediaContentFormat as MediaContentFormat -from tensorrt_llm.media.reference import MediaRef as MediaRef -from tensorrt_llm.media.reference import MediaRole as MediaRole -from tensorrt_llm.media.reference import reject_bare_refs as _reject_bare_refs +MediaRole = Literal["reference", "first_frame", "last_frame"] + +# Wire form of a reference's ``content``. Declared explicitly rather than +# sniffed: a bare string is otherwise ambiguous between a local path and +# base64, and guessing lets a mistyped path silently become base64 (or a +# malformed base64 silently become a filesystem read). +MediaContentFormat = Literal["path", "url", "base64", "bytes"] + + +@set_api_status("prototype") +class MediaRef(StrictBaseModel): + """A single media reference (image / video / audio). + + Carried by ``image_reference`` / ``video_reference`` / ``audio_reference``; + the field it sits in fixes the modality. ``role`` is required only when the + target model accepts that modality in more than one role (e.g. image first + + last frame); otherwise the pipeline knows the reference's meaning and + ``role`` may be omitted (video/audio are always the single ``reference``). + """ + + content: Union[str, bytes] = Field( + description="The reference payload, in the form declared by ``format``." + ) + format: MediaContentFormat = Field( + description=( + "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " + "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " + "loader), ``base64`` (a ``data:`` URI is also accepted), or ``bytes``." + ) + ) + role: Optional[MediaRole] = Field( + default=None, + description=( + "Which conditioning slot this reference fills. Required only when the " + "target model accepts this modality in more than one slot; omit it when " + "the model leaves no ambiguity." + ), + ) + + @model_validator(mode="after") + def _check_content_matches_format(self): + """Reject a ``content`` whose Python type contradicts ``format``. + + ``bytes`` is the only format carrying a binary payload; the other three + name a location or an encoding and are therefore strings. Checking the + pairing here fails at construction — an HTTP 422 or an immediate + ``ValueError`` — instead of deep in the engine's resolve step. + """ + if self.format == "bytes": + if not isinstance(self.content, bytes): + raise ValueError( + f"format='bytes' requires bytes content, got {type(self.content).__name__}." + ) + elif not isinstance(self.content, str): + raise ValueError( + f"format={self.format!r} requires string content, got " + f"{type(self.content).__name__}." + ) + return self + + +def _reject_bare_refs(value: Any) -> Any: + """Reject the bare path/bytes shorthand with an actionable message. + + Runs before coercion, so the caller sees what to do instead of a union + mismatch reported against an inner model. A bare string has nowhere to + declare its wire form, and guessing is what ``format`` exists to prevent. + """ + for x in value if isinstance(value, list) else [value]: + if isinstance(x, (str, bytes)): + raise ValueError( + "a reference must declare its wire form; a bare " + f"{type(x).__name__} is no longer accepted. Pass " + 'MediaRef(content=..., format="path"|"url"|"base64"|"bytes").' + ) + return value def _normalize_refs(value: Any) -> Optional[list]: @@ -315,3 +393,119 @@ def validate_visual_gen_params( return raise ValueError("Parameter validation failed:\n" + "\n".join(f" - {e}" for e in messages)) + + +def _read_reference_payload(reference: str) -> bytes: + """Decode one base64 (optionally ``data:`` URI) reference string to bytes. + + Payload size is deliberately not checked here: encoded size is not part + of the request-validity contract, and body limits belong to the + proxy/ASGI deployment layer (HTTP 413). Base64 decodes strictly so + malformed encodings — not sizes — are rejected. + """ + data = reference + if data.startswith("data:"): + comma = data.find(",") + if comma == -1: + raise ValueError("reference data: URI is malformed (missing comma).") + # Match the LLM loader: only base64 payloads are supported, and saying so + # beats letting a percent-encoded body fail as "not valid base64". + if "base64" not in data[:comma].split(";")[1:]: + raise ValueError("only base64 data: URIs are supported for references.") + data = data[comma + 1 :] + try: + return base64.b64decode(data, validate=True) + except ValueError as exc: + # binascii.Error subclasses ValueError. + raise ValueError("reference is not valid base64 data.") from exc + + +def _resolve_reference(content: Any, content_format: str) -> bytes: + """Resolve one reference to raw bytes using its declared wire form. + + Dispatch is on the caller-declared ``format``, never on the shape of the + value: a bare string is otherwise ambiguous between a local path and + base64, and guessing lets a mistyped path become base64 (or a malformed + base64 become a filesystem read). Fetch/read/decode failures become + ``ValueError`` so a bad reference is a client 400, not a server 500. + """ + if content_format == "bytes": + if not isinstance(content, bytes): + raise ValueError( + f"format='bytes' requires bytes content, got {type(content).__name__}." + ) + return content + if not isinstance(content, str): + raise ValueError( + f"format={content_format!r} requires string content, got {type(content).__name__}." + ) + if content_format == "url": + try: + return _safe_request_get(content).content + except Exception as exc: + raise ValueError(f"reference URL could not be fetched: {exc}") from exc + if content_format == "path": + return _safe_read_local_file(content) + if content_format == "base64": + return _read_reference_payload(content) + raise ValueError(f"unsupported reference format: {content_format!r}") + + +def _validate_reference_payload(payload: bytes, *, modality: str) -> None: + """Reject a payload whose container does not match the declared modality. + + HEIF/AVIF images are rejected on signature alone (Pillow support depends + on optional plugins the worker need not share). Video acceptance beyond the + container signature happens in the worker's NVDEC demux. + """ + if modality == "image": + if sniff_media_kind(payload) != "image": + raise ValueError( + "image_reference is not a recognized image; supported inputs are PNG/JPEG." + ) + if is_isobmff_image_bytes(payload): + raise ValueError( + "image_reference is a HEIF/AVIF image, which is not a supported " + "reference format; convert it to PNG or JPEG." + ) + elif modality == "video": + if sniff_media_kind(payload) != "video": + raise ValueError( + "video_reference is not a recognized media container; supported " + "inputs are MP4/AVI video." + ) + elif modality == "audio": + if sniff_media_kind(payload) != "audio": + raise ValueError( + "audio_reference is not a recognized audio container; supported " + "inputs are WAV/MP3/FLAC/OGG/M4A/AAC." + ) + + +def prepare_reference_slots(params: Any) -> None: + """Resolve every reference to raw bytes, in place. + + The single reference choke point, used by the engine (``generate_async``) + so serve and the standalone Python API share one path. Dispatch is on each + reference's declared ``format``, never on the shape of its content: the + declared form is resolved to bytes, content-validated against the slot's + modality, and written back with ``format`` set to ``"bytes"``. + + ``format`` is rewritten alongside ``content`` because the mutated params + object is what gets broadcast to the workers; a stale format would tell a + worker it is holding base64 when it is holding raw bytes. + + Bytes are the canonical form all the way to the pipeline, so a reference + never touches the filesystem: there is nothing to clean up afterwards, and + a worker needs no shared filesystem to read what the coordinator resolved. + + Runs before the coordinator broadcasts the request, so a bad reference + raises ``ValueError`` synchronously and serve keeps its immediate 400. + """ + for slot in ("image_reference", "video_reference", "audio_reference"): + modality = slot.split("_", 1)[0] + for ref in getattr(params, slot, None) or []: + data = _resolve_reference(ref.content, ref.format) + _validate_reference_payload(data, modality=modality) + ref.content = data + ref.format = "bytes" diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index ad595a44fab8..a6274d353406 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -31,9 +31,12 @@ from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema, RefSlotSpec from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm.visual_gen.args import VisualGenArgs -from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots from tensorrt_llm.visual_gen.output import VisualGenOutput -from tensorrt_llm.visual_gen.params import VisualGenParams, validate_visual_gen_params +from tensorrt_llm.visual_gen.params import ( + VisualGenParams, + prepare_reference_slots, + validate_visual_gen_params, +) __all__ = [ "VisualGen", diff --git a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py index 38b3656ebe11..551d3a1fbabc 100644 --- a/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py +++ b/tests/unittest/_torch/visual_gen/test_flux2_image_conditioning.py @@ -32,13 +32,6 @@ def test_load_reference_images_accepts_pil_and_bytes(tmp_path) -> None: assert [image.size for image in images] == [(64, 64), (64, 64)] -def test_load_reference_images_rejects_a_path() -> None: - """References reach the worker as bytes, so a path is a type error here, - not a filesystem read.""" - with pytest.raises(ValueError, match="PIL images or encoded bytes"): - Flux2Pipeline._load_reference_images(["/tmp/nope.png"]) - - def test_load_reference_images_drops_alpha_without_compositing() -> None: """RGBA is converted the way diffusers does it: the alpha channel is dropped, so a fully-transparent pixel keeps its RGB value instead of diff --git a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py index 2de382582c53..283ad37abb63 100644 --- a/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_qwen_image_pipeline.py @@ -324,17 +324,6 @@ def test_bytes_and_pil_both_load(self): assert all(image.mode == "RGB" for image in images) assert all(image.size == (8, 8) for image in images) - def test_a_path_is_a_type_error(self): - """References reach a pipeline as bytes; a path is not a filesystem read.""" - import pytest - - from tensorrt_llm._torch.visual_gen.models.qwen_image.pipeline_qwen_image_edit import ( - QwenImageEditPlusPipeline, - ) - - with pytest.raises(ValueError, match="PIL images or encoded bytes"): - QwenImageEditPlusPipeline._load_edit_images(["/tmp/ref.png"]) - def test_alpha_is_dropped_not_composited(self): """A transparent pixel keeps its stored RGB. 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 073b3d5fd5e7..48e48cf744d5 100644 --- a/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py +++ b/tests/unittest/_torch/visual_gen/test_trtllm_serve_endpoints.py @@ -35,9 +35,8 @@ from tensorrt_llm.serve.openai_server import _normalize_image_output from tensorrt_llm.serve.visual_gen_metrics import SERVER_TIMING_HEADER from tensorrt_llm.serve.visual_gen_utils import VIDEO_STORE -from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots from tensorrt_llm.visual_gen.output import VisualGenMetrics, VisualGenOutput -from tensorrt_llm.visual_gen.params import validate_visual_gen_params +from tensorrt_llm.visual_gen.params import prepare_reference_slots, validate_visual_gen_params pytestmark = pytest.mark.cpu_only 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 66fc23298b4f..0c655a6d9f7f 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -118,18 +118,6 @@ def test_image_reference_accepts_bytes(self): assert params.image_reference[0].content == b"\x89PNG" assert params.image_reference[0].format == "bytes" - def test_image_reference_accepts_list(self): - from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - - params = VisualGenParams( - image_reference=[ - MediaRef(content="/path/a.png", format="path"), - MediaRef(content=b"\x89PNG", format="bytes"), - ] - ) - assert len(params.image_reference) == 2 - assert params.image_reference[0].content == "/path/a.png" - def test_model_dump(self): from tensorrt_llm.visual_gen import VisualGenParams @@ -207,13 +195,6 @@ def test_content_type_must_match_format(self, content, content_format): with pytest.raises(ValidationError, match="requires (bytes|string) content"): MediaRef(content=content, format=content_format) - def test_content_type_matching_format_accepted(self): - from tensorrt_llm.visual_gen import MediaRef - - assert MediaRef(content=b"raw", format="bytes").content == b"raw" - for fmt in ("path", "url", "base64"): - assert MediaRef(content="x", format=fmt).format == fmt - def test_engine_rewrite_is_not_blocked_by_the_pairing_check(self): """``prepare_reference_slots`` rewrites content then format, so the intermediate state contradicts the pairing; assignment must not @@ -1046,24 +1027,10 @@ def test_num_frames_on_video_pipeline_ok(self): req = self._make_request(num_frames=81) self._merge_and_validate(executor, req) - def test_image_reference_on_i2v_pipeline_ok(self): - """image_reference is consumed by WanImageToVideoPipeline; validating - without ref_slot_specs should not raise.""" - from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_i2v import ( - WanImageToVideoPipeline, - ) - from tensorrt_llm.visual_gen.params import MediaRef - - executor = self._make_mock_executor(WanImageToVideoPipeline, _wan_mock(num_heads=12)) - req = self._make_request( - image_reference=MediaRef(content="/path/to/img.png", format="path") - ) - self._merge_and_validate(executor, req) - def test_ref_slot_required_vs_optional(self): """``min >= 1`` marks a required reference (clean error when absent); - ``min == 0`` leaves the slot optional; an undeclared absent slot is - fine, but an unsolicited one is rejected.""" + ``min == 0`` leaves the slot optional; an undeclared slot is fine while + absent but rejected when sent, as is a role the slot never declared.""" from tensorrt_llm._torch.visual_gen.pipeline import RefSlotSpec, RoleSpec from tensorrt_llm.visual_gen.params import ( MediaRef, @@ -1094,11 +1061,18 @@ def run(params, spec): run(VisualGenParams(), optional) # Required slot with the image present -> allowed. run(VisualGenParams(image_reference=MediaRef(content="a.png", format="path")), required) - # Undeclared slot left absent -> no spurious "not accepted". - run(VisualGenParams(), optional) # Undeclared slot actually sent -> rejected. with pytest.raises(ValueError, match=r"video_reference.*not accepted"): run(VisualGenParams(video_reference=MediaRef(content="v.mp4", format="path")), optional) + # A role the slot never declared -> rejected here, not carried to a worker + # that has no conditioning input to put it in. + with pytest.raises(ValueError, match=r"role 'last_frame' not supported"): + run( + VisualGenParams( + image_reference=MediaRef(content="a.png", format="path", role="last_frame") + ), + required, + ) def test_multi_role_slot_infers_single_required_role(self): """A role-less ref against a multi-role slot is inferred when only one diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 1a946258e2f8..473f41379ecc 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -31,7 +31,7 @@ parse_visual_gen_params, ) from tensorrt_llm.visual_gen import VisualGenParams -from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots +from tensorrt_llm.visual_gen.params import prepare_reference_slots def _parse_and_prepare(request, generator): @@ -301,22 +301,6 @@ def test_base64_image_reference_resolves_to_bytes(self, tmp_path): assert ref_path == buf.getvalue() assert params.image_reference[0].format == "bytes" - def test_image_reference_role_and_list(self, tmp_path): - generator = _StubVisualGen() - buf = BytesIO() - Image.new("RGB", (4, 4)).save(buf, format="PNG") - b64 = base64.b64encode(buf.getvalue()).decode() - request = VideoGenerationRequest( - prompt="x", - image_reference=[ - {"content": b64, "format": "base64"}, - {"content": b64, "format": "base64", "role": "last_frame"}, - ], - ) - params = _parse_and_prepare(request, generator) - assert [r.role for r in params.image_reference] == [None, "last_frame"] - assert [r.content for r in params.image_reference] == [buf.getvalue()] * 2 - _TEST_DATA = Path(__file__).parent / "test_data" @staticmethod @@ -388,48 +372,6 @@ def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): assert len(p.image_reference) == 1 assert p.video_reference is None # input_reference video dropped - def test_base64_video_reference_resolves_to_bytes(self, tmp_path): - # The JSON/base64 path carries video even though it has no content-type - # or filename; modality is declared by the field name. - generator = _StubVisualGen() - payload = self._mp4_bytes() - b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest( - prompt="x", video_reference={"content": b64, "format": "base64"} - ) - params = _parse_and_prepare(request, generator) - assert params.image_reference is None - assert params.video_reference[0].content == payload - - def test_video_reference_survives_real_specs(self, tmp_path): - """With the real cosmos3 specs loaded, the encoded payload is persisted - byte-identical — the boundary never transforms video content; the - worker decodes the conditioning window.""" - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS - - generator = _StubVisualGen(extra_param_specs=COSMOS3_EXTRA_SPECS) - payload = self._mp4_bytes() - b64 = base64.b64encode(payload).decode() - request = VideoGenerationRequest( - prompt="x", video_reference={"content": b64, "format": "base64"} - ) - params = _parse_and_prepare(request, generator) - assert params.video_reference[0].content == payload - - def test_multipart_image_reference_resolves_to_bytes(self, tmp_path): - # JPEG upload routed by field name to image_reference. The stored file - # has no type-suffix (PIL identifies by content, not name). - generator = _StubVisualGen() - img = Image.new("RGB", (4, 4), (10, 20, 30)) - buf = BytesIO() - img.save(buf, format="JPEG") - buf.seek(0) - upload = UploadFile(file=buf, filename="ref.jpg") - request = VideoGenerationRequest(prompt="x", image_reference=upload) - params = _parse_and_prepare(request, generator) - assert params.extra_params is None - assert isinstance(params.image_reference[0].content, bytes) - def test_wrong_modality_content_raises(self, tmp_path): # The field name declares modality; mismatched content is a client error. generator = _StubVisualGen() @@ -453,28 +395,6 @@ def test_wrong_modality_content_raises(self, tmp_path): ) assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - def test_undecodable_image_reference_raises_and_cleans_up(self, tmp_path): - generator = _StubVisualGen() - b64 = base64.b64encode(b"neither an image nor a video").decode() - request = VideoGenerationRequest( - prompt="x", image_reference={"content": b64, "format": "base64"} - ) - with pytest.raises(ValueError, match="not a recognized image"): - _parse_and_prepare(request, generator) - # Classification runs on the bytes; rejected content never touches disk. - assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - - def test_malformed_base64_reference_raises_and_cleans_up(self, tmp_path): - generator = _StubVisualGen() - # "ABC" has an invalid base64 length. The declared format is honored, - # so this is a decode error rather than a fallback to a filesystem read. - request = VideoGenerationRequest( - prompt="x", image_reference={"content": "ABC", "format": "base64"} - ) - with pytest.raises(ValueError, match="not valid base64"): - _parse_and_prepare(request, generator) - assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - def test_upload_stream_failure_cleans_up_tmp(self, tmp_path): generator = _StubVisualGen() @@ -490,88 +410,6 @@ def read(self, *args, **kwargs): # … and the payload read fails before any file is written, so nothing leaks. assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - def test_multi_reference_partial_failure_cleans_up(self, tmp_path): - # A later item's rejection removes the files earlier items already wrote, - # so a rejected multi-reference request leaves nothing on disk. - generator = _StubVisualGen() - buf = BytesIO() - Image.new("RGB", (4, 4)).save(buf, format="PNG") - good = base64.b64encode(buf.getvalue()).decode() - bad = base64.b64encode(b"neither an image nor a video").decode() - request = VideoGenerationRequest( - prompt="x", - image_reference=[ - {"content": good, "format": "base64"}, - {"content": bad, "format": "base64"}, - ], - ) - with pytest.raises(ValueError, match="not a recognized image"): - _parse_and_prepare(request, generator) - assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - - def test_file_uri_image_reference_is_read(self, tmp_path): - # format="path" also accepts a file:// URI, normalized before the read. - generator = _StubVisualGen() - src = tmp_path / "ref.png" - Image.new("RGB", (4, 4), (7, 8, 9)).save(src, format="PNG") - request = VideoGenerationRequest( - prompt="x", image_reference={"content": src.as_uri(), "format": "path"} - ) - params = _parse_and_prepare(request, generator) - assert params.image_reference[0].content == src.read_bytes() - assert params.image_reference[0].format == "bytes" - - def test_bare_path_image_reference_is_read(self, tmp_path): - # A path is read at the coordinator, so the worker needs no shared - # filesystem to see what the client named. - generator = _StubVisualGen() - src = tmp_path / "ref.png" - Image.new("RGB", (4, 4), (11, 22, 33)).save(src, format="PNG") - request = VideoGenerationRequest( - prompt="x", image_reference={"content": str(src), "format": "path"} - ) - params = _parse_and_prepare(request, generator) - assert params.image_reference[0].content == src.read_bytes() - assert params.image_reference[0].format == "bytes" - - def test_http_url_image_reference_is_fetched(self, tmp_path, monkeypatch): - # An http(s) reference is fetched through the guarded loader. - generator = _StubVisualGen() - buf = BytesIO() - Image.new("RGB", (4, 4)).save(buf, format="PNG") - png = buf.getvalue() - - class _FakeResp: - def __init__(self, content): - self.content = content - - monkeypatch.setattr( - "tensorrt_llm.visual_gen.media_refs._safe_request_get", - lambda url, **kwargs: _FakeResp(png), - ) - request = VideoGenerationRequest( - prompt="x", - image_reference={"content": "https://example.com/a.png", "format": "url"}, - ) - params = _parse_and_prepare(request, generator) - assert params.image_reference[0].content == png - - def test_http_url_fetch_failure_is_client_error(self, tmp_path, monkeypatch): - # A blocked/failed fetch (e.g. SSRF guard) is a client 400, not a 500, - # and leaves nothing on disk. - generator = _StubVisualGen() - - def _blocked(url, **kwargs): - raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") - - monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) - request = VideoGenerationRequest( - prompt="x", image_reference={"content": "http://10.0.0.1/a.png", "format": "url"} - ) - with pytest.raises(ValueError, match="reference URL could not be fetched"): - _parse_and_prepare(request, generator) - assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - def test_missing_file_uri_is_client_error(self, tmp_path): # A file:// path that does not exist is a client 400, not a server 500. generator = _StubVisualGen() @@ -579,32 +417,8 @@ def test_missing_file_uri_is_client_error(self, tmp_path): request = VideoGenerationRequest( prompt="x", image_reference={"content": missing, "format": "path"} ) - with pytest.raises(ValueError, match="reference file could not be read"): + with pytest.raises(ValueError, match="file could not be read"): _parse_and_prepare(request, generator) - assert list(tmp_path.iterdir()) == [] # nothing is ever written to disk - - def test_bare_reference_string_is_rejected(self): - # The bare-string shorthand is gone: a reference must declare its wire - # form rather than have it guessed from the shape of the value. - from pydantic import ValidationError - - buf = BytesIO() - Image.new("RGB", (4, 4)).save(buf, format="PNG") - b64 = base64.b64encode(buf.getvalue()).decode() - with pytest.raises(ValidationError): - VideoGenerationRequest(prompt="x", image_reference=b64) - with pytest.raises(ValidationError, match="a bare str is no longer accepted"): - VisualGenParams(image_reference=b64) - - def test_json_reference_cannot_declare_bytes(self): - # JSON cannot carry raw bytes; the HTTP schema says so instead of - # letting a str reach the engine claiming to be bytes. - from pydantic import ValidationError - - with pytest.raises(ValidationError, match="multipart/form-data"): - VideoGenerationRequest( - prompt="x", image_reference={"content": "abc", "format": "bytes"} - ) class TestMediaBytesProbes: @@ -882,14 +696,12 @@ def test_malformed_base64_is_a_client_error(self): class TestPrepareReferenceSlots: """The engine choke point: every declared form resolves to raw bytes.""" - def test_resolving_writes_nothing_to_disk(self, tmp_path, monkeypatch): - """References never touch the filesystem, so a worker needs no shared - filesystem to read what the coordinator resolved. The rewritten - ``format`` is what stops a worker being handed a stale spelling.""" + def test_resolving_rewrites_the_format_to_bytes(self): + """Every wire form leaves here as ``bytes``, which is what stops a + worker being handed a spelling it would have to resolve itself.""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + from tensorrt_llm.visual_gen.params import prepare_reference_slots - monkeypatch.setenv("TRTLLM_MEDIA_STORAGE_PATH", str(tmp_path)) buf = BytesIO() Image.new("RGB", (4, 4)).save(buf, format="PNG") params = VisualGenParams( @@ -898,21 +710,20 @@ def test_resolving_writes_nothing_to_disk(self, tmp_path, monkeypatch): ) ) prepare_reference_slots(params) - assert list(tmp_path.iterdir()) == [] + assert params.image_reference[0].format == "bytes" + assert params.image_reference[0].content == buf.getvalue() - def test_wrong_modality_is_rejected(self, tmp_path): - """Content is validated against the slot's modality before dispatch.""" + def test_audio_slot_rejects_non_audio(self): + """The audio slot is checked like the others: a container that is not + audio must not reach a worker that will try to decode it as one.""" from tensorrt_llm.visual_gen import MediaRef, VisualGenParams - from tensorrt_llm.visual_gen.media_refs import prepare_reference_slots + from tensorrt_llm.visual_gen.params import prepare_reference_slots buf = BytesIO() Image.new("RGB", (2, 2)).save(buf, format="PNG") - params = VisualGenParams( - image_reference=MediaRef(content=buf.getvalue(), format="bytes"), - video_reference=MediaRef(content=buf.getvalue(), format="bytes"), - ) - with pytest.raises(ValueError, match="video_reference is not a recognized"): + params = VisualGenParams(audio_reference=MediaRef(content=buf.getvalue(), format="bytes")) + with pytest.raises(ValueError, match="audio_reference is not a recognized"): prepare_reference_slots(params) @@ -926,7 +737,7 @@ def _png() -> bytes: return buf.getvalue() def test_every_format_resolves_to_the_same_bytes(self, tmp_path, monkeypatch): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference + from tensorrt_llm.visual_gen.params import _resolve_reference png = self._png() src = tmp_path / "ref.png" @@ -938,7 +749,7 @@ def __init__(self, content): self.content = content monkeypatch.setattr( - "tensorrt_llm.visual_gen.media_refs._safe_request_get", + "tensorrt_llm.visual_gen.params._safe_request_get", lambda url, **kwargs: _FakeResp(png), ) assert _resolve_reference(str(src), "path") == png @@ -948,14 +759,8 @@ def __init__(self, content): assert _resolve_reference(f"data:image/png;base64,{b64}", "base64") == png assert _resolve_reference(png, "bytes") == png - def test_missing_path_is_a_read_error(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference - - with pytest.raises(ValueError, match="reference file could not be read"): - _resolve_reference(str(tmp_path / "absent.png"), "path") - def test_base64_does_not_fall_back_to_a_disk_read(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference + from tensorrt_llm.visual_gen.params import _resolve_reference src = tmp_path / "ref.png" src.write_bytes(self._png()) @@ -963,38 +768,15 @@ def test_base64_does_not_fall_back_to_a_disk_read(self, tmp_path): _resolve_reference(str(src), "base64") def test_url_fetch_failure_is_a_client_error(self, monkeypatch): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference + from tensorrt_llm.visual_gen.params import _resolve_reference def _blocked(url, **kwargs): raise RuntimeError("URL resolves to a non-public address (10.0.0.1)") - monkeypatch.setattr("tensorrt_llm.visual_gen.media_refs._safe_request_get", _blocked) + monkeypatch.setattr("tensorrt_llm.visual_gen.params._safe_request_get", _blocked) with pytest.raises(ValueError, match="reference URL could not be fetched"): _resolve_reference("http://10.0.0.1/a.png", "url") - def test_non_base64_data_uri_is_rejected(self): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference - - with pytest.raises(ValueError, match="only base64 data: URIs"): - _resolve_reference("data:image/png,%89PNG", "base64") - with pytest.raises(ValueError, match="data: URI is malformed"): - _resolve_reference("data:image/png;base64", "base64") - - def test_content_type_must_match_the_declared_format(self): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference - - with pytest.raises(ValueError, match="requires bytes content"): - _resolve_reference("not bytes", "bytes") - for content_format in ("path", "url", "base64"): - with pytest.raises(ValueError, match="requires string content"): - _resolve_reference(b"raw bytes", content_format) - - def test_unknown_format_is_rejected(self): - from tensorrt_llm.visual_gen.media_refs import _resolve_reference - - with pytest.raises(ValueError, match="unsupported reference format"): - _resolve_reference("a.png", "filepath") - # ============================================================================= # reference transport (coordinator -> rank0) @@ -1018,50 +800,23 @@ def _request(*payloads: bytes): ) return DiffusionRequest(request_id=1, prompt=["x"], params=params) - def test_round_trip_is_byte_identical(self): - payloads = (b"\x89PNG\r\n\x1a\n" + os.urandom(4096), os.urandom(1024)) - req = self._request(*payloads) - - req.refs_to_shm() - req.refs_from_shm() - - assert tuple(r.content for r in req.params.image_reference) == payloads - assert all(r.format == "bytes" for r in req.params.image_reference) - - def test_payload_leaves_the_request_pickle(self): - """The handle is the transport, so the bytes must not also be pickled — - otherwise the hop still pays for a full copy of every reference.""" - payload = os.urandom(256 * 1024) - req = self._request(payload) - before = len(pickle.dumps(req)) - - req.refs_to_shm() - after = len(pickle.dumps(req)) - - assert after < before - len(payload) // 2 - assert payload not in pickle.dumps(req) - def test_survives_a_real_pickle_round_trip(self): """A handle is only useful if it still resolves after being serialized - and rebuilt, which is what the IPC queue does to it.""" - payload = os.urandom(8192) + and rebuilt, which is what the IPC queue does to it. The bytes must not + ride along in that pickle, or the hop still pays for a full copy.""" + payload = os.urandom(256 * 1024) req = self._request(payload) req.refs_to_shm() - received = pickle.loads(pickle.dumps(req)) + wire = pickle.dumps(req) + assert payload not in wire + assert len(wire) < len(payload) // 2 + + received = pickle.loads(wire) received.refs_from_shm() assert received.params.image_reference[0].content == payload - def test_no_references_costs_nothing(self): - """T2V/T2I requests carry no handle, so the hop is untouched for them.""" - from tensorrt_llm._torch.visual_gen import DiffusionRequest - from tensorrt_llm.visual_gen import VisualGenParams - - req = DiffusionRequest(request_id=1, prompt=["x"], params=VisualGenParams()) - req.refs_to_shm() - assert req.ref_handles is None - class TestReferenceBroadcastSplit: """Reference payloads leave the object before the rank0 -> N-rank hop. @@ -1198,57 +953,24 @@ def _png(tmp_path): return target def test_a_regular_file_reads(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + from tensorrt_llm.inputs.media_io import _safe_read_local_file target = self._png(tmp_path) assert _safe_read_local_file(str(target)) == target.read_bytes() assert _safe_read_local_file(target.as_uri()) == target.read_bytes() - def test_a_symlink_to_a_regular_file_reads(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file - - target = self._png(tmp_path) - link = tmp_path / "link.png" - link.symlink_to(target) - - assert _safe_read_local_file(str(link)) == target.read_bytes() - - @pytest.mark.parametrize("kind", ["chardev", "fifo", "directory"]) - def test_a_non_regular_file_is_refused(self, tmp_path, kind): - """Only a regular file has a size the read can trust: a character - device never reaches EOF and a FIFO blocks instead of returning.""" - import os - - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file - - if kind == "chardev": - target = "/dev/zero" - elif kind == "fifo": - target = str(tmp_path / "pipe") - os.mkfifo(target) - else: - target = str(tmp_path) - - with pytest.raises(ValueError, match="not a regular file"): - _safe_read_local_file(target) - - def test_a_symlink_to_a_device_is_refused(self, tmp_path): - """``stat`` follows the link, so the check sees what will be read.""" - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file + def test_a_non_regular_file_is_refused(self, tmp_path): + """Character devices, FIFOs and directories share one check, and + ``stat`` follows the link so it sees what will actually be read.""" + from tensorrt_llm.inputs.media_io import _safe_read_local_file link = tmp_path / "innocent.png" - link.symlink_to("/dev/zero") + link.symlink_to(tmp_path) with pytest.raises(ValueError, match="not a regular file"): _safe_read_local_file(str(link)) - def test_a_missing_file_is_a_client_error(self, tmp_path): - from tensorrt_llm.visual_gen.media_refs import _safe_read_local_file - - with pytest.raises(ValueError, match="could not be read"): - _safe_read_local_file(str(tmp_path / "nope.png")) - class TestLocalMediaPathCanBeDisallowed: """``format='path'`` reads server-side files, so a deployment can refuse it. @@ -1275,16 +997,6 @@ def test_path_is_refused_when_disallowed(self, monkeypatch): with pytest.raises(ValueError, match="is disallowed on this server"): parse_visual_gen_params(self._request(), _StubVisualGen()) - def test_disallowing_path_leaves_the_other_formats_alone(self, monkeypatch): - """The gate is about reading server-side files, not about references.""" - monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") - - params = parse_visual_gen_params( - self._request(fmt="base64", content="aGk="), _StubVisualGen() - ) - - assert params.image_reference[0].format == "base64" - def test_an_unrecognized_value_warns_and_stays_allowed(self, monkeypatch): """Silently reading a typo as "1" would break working deployments, and silently reading it as "0" would leave one that believes it is locked diff --git a/tests/unittest/llmapi/apps/test_media_io.py b/tests/unittest/llmapi/apps/test_media_io.py index f34cdc475a08..d139ea2cd4a6 100644 --- a/tests/unittest/llmapi/apps/test_media_io.py +++ b/tests/unittest/llmapi/apps/test_media_io.py @@ -5,13 +5,7 @@ import pytest from tensorrt_llm.inputs import MultimodalDataTracker -from tensorrt_llm.inputs.media_io import ( - AudioMediaIO, - BaseMediaIO, - ImageMediaIO, - VideoMediaIO, - convert_image_mode, -) +from tensorrt_llm.inputs.media_io import AudioMediaIO, BaseMediaIO, ImageMediaIO, VideoMediaIO from tensorrt_llm.serve.chat_utils import parse_chat_message_content_part pytestmark = pytest.mark.cpu_only @@ -97,83 +91,3 @@ def test_non_video_classes_use_plain_shallow_merge(self, media_io_cls): {"num_frames": 32}, ) assert merged == {"num_frames": 32, "fps": 1} - - -class TestImageAlphaHandling: - """RGBA -> RGB has two defensible semantics. - - The caller picks, and the default must stay what every existing caller - already gets. - """ - - @staticmethod - def _rgba_png() -> bytes: - """Build a 3-pixel RGBA fixture. - - One opaque, one half-transparent and one fully transparent pixel, all - sharing the same stored RGB so the two semantics are separable. - """ - from io import BytesIO - - from PIL import Image - - im = Image.new("RGBA", (3, 1)) - im.putpixel((0, 0), (200, 30, 30, 255)) - im.putpixel((1, 0), (200, 30, 30, 128)) - im.putpixel((2, 0), (200, 30, 30, 0)) - buf = BytesIO() - im.save(buf, format="PNG") - return buf.getvalue() - - def test_drop_alpha_matches_pil_and_diffusers(self): - """Match diffusers. - - Its load_image defaults to image.convert("RGB"), so pipelines ported - from diffusers need that exact behavior to stay aligned. - """ - from io import BytesIO - - from PIL import Image - - png = self._rgba_png() - reference = list(Image.open(BytesIO(png)).convert("RGB").getdata()) - got = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(png) - assert list(got.getdata()) == reference - - def test_default_composites_and_is_unchanged(self): - """Keep compositing onto white by default. - - No existing LLM or VLM caller may shift. - """ - got = ImageMediaIO(format="pil").load_bytes(self._rgba_png()) - assert list(got.getdata()) == [(200, 30, 30), (227, 142, 142), (255, 255, 255)] - - def test_semantics_coincide_for_opaque_images(self): - """Coincide on opaque media. - - Every committed golden uses opaque media, so the two semantics must be - bit-identical there. - """ - from io import BytesIO - - from PIL import Image - - buf = BytesIO() - Image.new("RGBA", (2, 1), (10, 20, 30, 255)).save(buf, format="PNG") - png = buf.getvalue() - composited = ImageMediaIO(format="pil").load_bytes(png) - dropped = ImageMediaIO(format="pil", drop_alpha=True).load_bytes(png) - assert list(composited.getdata()) == list(dropped.getdata()) - - def test_mode_rgba_preserves_alpha(self): - got = ImageMediaIO(format="pil", mode="RGBA").load_bytes(self._rgba_png()) - assert got.mode == "RGBA" - assert list(got.getdata())[2] == (200, 30, 30, 0) - - def test_convert_image_mode_default_is_unchanged(self): - """The shared helper is publicly exported; its default must not move.""" - from PIL import Image - - im = Image.new("RGBA", (1, 1), (200, 30, 30, 0)) - assert convert_image_mode(im, "RGB").getpixel((0, 0)) == (255, 255, 255) - assert convert_image_mode(im, "RGB", drop_alpha=True).getpixel((0, 0)) == (200, 30, 30) From a6143b4fe4274785944726850727b8ffa806fad6 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Wed, 26 Aug 2026 18:29:00 -0700 Subject: [PATCH 56/61] [TRTLLM-15277][fix] Drive the Cosmos3 transfer tests through video_reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The branch moved Cosmos3's V2V conditioning off `extra_params["video"]` and onto the typed `video_reference` slot, but this module still built requests the old way. `infer()` therefore saw no reference, skipped the source-header probe, and fell back to the 720p landscape defaults — six source-derived-default cases failed on size and frame rate. The request helper now routes `video=` into `params.video_reference` as a `format="bytes"` reference, which is the only spelling that reaches a worker; every call site is unchanged. Found by CI, not locally: the module is not in this branch's diff, so it never entered the hand-picked set of test files being run. The lesson is in the `write-test` skill — a full run of the directory that owns the changed code is the gate before pushing, and a targeted run cannot stand in for it. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py index 83a11058519e..7e4ded7fc773 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -53,7 +53,7 @@ from tensorrt_llm._torch.visual_gen.offloading import PipelineOffloader from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer from tensorrt_llm.media.decoding import VideoStreamInfo -from tensorrt_llm.visual_gen.params import VisualGenParams +from tensorrt_llm.visual_gen.params import MediaRef, VisualGenParams pytestmark = pytest.mark.cosmos3 @@ -428,7 +428,7 @@ class TestSourceDerivedDefaults: REFERENCE = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" - def _infer_req(self, _params=None, **extra): + def _infer_req(self, _params=None, *, video=None, **extra): # Executor-merged shape: num_frames/frame_rate carry pipeline defaults, # height/width are declared None, and nothing reads as caller intent. params = ( @@ -442,6 +442,10 @@ def _infer_req(self, _params=None, **extra): ) ) params.extra_params = dict(extra) + # The V2V reference reaches a worker resolved, so ``format="bytes"`` is + # the only spelling ``infer()`` can see. + if video is not None: + params.video_reference = [MediaRef(content=video, format="bytes")] return SimpleNamespace(params=params, prompt="a prompt") def _captured(self, req): From fe1c9e4241afb149038746d15ea76d6da66fb238 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 27 Aug 2026 08:50:25 +0000 Subject: [PATCH 57/61] [TRTLLM-15277][fix] Stop advertising a wire form JSON cannot carry MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review feedback: `MediaReferenceItem.format` and `input_reference_format` were annotated with the engine's four-value format type, so the generated OpenAPI schema offered `bytes` as a legal value while the request handlers rejected it. A generated client could construct a body that always came back 422. Both fields now declare `Literal["path", "url", "base64"]`, which is what a JSON body can actually express — raw bytes arrive as a multipart upload, and that transport carries no `format` of its own. The two validators that existed only to intercept `bytes` are gone with it: the field type refuses the value before either could run. `MediaContentFormat` had one use left after that and is inlined into `MediaRef`, matching how every other `Literal` in these two files is written. `MediaRole` stays a named alias — it is shared by `RefSlotSpec` in the pipeline layer. Not verified locally: the dev container on this node carries a different torch than the prebuilt libs, so `import tensorrt_llm` fails and the suite cannot run until the rebuild finishes. Static checks only. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- tensorrt_llm/serve/openai_protocol.py | 37 ++++++------------- tensorrt_llm/visual_gen/params.py | 11 ++---- .../visual_gen/test_visual_gen_params.py | 4 +- .../references/trtllm_serve_api.yaml | 4 +- 4 files changed, 21 insertions(+), 35 deletions(-) diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index acedd2bf64e4..0ef3530929f9 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -62,7 +62,7 @@ from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy -from tensorrt_llm.visual_gen.params import MediaContentFormat, MediaRole +from tensorrt_llm.visual_gen.params import MediaRole _LOGIT_BIAS_MIN = -100.0 _LOGIT_BIAS_MAX = 100.0 @@ -2027,12 +2027,12 @@ class MediaReferenceItem(OpenAIBaseModel): content: str = Field( description="The reference payload, in the form declared by ``format``." ) - format: MediaContentFormat = Field(description=( + format: Literal["path", "url", "base64"] = Field(description=( "Wire form of ``content``: ``path`` (a file readable by the server; a " "``file://`` URI is also accepted), ``url`` (``http(s)``, fetched " "through the SSRF-guarded loader), or ``base64`` (a ``data:`` URI is " - "also accepted). ``bytes`` cannot be carried in JSON — upload the file " - "as multipart/form-data instead. Distinct from the top-level " + "also accepted). Raw bytes reach the server as a multipart upload, " + "which carries no ``format`` of its own. Distinct from the top-level " "``format``, which selects the *output* encoding.")) role: Optional[MediaRole] = Field( default=None, @@ -2041,15 +2041,6 @@ class MediaReferenceItem(OpenAIBaseModel): "it when the model leaves no ambiguity.", ) - @field_validator("format") - @classmethod - def _reject_bytes_over_json(cls, v: str) -> str: - if v == "bytes": - raise ValueError( - "format='bytes' cannot be carried in JSON; upload the file as " - "multipart/form-data, or send format='base64'.") - return v - class VideoGenerationRequest(OpenAIBaseModel): """Video generation request (extended API). @@ -2114,7 +2105,7 @@ class VideoGenerationRequest(OpenAIBaseModel): "ignored whenever a typed ``image_reference`` or ``video_reference`` " "is provided. A string form requires ``input_reference_format``."), ) - input_reference_format: Optional[MediaContentFormat] = Field( + input_reference_format: Optional[Literal["path", "url", "base64"]] = Field( default=None, description=( "Deprecated, alongside ``input_reference``: the wire form of that " @@ -2203,17 +2194,13 @@ def _check_input_reference_format(self): field here would break exactly what the field is for. A multipart upload carries its own form and needs no sibling either way. """ - if isinstance(self.input_reference, str): - if self.input_reference_format is None: - logger.warning( - "'input_reference' without 'input_reference_format' is read as " - "base64; both are deprecated, use 'image_reference' / " - "'video_reference' with an explicit format.") - self.input_reference_format = "base64" - elif self.input_reference_format == "bytes": - raise ValueError( - "input_reference_format='bytes' cannot be carried in JSON; upload " - "the file as multipart/form-data, or send 'base64'") + if isinstance(self.input_reference, + str) and self.input_reference_format is None: + logger.warning( + "'input_reference' without 'input_reference_format' is read as " + "base64; both are deprecated, use 'image_reference' / " + "'video_reference' with an explicit format.") + self.input_reference_format = "base64" return self diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 2bf0a05d63b3..fb13809c8891 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -29,12 +29,6 @@ MediaRole = Literal["reference", "first_frame", "last_frame"] -# Wire form of a reference's ``content``. Declared explicitly rather than -# sniffed: a bare string is otherwise ambiguous between a local path and -# base64, and guessing lets a mistyped path silently become base64 (or a -# malformed base64 silently become a filesystem read). -MediaContentFormat = Literal["path", "url", "base64", "bytes"] - @set_api_status("prototype") class MediaRef(StrictBaseModel): @@ -50,7 +44,10 @@ class MediaRef(StrictBaseModel): content: Union[str, bytes] = Field( description="The reference payload, in the form declared by ``format``." ) - format: MediaContentFormat = Field( + # Declared rather than sniffed: a bare string is otherwise ambiguous between + # a local path and base64, and guessing lets a mistyped path silently become + # base64 (or a malformed base64 silently become a filesystem read). + format: Literal["path", "url", "base64", "bytes"] = Field( description=( "Wire form of ``content``: ``path`` (local file; a ``file://`` URI is " "also accepted), ``url`` (``http(s)``, fetched through the SSRF-guarded " 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 0c655a6d9f7f..e6ff0b446275 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1591,11 +1591,13 @@ def test_an_explicit_format_still_wins(self): assert request.input_reference_format == "path" def test_bytes_over_json_is_still_rejected(self): + """The wire form a JSON body can declare stops at ``base64``, so the + generated schema never advertises a value that is certain to 422.""" from pydantic import ValidationError from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest - with pytest.raises(ValidationError, match="cannot be carried in JSON"): + with pytest.raises(ValidationError, match="'path', 'url' or 'base64'"): VideoGenerationRequest( prompt="x", input_reference="aGk=", input_reference_format="bytes" ) diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index abd2f68b4ee2..782764059fdd 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1497,7 +1497,7 @@ models: required: false input_reference_format: kind: extension - type: Optional[MediaContentFormat] + type: Optional[Literal['path', 'url', 'base64']] default: null status: deprecated required: false @@ -1583,7 +1583,7 @@ models: required: true format: kind: extension - type: MediaContentFormat + type: Literal['path', 'url', 'base64'] default: null status: prototype required: true From 49b1c9512eef5111d412587f19f78575677766da Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:12:53 +0000 Subject: [PATCH 58/61] [TRTLLM-15277][fix] Stop growing the deprecated input_reference MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `input_reference_format` was a field this branch added, and `input_reference` is deprecated — the sibling widened a field on its way out to accept `path` and `url`, neither of which it has ever taken on main: input_reference: Optional[Union[str, UploadFile]] # "JSON requests carry base64 bytes; multipart requests upload the file." That is backwards. The field exists so callers written against the old API keep running, and such a caller cannot be using a field that does not exist upstream; offering it new wire forms only invites new code to reach for a deprecated one. The sibling is gone and the deprecated path is back to the two forms it has always had — a base64 string, or a multipart upload — so the transport decides the wire form and there is nothing left to declare. `_check_input_reference_format` went with it: it existed to fill in a field that no longer exists. This also settles the second half of the schema review: a field that is not there cannot advertise `bytes`. Tests: three cases covered the removed field and are gone; the deprecated field's own contract (a bare string is still accepted, a typed one still needs a format) is kept. Dropping `test_the_deprecated_field_is_gated_too` costs nothing real — the deprecated field has no `path` branch left to gate, and the typed fields keep their own gate test. Static checks only; the container on this node is mid-rebuild. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- examples/visual_gen/serve/README.md | 2 +- tensorrt_llm/serve/openai_protocol.py | 32 ++--------------- tensorrt_llm/serve/visual_gen_utils.py | 24 ++++--------- .../visual_gen/test_visual_gen_params.py | 36 +++---------------- .../visual_gen/test_visual_gen_utils.py | 18 ++-------- .../references/trtllm_serve_api.yaml | 6 ---- 6 files changed, 18 insertions(+), 100 deletions(-) diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index aad4bfd5ce22..bedec99e6df1 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -305,7 +305,7 @@ You can customize these by: ``` - **Supported formats**: PNG and JPEG images; MP4 and AVI video, with H.264 the tested codec and others best-effort. HEIF/AVIF are not supported. -- `input_reference` (deprecated): a single image or video reference, routed by content signature to I2V or V2V. Declares its wire form through the sibling `input_reference_format` field, and is ignored when `image_reference` / `video_reference` is also given. Prefer the typed fields. +- `input_reference` (deprecated): a single image or video reference, routed by content signature to I2V or V2V. A JSON request carries base64 bytes and a multipart request uploads the file; it is ignored when `image_reference` / `video_reference` is also given. Prefer the typed fields. - `extra_params`: model-specific overflow (see below) - `response_format`: `"file"` (default; `FileResponse` byte download) or `"path"` (server-side output path JSON, for co-located clients) - `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2). diff --git a/tensorrt_llm/serve/openai_protocol.py b/tensorrt_llm/serve/openai_protocol.py index 0ef3530929f9..f4d329e24b23 100644 --- a/tensorrt_llm/serve/openai_protocol.py +++ b/tensorrt_llm/serve/openai_protocol.py @@ -58,7 +58,6 @@ from tensorrt_llm.llmapi import (DisaggScheduleStyle, GuidedDecodingParams, SamplingParams) from tensorrt_llm.llmapi.reasoning_parser import ReasoningParserFactory -from tensorrt_llm.logger import logger from tensorrt_llm.sampling_params import (check_logprobs_limit, validate_thinking_token_budget) from tensorrt_llm.scheduling_params import AgentHierarchy @@ -2099,18 +2098,12 @@ class VideoGenerationRequest(OpenAIBaseModel): default=None, description= ("Deprecated. A single image or video reference, routed by content " - "signature to image-to-video or video-to-video. Kept for backward " + "signature to image-to-video or video-to-video. A JSON request carries " + "base64 bytes; a multipart request uploads the file. Kept for backward " "compatibility; prefer the typed ``image_reference`` / " "``video_reference`` fields, which take precedence — this field is " "ignored whenever a typed ``image_reference`` or ``video_reference`` " - "is provided. A string form requires ``input_reference_format``."), - ) - input_reference_format: Optional[Literal["path", "url", "base64"]] = Field( - default=None, - description=( - "Deprecated, alongside ``input_reference``: the wire form of that " - "field's value (``path`` / ``url`` / ``base64``). Required when " - "``input_reference`` is a string; implied for a multipart upload."), + "is provided."), ) # Resolution @@ -2184,25 +2177,6 @@ def _reject_removed_response_format(cls, value): raise ValueError(removed[value]) return value - @model_validator(mode="after") - def _check_input_reference_format(self): - """Fill in the deprecated ``input_reference``'s wire form when omitted. - - The typed fields require an explicit ``format`` — that is the point of - them. This one exists only so callers written against the old API keep - working, and those callers sent bare base64, so demanding a new sibling - field here would break exactly what the field is for. A multipart - upload carries its own form and needs no sibling either way. - """ - if isinstance(self.input_reference, - str) and self.input_reference_format is None: - logger.warning( - "'input_reference' without 'input_reference_format' is read as " - "base64; both are deprecated, use 'image_reference' / " - "'video_reference' with an explicit format.") - self.input_reference_format = "base64" - return self - class VideoJob(OpenAIBaseModel): """Metadata for an asynchronous video generation job. diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index e7fb9f81c285..0702470f6d4f 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -364,16 +364,15 @@ def _validate_image_edit_request_limits( def _apply_deprecated_input_reference( input_reference: str | UploadFile | None, params: VisualGenParams, - input_reference_format: Optional[str] = None, ) -> None: """Back-compat for the deprecated single ``input_reference``. Sniff-routes the payload to ``image_reference`` (image) or ``video_reference`` (video), preserving the pre-typed-fields behavior. Ignored when a typed image/video reference is already set — the typed fields take precedence. - Routing needs the bytes, so the payload is resolved here using the wire form - from the sibling ``input_reference_format`` (implied for an upload); the - resolved bytes are then handed to the engine like any other reference. + The field has only ever carried base64 in JSON or a multipart upload, so the + wire form follows from the transport; routing needs the bytes, so they are + read here and handed to the engine like any other reference. """ if input_reference is None: return @@ -382,20 +381,13 @@ def _apply_deprecated_input_reference( return from tensorrt_llm.visual_gen.params import MediaRef - if hasattr(input_reference, "file"): # multipart upload — form implied + if hasattr(input_reference, "file"): # multipart upload payload = input_reference.file.read() else: - if input_reference_format == "path" and local_media_path_is_disallowed(): - raise ValueError( - "reference format='path' is disallowed on this server " - "(TRTLLM_DISALLOW_LOCAL_MEDIA_PATH=1); it reads server-side " - "files and is only meaningful for co-located clients. Send the " - "file as base64 or upload it via multipart/form-data." - ) # Local import, for the reason given at the first one. - from tensorrt_llm.visual_gen.params import _resolve_reference + from tensorrt_llm.visual_gen.params import _read_reference_payload - payload = _resolve_reference(input_reference, input_reference_format) + payload = _read_reference_payload(input_reference) kind = sniff_media_kind(payload) if kind == "image": params.image_reference = [MediaRef(content=payload, format="bytes")] @@ -499,9 +491,7 @@ def parse_visual_gen_params( audio_refs = _build_reference_list(request.audio_reference) if audio_refs: params.audio_reference = audio_refs - _apply_deprecated_input_reference( - request.input_reference, params, request.input_reference_format - ) + _apply_deprecated_input_reference(request.input_reference, params) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) _decode_inline_media(request.extra_params, generator.extra_param_specs) 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 e6ff0b446275..bdc90d45c7f3 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1566,44 +1566,18 @@ def request_warmup_cache_key(req): class TestDeprecatedInputReferenceStaysCompatible: - """The deprecated field exists so old callers keep working. + """The deprecated field takes a bare string, the typed ones never do. - Callers written against the old API sent bare base64 with no sibling - format, so requiring one here would break exactly what the field is for. - The typed fields still demand an explicit format — that distinction is the - point. + Old callers sent base64 with nothing declaring it, and that has to keep + working; a typed reference has somewhere to say what it is, so it must. """ - def test_bare_base64_is_read_as_base64(self): + def test_a_bare_string_is_still_accepted(self): from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest - request = VideoGenerationRequest(prompt="x", input_reference="aGk=") - - assert request.input_reference_format == "base64" - - def test_an_explicit_format_still_wins(self): - from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest - - request = VideoGenerationRequest( - prompt="x", input_reference="/tmp/ref.png", input_reference_format="path" - ) - - assert request.input_reference_format == "path" - - def test_bytes_over_json_is_still_rejected(self): - """The wire form a JSON body can declare stops at ``base64``, so the - generated schema never advertises a value that is certain to 422.""" - from pydantic import ValidationError - - from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest - - with pytest.raises(ValidationError, match="'path', 'url' or 'base64'"): - VideoGenerationRequest( - prompt="x", input_reference="aGk=", input_reference_format="bytes" - ) + assert VideoGenerationRequest(prompt="x", input_reference="aGk=").input_reference == "aGk=" def test_the_typed_field_still_requires_a_format(self): - """Relaxing the deprecated alias must not relax the new API.""" from pydantic import ValidationError from tensorrt_llm.serve.openai_protocol import VideoGenerationRequest diff --git a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py index 473f41379ecc..b8d9aa865d60 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -338,17 +338,13 @@ def test_deprecated_input_reference_routes_by_sniff(self, tmp_path): vid_b64 = base64.b64encode(self._mp4_bytes()).decode() p = _parse_and_prepare( - VideoGenerationRequest( - prompt="x", input_reference=img_b64, input_reference_format="base64" - ), + VideoGenerationRequest(prompt="x", input_reference=img_b64), generator, ) assert len(p.image_reference) == 1 and p.video_reference is None p = _parse_and_prepare( - VideoGenerationRequest( - prompt="x", input_reference=vid_b64, input_reference_format="base64" - ), + VideoGenerationRequest(prompt="x", input_reference=vid_b64), generator, ) assert len(p.video_reference) == 1 and p.image_reference is None @@ -365,7 +361,6 @@ def test_input_reference_ignored_when_typed_reference_set(self, tmp_path): prompt="x", image_reference={"content": img_b64, "format": "base64"}, input_reference=vid_b64, - input_reference_format="base64", ), generator, ) @@ -1009,12 +1004,3 @@ def test_an_unrecognized_value_warns_and_stays_allowed(self, monkeypatch): assert params.image_reference[0].format == "path" assert any("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH" in w for w in warnings) - - def test_the_deprecated_field_is_gated_too(self, monkeypatch): - monkeypatch.setenv("TRTLLM_DISALLOW_LOCAL_MEDIA_PATH", "1") - request = VideoGenerationRequest( - prompt="x", input_reference="/tmp/ref.png", input_reference_format="path" - ) - - with pytest.raises(ValueError, match="is disallowed on this server"): - parse_visual_gen_params(request, _StubVisualGen()) diff --git a/tests/unittest/api_stability/references/trtllm_serve_api.yaml b/tests/unittest/api_stability/references/trtllm_serve_api.yaml index 782764059fdd..9e25b6abc66c 100644 --- a/tests/unittest/api_stability/references/trtllm_serve_api.yaml +++ b/tests/unittest/api_stability/references/trtllm_serve_api.yaml @@ -1495,12 +1495,6 @@ models: default: null status: deprecated required: false - input_reference_format: - kind: extension - type: Optional[Literal['path', 'url', 'base64']] - default: null - status: deprecated - required: false size: kind: extension type: Optional[str] From 602a148d2113d321fba7a3d1be3dc31ce5fb64aa Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Thu, 27 Aug 2026 09:13:23 +0000 Subject: [PATCH 59/61] [TRTLLM-15277][test] Cover the Wan I2V reference-role dispatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `WanImageToVideoPipeline` declares two roles for one slot, so `infer` picks the first frame and the last out of the reference list — the only pipeline on this branch that has to. Nothing exercised it: every Wan test drives `forward()` directly and so never crosses the reference layer at all. Choosing wrong here conditions the model on the wrong frame and still returns a video, so no caller would notice. The list in the first case is deliberately ordered last-frame-first, which is what an implementation that indexed by position instead of role would get away with. Placed beside `test_flux2_image_conditioning.py`, the same shape of test: `__new__` plus a mocked `forward`, no weights, no GPU. Six other migrated pipelines still have no test crossing that layer, but each takes a single reference and reads `refs[0]`, with none of the role dispatch that makes this one worth guarding. Static checks only; the container on this node is mid-rebuild, so this case has not been mutation-verified yet. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../test_wan_i2v_reference_conditioning.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 tests/unittest/_torch/visual_gen/test_wan_i2v_reference_conditioning.py diff --git a/tests/unittest/_torch/visual_gen/test_wan_i2v_reference_conditioning.py b/tests/unittest/_torch/visual_gen/test_wan_i2v_reference_conditioning.py new file mode 100644 index 000000000000..7b960b4a8a21 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_wan_i2v_reference_conditioning.py @@ -0,0 +1,54 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Unit tests for Wan I2V reference-role dispatch. + +The slot declares two roles, so ``infer`` has to decide which reference is the +first frame and which is the last. Getting that backwards conditions the model +on the wrong frame and still returns a video, so nothing downstream would say +anything was wrong. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +from tensorrt_llm._torch.visual_gen.models.wan.pipeline_wan_i2v import WanImageToVideoPipeline +from tensorrt_llm.visual_gen.params import MediaRef, VisualGenParams + + +def _request(*refs) -> SimpleNamespace: + """A request as the coordinator hands it over: references already bytes.""" + return SimpleNamespace( + prompt=["a prompt"], + params=VisualGenParams(image_reference=list(refs), seed=0), + ) + + +def _forward_kwargs(req) -> dict: + pipeline = WanImageToVideoPipeline.__new__(WanImageToVideoPipeline) + pipeline.forward = MagicMock(return_value=object()) + + pipeline.infer(req) + + return pipeline.forward.call_args.kwargs + + +def test_roles_select_the_frames() -> None: + kwargs = _forward_kwargs( + _request( + MediaRef(content=b"last", format="bytes", role="last_frame"), + MediaRef(content=b"first", format="bytes", role="first_frame"), + ) + ) + + # Declaration order must not decide: the roles do. + assert kwargs["image"] == b"first" + assert kwargs["last_image"] == b"last" + + +def test_an_omitted_role_is_the_first_frame() -> None: + """Plain I2V leaves ``role`` off, and the slot's other role is optional.""" + kwargs = _forward_kwargs(_request(MediaRef(content=b"only", format="bytes"))) + + assert kwargs["image"] == b"only" + assert kwargs["last_image"] is None From e8da4ed31e606b21a063bad187faf11ea45f4964 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Fri, 28 Aug 2026 08:27:52 +0000 Subject: [PATCH 60/61] [TRTLLM-15277][test] Drop the extra-param check this branch replaced MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `test_video_reference_must_be_bytes` arrived with Cosmos3 action generation (2938adaf3c) and asserts that `extra_params["video"]` rejects a server-local path, relying on the `ExtraParamSchema(type="bytes")` declaration to do it. This branch moves Cosmos3's V2V conditioning to the typed `video_reference` slot, so that declaration is gone and an undeclared key passes through unchecked — the case fails after the rebase for that reason, not for a merge error. The contract it guarded is now stronger and lives elsewhere: `MediaRef.format` makes the wire form explicit instead of inferring it from the payload type, the coordinator resolves every form to bytes before a worker sees it, and a server-local path additionally answers to `TRTLLM_DISALLOW_LOCAL_MEDIA_PATH`. `test_path_is_refused_when_disallowed` and `test_resolving_rewrites_the_format_to_bytes` cover those. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/test_visual_gen_params.py | 15 --------------- 1 file changed, 15 deletions(-) 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 bdc90d45c7f3..c48d0fa963d1 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_params.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_params.py @@ -1295,21 +1295,6 @@ def test_literal_extra_param_accepts_numeric_choice(self): extra_param_specs=COSMOS3_EXTRA_SPECS, ) - def test_video_reference_must_be_bytes(self): - """A server-local path must not reach the worker: the ``video`` contract - is encoded bytes, so a string (or anything else) fails preflight.""" - from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_EXTRA_SPECS - - req = self._make_request(extra_params={"video": "/server/local/path.mp4"}) - with pytest.raises(ValueError, match="expected type 'bytes'"): - from tensorrt_llm.visual_gen.params import validate_visual_gen_params - - validate_visual_gen_params( - req.params, - declared_defaults=None, - extra_param_specs=COSMOS3_EXTRA_SPECS, - ) - # ============================================================================= # Parameter validation — message content per category From cd27779678bdd40bb27d540ce63e9861f91f2ed3 Mon Sep 17 00:00:00 2001 From: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> Date: Mon, 31 Aug 2026 16:12:08 +0000 Subject: [PATCH 61/61] [TRTLLM-15277][chore] Drop a comment the rebase left over a wrong statement Resolving the Cosmos3 conflict against upstream's action-generation branch left the coordinator-choke-point note duplicated: the second copy landed above `is_action = extra_params.get("action_mode")`, which has nothing to do with the reference bytes, so it described the wrong statement. Reported by BowenFu in review. Signed-off-by: Yiyun Lu <55233584+luyiyun1021@users.noreply.github.com> --- .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index f316e01aa7f2..69a4434710a8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -767,8 +767,6 @@ def as_given(field_name): # Container and modality are already checked at the coordinator's # reference choke point, so the bytes reaching here are known video. video = refs_v[0].content if refs_v else None - # Container and modality are already checked at the coordinator's - # reference choke point, so the bytes reaching here are known video. is_action = extra_params.get("action_mode") is not None if is_action: # Action resolves its whole recipe in forward() -- the canvas from