[TRTLLM-12491][feat] Align VisualGen serve request schema with VisualGenParams

examples/visual_gen/serve/README.md

@@ -67,7 +67,7 @@ Demonstrates synchronous text-to-image generation using the OpenAI SDK. Supports
 
 **Features:**
 - Generates images from text prompts
-- Supports configurable model, image size, and quality
+- Supports configurable model and image size
 - Returns base64-encoded images or URLs
 - Saves generated images to disk
 
@@ -269,20 +269,52 @@ You can customize these by:
 ## Common Parameters
 
 ### Image Generation
-- `model`: Model identifier (e.g., "flux1", "flux2")
-- `prompt`: Text description
+- `prompt`: Text description (required)
 - `n`: Number of images to generate
-- `size`: Image dimensions (e.g., "512x512", "1024x1024")
-- `quality`: "standard" or "hd"
-- `response_format`: "b64_json" or "url"
+- `size`: Image dimensions in `WxH` format (e.g., `"512x512"`, `"1024x1024"`) — or use the structured pair `width` + `height` (both required when sent)
+- `seed`: Random seed; `null` / omitted means the engine draws a fresh seed
+- `num_inference_steps`, `guidance_scale`, `max_sequence_length`, `negative_prompt`: per-request denoise controls (override pipeline defaults when sent)
+- `extra_params`: model-specific overflow as a JSON object (see "Model-Specific `extra_params`" below). Unknown keys are rejected by the executor.
+- `response_format`: `"b64_json"` or `"url"`
+- `format`: Generation content encoding. Image encoders: `"png"`, `"webp"`, `"jpeg"`. Tensor formats: `"safetensors"`, `"pt"`.
+- Accept-and-warn OpenAI-shape fields (no engine semantic): `model`, `quality`, `style`, `user`. Sending `quality`/`style` logs a server-side WARNING; sending `model` warns on mismatch. None of these change generation behavior.
 
 ### Video Generation
-- `model`: Model identifier (e.g., "wan", "ltx2")
-- `prompt`: Text description
-- `size`: Video resolution (e.g., "256x256", "512x512", "1280x720")
-- `seconds`: Duration in seconds
-- `fps`: Frames per second
-- `input_reference`: Reference image file (for TI2V mode)
+- `prompt`: Text description (required)
+- `size` / `width` / `height`: same convention as image
+- `seconds`: Duration in seconds (engine multiplies by `frame_rate` to derive `num_frames` when the latter is absent)
+- `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 (TI2V mode); accepted as base64-encoded string in JSON or as a file in multipart form-data
+- `extra_params`: model-specific overflow (see below)
+- `response_format`: `"b64_json"` or `"url"`
+- `format`: Generation content encoding. Video encoders: `"mp4"`, `"avi"`, `"auto"`. Tensor formats: `"safetensors"`, `"pt"` (carries video + audio + scalar metadata in one payload for LTX-2).
+
+#### Tensor-format consumer contract
+
+When `format="safetensors"` or `format="pt"`, the payload bundles every populated media tensor (`image` / `video` / `audio`) and the scalar metadata (`frame_rate`, `audio_sample_rate`) into one file.
+
+- **`pt`**: `torch.load(buf, weights_only=True)` returns a dict with the tensor keys and the scalars as native Python values.
+- **`safetensors`**: `safetensors.torch.load(bytes)` returns a dict with the tensor keys and each scalar as a 0-d tensor under the same key — call `.item()` to unbox (e.g. `loaded["frame_rate"].item()`). The same scalars are also written to the safetensors file header as strings; `safe_open(path, framework="pt").metadata()` exposes them in that form for consumers that prefer header access.
+
+#### Unknown-field policy
+
+The visual-gen endpoints reject unknown top-level fields with HTTP 422 (`extra="forbid"`). Anything model-specific belongs inside `extra_params`. Sending `output_format`, top-level `guidance_rescale`, or — for video — top-level `n` returns 422 with the offending field named in the error body.
+
+#### Model-specific `extra_params`
+
+Use the Python API to discover accepted keys for a loaded pipeline:
+
+```python
+generator = VisualGen(model="...")
+print(generator.extra_param_specs)   # {key: ExtraParamSchema(type=..., range=..., default=..., description=...)}
+```
+
+Examples:
+- **LTX-2**: `stg_scale`, `stg_blocks`, `modality_scale`, `guidance_rescale`, `output_type`, ...
+- **Wan 2.2 A14B**: `guidance_scale_2`, `boundary_ratio`
+- **Wan 2.1 / Flux**: no model-specific `extra_params` declared
 
 > **Note:** LTX-2 generates video **with audio**. The `ltx2.yml` config must include
 > `text_encoder_path` pointing to a Gemma3 model (e.g., `google/gemma-3-12b-it`).

examples/visual_gen/serve/async_video_gen.py

@@ -33,7 +33,7 @@ def test_async_video_generation(
     fps: int = 24,
     size: str = "256x256",
     output_file: str = "output_async.mp4",
-    output_format: str = "auto",
+    format: str = "auto",
 ):
     """Test asynchronous video generation with OpenAI SDK.
 
@@ -77,7 +77,7 @@ def test_async_video_generation(
             "seconds": duration,
             "extra_body": {
                 "fps": fps,
-                "output_format": output_format,
+                "format": format,
             },
         }
 
@@ -131,12 +131,18 @@ def test_async_video_generation(
         # For binary content, use the underlying HTTP client
         content = client.videos.download_content(video_id, variant="video")
 
-        # Check content type to determine actual file extension
-        content_type = getattr(content.response, "headers", {}).get("content-type", "video/mp4")
-        if "x-msvideo" in content_type or "avi" in content_type:
-            actual_ext = ".avi"
+        # Determine the on-disk extension. Tensor formats are
+        # selected by the request and the server returns
+        # ``application/octet-stream``; encoder formats can be
+        # disambiguated from Content-Type (mp4 vs avi).
+        if format in ("safetensors", "pt"):
+            actual_ext = f".{format}"
         else:
-            actual_ext = ".mp4"
+            content_type = getattr(content.response, "headers", {}).get("content-type", "video/mp4")
+            if "x-msvideo" in content_type or "avi" in content_type:
+                actual_ext = ".avi"
+            else:
+                actual_ext = ".mp4"
 
         # Adjust output filename if extension doesn't match
         output_path = Path(output_file)
@@ -233,11 +239,15 @@ def test_async_video_generation(
     )
 
     parser.add_argument(
-        "--output-format",
+        "--format",
         type=str,
         default="auto",
-        choices=["mp4", "avi", "auto"],
-        help="Output video format: mp4 or avi or auto",
+        choices=["mp4", "avi", "auto", "safetensors", "pt"],
+        help=(
+            "Generation content encoding format. Encoders: mp4 / avi / auto. "
+            "Tensor formats safetensors / pt return raw tensor bytes for "
+            "programmatic post-processing."
+        ),
     )
 
     args = parser.parse_args()
@@ -264,7 +274,7 @@ def test_async_video_generation(
         fps=args.fps,
         size=args.size,
         output_file=args.output,
-        output_format=args.output_format,
+        format=args.format,
     )
 
     sys.exit(0 if success else 1)

examples/visual_gen/serve/sync_image_gen.py

@@ -28,11 +28,16 @@ def test_image_generation(
     prompt: str = "A lovely cat lying on a sofa",
     n: int = 1,
     size: str = "512x512",
-    quality: str = "standard",
+    format: str = "png",
     response_format: str = "b64_json",
     output_file: str = "output_generation.png",
 ):
-    """Test image generation endpoint."""
+    """Test image generation endpoint.
+
+    ``format`` selects the encoding for the returned bytes. Image encoders
+    are ``"png"``, ``"webp"``, ``"jpeg"``; tensor payloads are
+    ``"safetensors"`` and ``"pt"``.
+    """
     print("=" * 80)
     print("Testing Image Generation API (POST /v1/images/generations)")
     print("=" * 80)
@@ -44,30 +49,41 @@ def test_image_generation(
     print(f"   Model: {model}")
     print(f"   Prompt: {prompt}")
     print(f"   Size: {size}")
-    print(f"   Quality: {quality}")
+    print(f"   Format: {format}")
     print(f"   Number of images: {n}")
 
     try:
-        # Use OpenAI SDK's images.generate() method
+        # ``format`` is a trtllm-serve extension over the OpenAI image
+        # API; the SDK forwards it via ``extra_body``.
         response = client.images.generate(
             model=model,
             prompt=prompt,
             n=n,
             size=size,
-            quality=quality,
             response_format=response_format,
+            extra_body={"format": format},
         )
 
         print("\n✓ Image generated successfully!")
         print(f"   Number of images: {len(response.data)}")
 
+        # Choose the on-disk extension to match the requested format so
+        # the saved file's suffix reflects its actual contents.
+        ext_map = {
+            "png": ".png",
+            "webp": ".webp",
+            "jpeg": ".jpeg",
+            "safetensors": ".safetensors",
+            "pt": ".pt",
+        }
+        ext = ext_map[format]
+        stem = output_file.rsplit(".", 1)[0]
+
         # Save images
         for i, image in enumerate(response.data):
             if response_format == "b64_json":
-                # Decode base64 image
                 image_data = base64.b64decode(image.b64_json)
-                output = f"{output_file.rsplit('.', 1)[0]}_{i}.png" if n > 1 else output_file
-
+                output = f"{stem}_{i}{ext}" if n > 1 else f"{stem}{ext}"
                 with open(output, "wb") as f:
                     f.write(image_data)
 
@@ -116,6 +132,16 @@ def test_image_generation(
         default="512x512",
         help="Image size in WxH format (e.g., 512x512, 1024x1024)",
     )
+    parser.add_argument(
+        "--format",
+        type=str,
+        default="png",
+        choices=["png", "webp", "jpeg", "safetensors", "pt"],
+        help=(
+            "Generation content encoding format. Image encoders: png / "
+            "webp / jpeg. Tensor payloads: safetensors / pt."
+        ),
+    )
     parser.add_argument(
         "--output",
         type=str,
@@ -137,6 +163,7 @@ def test_image_generation(
         model=args.model,
         prompt=args.prompt,
         size=args.size,
+        format=args.format,
         output_file=args.output,
     )
 

examples/visual_gen/serve/sync_video_gen.py

@@ -32,6 +32,7 @@ def test_sync_video_generation(
     fps: int = 24,
     size: str = "256x256",
     output_file: str = "output_sync.mp4",
+    format: str = "auto",
 ):
     """Test synchronous video generation with direct HTTP requests.
 
@@ -79,6 +80,7 @@ def test_sync_video_generation(
                 "size": size,
                 "seconds": str(duration),
                 "fps": str(fps),
+                "format": format,
             }
 
             # Add the file
@@ -103,18 +105,25 @@ def test_sync_video_generation(
                     "size": size,
                     "seconds": duration,
                     "fps": fps,
+                    "format": format,
                 },
             )
 
         print(f"\nStatus code: {response_video.status_code}")
 
         if response_video.status_code == 200:
-            # Determine actual file extension from Content-Type header
-            content_type = response_video.headers.get("content-type", "video/mp4")
-            if "x-msvideo" in content_type or "avi" in content_type:
-                actual_ext = ".avi"
+            # Determine the on-disk extension. Tensor formats are
+            # selected by the request and the server returns
+            # ``application/octet-stream``; encoder formats can be
+            # disambiguated from Content-Type (mp4 vs avi).
+            if format in ("safetensors", "pt"):
+                actual_ext = f".{format}"
             else:
-                actual_ext = ".mp4"
+                content_type = response_video.headers.get("content-type", "video/mp4")
+                if "x-msvideo" in content_type or "avi" in content_type:
+                    actual_ext = ".avi"
+                else:
+                    actual_ext = ".mp4"
 
             # Adjust output filename if extension doesn't match
             output_path = Path(output_file)
@@ -214,6 +223,17 @@ def test_sync_video_generation(
         default="output_sync.mp4",
         help="Output video file path (extension may change based on server encoder: .mp4 or .avi)",
     )
+    parser.add_argument(
+        "--format",
+        type=str,
+        default="auto",
+        choices=["mp4", "avi", "auto", "safetensors", "pt"],
+        help=(
+            "Generation content encoding format. Video encoders: mp4 / "
+            "avi / auto. Tensor payloads: safetensors / pt carry video + "
+            "audio + scalar metadata in a single file."
+        ),
+    )
 
     args = parser.parse_args()
 
@@ -239,6 +259,7 @@ def test_sync_video_generation(
         fps=args.fps,
         size=args.size,
         output_file=args.output,
+        format=args.format,
     )
 
     sys.exit(0 if success else 1)