diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 5b6981bb5647..880f4b536d8e 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -6,6 +6,7 @@ Cosmos3 supports the following generation modes from a single checkpoint: - **T2I** — text-to-image (`prompts/t2i.json`); emits a still frame (use `--output_type image` / a non-video `--output_path`). - **I2V / TI2V** — image-conditioned video (`prompts/i2v.json`). Condition on a reference frame via the prompt file's `vision_path` or `--image_path`. The image may be a local path, a `file://` / `http(s)://` URL, or a `data:` URI. - **V2V** — video-conditioned video (`prompts/v2v.json`). Condition on a reference video via `--video_path` (a local MP4/AVI file). Only the first (or last, per `condition_video_keep`) `max(condition_video_latent_indexes) * 4 + 1` input frames condition the output (5 by default); the encoded bytes pass through and each worker decodes just that window on NVDEC (see [Media I/O dependencies](#media-io-dependencies)). Validated for Nano / Super only. +- **Transfer** — control-video conditioning (`edge`/`blur`/`depth`/`seg`/`wsm` hints via `--extra_params`). The control constrains structure frame by frame; the prompt supplies appearance. `edge` and `blur` are auto-computed from `--video_path`; any hint also accepts a precomputed control clip (`{"edge": "control.mp4"}` — the example reads it and sends encoded bytes, the same contract as the `video` reference). Multiple hints compose (each adds a full control-token copy of the video sequence); long videos run chunked (93 frames/chunk, stitched on overlap frames) — but only past the first chunk, so raise `num_frames` above the pipeline default to generate one: it bounds how many frames are decoded from the inputs, and so how long the output can be. A single-hint request picks up that hint's tuned sampling preset — guidance scale, control guidance and flow shift — for any of those the request leaves unset; requests with several hints fall back to the generic video defaults. The active hint names are also appended to the prompt as a one-sentence control-adherence directive; pass `"emphasize_control_in_prompt": false` to suppress it for clean baselines or ablations. - **T2AV** — text-to-video with synchronized audio (`prompts/t2av.json` with `enable_audio: true`, or pass `--enable_audio`). Combine with a `vision_path` for image-conditioned audio-video (TI2AV). ## Checkpoints @@ -38,6 +39,7 @@ export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 - Saving `.mp4` output requires the `ffmpeg` CLI on `PATH` (`apt-get install -y ffmpeg`); without it the encoder falls back to `.avi`. - Decoding MP4/AVI reference videos (V2V) happens in the worker processes on NVDEC via PyNvVideoCodec, a declared TensorRT-LLM dependency — nothing extra to install. Tested combinations: H.264 in MP4 and H.264 in AVI; other containers/codecs/profiles depend on the demuxer and the GPU's NVDEC capabilities and are best-effort. +- Transfer's `edge`/`blur` controls are derived on the GPU from the reference video — nothing extra to install. Precomputed controls (`depth`/`seg`/`wsm`, or a precomputed `edge`/`blur`) are decoded like any other reference video. ## Deployment configs @@ -90,8 +92,9 @@ python cosmos3.py --model nvidia/Cosmos3-Nano \ # V2V: video-conditioned video (continues the first frames of --video_path). # Best results when the prompt describes the input video — e.g. continue a -# T2V output reusing its original prompt. Output size is fixed (1280x720 -# default); inputs are center-cropped, not aspect-matched. +# T2V output reusing its original prompt. Output size follows the source's +# aspect ratio (the closest supported bucket) unless the request sets +# height/width; the reference is center-cropped to whatever size is chosen. python cosmos3.py --model /path/to/Cosmos3-Nano \ --prompt_file prompts/v2v.json \ --video_path /path/to/Cosmos3-Nano/assets/example_i2v_output.mp4 \ @@ -124,6 +127,34 @@ python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \ --image_path https://example.com/frame.jpg \ --output_path output.mp4 +# Transfer: control-video conditioning — structure from the control video, +# appearance from the prompt. edge/blur are computed from --video_path. +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt "The same scene rendered as a photorealistic video, sharp detail." \ + --video_path /path/to/reference.mp4 \ + --extra_params '{"edge": true}' \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# Transfer with a fully synthetic control (no assets): generate an edge-map +# video of a bouncing ball, then let the prompt paint it photoreal. +# Keep synthetic controls edge-style: the blur hint expects the low +# frequencies of natural video, and flat synthetic color fields degrade +# generation. Temporal exposure swings (e.g. pulsing global light) do not +# transfer — express lighting spatially or in the prompt instead. +python generate_bouncing_ball_control.py --out_dir ./ball_control +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt "A photorealistic beach ball with colorful panels bouncing between the walls of an enclosed room, studio lighting." \ + --extra_params '{"edge": "./ball_control/control.mp4"}' \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# Multi-hint transfer: edge pins the layout, blur pins the palette/lighting. +# Hints must describe the same underlying video as each other and the prompt. +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt "The same scene, ultra sharp, professional photography." \ + --video_path /path/to/reference.mp4 \ + --extra_params '{"edge": true, "blur": true}' \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # Cosmos3-Edge image-to-video (480p-native defaults: 832x480 x 121 frames). # Reproduces the model-card sample: the checkpoint ships a structured prompt and # its own negative prompt alongside the conditioning image. Fetch them with @@ -134,7 +165,7 @@ python cosmos3.py --model nvidia/Cosmos3-Edge \ --image_path Cosmos3-Edge/assets/example_i2v_input.jpg \ --output_path output.mp4 -# Inline prompt +# Inline prompt (--prompt or a JSON file path) python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt "A cute puppy playing with a ball in a park" \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py index 57acd701506f..9664940c71eb 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -15,10 +15,10 @@ # limitations under the License. """Cosmos3 Text(+Image/Video)-to-Video(+Audio) generation. -One checkpoint serves T2V, T2I, I2V/TI2V, V2V and T2AV; ``prompts/`` holds a -prompt file per mode and ``--help`` lists the flags. See ``README.md`` in this -directory for the checkpoints, guardrail setup, deployment configs, and a -worked command line per mode. +One checkpoint serves T2V, T2I, I2V/TI2V, V2V, Transfer and T2AV; +``prompts/`` holds a prompt file per mode and ``--help`` lists the flags. +See ``README.md`` in this directory for the checkpoints, guardrail setup, +deployment configs, and a worked command line per mode. """ import argparse @@ -28,6 +28,7 @@ from typing import Any, Dict, Optional from tensorrt_llm import VisualGen, VisualGenArgs +from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import TRANSFER_HINT_KEYS _SCRIPT_DIR = Path(__file__).resolve().parent @@ -45,6 +46,48 @@ def _resolve_path(path: str) -> str: return path +def _load_transfer_controls(extra_params: dict[str, Any]) -> None: + """Read precomputed transfer controls into ``control`` bytes, client-side. + + A hint may name a control file (``{"edge": "ctrl.mp4"}`` or + ``{"edge": {"control_path": "ctrl.mp4"}}``); the worker only accepts encoded + bytes, so the media is read here. + """ + for key in TRANSFER_HINT_KEYS: + hint = extra_params.get(key) + if isinstance(hint, str): + hint = {"control_path": hint} + if not isinstance(hint, dict): + continue + control_path = hint.pop("control_path", None) + if control_path is None: + continue + if not isinstance(control_path, str) or not control_path.strip(): + raise ValueError( + f"--extra_params {key}.control_path must be a non-empty file path, " + f"got {control_path!r}." + ) + hint["control"] = Path(_resolve_path(control_path)).read_bytes() + extra_params[key] = hint + + +def _json_object(text: str) -> dict[str, Any]: + """Argparse type for a JSON *object*. + + ``json.loads`` alone also accepts arrays, scalars and null, which then + either fail deep in the merge or, for ``[]``, succeed while doing nothing. + """ + try: + value = json.loads(text) + except json.JSONDecodeError as exc: + raise argparse.ArgumentTypeError(f"not valid JSON: {exc}") from exc + if not isinstance(value, dict): + raise argparse.ArgumentTypeError( + f"expected a JSON object, got {type(value).__name__}: {text!r}" + ) + return value + + def _is_prompt_file(value: str) -> bool: """Whether a ``--prompt``/``--negative_prompt`` value names an existing file.""" return bool(value) and os.path.isfile(_resolve_path(value)) @@ -117,14 +160,14 @@ def resolve_negative_prompt( def resolve_prompt_and_options( *, - prompt: Optional[str], - prompt_file: Optional[str], - image_path: Optional[str], + prompt: str | None, + prompt_file: str | None, + image_path: str | None, enable_audio: bool, output_type: str, -) -> tuple[str, Optional[str], bool, str]: +) -> tuple[str, str | None, bool, str]: """Merge CLI args with optional prompt-file defaults.""" - prompt_data: Dict[str, Any] = {} + prompt_data: dict[str, Any] = {} if prompt_file is not None: prompt_data = load_prompt_file(prompt_file) @@ -240,6 +283,19 @@ def main(): parser.add_argument( "--output_type", type=str, default="video", help="Output type (video, image)" ) + parser.add_argument( + "--extra_params", + type=_json_object, + default=None, + help=( + "Model-specific extra params as a JSON object, merged last (overrides " + "flag-derived values). Keys are validated against the pipeline's " + "extra_param_specs. Transfer example: " + '\'{"edge": true, "blur": true, "control_guidance": 1.5}\' with --video_path, ' + 'or \'{"edge": "/path/control.mp4"}\' for a precomputed control (read here and ' + "sent as encoded bytes)." + ), + ) # Guardrails parser.add_argument( @@ -282,6 +338,12 @@ def main(): if args.video_path is not None: params.extra_params["video"] = Path(args.video_path).read_bytes() + if args.extra_params: + # Merged last: explicit JSON wins over flag-derived values. + params.extra_params.update(args.extra_params) + # The pipeline fits the output to the reference's aspect when height/width + # are unset, so there is nothing to do client-side. + _load_transfer_controls(params.extra_params) params.negative_prompt = negative_prompt diff --git a/examples/visual_gen/models/cosmos3/generate_bouncing_ball_control.py b/examples/visual_gen/models/cosmos3/generate_bouncing_ball_control.py new file mode 100644 index 000000000000..20cc260186e8 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/generate_bouncing_ball_control.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 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. +r"""Generate a synthetic edge-map control video: a ball bouncing off walls. + +Draws white outlines on black — a room border plus a ball following simple +elastic-bounce physics — which is exactly what the Cosmos3 transfer ``edge`` +hint expects. No media assets required: the control is 30 lines of math, and +transfer turns it into a photorealistic video whose subject follows the +physics frame by frame. + +Generate the control, then run transfer with it: + + python generate_bouncing_ball_control.py --out_dir ./ball_control + + python cosmos3.py --model nvidia/Cosmos3-Nano \ + --visual_gen_args ../../configs/cosmos3-nano-1gpu.yaml \ + --prompt "A photorealistic beach ball with colorful panels bouncing \ + between the walls of an enclosed room, studio lighting." \ + --extra_params '{"edge": "./ball_control/control.mp4"}' \ + --output_path cosmos3_bouncing_ball.mp4 + +Tip: keep synthetic controls edge-style. The ``blur`` hint expects the low +frequencies of natural video; flat synthetic color fields are far from its +training distribution and degrade generation quality. +""" + +import argparse +from pathlib import Path + +import numpy as np +import PIL.Image +import PIL.ImageDraw +import torch + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Bouncing-ball edge-control generator") + parser.add_argument("--out_dir", default="./ball_control") + parser.add_argument("--width", type=int, default=1280) + parser.add_argument("--height", type=int, default=720) + parser.add_argument("--num_frames", type=int, default=49) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--radius", type=int, default=90) + parser.add_argument("--wall_inset", type=int, default=10, help="Room border inset in px") + parser.add_argument("--line_width", type=int, default=6) + parser.add_argument("--start", type=float, nargs=2, default=(300.0, 250.0)) + parser.add_argument( + "--velocity", + type=float, + nargs=2, + default=(26.0, 19.0), + help="px/frame; the defaults bounce a few times over 49 frames", + ) + parser.add_argument( + "--save_frames", action="store_true", help="Also write the individual PNG frames" + ) + return parser.parse_args() + + +def draw_frame(args: argparse.Namespace, x: float, y: float) -> PIL.Image.Image: + image = PIL.Image.new("RGB", (args.width, args.height), (0, 0, 0)) + draw = PIL.ImageDraw.Draw(image) + inset, r = args.wall_inset, args.radius + draw.rectangle( + [inset, inset, args.width - inset, args.height - inset], + outline=(255, 255, 255), + width=args.line_width, + ) + draw.ellipse([x - r, y - r, x + r, y + r], outline=(255, 255, 255), width=args.line_width) + return image + + +def main() -> None: + args = parse_args() + out_dir = Path(args.out_dir) + out_dir.mkdir(parents=True, exist_ok=True) + + lo_x, hi_x = args.wall_inset + args.radius, args.width - args.wall_inset - args.radius + lo_y, hi_y = args.wall_inset + args.radius, args.height - args.wall_inset - args.radius + x, y = args.start + vx, vy = args.velocity + + frames = [] + for i in range(args.num_frames): + frame = draw_frame(args, x, y) + frames.append(frame) + if args.save_frames: + frame.save(out_dir / f"frame_{i:03d}.png") + x, y = x + vx, y + vy + if x < lo_x or x > hi_x: + vx = -vx + x = max(lo_x, min(x, hi_x)) + if y < lo_y or y > hi_y: + vy = -vy + y = max(lo_y, min(y, hi_y)) + + # Same encoder the pipeline saves its outputs through, so this needs + # nothing installed beyond what running a generation already needs. + from tensorrt_llm.media.encoding import save_video + + video_path = out_dir / "control.mp4" + clip = torch.from_numpy(np.stack([np.asarray(frame) for frame in frames])) + save_video(clip, video_path, frame_rate=args.fps) + + print( + f"Wrote {video_path}" + (f" and {args.num_frames} PNG frames" if args.save_frames else "") + ) + + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/serve/README.md b/examples/visual_gen/serve/README.md index fa9c55c8ca70..e075f8080b2a 100644 --- a/examples/visual_gen/serve/README.md +++ b/examples/visual_gen/serve/README.md @@ -316,7 +316,37 @@ 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 -- **Cosmos3**: `condition_video_latent_indexes`, `condition_video_keep` (V2V conditioning), `flow_shift`, `use_system_prompt`, ... +- **Cosmos3**: `condition_video_latent_indexes`, `condition_video_keep` (V2V conditioning), `flow_shift`, `use_system_prompt`, and the transfer hints `edge`/`blur`/`depth`/`seg`/`wsm` with `control_guidance`, `control_guidance_interval`, `num_video_frames_per_chunk`, ... (see below) + +##### Cosmos3 transfer hints + +`extra_params` is JSON, so a control clip travels as a **base64-encoded** MP4/AVI +string under `.control`; the server decodes it at the HTTP boundary. Only +`edge` and `blur` can be auto-computed — pass `true` and supply a `video` +reference for them to derive from. `depth`/`seg`/`wsm` have no generator, so +they always need a control clip. + +```json +{ + "prompt": "a city street at dusk", + "extra_params": { + "video": "", + "edge": {"preset_edge_threshold": "medium"}, + "blur": {"preset_blur_strength": "medium"}, + "depth": {"control": ""}, + "control_guidance": 1.5 + } +} +``` + +`preset_edge_threshold` and `preset_blur_strength` accept +`none`/`very_low`/`low`/`medium`/`high`/`very_high` and default to `medium`; a +bare `true` (or `""`) is shorthand for the object form. Individual +values are validated before the job is queued, so a bad preset or an +unsupported frame count fails fast; combinations that only make sense together +— a transfer option with no hint selected, or `edge`/`blur` asked to +auto-compute with no `video` — are still reported by the worker, as a client +error, once the request is running. > **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`). diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 359907801466..1661adfe364c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -17,7 +17,8 @@ Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. """ -from typing import Dict, Iterable +from collections.abc import Mapping +from typing import Any, Dict, Iterable from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema from tensorrt_llm.inputs.media_io import sniff_media_kind @@ -26,9 +27,45 @@ # Constant tables # --------------------------------------------------------------------------- +# Cosmos3 output resolution buckets keyed by target level, then aspect ratio; +# each value is (width, height). A source frame maps onto the bucket whose +# aspect ratio is closest (see ``find_closest_target_size`` in ``transfer.py``). +VIDEO_RES_SIZE_INFO = { + "256": { + "1,1": (256, 256), + "4,3": (320, 256), + "3,4": (256, 320), + "16,9": (320, 192), + "9,16": (192, 320), + }, + "480": { + "1,1": (640, 640), + "4,3": (736, 544), + "3,4": (544, 736), + "16,9": (832, 480), + "9,16": (480, 832), + }, + "704": { + "1,1": (960, 960), + "4,3": (1088, 832), + "3,4": (832, 1088), + "16,9": (1280, 704), + "9,16": (704, 1280), + }, + "720": { + "1,1": (960, 960), + "4,3": (1104, 832), + "3,4": (832, 1104), + "16,9": (1280, 720), + "9,16": (720, 1280), + }, +} + +# The default video resolution is the 720p 16:9 bucket, ``(width, height)``. +_DEFAULT_VIDEO_W, _DEFAULT_VIDEO_H = VIDEO_RES_SIZE_INFO["720"]["16,9"] COSMOS3_720P_PARAMS = { - "height": 720, - "width": 1280, + "height": _DEFAULT_VIDEO_H, + "width": _DEFAULT_VIDEO_W, "num_inference_steps": 35, "guidance_scale": 6.0, "max_sequence_length": 4096, @@ -101,6 +138,110 @@ def _validate_video_reference(video) -> None: ) +# --------------------------------------------------------------------------- +# Transfer preflight validators. +# +# These mirror ``transfer.resolve_transfer_config``'s parsing rather than adding +# rules of their own: a validator that is stricter than the worker would reject +# requests the pipeline would have served. The worker keeps its own checks -- +# offline callers do not go through preflight -- so these exist purely to turn a +# deterministic client mistake into a 400 at enqueue instead of a failure deep +# in the pipeline, after the request has already been accepted with a 202. +# --------------------------------------------------------------------------- +def _transfer_hint_payload(value: Any) -> Mapping: + """Normalize a control hint to its object form, as the worker does. + + Normalization first, checks after, so a hint carried as bare bytes and the + same hint carried as ``{"control": }`` are held to one standard. + """ + if value is True: + payload: Mapping = {} + elif isinstance(value, bytes): + payload = {"control": value} + elif isinstance(value, bool): # False, having already excluded True + raise ValueError( + "control hint must be true, encoded MP4/AVI bytes, or an object; got false. " + "Omit the key entirely to leave the hint off." + ) + elif not isinstance(value, Mapping): + raise TypeError( + "control hint must be an object, encoded control bytes, or true; " + f"got {type(value).__name__}." + ) + else: + if value.get("control_path") is not None: + raise ValueError( + "control hint no longer accepts 'control_path'; pass the encoded control clip " + "as 'control' bytes (Path(control).read_bytes())." + ) + payload = value + + control = payload.get("control") + if control is not None: + if not isinstance(control, bytes): + raise TypeError( + "control hint 'control' must be encoded MP4/AVI bytes, got " + f"{type(control).__name__}." + ) + # Same bar the `video` reference is held to: undecodable bytes fail at + # decode anyway, so name the problem now rather than mid-request. + if sniff_media_kind(control) != "video": + raise ValueError( + "control hint bytes are not a recognized video container (supported: MP4/AVI)." + ) + return payload + + +def _validate_edge_hint(value: Any) -> None: + # Imported at call time: transfer imports this module, so a module-level + # import would be circular. Specs pickle by name, so this stays picklable. + from .transfer import EDGE_PRESETS + + payload = _transfer_hint_payload(value) + # `or "medium"` matches the worker, which treats empty/None as the default. + preset = str(payload.get("preset_edge_threshold") or "medium").lower() + if preset not in EDGE_PRESETS: + raise ValueError( + f"unsupported preset_edge_threshold {preset!r}; expected one of {sorted(EDGE_PRESETS)}." + ) + + +def _validate_blur_hint(value: Any) -> None: + from .transfer import BLUR_PRESETS + + payload = _transfer_hint_payload(value) + preset = str(payload.get("preset_blur_strength") or "medium").lower() + if preset not in BLUR_PRESETS: + raise ValueError( + f"unsupported preset_blur_strength {preset!r}; expected one of {sorted(BLUR_PRESETS)}." + ) + + +def _validate_precomputed_control_hint(value: Any) -> None: + """For hints with no on-the-fly generator: a control clip is mandatory.""" + if _transfer_hint_payload(value).get("control") is None: + raise ValueError( + "this control has no on-the-fly generator, so it requires a precomputed clip " + "as encoded MP4/AVI bytes; only 'edge' and 'blur' accept true." + ) + + +def _validate_control_guidance_interval(value: Any) -> None: + from .transfer import _as_interval + + _as_interval(value) # reused outright, so preflight cannot drift from the worker + + +def _validate_positive_frames(value: Any) -> None: + if int(value) <= 0: + raise ValueError(f"must be a positive frame count, got {value}.") + + +def _validate_non_negative_frames(value: Any) -> None: + if int(value) < 0: + raise ValueError(f"must be a non-negative frame count, got {value}.") + + # Text-to-image (``output_type="image"``) defaults; resolved in ``infer()``. COSMOS3_T2I_PARAMS = { "height": 1024, @@ -174,6 +315,7 @@ def _validate_video_reference(video) -> None: "nemotron_dense": COSMOS3_EDGE_ENVELOPE, } +COSMOS3_V2V_DEFAULT_FLOW_SHIFT = 10.0 COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( @@ -245,4 +387,108 @@ def _validate_video_reference(video) -> None: ), validator=_validate_video_reference, ), + # Transfer + "edge": ExtraParamSchema( + type="bool_or_bytes_or_dict", + default=None, + description=( + "Canny-edge control. true auto-computes it from the `video` extra param; " + "or pass the encoded MP4/AVI bytes of a precomputed control clip, or an " + 'object {"control": , "preset_edge_threshold": "medium"}.' + ), + validator=_validate_edge_hint, + ), + "blur": ExtraParamSchema( + type="bool_or_bytes_or_dict", + default=None, + description=( + "Low-frequency (color/lighting) control. true auto-computes it from the " + "`video` extra param; or pass encoded control bytes, or an object " + '{"control": , "preset_blur_strength": "medium"}.' + ), + validator=_validate_blur_hint, + ), + "depth": ExtraParamSchema( + type="bool_or_bytes_or_dict", + default=None, + description=( + "Depth control. Requires precomputed encoded MP4/AVI bytes (no auto-computation)." + ), + validator=_validate_precomputed_control_hint, + ), + "seg": ExtraParamSchema( + type="bool_or_bytes_or_dict", + default=None, + description="Semantic-segmentation control. Requires precomputed encoded MP4/AVI bytes.", + validator=_validate_precomputed_control_hint, + ), + "wsm": ExtraParamSchema( + type="bool_or_bytes_or_dict", + default=None, + description=( + "World-scenario-model control. Requires precomputed encoded MP4/AVI bytes; " + "runs 101-frame chunks at 10 fps." + ), + validator=_validate_precomputed_control_hint, + ), + "control_guidance": ExtraParamSchema( + type="float", + default=None, + description="Transfer control-guidance scale (CFG for the control branch).", + ), + "control_guidance_interval": ExtraParamSchema( + type="list", + default=None, + description=( + "[lo, hi] window where control guidance is active, in raw scheduler " + "timesteps (typically 0-1000, counting down) rather than as a fraction " + "of the schedule: [0.0, 0.8] gates on the final step, not the first " + "80%. To cover most of the schedule use e.g. [200, 1000]." + ), + validator=_validate_control_guidance_interval, + ), + "num_video_frames_per_chunk": ExtraParamSchema( + type="int", + default=None, + description="Transfer chunk length in frames (default 93; 101 for wsm).", + validator=_validate_positive_frames, + ), + "num_conditional_frames": ExtraParamSchema( + type="int", + default=None, + description="Overlap frames pinned from the previous chunk when stitching.", + validator=_validate_non_negative_frames, + ), + "num_first_chunk_conditional_frames": ExtraParamSchema( + type="int", + default=None, + description="Input-video frames pinned at the start of the first chunk.", + validator=_validate_non_negative_frames, + ), + "max_frames": ExtraParamSchema( + type="int", + default=None, + description="Cap on frames decoded from transfer inputs/controls.", + validator=_validate_positive_frames, + ), + "show_control_condition": ExtraParamSchema( + type="bool", default=False, description="Concatenate the control video beside the output." + ), + "show_input": ExtraParamSchema( + type="bool", default=False, description="Concatenate the input video beside the output." + ), + "share_vision_temporal_positions": ExtraParamSchema( + type="bool", + default=None, + description="Controls share the target frames' temporal mRoPE positions.", + ), + "emphasize_control_in_prompt": ExtraParamSchema( + type="bool", + default=None, + description=( + "Append a one-sentence control-adherence directive naming the active " + "hints to the user prompt (default true). Set false for clean " + "baselines / ablations. The system prompt is unchanged." + ), + ), } 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 6ecd12ae8884..c7ff75392046 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -27,6 +27,14 @@ from diffusers.video_processor import VideoProcessor from transformers import AutoTokenizer +try: + from tqdm.auto import tqdm +except ImportError: # pragma: no cover - tqdm is optional at runtime. + + def tqdm(iterable, **kwargs): + return iterable + + 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_registry import PipelineComponent, register_pipeline @@ -38,19 +46,31 @@ from tensorrt_llm._utils import nvtx_range from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from tensorrt_llm.media.decoding import decode_video_reference_window +from tensorrt_llm.media.decoding import decode_video_reference_window, video_stream_info from .defaults import ( + COSMOS3_720P_PARAMS, COSMOS3_ENVELOPES, COSMOS3_EXTRA_SPECS, COSMOS3_GENERATION_DEFAULTS, + COSMOS3_V2V_DEFAULT_FLOW_SHIFT, _normalize_condition_video_keep, _normalize_condition_video_latent_indexes, ) from .guardrails import check_video_safety, download_guardrail_checkpoint from .negative_prompt import COSMOS3_VIDEO_NEGATIVE_PROMPT -from .sampling import Cosmos3SamplingPolicy, load_scheduler +from .sampling import DISTILLED_GUIDANCE_SCALE, Cosmos3SamplingPolicy, load_scheduler from .sound_tokenizer import LatentAutoEncoderV2 +from .transfer import ( + TRANSFER_DEFAULTS, + Cosmos3TransferConfig, + decode_media_to_uint8_cthw, + find_closest_target_size, + load_or_compute_control_frames, + pad_temporal_frames, + resolve_transfer_config, + uint8_cthw_to_normalized_5d, +) from .transformer_cosmos3 import NEMOTRON_DENSE_RECIPE, Cosmos3VFMTransformer, resolve_arch_recipe # Image modes declare no negative prompt in the reference @@ -89,6 +109,11 @@ def default_negative_prompt(output_type: str) -> str: "You are a helpful assistant who will generate images from a give prompt." ) COSMOS3_V2V_FLOW_SHIFT = 10.0 +# Fraction of a transfer hint that may be mirror-padding before it is worth a +# warning. A few frames of tail ping-pong is normal when clips differ slightly; +# beyond this the control is mostly invented and the client likely sent clips of +# different videos. +CONTROL_LENGTH_MISMATCH_RATIO = 0.10 COSMOS3_DURATION_TEMPLATE = "The video is {duration:.1f} seconds long and is of {fps:.0f} FPS." COSMOS3_DEFAULT_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." @@ -541,13 +566,19 @@ def _cacheable_flow_shifts(self) -> frozenset: Derived rather than fixed so a new family or mode table is picked up automatically: the per-mode generation tables, the checkpoint's own - shift, and V2V's stronger shift. Entries are still created lazily, so a - mode that is never served never builds one. + shift, V2V's stronger shift, and the per-hint transfer presets. Entries + are still created lazily, so a mode that is never served never builds + one. + + Transfer's shifts are included even though every hint currently declares + the same value as V2V: tuning one off that value would otherwise drop it + out of the cache silently and rebuild a scheduler on every request for + that hint. """ cached = getattr(self, "_cacheable_flow_shifts_cache", None) if cached is not None: return cached - shifts = {COSMOS3_V2V_FLOW_SHIFT} + shifts = {COSMOS3_V2V_FLOW_SHIFT, COSMOS3_V2V_DEFAULT_FLOW_SHIFT} checkpoint_shift = getattr(self.sampling, "checkpoint_flow_shift", None) if checkpoint_shift is not None: shifts.add(float(checkpoint_shift)) @@ -556,6 +587,10 @@ def _cacheable_flow_shifts(self) -> frozenset: mode_shift = table.get("flow_shift") if mode_shift is not None: shifts.add(float(mode_shift)) + for hint_defaults in TRANSFER_DEFAULTS.values(): + hint_shift = hint_defaults.get("flow_shift") + if hint_shift is not None: + shifts.add(float(hint_shift)) cached = self._cacheable_flow_shifts_cache = frozenset(shifts) return cached @@ -590,6 +625,7 @@ def infer(self, req): extra_params = req.params.extra_params or {} output_type = extra_params.get("output_type", "video") is_t2i = str(output_type).lower() == "image" + transfer_config = resolve_transfer_config(extra_params, req.params, req.prompt) # Caller-assigned values win. Anything still carrying a pipeline # default — unset, or merged by the executor from the video table — @@ -600,10 +636,49 @@ 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) + # Source-derived sizes go in as non-None, so they win over the mode + # table: _resolve_generation_params only fills what is still None. + height, width = as_given("height"), as_given("width") + frame_rate = req.params.frame_rate + wants_source_size = "height" not in specified and "width" not in specified + # Timing resolves as a unit. A caller who pinned `num_frames` -- directly, + # or via `seconds`, which the serving layer already converted at the + # default rate -- has fixed the output duration; adopting a different + # frame rate underneath that would silently stretch or compress it. + wants_source_fps = "frame_rate" not in specified and "num_frames" not in specified + if wants_source_size or wants_source_fps: + source = video + if source is None and transfer_config is not None: + # Transfer can run on precomputed controls alone; then the + # control clip is what defines the structure to match. + source = next( + (hint.control for hint in transfer_config.ordered_hints if hint.control), + None, + ) + # One header read serves both: no GPU, no frame decoded. + info = video_stream_info(source) if source is not None else None + if info is not None: + if wants_source_size: + # Follow the reference's aspect, so a portrait or square + # source is not center-cropped into the default landscape + # bucket. Skipped when either dimension was named: that + # states an intent, and overriding half of it would be + # worse than leaving it alone. + width, height = find_closest_target_size(info.height, info.width, 720) + if wants_source_fps and info.frame_rate is not None: + # Emitting an 8 fps source at the 24 fps default would play + # it back at 3x speed and misreport its duration to the + # text conditioning. + frame_rate = info.frame_rate + if self.rank == 0: + logger.info( + f"Cosmos3 following the source: {width}x{height} (WxH) @ {frame_rate} fps" + ) resolved = self._resolve_generation_params( "image" if is_t2i else "video", - height=as_given("height"), - width=as_given("width"), + height=height, + width=width, num_inference_steps=as_given("num_inference_steps"), guidance_scale=as_given("guidance_scale"), ) @@ -611,7 +686,6 @@ def as_given(field_name): width = resolved["width"] num_inference_steps = resolved["num_inference_steps"] guidance_scale = resolved["guidance_scale"] - video = extra_params.get("video") # encoded MP4/AVI bytes (the extra-param contract) return self.forward( prompt=req.prompt, @@ -624,7 +698,7 @@ def as_given(field_name): guidance_scale=guidance_scale, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, - frame_rate=req.params.frame_rate, + frame_rate=frame_rate, use_duration_template=extra_params.get( "use_duration_template", COSMOS3_EXTRA_SPECS["use_duration_template"].default, @@ -642,6 +716,7 @@ def as_given(field_name): condition_video_latent_indexes=extra_params.get("condition_video_latent_indexes"), condition_video_keep=extra_params.get("condition_video_keep"), flow_shift=extra_params.get("flow_shift"), + transfer_config=transfer_config, ) def _apply_metadata_templates( @@ -781,7 +856,8 @@ def _tokenize_prompt( token_ids.append(self.tokenizer.convert_tokens_to_ids("<|vision_start|>")) # 151652 seq_len = len(token_ids) - # Pad to max_sequence_length + # Pad to max_sequence_length. TRT's shared denoiser concatenates CFG + # prompt tensors, so cond/uncond sequence lengths must match. pad_len = max_sequence_length - seq_len attention_mask = [1] * seq_len + [0] * pad_len token_ids = token_ids + [self.tokenizer.pad_token_id or 0] * pad_len @@ -923,8 +999,7 @@ def post_step_fn(latents: torch.Tensor) -> torch.Tensor: # VAE decode # ========================================================================= - @nvtx_range("_decode_latents", color="blue") - def _decode_latents(self, latents): + def _decode_latents_raw(self, latents): latents = latents.to(self.vae.dtype) if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): @@ -944,9 +1019,11 @@ def _decode_latents(self, latents): scaling_factor = self.vae.config.get("scaling_factor", 1.0) latents = latents / scaling_factor - video = self.vae.decode(latents, return_dict=False)[0] - video = postprocess_video_tensor(video) - return video + return self.vae.decode(latents, return_dict=False)[0] + + @nvtx_range("_decode_latents", color="blue") + def _decode_latents(self, latents): + return postprocess_video_tensor(self._decode_latents_raw(latents)) # ========================================================================= # Audio generation @@ -1080,6 +1157,39 @@ def _prepare_latents_v2v( velocity_mask = 1.0 - condition_mask return latents, velocity_mask, condition_latents + # ========================================================================= + # Transfer + # ========================================================================= + + def _prepare_transfer_latents( + self, + target_video: torch.Tensor, + current_conditional_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + condition_latents = self._encode_video_tensor(target_video) + noise = randn_tensor( + condition_latents.shape, + generator=generator, + device=self.device, + dtype=self.dtype, + ) + condition_mask = torch.zeros( + 1, + 1, + condition_latents.shape[2], + 1, + 1, + device=self.device, + dtype=self.dtype, + ) + if current_conditional_frames > 0: + latent_frames = (current_conditional_frames - 1) // self.vae_scale_factor_temporal + 1 + condition_mask[:, :, :latent_frames] = 1.0 + latents = condition_mask * condition_latents + (1.0 - condition_mask) * noise + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, condition_mask * condition_latents + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -1109,6 +1219,7 @@ def forward( condition_video_latent_indexes: Iterable[int] | None = None, condition_video_keep: str | None = None, flow_shift: Optional[float] = None, + transfer_config: Optional[Cosmos3TransferConfig] = None, ): """Run one generation. ``infer()`` is the resolved entry point. @@ -1179,9 +1290,29 @@ def forward( is_v2v = video is not None and not is_t2i if use_system_prompt is None: # V2V always wants it; otherwise the checkpoint declares the default. - use_system_prompt = is_v2v or self.default_use_system_prompt + # Transfer opts out for reference parity (vllm-omni + # `_forward_transfer` defaults it False). + if transfer_config is not None: + use_system_prompt = False + else: + use_system_prompt = is_v2v or self.default_use_system_prompt else: use_system_prompt = bool(use_system_prompt) + if transfer_config is not None: + if is_t2i: + raise ValueError("Cosmos3 transfer inference is supported only for video outputs.") + if enable_audio: + raise ValueError( + "Cosmos3 transfer inference cannot be combined with sound generation." + ) + if image is not None: + # _forward_transfer takes no image: structure comes from the + # control hints and the first chunk conditions on `video`. Say + # so rather than dropping the reference silently. + raise ValueError( + "Cosmos3 transfer inference cannot be combined with an image reference; " + "pass the conditioning clip as the 'video' extra param instead." + ) guidance_interval = None if is_t2i: @@ -1208,9 +1339,14 @@ def forward( else: target_shift = mode_shift if flow_shift is None else flow_shift target_karras = None - self.scheduler = self._scheduler_for(target_shift, target_karras) - if getattr(self, "audio_scheduler", None) is not None: - self.audio_scheduler = self._scheduler_for(target_shift, target_karras, stream="audio") + if transfer_config is None: + # Transfer applies the hint's own shift inside _forward_transfer, so + # rebuilding the schedulers here would only be undone a few lines later. + self.scheduler = self._scheduler_for(target_shift, target_karras) + if getattr(self, "audio_scheduler", None) is not None: + self.audio_scheduler = self._scheduler_for( + target_shift, target_karras, stream="audio" + ) if self.rank == 0: logger.info( @@ -1265,6 +1401,27 @@ def forward( if text_blocked.item(): return PipelineOutput() + if transfer_config is not None: + return self._forward_transfer( + prompt=prompt, + negative_prompt=negative_prompt, + height=height, + width=width, + max_frames=transfer_config.max_frames, + num_inference_steps=num_inference_steps, + max_sequence_length=max_sequence_length, + use_system_prompt=use_system_prompt, + use_duration_template=False, + use_resolution_template=False, + seed=seed, + frame_rate=frame_rate, + num_frames=num_frames, + use_guardrails=use_guardrails, + timer=timer, + transfer_config=transfer_config, + video=video, + ) + generator = torch.Generator(device=self.device).manual_seed(seed) if negative_prompt is None: @@ -1590,3 +1747,593 @@ def post_step_fn(step_latents): else None, ) ) + + @staticmethod + def _get_transfer_num_chunks( + total_frames: int, + frames_per_chunk: int, + conditional_frames: int, + ) -> tuple[int, int]: + if frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer frames_per_chunk must be positive.") + if total_frames <= frames_per_chunk: + return 1, frames_per_chunk + stride = frames_per_chunk - conditional_frames + if stride <= 0: + raise ValueError( + "Cosmos3 transfer num_conditional_frames must be smaller than num_video_frames_per_chunk." + ) + remaining = total_frames - frames_per_chunk + extra_chunks = remaining // stride + (1 if remaining % stride else 0) + return 1 + extra_chunks, stride + + def _warn_on_control_length_mismatch( + self, per_hint_frames: dict[str, torch.Tensor], total_frames: int + ) -> None: + """Flag hints short enough that mirror-padding invents most of their control. + + Measured as padding actually applied rather than as disagreement between + hints, because that is what reaches the model: a request that pins + ``num_frames`` down to the shortest clip pads nothing and stays silent, + while a few frames of tail ping-pong is well under the threshold. A + large gap is almost always a client that sent clips of different videos. + """ + if self.rank != 0 or total_frames <= 0: + return + short = { + key: frames.shape[1] + for key, frames in per_hint_frames.items() + if (total_frames - frames.shape[1]) / total_frames > CONTROL_LENGTH_MISMATCH_RATIO + } + if not short: + return + counts = ", ".join(f"{key}: {count}" for key, count in sorted(short.items())) + logger.warning( + f"Cosmos3 transfer control length mismatch: {counts} against {total_frames} frames " + f"generated. The short hints are mirror-padded up to that length, so the output is " + f"conditioned on control content that does not exist in them. Supply clips of " + f"matching length, or set num_frames to the shortest one." + ) + + @staticmethod + def _positive_float(value: Optional[float]) -> Optional[float]: + if value is None: + return None + try: + value = float(value) + except (TypeError, ValueError): + return None + if value <= 0: + return None + return value + + @staticmethod + def _transfer_active_at( + timestep: torch.Tensor, + interval: tuple[float, float] | None, + ) -> bool: + """Is guidance active at ``timestep``? + + ``interval`` is in the scheduler's own timestep units (raw, typically + 0-1000 counting down), not a fraction of the schedule -- ``[0.0, 0.8]`` + selects the last step rather than the first 80%. + """ + if interval is None: + return True + t_scalar = float(timestep.item()) if torch.is_tensor(timestep) else float(timestep) + lo, hi = interval + return float(lo) <= t_scalar <= float(hi) + + @staticmethod + def _combine_transfer_predictions( + *, + cond_full: torch.Tensor, + cond_no_control: torch.Tensor | None, + uncond_full: torch.Tensor | None, + guidance_scale: float, + control_guidance: float, + ) -> torch.Tensor: + needs_control_cfg = cond_no_control is not None and control_guidance != 1.0 + needs_text_cfg = uncond_full is not None and guidance_scale > 1.0 + + if needs_control_cfg and needs_text_cfg: + control_cond = cond_no_control + control_guidance * (cond_full - cond_no_control) + return uncond_full + guidance_scale * (control_cond - uncond_full) + if needs_control_cfg: + return cond_no_control + control_guidance * (cond_full - cond_no_control) + if needs_text_cfg: + return uncond_full + guidance_scale * (cond_full - uncond_full) + return cond_full + + def diffuse_transfer( + self, + *, + latents: torch.Tensor, + timesteps: torch.Tensor, + cond_ids: torch.Tensor, + cond_mask: torch.Tensor, + uncond_ids: torch.Tensor, + uncond_mask: torch.Tensor, + guidance_scale: float, + control_guidance: float, + control_guidance_interval: tuple[float, float] | None, + control_latents: list[torch.Tensor], + shared_kwargs: dict[str, Any], + velocity_mask: torch.Tensor, + condition_latents: torch.Tensor, + generator: torch.Generator, + guidance_interval: tuple[float, float] | None = None, + ) -> torch.Tensor: + """Run Cosmos3 transfer denoising with sequential control/text CFG branches.""" + + branch_caches: dict[str, tuple[Any, Any]] = {} + + def run_branch( + cache_key: str, + *, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + branch_control_latents: list[torch.Tensor] | None, + timestep: torch.Tensor, + ) -> torch.Tensor: + self.transformer.cached_kv, self.transformer.cached_freqs_gen = branch_caches.get( + cache_key, + (None, None), + ) + result = self.transformer( + hidden_states=latents, + timestep=timestep / self.scheduler.config.num_train_timesteps, + raw_timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + control_latents=branch_control_latents, + **shared_kwargs, + ) + branch_caches[cache_key] = ( + self.transformer.cached_kv, + self.transformer.cached_freqs_gen, + ) + if result.video is None: + raise ValueError("Cosmos3 transfer diffusion expects video predictions.") + return result.video + + self.transformer.reset_cache() + try: + transfer_steps = tqdm( + timesteps, + total=len(timesteps), + desc="Transfer denoising", + disable=self.rank != 0, + dynamic_ncols=True, + ) + for t in transfer_steps: + timestep = t.expand(latents.shape[0]) + step_guidance = ( + float(guidance_scale) if self._transfer_active_at(t, guidance_interval) else 1.0 + ) + step_control = ( + float(control_guidance) + if self._transfer_active_at(t, control_guidance_interval) + else 1.0 + ) + + cond_full = run_branch( + "transfer_cond_full", + text_ids=cond_ids, + text_mask=cond_mask, + branch_control_latents=control_latents, + timestep=timestep, + ) + cond_no_control = None + if step_control != 1.0: + cond_no_control = run_branch( + "transfer_cond_no_control", + text_ids=cond_ids, + text_mask=cond_mask, + branch_control_latents=None, + timestep=timestep, + ) + + uncond_full = None + if step_guidance > 1.0: + uncond_full = run_branch( + "transfer_uncond_full", + text_ids=uncond_ids, + text_mask=uncond_mask, + branch_control_latents=control_latents, + timestep=timestep, + ) + + noise_pred = self._combine_transfer_predictions( + cond_full=cond_full, + cond_no_control=cond_no_control, + uncond_full=uncond_full, + guidance_scale=step_guidance, + control_guidance=step_control, + ) + noise_pred = noise_pred * velocity_mask + latents = self.scheduler.step( + noise_pred, + t, + latents, + return_dict=False, + **self.sampling.scheduler_step_kwargs(generator), + )[0] + latents = velocity_mask * latents + (1.0 - velocity_mask) * condition_latents + finally: + self.transformer.reset_cache() + + return latents + + def _forward_transfer( + self, + *, + prompt: Union[str, List[str]], + negative_prompt: Optional[str], + height: int, + width: int, + max_frames: int, + num_inference_steps: int, + max_sequence_length: int, + use_system_prompt: bool, + use_duration_template: bool, + use_resolution_template: bool, + seed: int, + frame_rate: float, + num_frames: int, + use_guardrails: bool, + timer: CudaPhaseTimer, + transfer_config: Cosmos3TransferConfig, + video: Optional[bytes], + ) -> PipelineOutput: + if self.rank == 0: + logger.info(f"Cosmos3 transfer target={width}x{height} (WxH)") + + # Decode the input video and every precomputed control on this rank's + # NVDEC. A decode can fail non-uniformly across ranks (decoder init, + # corrupt stream, allocation), so converge all ranks on one outcome + # before any model collective rather than hanging the healthy ones. + # The decoder sizes its retention ring from the requested window, so + # asking for `max_frames` (5000 by default) would reserve ~14 GB at 720p + # before a single frame lands. Bound it by what is actually generated: + # the output is `num_frames` long, and `max_frames` stays the ceiling. + decode_frames = min(int(max_frames), int(transfer_config.num_frames or num_frames)) + if self.rank == 0: + logger.info(f"Cosmos3 transfer decoding up to {decode_frames} frames per reference") + + input_frames = None + per_hint_frames: dict[str, torch.Tensor] = {} + prepare_error: Optional[Exception] = None + try: + if video is not None: + input_frames = decode_media_to_uint8_cthw( + video, + height=height, + width=width, + max_frames=decode_frames, + device=self.device, + ) + + for hint in transfer_config.ordered_hints: + frames = load_or_compute_control_frames( + hint, + height=height, + width=width, + max_frames=decode_frames, + input_frames=input_frames, + device=self.device, + ) + if frames.shape[1] < 1: + raise ValueError(f"Cosmos3 transfer hint '{hint.key}' produced no frames.") + per_hint_frames[hint.key] = frames + if not per_hint_frames: + raise ValueError("Cosmos3 transfer requires at least one control hint.") + except Exception as exc: + prepare_error = exc + synchronize_media_prepare_status(prepare_error) + + # The longest hint sets the length. Taking the first hint's instead would + # make the result depend on TRANSFER_HINT_KEYS order rather than on the + # request: with the same two clips, the earlier modality would win, so a + # shorter first hint silently truncated the longer one (the chunk loop + # never reads past total_frames) while a longer one padded it. No + # further clamp is needed -- every hint was decoded under decode_frames, + # which already bounds them by num_frames. + total_frames = max(1, max(frames.shape[1] for frames in per_hint_frames.values())) + self._warn_on_control_length_mismatch(per_hint_frames, total_frames) + per_hint_frames = { + key: pad_temporal_frames(frames, total_frames) + for key, frames in per_hint_frames.items() + } + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, total_frames) + + temporal_compression = self.vae_scale_factor_temporal + chunk_frames = 1 if total_frames == 1 else transfer_config.num_video_frames_per_chunk + chunk_frames = ( + math.ceil((chunk_frames - 1) / temporal_compression) * temporal_compression + 1 + ) + num_chunks, stride = self._get_transfer_num_chunks( + total_frames, + chunk_frames, + transfer_config.num_conditional_frames, + ) + padded_frames = max(total_frames, chunk_frames) + per_hint_frames = { + key: pad_temporal_frames(frames, padded_frames) + for key, frames in per_hint_frames.items() + } + if input_frames is not None: + input_frames = pad_temporal_frames(input_frames, padded_frames) + + # The encoded reference carries no frame rate across the extra-param + # boundary, so the hint's configured fps (wsm's 10) wins over the + # request's, which falls back to the mode default. + frame_rate = ( + self._positive_float(transfer_config.fps) or self._positive_float(frame_rate) or 24.0 + ) + num_inference_steps = num_inference_steps or COSMOS3_720P_PARAMS["num_inference_steps"] + guidance_scale = ( + float(transfer_config.guidance_scale) + if transfer_config.guidance_scale is not None + else COSMOS3_720P_PARAMS["guidance_scale"] + ) + if self.sampling.is_distilled: + # A distilled checkpoint runs one schedule and bakes guidance into + # the weights, so the per-hint guidance presets cannot apply. Say so + # rather than silently sampling off-distribution. + num_inference_steps = self.sampling.num_steps(num_inference_steps) + if guidance_scale != DISTILLED_GUIDANCE_SCALE: + logger.warning( + f"Cosmos3 transfer on a distilled checkpoint: overriding " + f"guidance_scale {guidance_scale} with the mandated " + f"{DISTILLED_GUIDANCE_SCALE}; per-hint guidance presets do not apply." + ) + guidance_scale = DISTILLED_GUIDANCE_SCALE + flow_shift_target = float( + transfer_config.flow_shift + if transfer_config.flow_shift is not None + else COSMOS3_V2V_DEFAULT_FLOW_SHIFT + ) + max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] + self._guidance_scale = guidance_scale + self._num_timesteps = num_inference_steps + # forward() skips its own scheduler rebuild for transfer, so this is the + # only setup on this path: assigning is what keeps a previous request's + # scheduler from carrying over on a reused worker. + self.scheduler = self._scheduler_for(flow_shift_target, use_karras_sigmas=False) + + generator = torch.Generator(device=self.device).manual_seed(seed) + + if negative_prompt is None: + negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT + + # Transfer prompts are already upsampled by the benchmark/config path. + # Keep them verbatim; duration/resolution templates would change parity. + prompt = [prompt] if isinstance(prompt, str) else list(prompt) + prompt = prompt[0] + prompt = transfer_config.emphasized_prompt(prompt) + if self.rank == 0: + logger.info(f"Transfer prompt: '{prompt}'") + + # 1. Tokenize prompts (no separate text encoder — transformer embeds internally) + logger.info("Tokenizing prompts...") + system_prompt = COSMOS3_DEFAULT_SYSTEM_PROMPT + cond_ids, cond_mask = self._tokenize_prompt( + prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt + ) + uncond_ids, uncond_mask = self._tokenize_prompt( + negative_prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt + ) + + # Finished frames land in host memory as they are produced rather than + # stacking on the GPU for the whole run: at 720p a decoded chunk is + # ~0.5 GB, so a long generation would pile several GB of *completed* + # output on top of the denoise working set -- and `torch.cat` would then + # briefly double it. Only the next chunk's conditioning has to stay + # resident. + # + # Pinned, because the copy is then asynchronous and overlaps the next + # chunk's denoise; into pageable memory the same copy blocks and runs + # ~40x slower. Allocated per request rather than held, since torch's + # caching host allocator keeps the block and only the first request of a + # given size reaches cudaHostAlloc -- so deployments that never serve + # transfer page-lock nothing. + # + # Assembled on rank 0 alone: every rank decodes each chunk because the + # next one conditions on it, but the executor sends only rank 0's + # response, so assembling on the others allocates a pinned buffer per + # rank to build a video nobody reads. + assembles_output = self.rank == 0 + show_controls = transfer_config.show_control_condition + show_input_panel = transfer_config.show_input and input_frames is not None + panels = (len(per_hint_frames) if show_controls else 0) + (1 if show_input_panel else 0) + host_output = ( + torch.empty( + (1, total_frames, height, width * (panels + 1), 3), + dtype=torch.uint8, + # Page-locking needs a CUDA context; only the CPU-only unit + # tests run without one, and there is nothing to overlap there. + pin_memory=torch.cuda.is_available(), + ) + if assembles_output + else None + ) + frames_written = 0 + previous_output: torch.Tensor | None = None + + # Every chunk's decode is part of generation here (it feeds the next + # chunk), so the whole loop counts as denoise; post covers assembly. + timer.mark_denoise_start() + for chunk_id in range(num_chunks): + start_frame = chunk_id * stride + end_frame = min(start_frame + chunk_frames, total_frames) + control_norms = { + key: uint8_cthw_to_normalized_5d( + pad_temporal_frames(frames[:, start_frame:end_frame], chunk_frames), + dtype=self.dtype, + ) + for key, frames in per_hint_frames.items() + } + target_norm = torch.zeros_like(next(iter(control_norms.values()))) + current_conditional_frames = 0 + + if chunk_id == 0 and transfer_config.num_first_chunk_conditional_frames > 0: + if input_frames is None: + raise ValueError( + "Cosmos3 transfer num_first_chunk_conditional_frames > 0 requires a video input." + ) + current_conditional_frames = min( + transfer_config.num_first_chunk_conditional_frames, + input_frames.shape[1], + chunk_frames, + ) + if current_conditional_frames > 0: + input_cond = uint8_cthw_to_normalized_5d( + input_frames[:, :current_conditional_frames], + dtype=self.dtype, + ) + target_norm[:, :, :current_conditional_frames] = input_cond + if current_conditional_frames < chunk_frames: + fill = target_norm[ + :, :, current_conditional_frames - 1 : current_conditional_frames + ] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + elif chunk_id > 0 and previous_output is not None: + current_conditional_frames = min( + transfer_config.num_conditional_frames, + previous_output.shape[2], + chunk_frames, + ) + if current_conditional_frames > 0: + target_norm[:, :, :current_conditional_frames] = previous_output[ + :, :, -current_conditional_frames: + ].to(target_norm) + if current_conditional_frames < chunk_frames: + fill = target_norm[ + :, :, current_conditional_frames - 1 : current_conditional_frames + ] + target_norm[:, :, current_conditional_frames:] = fill.expand( + -1, + -1, + chunk_frames - current_conditional_frames, + -1, + -1, + ) + + control_latents = [ + self._encode_video_tensor(control) for control in control_norms.values() + ] + latents, velocity_mask, condition_latents = self._prepare_transfer_latents( + target_norm, + current_conditional_frames, + generator, + ) + video_shape = (latents.shape[2], latents.shape[3], latents.shape[4]) + shared_kwargs = dict( + video_shape=video_shape, + fps=frame_rate, + noisy_frame_mask=velocity_mask, + transfer_share_vision_temporal_positions=transfer_config.share_vision_temporal_positions, + ) + + self.sampling.set_timesteps(self.scheduler, num_inference_steps, device=self.device) + latents = self.diffuse_transfer( + latents=latents, + timesteps=self.scheduler.timesteps, + cond_ids=cond_ids, + cond_mask=cond_mask, + uncond_ids=uncond_ids, + uncond_mask=uncond_mask, + guidance_scale=guidance_scale, + control_guidance=transfer_config.control_guidance, + control_guidance_interval=transfer_config.control_guidance_interval, + control_latents=control_latents, + shared_kwargs=shared_kwargs, + velocity_mask=velocity_mask, + condition_latents=condition_latents, + generator=generator, + ) + # Deliberately the raw decode rather than `decode_latents()`: the + # decoded chunk is this loop's *input* as well as its output — the + # next chunk conditions on its tail — so every rank needs it. Owning + # one rank's decode and broadcasting would ship ~0.5 GB per chunk + # and serialize the loop behind a collective, which costs more than + # recomputing the decode locally. + output_video = self._decode_latents_raw(latents).clamp(-1, 1) + + # Chunk 0 keeps every frame; later chunks drop the overlap they were + # conditioned on. The tail trim that used to follow the final concat + # happens here instead, as a bound on what each chunk contributes. + skip = 0 if chunk_id == 0 else current_conditional_frames + take = min(output_video.shape[2] - skip, total_frames - frames_written) + if assembles_output and take > 0: + chunk = output_video[:, :, skip : skip + take] + panel_tensors = [] + if show_controls: + panel_tensors.extend( + control_norms[key][:, :, skip : skip + take] for key in per_hint_frames + ) + if show_input_panel: + panel_tensors.append( + uint8_cthw_to_normalized_5d( + input_frames[:, frames_written : frames_written + take], + dtype=torch.float32, + ) + ) + if panel_tensors: + chunk = torch.cat([p.to(chunk) for p in panel_tensors] + [chunk], dim=-1) + # Post-processing is elementwise, so doing it per chunk is + # bitwise identical to doing it on the assembled clip -- and it + # halves the transfer, since uint8 crosses instead of bf16. + host_output[:, frames_written : frames_written + take].copy_( + postprocess_video_tensor(chunk), non_blocking=True + ) + frames_written += take + + # Only the tail is read as the next chunk's conditioning (one frame + # by default), so keep that slice rather than pinning the whole + # ~0.5 GB chunk across the next denoise. + keep = min(transfer_config.num_conditional_frames, output_video.shape[2]) + previous_output = output_video[:, :, -keep:].clone() if keep > 0 else None + + timer.mark_post_start() + + # The chunk copies are asynchronous, so the host buffer is not readable + # until the stream drains. The upload below would be correctly ordered + # without this, but the guardrail reads the frames on the CPU. + # + # Only false under the CPU-only unit tests, where the copies above were + # host-to-host and there is no stream to wait on (``CudaPhaseTimer`` + # disables itself on the same condition). + if torch.cuda.is_available(): + torch.cuda.current_stream().synchronize() + + # Same screening the other modes get: rank 0, post-processed frames, + # and a None result means the guardrail rejected the clip. With the + # debug panels enabled the caller's own control frames ride along in + # the same tensor and are screened too, which errs safe. Screening the + # host copy rather than a device one removes a full round trip: + # ``check_video_safety`` moves to CPU internally and returns on the + # input tensor's device. + video = host_output + if self.rank == 0 and use_guardrails and self.safety_checker is not None: + video = check_video_safety(video, self.safety_checker) + # Back to the device the other pipelines hand back, now that denoising + # has released its working set. + if video is not None: + video = video.to(self.device) + + timer.mark_end() + return timer.fill( + PipelineOutput( + video=video, + frame_rate=frame_rate, + ) + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py new file mode 100644 index 000000000000..902786e850ea --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transfer.py @@ -0,0 +1,558 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cosmos3 transfer inference helpers.""" + +from __future__ import annotations + +from collections.abc import Mapping +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import torch + +from tensorrt_llm._torch.visual_gen.triton_kernels import ( + bilateral_filter, + canny_edges, + resize_area_u8, + resize_cubic_u8, + resize_linear_u8, +) +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.media.decoding import decode_video_reference_window + +from .defaults import VIDEO_RES_SIZE_INFO + + +def find_closest_target_size(h: int, w: int, resolution: str | int) -> tuple[int, int]: + """Pick the ``resolution`` bucket whose aspect ratio best matches ``h/w``. + + Returns ``(target_w, target_h)`` so a source frame maps onto a supported + output size without distorting its aspect ratio. + """ + key = str(resolution) + if key not in VIDEO_RES_SIZE_INFO: + raise ValueError( + f"Unknown Cosmos3 transfer resolution={resolution!r}; " + f"expected one of {sorted(VIDEO_RES_SIZE_INFO)}." + ) + input_ratio = h / w + best_size = None + best_diff = float("inf") + for cand_w, cand_h in VIDEO_RES_SIZE_INFO[key].values(): + diff = abs(input_ratio - cand_h / cand_w) + if diff < best_diff: + best_diff = diff + best_size = (cand_w, cand_h) + assert best_size is not None + return best_size + + +# --------------------------------------------------------------------------- +# Transfer sampling + per-control-hint defaults, ported from vllm-omni's +# Cosmos3 transfer implementation; the guidance and edge/blur presets are +# empirically tuned per control modality: +# https://github.com/vllm-project/vllm-omni/blob/main/vllm_omni/diffusion/models/cosmos3/transfer.py +# --------------------------------------------------------------------------- + +# Supported control hints, in application order. +TRANSFER_HINT_KEYS: tuple[str, ...] = ("edge", "blur", "depth", "seg", "wsm") + +# Chunked-sampling defaults for long-video transfer. +TRANSFER_SAMPLE_DEFAULTS: dict[str, Any] = { + "num_video_frames_per_chunk": 93, # frames per autoregressive chunk (4k+1 form) + "num_conditional_frames": 1, # overlap frames reused from the previous chunk + "max_frames": 5000, # hard cap on total input frames + "show_control_condition": False, # debug: also emit the control frames + "show_input": False, # debug: also emit the input frames + "num_first_chunk_conditional_frames": 0, # chunk 0 has no prior chunk to condition on + "share_vision_temporal_positions": True, # control tokens reuse the video patches' positions + "emphasize_control_in_prompt": True, # name the active hints in the user prompt +} + +# Appended to the user prompt when ``emphasize_control_in_prompt`` is on, naming +# the active hint modalities so the model is told which control it is following. +# The system prompt is left alone, which keeps the text in the training +# distribution. +CONTROL_DIRECTIVE_TEMPLATE = ( + " Follow the {hints} control video precisely: shape, contour, silhouette," + " position, and motion of every visible structure must align with the {hints}" + " signal at every frame." +) + +# Per-hint guidance tuning: text guidance_scale, control_guidance, and flow_shift, +# chosen empirically per control modality. +TRANSFER_DEFAULTS: dict[str, dict[str, Any]] = { + "edge": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "blur": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "depth": {"guidance_scale": 3.0, "control_guidance": 1.5, "flow_shift": 10.0}, + "seg": { + "guidance_scale": 3.0, + "control_guidance": 2.0, + "flow_shift": 10.0, + }, # leans harder on control + # Precomputed control; control-only guidance (text off), 10 fps / 101-frame chunks. + "wsm": { + "guidance_scale": 1.0, + "control_guidance": 3.0, + "flow_shift": 10.0, + "num_frames": 101, + "fps": 10, + "num_video_frames_per_chunk": 101, + }, +} + +# Canny (lower, upper) hysteresis thresholds per edge-strength preset; +# higher preset = higher thresholds = sparser edges. +EDGE_PRESETS: dict[str, tuple[int, int]] = { + "none": (20, 50), + "very_low": (20, 50), + "low": (50, 100), + "medium": (100, 200), + "high": (200, 300), + "very_high": (300, 400), +} + +# Bilateral-blur strength presets: ``pre_blur_downscale`` (downscale applied +# before the bilateral filter) and ``downup`` (down/up-sample factor after); +# larger = blurrier. +BLUR_PRESETS: dict[str, dict[str, int]] = { + "none": {"pre_blur_downscale": 1, "downup": 1}, + "very_low": {"pre_blur_downscale": 1, "downup": 4}, + "low": {"pre_blur_downscale": 4, "downup": 4}, + "medium": {"pre_blur_downscale": 2, "downup": 10}, + "high": {"pre_blur_downscale": 1, "downup": 16}, + "very_high": {"pre_blur_downscale": 4, "downup": 16}, +} + +# Bilateral-filter parameters for the blur control, tuned at a 720p reference. +BILATERAL_REFERENCE_RESOLUTION = 720 # resolution the params below are tuned for +BILATERAL_D = 30 # filter diameter in pixels +BILATERAL_SIGMA_COLOR = 150 # color-space sigma +BILATERAL_SIGMA_SPACE = 100 # coordinate-space sigma +BILATERAL_ITERATIONS = 1 # number of bilateral passes + +# Control generation is batched over frames, so unwindowed its scratch scales +# with the whole clip: a 189-frame 704p clip peaked at 7.5 GiB for edge and +# 10.2 GiB for blur before denoising had allocated anything. No kernel reaches +# across the temporal axis, so a bounded window is bitwise identical and makes +# the scratch O(window) rather than O(clip). Measured at 704p: 32 frames holds +# both under 3.4 GiB for +5% (edge) / +0.7% (blur) wall time, and still gives +# the kernels enough parallelism to saturate the GPU. +CONTROL_FRAME_WINDOW = 32 + + +@dataclass +class Cosmos3TransferHint: + key: str + control: bytes | None = None + """Precomputed control clip as encoded MP4/AVI bytes, or None to auto-compute.""" + preset_edge_threshold: str = "medium" + preset_blur_strength: str = "medium" + + +@dataclass +class Cosmos3TransferConfig: + hints: dict[str, Cosmos3TransferHint] = field(default_factory=dict) + guidance_scale: float | None = None + control_guidance: float = 1.0 + control_guidance_interval: tuple[float, float] | None = None + flow_shift: float | None = None + num_video_frames_per_chunk: int = 93 + num_conditional_frames: int = 1 + max_frames: int = 5000 + show_control_condition: bool = False + show_input: bool = False + num_first_chunk_conditional_frames: int = 0 + share_vision_temporal_positions: bool = True + num_frames: int | None = None + fps: float | None = None + emphasize_control_in_prompt: bool = True + + @property + def ordered_hints(self) -> list[Cosmos3TransferHint]: + return [self.hints[key] for key in TRANSFER_HINT_KEYS if key in self.hints] + + def emphasized_prompt(self, prompt: str) -> str: + """The user prompt with the control-adherence directive appended. + + Returned unchanged when the directive is off or no hint is active. + """ + if not self.emphasize_control_in_prompt or not self.hints: + return prompt + hint_names = ", ".join(hint.key for hint in self.ordered_hints) + return prompt.rstrip() + CONTROL_DIRECTIVE_TEMPLATE.format(hints=hint_names) + + +def _as_interval(value: Any) -> tuple[float, float] | None: + if value is None: + return None + if isinstance(value, str): + value = [item.strip() for item in value.split(",") if item.strip()] + if not isinstance(value, (list, tuple)) or len(value) != 2: + raise ValueError( + "Cosmos3 transfer control_guidance_interval must contain exactly two values." + ) + lo, hi = float(value[0]), float(value[1]) + if lo > hi: + raise ValueError( + f"Cosmos3 transfer control_guidance_interval must be ordered as [lo, hi], got {(lo, hi)}." + ) + return lo, hi + + +def _extra_or_default(extra_params: dict, key: str, default: Any = None) -> Any: + value = extra_params.get(key, None) + return default if value is None else value + + +def resolve_transfer_config( + extra_params: dict, req_params: Any, prompt_data: Any = None +) -> Cosmos3TransferConfig | None: + hints: dict[str, Cosmos3TransferHint] = {} + for key in TRANSFER_HINT_KEYS: + raw = extra_params.get(key, None) + if raw is None: + continue + if raw is True: + raw = {} + elif isinstance(raw, bytes): + raw = {"control": raw} + if isinstance(raw, str | Path): + raise ValueError( + f"Cosmos3 transfer hint '{key}' must carry encoded control bytes, not a path " + f"({raw!r}); read the file client-side (Path(control).read_bytes()), as the " + "'video' extra param does." + ) + if not isinstance(raw, Mapping): + # ValueError, not TypeError: the worker's error classifier maps + # ValueError to a client error (400) and anything else to an + # unclassified server fault (500). A malformed hint is the caller's + # to fix, so it must not be reported as our failure. + raise ValueError( + f"Cosmos3 transfer hint '{key}' must be an object, encoded control bytes, or " + f"true; got {type(raw)!r}." + ) + if raw.get("control_path") is not None: + raise ValueError( + f"Cosmos3 transfer hint '{key}' no longer accepts 'control_path'; pass the " + "encoded control clip as 'control' bytes (Path(control).read_bytes())." + ) + control = raw.get("control") + if control is not None and not isinstance(control, bytes): + raise ValueError( + f"Cosmos3 transfer hint '{key}' control must be encoded MP4/AVI bytes, got " + f"{type(control)!r}." + ) + hints[key] = Cosmos3TransferHint( + key=key, + control=control, + preset_edge_threshold=str(raw.get("preset_edge_threshold") or "medium").lower(), + preset_blur_strength=str(raw.get("preset_blur_strength") or "medium").lower(), + ) + + if not hints: + transfer_only = ( + "control_guidance", + "control_guidance_interval", + "num_video_frames_per_chunk", + "num_conditional_frames", + "num_first_chunk_conditional_frames", + "max_frames", + "show_control_condition", + "show_input", + "share_vision_temporal_positions", + "emphasize_control_in_prompt", + ) + if any(extra_params.get(key, None) for key in transfer_only): + raise ValueError( + "Cosmos3 transfer options were provided, but no transfer hint was selected." + ) + return None + + # `guidance_scale`, `frame_rate` and `num_frames` are advertised defaults the + # executor merges into every request, so only `model_fields_set` distinguishes + # a caller's value from a merged one. + specified = getattr(req_params, "model_fields_set", frozenset()) + + # Stays None unless the caller asked for a value, so the single-hint preset + # below can apply. Reading `req_params.guidance_scale` unconditionally would + # capture the executor-merged default and make every request look explicit, + # which pins transfer to the generic scale and never reaches the tuned + # per-hint presets. With no preset (multi-hint), `_forward_transfer` falls + # back to the generic video default -- matching the reference, where the + # per-task table applies only to a single hint. + guidance_scale_user_set = ( + "guidance_scale" in specified or extra_params.get("guidance_scale", None) is not None + ) + request_guidance_scale = ( + getattr(req_params, "guidance_scale", None) if guidance_scale_user_set else None + ) + + config = Cosmos3TransferConfig( + hints=hints, + guidance_scale=request_guidance_scale, + control_guidance=_extra_or_default(extra_params, "control_guidance", 1.0), + control_guidance_interval=_as_interval( + _extra_or_default(extra_params, "control_guidance_interval", None) + ), + flow_shift=_extra_or_default(extra_params, "flow_shift", None), + num_video_frames_per_chunk=_extra_or_default( + extra_params, + "num_video_frames_per_chunk", + TRANSFER_SAMPLE_DEFAULTS["num_video_frames_per_chunk"], + ), + num_conditional_frames=_extra_or_default( + extra_params, + "num_conditional_frames", + TRANSFER_SAMPLE_DEFAULTS["num_conditional_frames"], + ), + max_frames=_extra_or_default( + extra_params, "max_frames", TRANSFER_SAMPLE_DEFAULTS["max_frames"] + ), + show_control_condition=_extra_or_default( + extra_params, + "show_control_condition", + TRANSFER_SAMPLE_DEFAULTS["show_control_condition"], + ), + show_input=_extra_or_default( + extra_params, "show_input", TRANSFER_SAMPLE_DEFAULTS["show_input"] + ), + num_first_chunk_conditional_frames=_extra_or_default( + extra_params, + "num_first_chunk_conditional_frames", + TRANSFER_SAMPLE_DEFAULTS["num_first_chunk_conditional_frames"], + ), + share_vision_temporal_positions=_extra_or_default( + extra_params, + "share_vision_temporal_positions", + TRANSFER_SAMPLE_DEFAULTS["share_vision_temporal_positions"], + ), + emphasize_control_in_prompt=_extra_or_default( + extra_params, + "emphasize_control_in_prompt", + TRANSFER_SAMPLE_DEFAULTS["emphasize_control_in_prompt"], + ), + # `num_frames` and `frame_rate` are request fields, not extra params, so + # they are read from `req_params` alone: `extra_params` spellings of them + # would be a second name for the same knob with different precedence, and + # `validate_visual_gen_params` rejects them as undeclared keys anyway. + num_frames=getattr(req_params, "num_frames", None), + # Only a caller-supplied frame_rate seeds this. Taking the request's + # value unconditionally would capture the executor-merged default, + # which then wins in _forward_transfer over a rate the pipeline + # inferred from the source -- silently pinning every transfer to 24. + fps=getattr(req_params, "frame_rate", None) if "frame_rate" in specified else None, + ) + + if len(hints) == 1: + hint_key = next(iter(hints)) + for field_name, default_value in TRANSFER_DEFAULTS[hint_key].items(): + if field_name == "guidance_scale": + user_set = guidance_scale_user_set + elif field_name == "flow_shift": + user_set = extra_params.get("flow_shift", None) is not None + elif field_name == "fps": + user_set = "frame_rate" in specified + elif field_name == "num_frames": + user_set = "num_frames" in specified + else: + user_set = extra_params.get(field_name, None) is not None + if not user_set: + setattr(config, field_name, default_value) + + if config.num_video_frames_per_chunk <= 0: + raise ValueError("Cosmos3 transfer num_video_frames_per_chunk must be positive.") + if config.num_conditional_frames < 0: + raise ValueError("Cosmos3 transfer num_conditional_frames must be non-negative.") + if config.max_frames <= 0: + raise ValueError("Cosmos3 transfer max_frames must be positive.") + if config.num_first_chunk_conditional_frames < 0: + raise ValueError( + "Cosmos3 transfer num_first_chunk_conditional_frames must be non-negative." + ) + for hint in hints.values(): + if hint.key == "edge" and hint.preset_edge_threshold not in EDGE_PRESETS: + raise ValueError(f"Unsupported Cosmos3 edge preset: {hint.preset_edge_threshold!r}.") + if hint.key == "blur" and hint.preset_blur_strength not in BLUR_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {hint.preset_blur_strength!r}.") + return config + + +def decode_media_to_uint8_cthw( + data: bytes, *, height: int, width: int, max_frames: int, device: torch.device +) -> torch.Tensor: + """Decode encoded MP4/AVI ``data`` to uint8 ``[3, T, H, W]`` frames on ``device``. + + The decoder resizes to ``(height, width)`` before retaining each frame, so a + high-resolution control never materializes at full size. + """ + if not isinstance(data, bytes): + raise ValueError( + f"Cosmos3 transfer media must be encoded MP4/AVI bytes, got {type(data)!r}." + ) + max_frames = int(max_frames) + if max_frames < 1: + raise ValueError(f"Cosmos3 transfer max_frames must be positive, got {max_frames}.") + frames_thwc = decode_video_reference_window( + data, + first_frame=0, + last_frame=max_frames - 1, + target_h=int(height), + target_w=int(width), + device=device, + ) + return frames_thwc.permute(3, 0, 1, 2).contiguous() + + +def uint8_cthw_to_normalized_5d(frames: torch.Tensor, *, dtype: torch.dtype) -> torch.Tensor: + """Normalize uint8 control frames into the transfer model's input tensor. + + ``frames``: uint8 ``[3, T, H, W]`` -> ``[1, 3, T, H, W]`` in ``[-1, 1]`` + (``/127.5 - 1``), the batched normalized form the control encoder consumes. + """ + if frames.ndim != 4 or frames.shape[0] != 3: + raise ValueError( + f"Cosmos3 transfer frames must have shape [3, T, H, W], got {tuple(frames.shape)}." + ) + return frames.to(dtype=dtype).div(127.5).sub(1.0).unsqueeze(0).contiguous() + + +@nvtx_range("make_edge_control", color="blue") +def make_edge_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + """Canny edge control: uint8 ``[3, T, H, W]`` CUDA -> the same shape. + + The single-channel edge map is broadcast across RGB, since the control + encoder consumes three channels. + """ + try: + lower, upper = EDGE_PRESETS[preset] + except KeyError as exc: + raise ValueError(f"Unsupported Cosmos3 edge preset: {preset!r}.") from exc + edges = torch.empty_like(frames) + for start in range(0, frames.shape[1], CONTROL_FRAME_WINDOW): + stop = min(start + CONTROL_FRAME_WINDOW, frames.shape[1]) + # assigning [t, H, W] into [3, t, H, W] broadcasts the map across RGB + edges[:, start:stop] = canny_edges(frames[:, start:stop], lower, upper) + return edges + + +def _scale_for_bilateral_resolution(value: float, longest_side: int) -> float: + if longest_side <= 0: + return value + return value * (longest_side / BILATERAL_REFERENCE_RESOLUTION) + + +def _scaled_bilateral_params(height: int, width: int) -> tuple[int, float, float]: + longest_side = int(max(height, width)) + diameter = max(1, int(round(_scale_for_bilateral_resolution(float(BILATERAL_D), longest_side)))) + if diameter % 2 == 0: + diameter += 1 + sigma_color = max( + 1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_COLOR), longest_side) + ) + sigma_space = max( + 1.0, _scale_for_bilateral_resolution(float(BILATERAL_SIGMA_SPACE), longest_side) + ) + return diameter, sigma_color, sigma_space + + +@nvtx_range("make_blur_control", color="blue") +def make_blur_control(frames: torch.Tensor, preset: str) -> torch.Tensor: + """Bilateral-blur control: uint8 ``[3, T, H, W]`` CUDA -> the same shape. + + Edge-preserving blur at ``pre_blur_downscale`` resolution, then a + ``downup`` round trip that discards high-frequency detail. + """ + preset = preset.lower() + if preset not in BLUR_PRESETS: + raise ValueError(f"Unsupported Cosmos3 blur preset: {preset!r}.") + if preset == "none": + return frames.clone() + + _, t, h, w = frames.shape + blur_params = BLUR_PRESETS[preset] + pre_blur_factor = max(1, int(blur_params["pre_blur_downscale"])) + downup_factor = max(1, int(blur_params["downup"])) + + blurred = torch.empty_like(frames) + for start in range(0, t, CONTROL_FRAME_WINDOW): + stop = min(start + CONTROL_FRAME_WINDOW, t) + # The kernels are channels-last, so permute once per window rather than + # per frame; windowing the whole chain also bounds the tensors handed + # between its stages, not just each stage's own scratch. + result = frames[:, start:stop].permute(1, 2, 3, 0).contiguous() + if pre_blur_factor > 1: + result = resize_area_u8(result, pre_blur_factor) + diameter, sigma_color, sigma_space = _scaled_bilateral_params( + result.shape[1], result.shape[2] + ) + for _ in range(BILATERAL_ITERATIONS): + result = bilateral_filter(result, diameter, sigma_color, sigma_space) + if pre_blur_factor > 1: + result = resize_linear_u8(result, w, h) + if downup_factor > 1: + result = resize_cubic_u8(result, max(1, w // downup_factor), max(1, h // downup_factor)) + result = resize_cubic_u8(result, w, h) + blurred[:, start:stop] = result.permute(3, 0, 1, 2) + return blurred + + +def load_or_compute_control_frames( + hint: Cosmos3TransferHint, + *, + height: int, + width: int, + max_frames: int, + input_frames: torch.Tensor | None, + device: torch.device, +) -> torch.Tensor: + """Decode a hint's precomputed control, or derive one from the input video.""" + if hint.control is not None: + return decode_media_to_uint8_cthw( + hint.control, height=height, width=width, max_frames=max_frames, device=device + ) + # Generated controls stay on the input frames' device, which is where the + # decoded ones already are, so hints never mix devices. + if hint.key == "edge": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'edge' requires either a video input for on-the-fly " + "control generation or precomputed control bytes." + ) + return make_edge_control(input_frames[:, :max_frames], hint.preset_edge_threshold) + if hint.key == "blur": + if input_frames is None: + raise ValueError( + "Cosmos3 transfer hint 'blur' requires either a video input for on-the-fly " + "control generation or precomputed control bytes." + ) + return make_blur_control(input_frames[:, :max_frames], hint.preset_blur_strength) + raise ValueError( + f"Cosmos3 transfer hint '{hint.key}' requires precomputed control bytes; " + "on-the-fly generation is supported only for edge and blur." + ) + + +def pad_temporal_frames(frames: torch.Tensor, target_frames: int) -> torch.Tensor: + if frames.ndim != 4: + raise ValueError( + f"Cosmos3 transfer frames must have shape [C, T, H, W], got {tuple(frames.shape)}." + ) + target_frames = int(target_frames) + if target_frames <= 0: + raise ValueError("Cosmos3 transfer target frame count must be positive.") + if frames.shape[1] >= target_frames: + return frames + if frames.shape[1] == 0: + raise ValueError("Cannot pad an empty Cosmos3 transfer frame tensor.") + padded = frames + while padded.shape[1] < target_frames: + pad_len = min(padded.shape[1] - 1, target_frames - padded.shape[1]) + if pad_len <= 0: + pad_frame = padded[:, -1:].repeat(1, target_frames - padded.shape[1], 1, 1) + padded = torch.cat([padded, pad_frame], dim=1) + break + padded = torch.cat([padded, padded.flip(dims=[1])[:, :pad_len]], dim=1) + return padded diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py index dbce5ca6ea45..ebe83e605626 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1109,6 +1109,8 @@ def _compute_rope_freqs( fps: float | None, device: torch.device, dtype: torch.dtype, + num_vision_items: int = 1, + share_vision_temporal_positions: bool = False, ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: """Compute mRoPE cos/sin for UND (text) and GEN (visual) pathways.""" B = text_mask.shape[0] @@ -1121,16 +1123,39 @@ def _compute_rope_freqs( for b in range(B): real_len = int(text_lengths[b].item()) t_pos, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) - v_pos, _ = compute_mrope_position_ids_vision( - T, - Hp, - Wp, - temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, - fps=effective_fps, - base_fps=self.base_fps, - temporal_compression_factor=self.temporal_compression_factor, - enable_fps_modulation=self.enable_fps_modulation, - ) + media_temporal_offset = t_offset + self.unified_3d_mrope_temporal_modality_margin + gen_positions = [] + if num_vision_items == 1 or share_vision_temporal_positions: + v_pos, _ = compute_mrope_position_ids_vision( + T, + Hp, + Wp, + temporal_offset=media_temporal_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.extend([v_pos] * num_vision_items) + else: + vision_offset: int | float = media_temporal_offset + for _ in range(num_vision_items): + v_pos, vision_offset = compute_mrope_position_ids_vision( + T, + Hp, + Wp, + temporal_offset=vision_offset, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + gen_positions.append(v_pos) + + pos_dtype = gen_positions[0].dtype + for pos in gen_positions[1:]: + pos_dtype = torch.promote_types(pos_dtype, pos.dtype) + v_pos = torch.cat([pos.to(pos_dtype) for pos in gen_positions], dim=1) if real_len < S_text: t_pos = torch.cat( [t_pos, torch.zeros(3, S_text - real_len, dtype=t_pos.dtype)], dim=1 @@ -1218,6 +1243,8 @@ def forward( fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, audio_latents: Optional[torch.Tensor] = None, + control_latents: list[torch.Tensor] | tuple[torch.Tensor, ...] | torch.Tensor | None = None, + transfer_share_vision_temporal_positions: bool = True, **kwargs, ) -> "TransformerOutput": """ @@ -1241,6 +1268,14 @@ def forward( When provided, audio tokens are appended to the generation sequence and an audio velocity is returned alongside the video velocity. Requires ``audio_gen=True`` in the pretrained config. + control_latents: Optional transfer-control latents. Controls are + clean (un-noised) vision context and are packed before the + noisy target; their outputs are discarded. + transfer_share_vision_temporal_positions: When True (the default), + control tokens reuse the target frames' mRoPE temporal + coordinates, so a control patch sits at zero displacement from + the output patch it governs and the shared coordinate carries + the correspondence. When False they take their own positions. Returns: TransformerOutput with video (and image alias) always set. @@ -1263,6 +1298,13 @@ def forward( time_embed = self.time_embedder((raw_timestep * self.timestep_scale)) time_embed = time_embed.to(hidden_states.dtype) + if control_latents is None: + control_lantent_list: list[torch.Tensor] = [] + elif isinstance(control_latents, torch.Tensor): + control_lantent_list = [control_latents] + else: + control_lantent_list = list(control_latents) + if noisy_frame_mask is not None: # Build per-token mask from per-frame mask. # noisy_frame_mask: [B, 1, T, 1, 1] → token mask: [B, T*Hp*Wp, 1] @@ -1286,6 +1328,8 @@ def forward( fps, hidden_states.device, hidden_states.dtype, + num_vision_items=len(control_lantent_list) + 1, + share_vision_temporal_positions=transfer_share_vision_temporal_positions, ) cached_kv_full = self.language_model( text_ids, @@ -1317,6 +1361,15 @@ def forward( # --- Audio token injection ------------------------------------------------- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp T_audio = 0 + T_control = 0 + hidden_controls: list[torch.Tensor] = [] + has_control = len(control_lantent_list) > 0 + + if has_control and audio_latents is not None: + raise ValueError( + "Cosmos3 transfer control latents cannot be combined with sound latents" + ) + if audio_latents is not None and self.audio_gen: T_audio = audio_latents.shape[2] hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) @@ -1336,6 +1389,22 @@ def forward( torch.cat([cos_v, cos_a], dim=1), torch.cat([sin_v, sin_a], dim=1), ) + elif has_control: + for idx, control in enumerate(control_lantent_list): + if control.shape != hidden_states.shape: + raise ValueError( + "Cosmos3 transfer control latent shape must match target latent shape: " + f"control[{idx}]={tuple(control.shape)}, target={tuple(hidden_states.shape)}." + ) + hidden_control = self.vae2llm( + self.patchify( + control.to(device=hidden_gen.device, dtype=hidden_gen.dtype), T, H, W + ) + ) + hidden_controls.append(hidden_control) + T_control += hidden_control.shape[1] + hidden_gen = torch.cat([*hidden_controls, hidden_gen], dim=1) + freqs_gen_combined = self.cached_freqs_gen else: freqs_gen_combined = self.cached_freqs_gen # -------------------------------------------------------------------------- @@ -1374,9 +1443,13 @@ def forward( hidden_gen = self.norm_moe_gen(hidden_gen) # --- Decode video velocity ------------------------------------------------ - video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) + video_vel = self.unpatchify( + self.llm2vae(hidden_gen[:, T_control : T_control + T_vid_tokens]), T, H, W + ) # --- Decode audio velocity (if requested) --------------------------------- + # Control latents and audio are mutually exclusive (guarded above), so the + # audio span always starts right after the video tokens. audio_vel = None if T_audio > 0 and audio_latents is not None and self.audio_gen: # hidden_gen[:, T_vid_tokens:] → [B, T_audio, hidden_size] diff --git a/tensorrt_llm/_torch/visual_gen/triton_kernels/__init__.py b/tensorrt_llm/_torch/visual_gen/triton_kernels/__init__.py new file mode 100644 index 000000000000..48423e29f4cc --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/triton_kernels/__init__.py @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU image-processing kernels for VisualGen control preprocessing. + +Implementations of published algorithms -- Canny edge detection (Canny, +1986), the bilateral filter (Tomasi & Manduchi, 1998), and bilinear / +area-average / bicubic interpolation in their standard fixed-point forms -- +written so control frames can be derived on the GPU alongside the pipeline +that consumes them. + +``reference`` holds torch-op implementations of the same arithmetic, used as +the executable specification the kernels are asserted against; it is not +imported here because nothing in the inference path should reach for it. +""" + +from .bilateral import bilateral_filter +from .canny import canny_edges +from .resize import resize_area_u8, resize_cubic_u8, resize_linear_u8 + +__all__ = [ + "bilateral_filter", + "canny_edges", + "resize_area_u8", + "resize_cubic_u8", + "resize_linear_u8", +] diff --git a/tensorrt_llm/_torch/visual_gen/triton_kernels/bilateral.py b/tensorrt_llm/_torch/visual_gen/triton_kernels/bilateral.py new file mode 100644 index 000000000000..0a231e081e4e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/triton_kernels/bilateral.py @@ -0,0 +1,166 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Bilateral filter (Tomasi & Manduchi, 1998) for uint8 RGB clips. + +Contract, in the widely-used formulation this reproduces: + +- ``radius = max(d // 2, 1)``; the support is the CIRCLE + ``sqrt(i^2 + j^2) <= radius``, not the enclosing square. +- Space weight ``exp(r^2 * -0.5 / sigma_space^2)``; colour weight is a lookup + over the L1 channel distance ``|db| + |dg| + |dr|``, i.e. + ``exp(k^2 * -0.5 / sigma_color^2)``. +- Borders reflect without repeating the edge pixel. +- Accumulation and the final divide are float32, with taps visited in + row-major circular order so the sums associate identically to the reference. + +Unlike the resize and Canny paths, this op cannot be made bitwise reproducible +against an arbitrary CPU SIMD implementation: float32 tap-summation order +decides pixels that land on a .5 rounding tie, so two vectorisations of the +same formula legitimately differ by 1 LSB on a small fraction of pixels. The +enforced contract is therefore bitwise equality with the torch reference in +``reference.py``, which fixes the summation order. +""" + +import math + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl +from triton.language.extra import libdevice + +from .resize import _check_frames + + +def circle_offsets(radius: int, gauss_space_coeff: float): + """Taps of the circular support in row-major order, with space weights. + + Shared with the reference implementation so both see identical tables and + any mismatch is attributable to the arithmetic, not the setup. + """ + dys, dxs, sws = [], [], [] + for i in range(-radius, radius + 1): + for j in range(-radius, radius + 1): + r = math.sqrt(float(i * i + j * j)) + if r > radius: + continue + dys.append(i) + dxs.append(j) + # evaluated in double, stored as float32 + sws.append(math.exp(r * r * gauss_space_coeff)) + return dys, dxs, sws + + +def color_lut(channels: int, gauss_color_coeff: float, device: torch.device) -> torch.Tensor: + """Colour weight lookup over the L1 channel distance, float32.""" + return ( + (torch.arange(256 * channels, dtype=torch.float64, device=device) ** 2 * gauss_color_coeff) + .exp() + .to(torch.float32) + ) + + +def reflect_pad(frames: torch.Tensor, radius: int) -> torch.Tensor: + """uint8 ``[T, H, W, C]`` -> float32 ``[T, H+2r, W+2r, C]``, reflected.""" + x = frames.permute(0, 3, 1, 2).to(torch.float32) + return F.pad(x, (radius,) * 4, mode="reflect").permute(0, 2, 3, 1).contiguous() + + +@triton.jit +def _bilateral_kernel( + src_ptr, + dst_ptr, + lut_ptr, + dy_ptr, + dx_ptr, + sw_ptr, + T, + H, + W, + Hp, + Wp, + K, + radius, + BLOCK: tl.constexpr, +): + # (x-block, y, t) grid: no runtime integer division in the hot path + y = tl.program_id(1) + t = tl.program_id(2) + x = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = x < W + offs = (t * H + y) * W + x + + cbase = ((t * Hp + y + radius) * Wp + x + radius) * 3 + cb = tl.load(src_ptr + cbase + 0, mask=m, other=0.0) + cg = tl.load(src_ptr + cbase + 1, mask=m, other=0.0) + cr = tl.load(src_ptr + cbase + 2, mask=m, other=0.0) + + sb = tl.zeros((BLOCK,), dtype=tl.float32) + sg = tl.zeros((BLOCK,), dtype=tl.float32) + sr = tl.zeros((BLOCK,), dtype=tl.float32) + ws = tl.zeros((BLOCK,), dtype=tl.float32) + + for k in range(K): # K is runtime, so this stays a serial loop in tap order + dy = tl.load(dy_ptr + k) + dx = tl.load(dx_ptr + k) + sw = tl.load(sw_ptr + k) + nbase = ((t * Hp + y + radius + dy) * Wp + x + radius + dx) * 3 + nb = tl.load(src_ptr + nbase + 0, mask=m, other=0.0) + ng = tl.load(src_ptr + nbase + 1, mask=m, other=0.0) + nr = tl.load(src_ptr + nbase + 2, mask=m, other=0.0) + dist = tl.abs(nb - cb) + tl.abs(ng - cg) + tl.abs(nr - cr) + w = tl.load(lut_ptr + dist.to(tl.int32), mask=m, other=0.0) * sw + # mul_rn/add_rn keep the accumulation unfused: ptxas would contract + # nb*w + sb into an FMA, which rounds differently from the reference's + # separate mul and add on .5-tie pixels. + sb = libdevice.add_rn(libdevice.mul_rn(nb, w), sb) + sg = libdevice.add_rn(libdevice.mul_rn(ng, w), sg) + sr = libdevice.add_rn(libdevice.mul_rn(nr, w), sr) + ws = libdevice.add_rn(w, ws) # += would let ptxas fuse the lut*sw mul into an fma + + obase = offs * 3 + # div_rn: Triton's `/` is not guaranteed IEEE correctly-rounded on f32, + # and a 1-ulp quotient difference flips torch.round on .5-tie pixels. + tl.store(dst_ptr + obase + 0, libdevice.div_rn(sb, ws), mask=m) + tl.store(dst_ptr + obase + 1, libdevice.div_rn(sg, ws), mask=m) + tl.store(dst_ptr + obase + 2, libdevice.div_rn(sr, ws), mask=m) + + +def bilateral_filter( + frames: torch.Tensor, d: int, sigma_color: float, sigma_space: float +) -> torch.Tensor: + """Bilateral filter over a clip: uint8 ``[T, H, W, 3]`` CUDA -> same shape.""" + _check_frames(frames, "bilateral_filter") + if frames.shape[-1] != 3: + raise ValueError(f"bilateral_filter expects [T, H, W, 3], got shape={tuple(frames.shape)}") + T, H, W, C = frames.shape + dev = frames.device + radius = max(d // 2, 1) + + lut = color_lut(C, -0.5 / (sigma_color * sigma_color), dev) + dys, dxs, sws = circle_offsets(radius, -0.5 / (sigma_space * sigma_space)) + dy = torch.tensor(dys, dtype=torch.int32, device=dev) + dx = torch.tensor(dxs, dtype=torch.int32, device=dev) + sw = torch.tensor(sws, dtype=torch.float32, device=dev) + + src = reflect_pad(frames, radius) + dst = torch.empty(T * H * W * 3, dtype=torch.float32, device=dev) + BLOCK = 256 + _bilateral_kernel[(triton.cdiv(W, BLOCK), H, T)]( + src.reshape(-1), + dst, + lut, + dy, + dx, + sw, + T, + H, + W, + src.shape[1], + src.shape[2], + len(dys), + radius, + BLOCK=BLOCK, + ) + return torch.round(dst.reshape(T, H, W, 3)).clamp(0, 255).to(torch.uint8) diff --git a/tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py b/tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py new file mode 100644 index 000000000000..463e42377e5f --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/triton_kernels/canny.py @@ -0,0 +1,197 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Canny edge detection (Canny, 1986) over a clip of uint8 frames. + +Two elementwise kernels plus a hysteresis fixpoint loop: + +1. ``_grad_kernel``: Sobel 3x3 with replicate borders per channel, then + per-pixel channel selection by max ``|dx| + |dy|`` (strict ``>``, so the + lowest channel index wins ties); writes int32 dx/dy/magnitude maps. +2. ``_nms_kernel``: direction binning in Q15 fixed point against + ``tan(22.5deg)``, non-maximum suppression with per-direction tie-breaks, + then the double threshold; writes a uint8 map (2 = strong, 1 = weak). +3. Hysteresis grows strong seeds through weak pixels, 8-connected. This stays + a torch max-pool fixpoint loop: it is data-dependent global propagation, + and it dominates the op's cost, so it is the next optimization target. + +The gradient magnitude is L1 (no L2 gradient), and every stage is integer +arithmetic, so bit-exactness here has no floating-point caveats. +""" + +import torch +import torch.nn.functional as F +import triton +import triton.language as tl + +from .resize import _check_frames + +_BLOCK = 256 +# tl.constexpr, not a bare int: a @triton.jit kernel may only read globals that +# are declared this way, and the NMS kernel uses this directly. +_TG22 = tl.constexpr(13573) # round(tan(22.5deg) * 2**15) +# Output rows per NMS program. The launch grid divides by this, so the two +# must not drift apart -- a mismatch silently mis-strides the row band. +_NMS_ROWS = 4 + + +@triton.jit +def _grad_kernel(src, bdx, bdy, mag, SC, H, W, C: tl.constexpr, BLOCK: tl.constexpr): + # (x-block, y, t) grid: no runtime integer division (idiv dominated the SM + # in the flat-index variant) + y = tl.program_id(1) + t = tl.program_id(2) + x = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = x < W + offs = (t * H + y) * W + x + + ym1 = tl.maximum(y - 1, 0) + yp1 = tl.minimum(y + 1, H - 1) + xm1 = tl.maximum(x - 1, 0) + xp1 = tl.minimum(x + 1, W - 1) + + best = tl.zeros((BLOCK,), dtype=tl.int32) - 1 # any real mag (>=0) beats it + dx = tl.zeros((BLOCK,), dtype=tl.int32) + dy = tl.zeros((BLOCK,), dtype=tl.int32) + # SC is dim 0's stride rather than a frame count, so a clip windowed along T + # is read in place. c is a static_range constant, so c*SC folds to 0/SC/2*SC + # -- the same address arithmetic the dense form compiled to. + for c in tl.static_range(C): + p = src + c * SC + t * H * W + v00 = tl.load(p + ym1 * W + xm1, mask=m, other=0).to(tl.int32) + v01 = tl.load(p + ym1 * W + x, mask=m, other=0).to(tl.int32) + v02 = tl.load(p + ym1 * W + xp1, mask=m, other=0).to(tl.int32) + v10 = tl.load(p + y * W + xm1, mask=m, other=0).to(tl.int32) + v12 = tl.load(p + y * W + xp1, mask=m, other=0).to(tl.int32) + v20 = tl.load(p + yp1 * W + xm1, mask=m, other=0).to(tl.int32) + v21 = tl.load(p + yp1 * W + x, mask=m, other=0).to(tl.int32) + v22 = tl.load(p + yp1 * W + xp1, mask=m, other=0).to(tl.int32) + dx_c = (v02 + 2 * v12 + v22) - (v00 + 2 * v10 + v20) + dy_c = (v20 + 2 * v21 + v22) - (v00 + 2 * v01 + v02) + mag_c = tl.abs(dx_c) + tl.abs(dy_c) + take = mag_c > best + best = tl.where(take, mag_c, best) + dx = tl.where(take, dx_c, dx) + dy = tl.where(take, dy_c, dy) + + # int32 intermediates on purpose: int16 halves the DRAM traffic but Triton + # unpacks every sub-word load through PRMT, which costs more ALU time than + # the saved bytes buys (91 -> 114 us measured) -- nms is issue-bound, not + # bandwidth-bound. + tl.store(bdx + offs, dx, mask=m) + tl.store(bdy + offs, dy, mask=m) + tl.store(mag + offs, best, mask=m) + + +@triton.jit +def _nms_kernel(bdx, bdy, mag, out, lo, hi, H, W, R: tl.constexpr, BLOCK: tl.constexpr): + # R output rows per program with register rotation of the 3-row magnitude + # band: 3*(R+1) mag loads instead of 9*R -- nms is issue-bound (ALU pipe == + # SM throughput), so fewer load issues is the lever that matters. + y0 = tl.program_id(1) * R + t = tl.program_id(2) + x = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = x < W + ok_l = m & (x > 0) + ok_r = m & (x < W - 1) + base = (t * H + y0) * W + x + + # masked-out taps read 0, which is exactly the zero-padded border + up_ok = y0 > 0 + pC = tl.load(mag + base - W, mask=m & up_ok, other=0) + pL = tl.load(mag + base - W - 1, mask=ok_l & up_ok, other=0) + pR = tl.load(mag + base - W + 1, mask=ok_r & up_ok, other=0) + cC = tl.load(mag + base, mask=m, other=0) + cL = tl.load(mag + base - 1, mask=ok_l, other=0) + cR = tl.load(mag + base + 1, mask=ok_r, other=0) + + # software pipeline: row r's down-band and gradients load one iteration + # ahead, so the current iteration's compare chain covers their latency + # (without this, load -> compute -> load serializes: 61% long_scoreboard) + dn0 = y0 < H - 1 + nC = tl.load(mag + base + W, mask=m & dn0, other=0) + nL = tl.load(mag + base + W - 1, mask=ok_l & dn0, other=0) + nR = tl.load(mag + base + W + 1, mask=ok_r & dn0, other=0) + dx = tl.load(bdx + base, mask=m, other=0) + dy = tl.load(bdy + base, mask=m, other=0) + + for r in tl.static_range(R): + yr = y0 + r + offs = base + r * W + mr = m & (yr < H) + if r < R - 1: + dn_ok = yr + 1 < H - 1 + mr1 = m & (yr + 1 < H) + fC = tl.load(mag + offs + 2 * W, mask=m & dn_ok, other=0) + fL = tl.load(mag + offs + 2 * W - 1, mask=ok_l & dn_ok, other=0) + fR = tl.load(mag + offs + 2 * W + 1, mask=ok_r & dn_ok, other=0) + fdx = tl.load(bdx + offs + W, mask=mr1, other=0) + fdy = tl.load(bdy + offs + W, mask=mr1, other=0) + ax = tl.abs(dx) + y15 = tl.abs(dy) << 15 + tg22 = ax * _TG22 + tg67 = tg22 + (ax << 16) + horiz = y15 < tg22 + vert = (~horiz) & (y15 > tg67) + diag = ~(horiz | vert) + s_pos = (dx ^ dy) >= 0 + + # select the direction's two neighbours, then compare once: + # horiz: cC > cL && cC >= cR vert: cC > pC && cC >= nC + # diag+: cC > pL && cC > nR diag-: cC > pR && cC > nL + # (c == nb only survives for the non-diagonal >= comparisons) + na = tl.where(horiz, cL, tl.where(vert, pC, tl.where(s_pos, pL, pR))) + nb = tl.where(horiz, cR, tl.where(vert, nC, tl.where(s_pos, nR, nL))) + keep = (cC > na) & ((cC > nb) | ((~diag) & (cC == nb))) + res = tl.where(keep & (cC > hi), 2, tl.where(keep & (cC > lo), 1, 0)) + tl.store(out + offs, res.to(tl.uint8), mask=mr) + + if r < R - 1: + pL, pC, pR = cL, cC, cR + cL, cC, cR = nL, nC, nR + nL, nC, nR = fL, fC, fR + dx, dy = fdx, fdy + + +def _hysteresis(strong: torch.Tensor, weak: torch.Tensor) -> torch.Tensor: + """Grow strong seeds through weak pixels to a fixpoint, 8-connected. + + Several growth steps per convergence check keep the host syncs down while + staying exact, since the growth is monotone. + """ + out = strong + while True: + prev = out + for _ in range(4): + out = torch.minimum(F.max_pool2d(out, 3, 1, 1), weak) + if torch.equal(out, prev): + break + return out + + +def canny_edges(frames: torch.Tensor, low: int, high: int) -> torch.Tensor: + """Canny over a clip: uint8 ``[C, T, H, W]`` CUDA -> uint8 ``[T, H, W]``. + + Dim 0 may be strided, so a caller windowing a longer clip along T can pass + ``frames[:, start:stop]`` directly instead of materializing it. + """ + _check_frames(frames, "canny_edges", layout="[C, T, H, W]", allow_outer_stride=True) + C, T, H, W = frames.shape + dev = frames.device + + bdx = torch.empty(T, H, W, dtype=torch.int32, device=dev) + bdy = torch.empty_like(bdx) + mag = torch.empty_like(bdx) + _grad_kernel[(triton.cdiv(W, _BLOCK), H, T)]( + frames, bdx, bdy, mag, frames.stride(0), H, W, C=C, BLOCK=_BLOCK + ) + + cmap = torch.empty(T, H, W, dtype=torch.uint8, device=dev) + _nms_kernel[(triton.cdiv(W, _BLOCK), triton.cdiv(H, 4), T)]( + bdx, bdy, mag, cmap, low, high, H, W, R=4, BLOCK=_BLOCK + ) + + out = _hysteresis( + (cmap == 2).to(torch.float32)[:, None], (cmap >= 1).to(torch.float32)[:, None] + ) + return (out[:, 0] > 0).to(torch.uint8) * 255 diff --git a/tensorrt_llm/_torch/visual_gen/triton_kernels/reference.py b/tensorrt_llm/_torch/visual_gen/triton_kernels/reference.py new file mode 100644 index 000000000000..f1c1446cde1a --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/triton_kernels/reference.py @@ -0,0 +1,180 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Torch-op reference implementations of the control-generation kernels. + +These are the executable specification of the arithmetic the Triton kernels +implement: readable, obviously-correct-by-inspection, and slow. Nothing in +the inference path calls them -- they exist so +``tests/unittest/_torch/visual_gen/test_control_kernels.py`` can assert the +kernels bitwise, and so a future change to a kernel has something precise to +be checked against. + +The axis tables, tap offsets and lookup tables are imported from the kernel +modules rather than rebuilt here: a mismatch should point at the arithmetic, +not at two subtly different table setups. For the same reason Canny's +hysteresis is shared, since it is a torch loop in both paths. +""" + +import numpy as np +import torch +import torch.nn.functional as F + +from .bilateral import circle_offsets, color_lut, reflect_pad +from .canny import _hysteresis +from .resize import _COEF_BITS, _COEF_SCALE, _check_frames, _cubic_axis, _linear_axis + +_SHIFT = 15 +_TG22 = 13573 # round(tan(22.5deg) * 2**15) +_SOBEL_X = [[-1.0, 0.0, 1.0], [-2.0, 0.0, 2.0], [-1.0, 0.0, 1.0]] +_SOBEL_Y = [[-1.0, -2.0, -1.0], [0.0, 0.0, 0.0], [1.0, 2.0, 1.0]] + + +def _dev_idx(a: np.ndarray, device: torch.device) -> torch.Tensor: + return torch.from_numpy(a).to(device=device, dtype=torch.int32) + + +def canny_edges(frames: torch.Tensor, low: int, high: int) -> torch.Tensor: + """Canny over a clip: uint8 ``[C, T, H, W]`` CUDA -> uint8 ``[T, H, W]``.""" + C, T, H, W = frames.shape + dev = frames.device + + # Sobel values fit float32 exactly (|dx| <= 1020), then integer thereafter. + x = frames.reshape(C * T, 1, H, W).to(torch.float32) + xp = F.pad(x, (1, 1, 1, 1), mode="replicate") + kx = torch.tensor(_SOBEL_X, device=dev)[None, None] + ky = torch.tensor(_SOBEL_Y, device=dev)[None, None] + dx = F.conv2d(xp, kx).to(torch.int32).reshape(C, T, H, W) + dy = F.conv2d(xp, ky).to(torch.int32).reshape(C, T, H, W) + + bdx, bdy = dx[0], dy[0] + if C > 1: + best = bdx.abs() + bdy.abs() + for c in range(1, C): + mag_c = dx[c].abs() + dy[c].abs() + take = mag_c > best # strict: the first channel wins ties + best = torch.where(take, mag_c, best) + bdx = torch.where(take, dx[c], bdx) + bdy = torch.where(take, dy[c], bdy) + + mag = bdx.abs() + bdy.abs() + ax = bdx.abs() + y15 = bdy.abs() << _SHIFT + tg22x = ax * _TG22 + tg67x = tg22x + (ax << 16) + horiz = y15 < tg22x + vert = (~horiz) & (y15 > tg67x) + diag = ~(horiz | vert) + s_pos = (bdx ^ bdy) >= 0 + + m = F.pad(mag, (1, 1, 1, 1)) # zeros outside, matching the map borders + c = m[:, 1 : H + 1, 1 : W + 1] + + def nb(dr: int, dc: int) -> torch.Tensor: + return m[:, 1 + dr : H + 1 + dr, 1 + dc : W + 1 + dc] + + keep = horiz & (c > nb(0, -1)) & (c >= nb(0, 1)) + keep |= vert & (c > nb(-1, 0)) & (c >= nb(1, 0)) + keep |= diag & s_pos & (c > nb(-1, -1)) & (c > nb(1, 1)) + keep |= diag & ~s_pos & (c > nb(-1, 1)) & (c > nb(1, -1)) + + strong = (keep & (mag > high)).to(torch.float32)[:, None] + weak = (keep & (mag > low)).to(torch.float32)[:, None] + return (_hysteresis(strong, weak)[:, 0] > 0).to(torch.uint8) * 255 + + +def bilateral_filter( + frames: torch.Tensor, d: int, sigma_color: float, sigma_space: float +) -> torch.Tensor: + """Bilateral filter over a clip: uint8 ``[T, H, W, 3]`` CUDA -> same shape. + + One full-size temporary per tap, so it is O(taps) in both time and + bandwidth -- fine for test sizes, hopeless at production resolution. + """ + T, H, W, C = frames.shape + dev = frames.device + radius = max(d // 2, 1) + + lut = color_lut(C, -0.5 / (sigma_color * sigma_color), dev) + dys, dxs, sws = circle_offsets(radius, -0.5 / (sigma_space * sigma_space)) + sw = torch.tensor(sws, dtype=torch.float32, device=dev) + + src = reflect_pad(frames, radius) + center = src[:, radius : radius + H, radius : radius + W, :] + total = torch.zeros(T, H, W, C, dtype=torch.float32, device=dev) + wsum = torch.zeros(T, H, W, 1, dtype=torch.float32, device=dev) + for k, (i, j) in enumerate(zip(dys, dxs)): # row-major circular tap order + nb = src[:, radius + i : radius + i + H, radius + j : radius + j + W, :] + dist = (nb - center).abs().sum(-1).long() + w = (lut[dist] * sw[k]).unsqueeze(-1) + total = nb * w + total # unfused mul then add, as the kernel forces + wsum = w + wsum + return torch.round(total / wsum).clamp(0, 255).to(torch.uint8) + + +def resize_linear_u8(frames: torch.Tensor, dst_w: int, dst_h: int) -> torch.Tensor: + """Bilinear resize: uint8 ``[T, H, W, C]`` CUDA -> ``[T, dst_h, dst_w, C]``.""" + _check_frames(frames, "resize_linear_u8") + T, H, W, C = frames.shape + dev = frames.device + sx0, sx1, ax0, ax1 = _linear_axis(dst_w, W, clamp_coeffs=True) + sy0, sy1, ay0, ay1 = _linear_axis(dst_h, H, clamp_coeffs=False) + + f = frames.to(torch.int32) + tx0 = _dev_idx(ax0, dev)[None, None, :, None] + tx1 = _dev_idx(ax1, dev)[None, None, :, None] + h = f.index_select(2, _dev_idx(sx0, dev)) * tx0 + f.index_select(2, _dev_idx(sx1, dev)) * tx1 + ty0 = _dev_idx(ay0, dev)[None, :, None, None] + ty1 = _dev_idx(ay1, dev)[None, :, None, None] + v = ((h.index_select(1, _dev_idx(sy0, dev)) >> 4) * ty0 >> 16) + ( + (h.index_select(1, _dev_idx(sy1, dev)) >> 4) * ty1 >> 16 + ) + return ((v + 2) >> 2).clamp(0, 255).to(torch.uint8) + + +def resize_area_u8(frames: torch.Tensor, factor: int) -> torch.Tensor: + """Area-average downscale by an integer ``factor`` of 2 or 4.""" + _check_frames(frames, "resize_area_u8") + T, H, W, C = frames.shape + dh, dw = H // factor, W // factor + s = frames.to(torch.int32).reshape(T, dh, factor, dw, factor, C).sum(dim=(2, 4)) + if factor == 2 and C in (1, 3, 4): + out = (s + 2) >> 2 + else: + bits = 2 if factor == 2 else 4 + half = 1 << (bits - 1) + # branch-free round-half-even for division by 2**bits + out = (s + half - 1 + ((s >> bits) & 1)) >> bits + return out.clamp(0, 255).to(torch.uint8) + + +def resize_cubic_u8(frames: torch.Tensor, dst_w: int, dst_h: int) -> torch.Tensor: + """Bicubic resize: uint8 ``[T, H, W, C]`` CUDA -> ``[T, dst_h, dst_w, C]``.""" + _check_frames(frames, "resize_cubic_u8") + T, H, W, C = frames.shape + dev = frames.device + xtaps, xcoef = _cubic_axis(dst_w, W) + ytaps, ycoef = _cubic_axis(dst_h, H) + + f = frames.to(torch.int32) + h = sum( + f.index_select(2, _dev_idx(t, dev)) * _dev_idx(a, dev)[None, None, :, None] + for t, a in zip(xtaps, xcoef) + ) + rows = [h.index_select(1, _dev_idx(t, dev)) for t in ytaps] + betas = [_dev_idx(a, dev)[None, :, None, None] for a in ycoef] + + v_int = sum(r * b for r, b in zip(rows, betas)) + out_int = (v_int + (1 << (2 * _COEF_BITS - 1))) >> (2 * _COEF_BITS) + + inv = np.float32(1.0 / (_COEF_SCALE * _COEF_SCALE)) # 2^-22, exact + t = rows[3].to(torch.float32) * (betas[3].to(torch.float32) * inv) + for k in (2, 1, 0): + t = rows[k].to(torch.float32) * (betas[k].to(torch.float32) * inv) + t + out_float = torch.round(t).to(torch.int32) + + n = dst_w * C + covered = (n // 8) * 8 + elem = torch.arange(n, device=dev, dtype=torch.int32).reshape(dst_w, C) + out = torch.where(elem[None, None] < covered, out_float, out_int) + return out.clamp(0, 255).to(torch.uint8) diff --git a/tensorrt_llm/_torch/visual_gen/triton_kernels/resize.py b/tensorrt_llm/_torch/visual_gen/triton_kernels/resize.py new file mode 100644 index 000000000000..eeb568ba4636 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/triton_kernels/resize.py @@ -0,0 +1,608 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Fixed-point uint8 image resizing: bilinear, area-average, and bicubic. + +These follow the standard fixed-point formulations used by image-processing +libraries, reproduced exactly so control frames are reproducible across +backends. The arithmetic details below are load-bearing -- each was traced +back from a 1-LSB mismatch -- so treat them as the contract, not as +implementation freedom: + +- Source coordinates ``(dx + 0.5) * scale - 0.5`` are narrowed to float32 + BEFORE the floor, and coefficients are quantised to 11-bit fixed point with + round-half-even on the float32 product. +- Horizontal passes accumulate exactly in int32. +- The vertical descale differs per mode and is NOT a single fused rounding; + see each function's docstring. + +Layout is ``[T, H, W, C]`` uint8 on CUDA throughout. One fused kernel per +mode: every output element gathers its source taps and runs the whole +fixed-point pipeline in registers, rather than materialising int32 +intermediates between a horizontal and a vertical pass. + +Bit-exactness notes: bilinear and area are pure integer, so nothing can +drift. Bicubic's float32 lane uses ``libdevice`` ``mul_rn``/``add_rn`` so +ptxas cannot contract mul+add into an FMA (which would round differently from +the unfused eager reference), and ``float2int_rn`` for round-half-even. +""" + +import functools + +import numpy as np +import torch +import triton +import triton.language as tl +from triton.language.extra import libdevice + +_BLOCK = 256 +_COEF_BITS = 11 +_COEF_SCALE = 1 << _COEF_BITS + +# Axis tables are cached per (src, dst, C, device): a clip reuses the same +# handful of sizes for every frame, and rebuilding them on the host dominated +# the resize wall time (~3x the kernel time when measured). The cache is +# bounded because the key carries caller-supplied source dimensions -- a +# long-lived worker serving many distinct resolutions would otherwise retain a +# GPU table for every one it had ever seen. A transfer chain touches a handful +# of sizes, so this evicts only across unrelated requests. +_TABLE_CACHE_ENTRIES = 32 + + +def _linear_axis(dst_n: int, src_n: int, clamp_coeffs: bool): + d = np.arange(dst_n, dtype=np.float64) + scale = np.float64(src_n) / np.float64(dst_n) + f32 = ((d + 0.5) * scale - 0.5).astype(np.float32) + s = np.floor(f32).astype(np.int64) + f = (f32 - s.astype(np.float32)).astype(np.float32) + if clamp_coeffs: + # Only the horizontal axis clamps coefficients at the borders (fx=0, + # sx pinned); the vertical axis keeps raw (sy, fy) and clips just the + # row indices. Clamping both costs 1 LSB on border rows. + f = np.where(s < 0, np.float32(0), f).astype(np.float32) + s = np.where(s < 0, 0, s) + f = np.where(s >= src_n - 1, np.float32(0), f).astype(np.float32) + s = np.where(s >= src_n - 1, src_n - 1, s) + a0 = np.rint((np.float32(1.0) - f).astype(np.float32) * np.float32(_COEF_SCALE)) + a1 = np.rint(f * np.float32(_COEF_SCALE)) + s0 = np.clip(s, 0, src_n - 1) + s1 = np.clip(s + 1, 0, src_n - 1) + return s0, s1, a0.astype(np.int64), a1.astype(np.int64) + + +def _cubic_axis(dst_n: int, src_n: int): + d = np.arange(dst_n, dtype=np.float64) + scale = np.float64(src_n) / np.float64(dst_n) + f32 = ((d + 0.5) * scale - 0.5).astype(np.float32) + s = np.floor(f32).astype(np.int64) + x = (f32 - s.astype(np.float32)).astype(np.float32) + # Catmull-Rom-style cubic convolution kernel with A = -0.75, evaluated + # op-for-op in float32. + A = np.float32(-0.75) + one = np.float32(1.0) + xp1 = (x + one).astype(np.float32) + c0 = ((A * xp1 - np.float32(5) * A) * xp1 + np.float32(8) * A) * xp1 - np.float32(4) * A + c1 = ((A + 2) * x - (A + 3)) * x * x + one + y = (one - x).astype(np.float32) + c2 = ((A + 2) * y - (A + 3)) * y * y + one + c3 = one - c0 - c1 - c2 + coef = [ + np.rint(c.astype(np.float32) * np.float32(_COEF_SCALE)).astype(np.int64) + for c in (c0, c1, c2, c3) + ] + # Bicubic never clamps coefficients at the borders, only the tap indices. + taps = [np.clip(s - 1 + j, 0, src_n - 1) for j in range(4)] + return taps, coef + + +def _check_frames( + frames: torch.Tensor, + op: str, + *, + layout: str = "[T, H, W, C]", + allow_outer_stride: bool = False, +) -> None: + """Reject anything the kernels cannot address. + + Explicit raises rather than ``assert``: assertions vanish under ``python -O``, + and only ``ValueError`` is classified as a client error by the worker, so an + ``AssertionError`` would surface as a server fault. + + Most kernels index raw storage as one dense block and never consult strides, + so a strided view would be read as if it were dense -- wrong pixels, no + error. ``allow_outer_stride`` relaxes that for kernels that take dim 0's + stride as an argument: dim 0 may then sit anywhere, but everything inside it + must still be dense. Either way we reject rather than call ``.contiguous()``, + since a silent full-clip copy does not belong on an inference path. + """ + if not frames.is_cuda: + raise ValueError(f"{op} requires a CUDA tensor, got device={frames.device}") + if frames.dtype != torch.uint8: + raise TypeError(f"{op} requires uint8 frames, got dtype={frames.dtype}") + if frames.ndim != 4: + raise ValueError(f"{op} expects {layout}, got shape={tuple(frames.shape)}") + if allow_outer_stride: + if not frames[0].is_contiguous(): + raise ValueError( + f"{op} takes dim 0's stride but addresses the rest densely, so every " + f"slice along dim 0 must be contiguous. Got shape={tuple(frames.shape)} " + f"strides={tuple(frames.stride())}." + ) + elif not frames.is_contiguous(): + raise ValueError( + f"{op} requires a contiguous tensor; the kernels address storage densely and " + f"would misread a strided view. Got shape={tuple(frames.shape)} " + f"strides={tuple(frames.stride())} -- call .contiguous() first." + ) + + +@functools.lru_cache(maxsize=_TABLE_CACHE_ENTRIES) +def _linear_tables_x(src_w, dst_w, C, dev): + """Horizontal tables re-indexed per output *byte* (k = x*C + c), so the + kernel loads them contiguously instead of gathering per-x through a + runtime div/mod.""" + xs0, xs1, xa0, xa1 = _linear_axis(dst_w, src_w, clamp_coeffs=True) + ch = np.tile(np.arange(C, dtype=np.int64), dst_w) + tabs = ( + np.repeat(xs0.astype(np.int64) * C, C) + ch, + np.repeat(xs1.astype(np.int64) * C, C) + ch, + np.repeat(xa0, C), + np.repeat(xa1, C), + ) + return tuple(torch.from_numpy(a).to(device=dev, dtype=torch.int32) for a in tabs) + + +@functools.lru_cache(maxsize=_TABLE_CACHE_ENTRIES) +def _linear_tables_y(src_h, dst_h, dev): + return tuple( + torch.from_numpy(a).to(device=dev, dtype=torch.int32) + for a in _linear_axis(dst_h, src_h, clamp_coeffs=False) + ) + + +# --------------------------------------------------------------------------- +# Bilinear +# --------------------------------------------------------------------------- +# Grid scheme (all kernels here): (x-block, y, t) 3D grid + constexpr C, so +# there is NO runtime integer division in the hot path -- profiling showed the +# flat-index variant spending ~80% SM on idiv at 1-15% DRAM. +@triton.jit +def _linear_kernel( + src, + dst, + xo0, + xo1, + aa0, + aa1, + sy0, + sy1, + ay0, + ay1, + SH, + SW, + DH, + DW, + C: tl.constexpr, + BLOCK: tl.constexpr, +): + # Two output rows per program: the x-tables are loaded once for both, and + # the second row's independent gathers give the scheduler work to issue + # while the first row's loads are in flight (the kernel is issue-bound). + yA = tl.program_id(1) * 2 + t = tl.program_id(2) + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) # x*C + c within the row + m = offs < DW * C + yB = yA + 1 + mB = yB < DH + + x0 = tl.load(xo0 + offs, mask=m, other=0) # sx*C + c, contiguous + x1 = tl.load(xo1 + offs, mask=m, other=0) + a0 = tl.load(aa0 + offs, mask=m, other=0) + a1 = tl.load(aa1 + offs, mask=m, other=0) + + base = t * SH * SW * C + out = dst + (t * DH + yA) * DW * C + offs + + y0 = tl.load(sy0 + yA) + y1 = tl.load(sy1 + yA) + b0 = tl.load(ay0 + yA) + b1 = tl.load(ay1 + yA) + r0 = base + y0 * SW * C + r1 = base + y1 * SW * C + p00 = tl.load(src + r0 + x0, mask=m, other=0).to(tl.int32) + p01 = tl.load(src + r0 + x1, mask=m, other=0).to(tl.int32) + p10 = tl.load(src + r1 + x0, mask=m, other=0).to(tl.int32) + p11 = tl.load(src + r1 + x1, mask=m, other=0).to(tl.int32) + h0 = p00 * a0 + p01 * a1 + h1 = p10 * a0 + p11 * a1 + v = (((h0 >> 4) * b0) >> 16) + (((h1 >> 4) * b1) >> 16) + r = (v + 2) >> 2 + r = tl.minimum(tl.maximum(r, 0), 255) + tl.store(out, r.to(tl.uint8), mask=m) + + y0 = tl.load(sy0 + yB, mask=mB, other=0) + y1 = tl.load(sy1 + yB, mask=mB, other=0) + b0 = tl.load(ay0 + yB, mask=mB, other=0) + b1 = tl.load(ay1 + yB, mask=mB, other=0) + r0 = base + y0 * SW * C + r1 = base + y1 * SW * C + p00 = tl.load(src + r0 + x0, mask=m & mB, other=0).to(tl.int32) + p01 = tl.load(src + r0 + x1, mask=m & mB, other=0).to(tl.int32) + p10 = tl.load(src + r1 + x0, mask=m & mB, other=0).to(tl.int32) + p11 = tl.load(src + r1 + x1, mask=m & mB, other=0).to(tl.int32) + h0 = p00 * a0 + p01 * a1 + h1 = p10 * a0 + p11 * a1 + v = (((h0 >> 4) * b0) >> 16) + (((h1 >> 4) * b1) >> 16) + r = (v + 2) >> 2 + r = tl.minimum(tl.maximum(r, 0), 255) + tl.store(out + DW * C, r.to(tl.uint8), mask=m & mB) + + +def resize_linear_u8(frames: torch.Tensor, dst_w: int, dst_h: int) -> torch.Tensor: + """Bilinear resize: uint8 ``[T, H, W, C]`` CUDA -> ``[T, dst_h, dst_w, C]``. + + The uint8 vertical pass truncates PER TERM rather than as one fused + descale:: + + dst = (((b0 * (h0 >> 4)) >> 16) + ((b1 * (h1 >> 4)) >> 16) + 2) >> 2 + """ + _check_frames(frames, "resize_linear_u8") + T, H, W, C = frames.shape + dev = frames.device + xo0, xo1, aa0, aa1 = _linear_tables_x(W, dst_w, C, dev) + ys0, ys1, ya0, ya1 = _linear_tables_y(H, dst_h, dev) + dst = torch.empty(T, dst_h, dst_w, C, dtype=torch.uint8, device=dev) + _linear_kernel[(triton.cdiv(dst_w * C, _BLOCK), triton.cdiv(dst_h, 2), T)]( + frames, dst, xo0, xo1, aa0, aa1, ys0, ys1, ya0, ya1, H, W, dst_h, dst_w, C=C, BLOCK=_BLOCK + ) + return dst + + +# --------------------------------------------------------------------------- +# Area average (integer factors 2 / 4) +# --------------------------------------------------------------------------- +@triton.jit +def _area_kernel( + src, + dst, + SH, + SW, + DH, + DW, + C: tl.constexpr, + F: tl.constexpr, + HALF_UP: tl.constexpr, + BLOCK: tl.constexpr, +): + y = tl.program_id(1) + t = tl.program_id(2) + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = offs < DW * C + x = offs // C + c = offs % C + + base = t * SH * SW * C + c + s = tl.zeros((BLOCK,), dtype=tl.int32) + for i in tl.static_range(F): + row = base + (y * F + i) * SW * C + for j in tl.static_range(F): + s += tl.load(src + row + (x * F + j) * C, mask=m, other=0).to(tl.int32) + if HALF_UP: + r = (s + 2) >> 2 + else: + BITS: tl.constexpr = 2 if F == 2 else 4 + HALF: tl.constexpr = 1 << (BITS - 1) + r = (s + HALF - 1 + ((s >> BITS) & 1)) >> BITS + r = tl.minimum(tl.maximum(r, 0), 255) + tl.store(dst + (t * DH + y) * DW * C + offs, r.to(tl.uint8), mask=m) + + +@triton.jit +def _area3_kernel( + src, + dst, + NLANES, + SW: tl.constexpr, + DW: tl.constexpr, + F: tl.constexpr, + HALF_UP: tl.constexpr, + BLOCK: tl.constexpr, +): + # SW/DW constexpr: the lane->(row, p) split and row-stride math strength- + # reduce to mul/shift, off the critical path in front of the loads. One + # compile per (source, dest) size, which a clip reuses for every frame. + # C=3 fast path: a lane owns 24 contiguous, word-aligned source bytes per + # row (6 int32 loads) -> 4 output pixels for F=2 or 2 for F=4, instead of + # the generic kernel's per-byte gathers (whose load latency was 63% of + # stall time at 24% DRAM). F=2 stores 3 full aligned words per lane; an + # earlier pixel-pair variant's stride-6 u16 stores wasted write sectors. + # One flat 1D grid over all lanes: per-row blocks were too small and + # short-lived to hide latency (32% occupancy, 74% long_scoreboard). + # Since SH = DH*F, the source row is F*row + i with row the flattened + # (t, y) index -- no per-lane t/y split needed. + lane = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = lane < NLANES + P: tl.constexpr = 8 // F # output pixels per lane + NPL = DW // P + row = lane // NPL + p = lane % NPL + src32 = src.to(tl.pointer_type(tl.int32)) + + if F == 2: + # both rows' 12 loads issued before any extraction: interleaving the + # ~40 byte-extraction ops between 6-load batches halved the memory- + # level parallelism and left DRAM under 80% + rw = (2 * row) * (SW * 3 // 4) + p * 6 + # single-touch streams: evict_first keeps L2 free for write coalescing + ep: tl.constexpr = "evict_first" + w0 = tl.load(src32 + rw, mask=m, other=0, eviction_policy=ep) + w1 = tl.load(src32 + rw + 1, mask=m, other=0, eviction_policy=ep) + w2 = tl.load(src32 + rw + 2, mask=m, other=0, eviction_policy=ep) + w3 = tl.load(src32 + rw + 3, mask=m, other=0, eviction_policy=ep) + w4 = tl.load(src32 + rw + 4, mask=m, other=0, eviction_policy=ep) + w5 = tl.load(src32 + rw + 5, mask=m, other=0, eviction_policy=ep) + rw2 = rw + SW * 3 // 4 + x0 = tl.load(src32 + rw2, mask=m, other=0, eviction_policy=ep) + x1 = tl.load(src32 + rw2 + 1, mask=m, other=0, eviction_policy=ep) + x2 = tl.load(src32 + rw2 + 2, mask=m, other=0, eviction_policy=ep) + x3 = tl.load(src32 + rw2 + 3, mask=m, other=0, eviction_policy=ep) + x4 = tl.load(src32 + rw2 + 4, mask=m, other=0, eviction_policy=ep) + x5 = tl.load(src32 + rw2 + 5, mask=m, other=0, eviction_policy=ep) + o0 = (w0 & 255) + ((w0 >> 24) & 255) + (x0 & 255) + ((x0 >> 24) & 255) + o1 = ((w0 >> 8) & 255) + (w1 & 255) + ((x0 >> 8) & 255) + (x1 & 255) + o2 = ((w0 >> 16) & 255) + ((w1 >> 8) & 255) + ((x0 >> 16) & 255) + ((x1 >> 8) & 255) + o3 = ((w1 >> 16) & 255) + ((w2 >> 8) & 255) + ((x1 >> 16) & 255) + ((x2 >> 8) & 255) + o4 = ((w1 >> 24) & 255) + ((w2 >> 16) & 255) + ((x1 >> 24) & 255) + ((x2 >> 16) & 255) + o5 = (w2 & 255) + ((w2 >> 24) & 255) + (x2 & 255) + ((x2 >> 24) & 255) + o6 = (w3 & 255) + ((w3 >> 24) & 255) + (x3 & 255) + ((x3 >> 24) & 255) + o7 = ((w3 >> 8) & 255) + (w4 & 255) + ((x3 >> 8) & 255) + (x4 & 255) + o8 = ((w3 >> 16) & 255) + ((w4 >> 8) & 255) + ((x3 >> 16) & 255) + ((x4 >> 8) & 255) + o9 = ((w4 >> 16) & 255) + ((w5 >> 8) & 255) + ((x4 >> 16) & 255) + ((x5 >> 8) & 255) + o10 = ((w4 >> 24) & 255) + ((w5 >> 16) & 255) + ((x4 >> 24) & 255) + ((x5 >> 16) & 255) + o11 = (w5 & 255) + ((w5 >> 24) & 255) + (x5 & 255) + ((x5 >> 24) & 255) + else: + o0 = tl.zeros((BLOCK,), dtype=tl.int32) + o1 = tl.zeros((BLOCK,), dtype=tl.int32) + o2 = tl.zeros((BLOCK,), dtype=tl.int32) + o3 = tl.zeros((BLOCK,), dtype=tl.int32) + o4 = tl.zeros((BLOCK,), dtype=tl.int32) + o5 = tl.zeros((BLOCK,), dtype=tl.int32) + for i in tl.static_range(F): + # word offset of this lane's 24 input bytes in source row F*row + i + rw = (F * row + i) * (SW * 3 // 4) + p * 6 + w0 = tl.load(src32 + rw, mask=m, other=0) + w1 = tl.load(src32 + rw + 1, mask=m, other=0) + w2 = tl.load(src32 + rw + 2, mask=m, other=0) + w3 = tl.load(src32 + rw + 3, mask=m, other=0) + w4 = tl.load(src32 + rw + 4, mask=m, other=0) + w5 = tl.load(src32 + rw + 5, mask=m, other=0) + o0 += (w0 & 255) + ((w0 >> 24) & 255) + ((w1 >> 16) & 255) + ((w2 >> 8) & 255) + o1 += ((w0 >> 8) & 255) + (w1 & 255) + ((w1 >> 24) & 255) + ((w2 >> 16) & 255) + o2 += ((w0 >> 16) & 255) + ((w1 >> 8) & 255) + (w2 & 255) + ((w2 >> 24) & 255) + o3 += (w3 & 255) + ((w3 >> 24) & 255) + ((w4 >> 16) & 255) + ((w5 >> 8) & 255) + o4 += ((w3 >> 8) & 255) + (w4 & 255) + ((w4 >> 24) & 255) + ((w5 >> 16) & 255) + o5 += ((w3 >> 16) & 255) + ((w4 >> 8) & 255) + (w5 & 255) + ((w5 >> 24) & 255) + + BITS: tl.constexpr = 2 if F == 2 else 4 + HALF: tl.constexpr = 1 << (BITS - 1) + if HALF_UP: + r0 = (o0 + 2) >> 2 + r1 = (o1 + 2) >> 2 + r2 = (o2 + 2) >> 2 + r3 = (o3 + 2) >> 2 + r4 = (o4 + 2) >> 2 + r5 = (o5 + 2) >> 2 + else: + r0 = (o0 + HALF - 1 + ((o0 >> BITS) & 1)) >> BITS + r1 = (o1 + HALF - 1 + ((o1 >> BITS) & 1)) >> BITS + r2 = (o2 + HALF - 1 + ((o2 >> BITS) & 1)) >> BITS + r3 = (o3 + HALF - 1 + ((o3 >> BITS) & 1)) >> BITS + r4 = (o4 + HALF - 1 + ((o4 >> BITS) & 1)) >> BITS + r5 = (o5 + HALF - 1 + ((o5 >> BITS) & 1)) >> BITS + + if F == 2: + r6 = (o6 + 2) >> 2 + r7 = (o7 + 2) >> 2 + r8 = (o8 + 2) >> 2 + r9 = (o9 + 2) >> 2 + r10 = (o10 + 2) >> 2 + r11 = (o11 + 2) >> 2 + dst32 = dst.to(tl.pointer_type(tl.int32)) + ow = row * (DW * 3) // 4 + p * 3 + tl.store(dst32 + ow, r0 | (r1 << 8) | (r2 << 16) | (r3 << 24), mask=m) + tl.store(dst32 + ow + 1, r4 | (r5 << 8) | (r6 << 16) | (r7 << 24), mask=m) + tl.store(dst32 + ow + 2, r8 | (r9 << 8) | (r10 << 16) | (r11 << 24), mask=m) + else: + dst16 = dst.to(tl.pointer_type(tl.uint16)) + ob = row * (DW * 3) // 2 + p * 3 + tl.store(dst16 + ob, (r0 | (r1 << 8)).to(tl.uint16), mask=m) + tl.store(dst16 + ob + 1, (r2 | (r3 << 8)).to(tl.uint16), mask=m) + tl.store(dst16 + ob + 2, (r4 | (r5 << 8)).to(tl.uint16), mask=m) + + +def resize_area_u8(frames: torch.Tensor, factor: int) -> torch.Tensor: + """Area-average downscale by an integer ``factor`` of 2 or 4. + + Rounds half-UP for factor 2 with 1/3/4 channels and half-EVEN otherwise, + matching the integer fast path and the generic ``sum * 1/area`` path + respectively. Non-divisible dimensions would need fractional area + weights, which are deliberately unimplemented: every supported output + bucket is a multiple of 16. + """ + _check_frames(frames, "resize_area_u8") + T, H, W, C = frames.shape + if factor not in (2, 4): + raise ValueError(f"resize_area_u8 factor={factor}, expected one of (2, 4)") + if H % factor or W % factor: + raise ValueError(f"resize_area_u8 source {W}x{H} not divisible by factor={factor}") + dh, dw = H // factor, W // factor + dst = torch.empty(T, dh, dw, C, dtype=torch.uint8, device=frames.device) + half_up = factor == 2 and C in (1, 3, 4) + pix = 8 // factor # output pixels per lane (24 source bytes either way) + # data_ptr alignment matters: the C=3 path reinterprets the buffer as + # int32, and a contiguous *view* can still start on an odd byte, which + # faults the device rather than returning wrong data. + word_aligned = frames.data_ptr() % 4 == 0 and (W * 3) % 4 == 0 + if C == 3 and dw % pix == 0 and word_aligned and frames.is_contiguous(): + nlanes = T * dh * (dw // pix) + # one lane per thread: this kernel streams (DRAM-bound), so resident + # warps matter more than per-thread ILP (2 lanes/thread measured + # slower for both factors) + _area3_kernel[(triton.cdiv(nlanes, 256),)]( + frames, dst, nlanes, SW=W, DW=dw, F=factor, HALF_UP=half_up, BLOCK=256, num_warps=8 + ) + else: + _area_kernel[(triton.cdiv(dw * C, _BLOCK), dh, T)]( + frames, dst, H, W, dh, dw, C=C, F=factor, HALF_UP=half_up, BLOCK=_BLOCK + ) + return dst + + +# --------------------------------------------------------------------------- +# Bicubic +# --------------------------------------------------------------------------- +@functools.lru_cache(maxsize=_TABLE_CACHE_ENTRIES) +def _cubic_tables_x(src_w, dst_w, C, dev): + """Horizontal tap/coeff tables re-indexed per output byte, like bilinear.""" + xtaps, xcoef = _cubic_axis(dst_w, src_w) + ch = np.tile(np.arange(C, dtype=np.int64), dst_w) + txb = np.repeat(np.stack(xtaps).astype(np.int64) * C, C, axis=1) + ch[None, :] + cxb = np.repeat(np.stack(xcoef), C, axis=1) + return ( + torch.from_numpy(txb).to(device=dev, dtype=torch.int32).contiguous(), + torch.from_numpy(cxb).to(device=dev, dtype=torch.int32).contiguous(), + ) + + +@functools.lru_cache(maxsize=_TABLE_CACHE_ENTRIES) +def _cubic_tables_y(src_h, dst_h, dev): + ytaps, ycoef = _cubic_axis(dst_h, src_h) + return ( + torch.from_numpy(np.stack(ytaps)).to(device=dev, dtype=torch.int32), + torch.from_numpy(np.stack(ycoef)).to(device=dev, dtype=torch.int32), + ) + + +@triton.jit +def _cubic_kernel( + src, + dst, + xt, + xc, + yt, + yc, + SH, + SW, + DH, + DW, + C: tl.constexpr, + covered, + TAIL: tl.constexpr, + R: tl.constexpr, + BLOCK: tl.constexpr, +): + # TAIL=False when DW*C is a multiple of 8: the hybrid vertical pass then + # takes the float lane for EVERY element and the integer accumulators are + # dead work (~15% of the issue-bound instruction stream) -- skip them at + # compile time. All production bucket sizes are tail-free. + # R=2: two output rows per program, sharing the 8 x-table loads per tap + # column and doubling the independent gathers in flight -- pays off on + # wide rows (upscale). R=1 for narrow rows (deep downscale), where the + # doubled accumulator set costs occupancy instead. + yA = tl.program_id(1) * R + t = tl.program_id(2) + offs = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK) + m = offs < DW * C + yB = yA + 1 + mB = yB < DH + + base = t * SH * SW * C + vA_int = tl.zeros((BLOCK,), dtype=tl.int32) + vA_flt = tl.zeros((BLOCK,), dtype=tl.float32) + vB_int = tl.zeros((BLOCK,), dtype=tl.int32) + vB_flt = tl.zeros((BLOCK,), dtype=tl.float32) + INV: tl.constexpr = 1.0 / 4194304.0 # 2^-22, exact + # k walks the vertical taps last-to-first so the float chain associates + # b3-first, matching the reference. + for k in tl.static_range(3, -1, -1): + tyA = tl.load(yt + k * DH + yA) + byA = tl.load(yc + k * DH + yA) + rowA = base + tyA * SW * C + hA = tl.zeros((BLOCK,), dtype=tl.int32) + if R == 2: + tyB = tl.load(yt + k * DH + yB, mask=mB, other=0) + byB = tl.load(yc + k * DH + yB, mask=mB, other=0) + rowB = base + tyB * SW * C + hB = tl.zeros((BLOCK,), dtype=tl.int32) + for j in tl.static_range(4): + tx = tl.load(xt + j * DW * C + offs, mask=m, other=0) # sx*C + c + ax = tl.load(xc + j * DW * C + offs, mask=m, other=0) + hA += tl.load(src + rowA + tx, mask=m, other=0).to(tl.int32) * ax + if R == 2: + hB += tl.load(src + rowB + tx, mask=m & mB, other=0).to(tl.int32) * ax + if TAIL: + vA_int += hA * byA + vA_flt = libdevice.add_rn( + libdevice.mul_rn(hA.to(tl.float32), byA.to(tl.float32) * INV), vA_flt + ) + if R == 2: + if TAIL: + vB_int += hB * byB + vB_flt = libdevice.add_rn( + libdevice.mul_rn(hB.to(tl.float32), byB.to(tl.float32) * INV), vB_flt + ) + + out = dst + (t * DH + yA) * DW * C + offs + rA = libdevice.float2int_rn(vA_flt) + if TAIL: + rA_int = (vA_int + 2097152) >> 22 # fixed-point cast, 22 fractional bits + rA = tl.where(offs < covered, rA, rA_int) + rA = tl.minimum(tl.maximum(rA, 0), 255) + tl.store(out, rA.to(tl.uint8), mask=m) + if R == 2: + rB = libdevice.float2int_rn(vB_flt) + if TAIL: + rB_int = (vB_int + 2097152) >> 22 + rB = tl.where(offs < covered, rB, rB_int) + rB = tl.minimum(tl.maximum(rB, 0), 255) + tl.store(out + DW * C, rB.to(tl.uint8), mask=m & mB) + + +def resize_cubic_u8(frames: torch.Tensor, dst_w: int, dst_h: int) -> torch.Tensor: + """Bicubic resize: uint8 ``[T, H, W, C]`` CUDA -> ``[T, dst_h, dst_w, C]``. + + The uint8 vertical pass is hybrid per output row of ``n = dst_w * C`` + elements: the vectorised lane covers ``[0, 8 * (n // 8))`` in float32 + (b3-first mul/add chain, round-half-even) and the row tail uses the + scalar integer fixed-point cast ``(v + 2^21) >> 22``. + """ + _check_frames(frames, "resize_cubic_u8") + T, H, W, C = frames.shape + dev = frames.device + xt, xc = _cubic_tables_x(W, dst_w, C, dev) + yt, yc = _cubic_tables_y(H, dst_h, dev) + dst = torch.empty(T, dst_h, dst_w, C, dtype=torch.uint8, device=dev) + covered = (dst_w * C // 8) * 8 + # R=2 on wide rows (ILP is what feeds this issue-bound kernel); narrow + # deep-downscale rows keep R=1, where the doubled accumulators would cost + # occupancy instead. + R = 2 if dst_w * C >= 1024 else 1 + BLOCK = 256 if R == 2 else 128 + _cubic_kernel[(triton.cdiv(dst_w * C, BLOCK), triton.cdiv(dst_h, R), T)]( + frames, + dst, + xt, + xc, + yt, + yc, + H, + W, + dst_h, + dst_w, + C=C, + covered=covered, + TAIL=(dst_w * C) % 8 != 0, + R=R, + BLOCK=BLOCK, + ) + return dst diff --git a/tensorrt_llm/media/decoding.py b/tensorrt_llm/media/decoding.py index 9fa72833bc0d..71708229cc38 100644 --- a/tensorrt_llm/media/decoding.py +++ b/tensorrt_llm/media/decoding.py @@ -27,6 +27,7 @@ import functools import math +from typing import NamedTuple import torch @@ -107,6 +108,59 @@ def resize_center_crop_uint8(frames: torch.Tensor, target_h: int, target_w: int) return x.round_().clamp_(0, 255).to(torch.uint8).permute(0, 2, 3, 1).contiguous() +class VideoStreamInfo(NamedTuple): + """What a container header reports about its video stream.""" + + height: int + width: int + frame_rate: float | None # None when the header reports nothing usable + + +def video_stream_info(data: bytes) -> VideoStreamInfo | None: + """Read a clip's dimensions and frame rate from its container header. + + Demuxing is CPU-side FFmpeg inside PyNvVideoCodec, so this costs no GPU and + decodes no frame — everything here comes straight off the header, in one + open, so a caller wanting both does not pay for two. + + The dimensions are the *coded* ones. A container may additionally carry a + display matrix (a phone shooting portrait usually records landscape frames + plus a 90-degree rotation); the demuxer does not expose it and the decode + path does not apply it, so a clip carrying that metadata decodes + pixel-identically to the same clip without it. Coded dimensions therefore + describe the frames a caller actually receives, which is what a caller + sizing its output against them needs. + + Returns ``None`` when the header cannot be read or reports no usable + dimensions, leaving the caller on its own defaults: this is a convenience + probe, and a genuinely unreadable stream still fails with a proper error at + decode. + """ + try: + import PyNvVideoCodec as nvc + except ImportError: + return None + + position = 0 + + def _read(buf: bytearray) -> int: + nonlocal position + chunk = data[position : position + len(buf)] + buf[: len(chunk)] = chunk + position += len(chunk) + return len(chunk) + + try: + demuxer = nvc.CreateDemuxer(_read) + height, width = int(demuxer.Height()), int(demuxer.Width()) + frame_rate = float(demuxer.FrameRate()) + except nvc.PyNvVCException: + return None + if height <= 0 or width <= 0: + return None + return VideoStreamInfo(height, width, frame_rate if frame_rate > 0 else None) + + def decode_video_reference_window( data: bytes, *, diff --git a/tensorrt_llm/serve/visual_gen_utils.py b/tensorrt_llm/serve/visual_gen_utils.py index 49f093a2297e..be916bcc8ba0 100644 --- a/tensorrt_llm/serve/visual_gen_utils.py +++ b/tensorrt_llm/serve/visual_gen_utils.py @@ -4,6 +4,7 @@ import base64 import binascii import os +from collections.abc import Mapping from io import BytesIO from typing import TYPE_CHECKING, Any, Dict, List, Optional from urllib.parse import urlparse @@ -123,6 +124,39 @@ def _read_reference_payload(reference) -> bytes: return reference.file.read() +def _decode_inline_media(extra_params: dict | None, specs) -> None: + """Turn base64 strings into bytes for extra params declared as media. + + JSON has no byte type, so a client can only inline binary as base64. Any + extra param whose spec accepts ``bytes`` is decoded here, at the HTTP + boundary, so pipelines keep a bytes-only contract and never parse + transport encodings. Values that already arrived as bytes (multipart) + pass through. + """ + if not extra_params: + return + for key, value in list(extra_params.items()): + spec = specs.get(key) if specs else None + if spec is None or "bytes" not in getattr(spec, "type", ""): + continue + if isinstance(value, Mapping): + inner = value.get("control") + if isinstance(inner, str): + extra_params[key] = {**value, "control": _b64(key, inner)} + elif isinstance(value, str): + extra_params[key] = _b64(key, value) + + +def _b64(key: str, value: str) -> bytes: + try: + return base64.b64decode(value, validate=True) + except ValueError as exc: # binascii.Error subclasses ValueError + raise ValueError( + f"extra_params['{key}'] must be base64-encoded media bytes; " + "it is not valid base64 data." + ) from exc + + def _decode_base64_media(value: str) -> Optional[bytes]: payload = value if value.startswith("data:"): @@ -430,6 +464,7 @@ def parse_visual_gen_params( ) _warn_if_set_with_no_semantic(request, getattr(generator, "model", None)) + _decode_inline_media(request.extra_params, generator.extra_param_specs) _merge_extra_params(params, request.extra_params, generator.extra_param_specs) return params diff --git a/tensorrt_llm/visual_gen/params.py b/tensorrt_llm/visual_gen/params.py index 874698dfc425..a3145ebbe81f 100644 --- a/tensorrt_llm/visual_gen/params.py +++ b/tensorrt_llm/visual_gen/params.py @@ -32,6 +32,14 @@ class VisualGenParams(StrictBaseModel): ``guidance_scale_2``) should be passed via ``extra_params``. Use ``VisualGen.extra_param_specs`` to discover valid keys for the loaded pipeline. + + **``model_fields_set`` carries caller intent.** Defaults are merged in + before a pipeline sees the request, so a non-``None`` field says nothing + about who chose it: the merge assigns the pipeline default and then + ``discard``s that field, leaving only what the caller supplied. To + distinguish the two, test ``"frame_rate" in params.model_fields_set`` + rather than ``params.frame_rate is not None``. The set is live state -- + assigning a field re-marks it as caller intent. """ # Core — None means "use model default" @@ -93,6 +101,7 @@ class VisualGenParams(StrictBaseModel): "str": (str,), "list": (list,), "bytes": (bytes,), + "bool_or_bytes_or_dict": (bool, bytes, dict), } # Generation config fields that pipelines declare defaults for. If a user diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index 4e0edff9f9a0..b8f6a2b39e4d 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -299,6 +299,11 @@ def default_params(self) -> "VisualGenParams": pipeline's defaults. All declared ``extra_params`` keys are included with their defaults (``None`` for params without one). + Fields carrying a pipeline default are reported as unset, so a + round trip through this object does not read as caller intent + and request-dependent defaults stay resolvable; assigning any of + them marks it as yours. + Use this to inspect what the model will use, then modify and pass to ``generate()``:: diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index e1acb231f60a..4ab12555a42e 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -263,6 +263,8 @@ l0_b200: - unittest/_torch/visual_gen/test_cosmos3_transformer.py - unittest/_torch/visual_gen/test_cosmos3_pipeline.py - unittest/_torch/visual_gen/test_cosmos3_distilled.py + - unittest/_torch/visual_gen/test_cosmos3_transfer.py + - unittest/_torch/visual_gen/test_control_kernels.py - unittest/_torch/visual_gen/test_hunyuan_video1_5_transformer.py - unittest/_torch/visual_gen/test_hunyuan_video1_5_pipeline.py - examples/visual_gen/test_visual_gen_wan.py::test_wan_t2v_example diff --git a/tests/integration/test_lists/waives.txt b/tests/integration/test_lists/waives.txt index d664cbfd06d9..5397e2d75857 100644 --- a/tests/integration/test_lists/waives.txt +++ b/tests/integration/test_lists/waives.txt @@ -112,6 +112,7 @@ disaggregated/test_workers.py::test_workers_kv_cache_aware_router_eviction[TinyL examples/test_ad_speculative_decoding.py::test_autodeploy_eagle3_one_model_acceptance_rate[trtllm-torch-cudagraph] SKIP (https://nvbugs/6426841) examples/test_ad_speculative_decoding.py::test_nemotron_mtp_model_with_weights SKIP (https://nvbugs/6630699) examples/test_ray.py::test_ray_disaggregated_serving[tp2] SKIP (https://nvbugs/6632606) +examples/test_ray.py::test_ray_disaggregated_serving_python[tp2] SKIP (https://nvbugs/6601574) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_feature_accuracy_against_golden[nvfp4] SKIP (https://nvbugs/6572800) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2i_lpips_against_golden SKIP (https://nvbugs/6418815) examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden SKIP (https://nvbugs/6437341) diff --git a/tests/unittest/_torch/visual_gen/test_control_kernels.py b/tests/unittest/_torch/visual_gen/test_control_kernels.py new file mode 100644 index 000000000000..5197ff0a8b37 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_control_kernels.py @@ -0,0 +1,295 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""GPU tests for the VisualGen control-generation kernels. + +Every kernel is asserted **bitwise** against the torch reference in +``triton_kernels/reference.py``. Bitwise, not approximate: these produce +control frames that condition a diffusion model, and the fixed-point +arithmetic they reproduce has no tolerance band -- a 1-LSB drift means one of +the two implementations is simply wrong. The reference shares the kernels' +axis/tap tables, so a failure points at the arithmetic rather than at setup. + +The last class covers the Cosmos3 transfer entry points that compose them. +""" + +import os + +import pytest +import torch + +os.environ["TLLM_DISABLE_MPI"] = "1" + +from tensorrt_llm._torch.visual_gen.models.cosmos3 import transfer as transfer_module +from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import ( + BLUR_PRESETS, + EDGE_PRESETS, + make_blur_control, + make_edge_control, +) +from tensorrt_llm._torch.visual_gen.triton_kernels import ( + bilateral_filter, + canny_edges, + reference, + resize_area_u8, + resize_cubic_u8, + resize_linear_u8, +) +from tensorrt_llm._torch.visual_gen.triton_kernels import resize as resize_module + +pytestmark = [ + pytest.mark.cosmos3, + pytest.mark.skipif(not torch.cuda.is_available(), reason="control kernels require CUDA"), +] + + +def _clip(h: int, w: int, t: int = 3, c: int = 3, *, seed: int = 0) -> torch.Tensor: + """Smoothly-varying uint8 ``[T, H, W, C]``: video-like, so Canny's + hysteresis and the bilateral colour lookup see realistic neighbourhoods + instead of white noise.""" + g = torch.Generator(device="cuda").manual_seed(seed) + x = torch.randint(0, 256, (t, h, w, c), dtype=torch.uint8, device="cuda", generator=g) + f = x.to(torch.float32) + return ((f + f.roll(1, 1) + f.roll(1, 2) + f.roll(2, 2)) / 4).to(torch.uint8) + + +class TestResizeKernels: + @pytest.mark.parametrize("src", [(90, 160), (181, 321)]) + @pytest.mark.parametrize("scale", ["up", "down", "fractional", "near_identity"]) + def test_linear_matches_reference(self, src, scale): + h, w = src + dst = { + "up": (w * 2, h * 2), + "down": (w // 2, h // 2), + "fractional": (int(w * 1.7), int(h * 1.3)), + "near_identity": (w + 1, h - 1), + }[scale] + frames = _clip(h, w) + got = resize_linear_u8(frames, *dst) + assert got.shape == (3, dst[1], dst[0], 3) and got.dtype == torch.uint8 + assert torch.equal(got, reference.resize_linear_u8(frames, *dst)) + + @pytest.mark.parametrize("src", [(704, 1280), (92, 164)]) + @pytest.mark.parametrize("channels", [1, 2, 3, 4]) + @pytest.mark.parametrize("factor", [2, 4]) + def test_area_matches_reference(self, src, channels, factor): + # Channel count matters: the C=3 path takes a word-load fast path, and + # factor 2 with 1/3/4 channels rounds half-up where the rest round + # half-even. + h, w = src + frames = _clip(h, w, c=channels) + got = resize_area_u8(frames, factor) + assert got.shape == (3, h // factor, w // factor, channels) + assert torch.equal(got, reference.resize_area_u8(frames, factor)) + + @pytest.mark.parametrize("src", [(90, 160), (704, 1280), (180, 320)]) + @pytest.mark.parametrize("scale", ["up", "down", "fractional", "decimate", "near_identity"]) + def test_cubic_matches_reference(self, src, scale): + # "decimate" and "near_identity" hit the row-tail path where the + # vertical pass falls back to integer fixed point. + h, w = src + dst = { + "up": (w * 2, h * 2), + "down": (w // 2, h // 2), + "fractional": (int(w * 1.7), int(h * 1.3)), + "decimate": (max(1, w // 10), max(1, h // 10)), + "near_identity": (w + 1, h - 1), + }[scale] + frames = _clip(h, w) + got = resize_cubic_u8(frames, *dst) + assert got.shape == (3, dst[1], dst[0], 3) + assert torch.equal(got, reference.resize_cubic_u8(frames, *dst)) + + @pytest.mark.parametrize("byte_offset", [0, 1, 2, 3]) + def test_area_handles_unaligned_input(self, byte_offset): + # The C=3 fast path reinterprets the buffer as int32. A contiguous + # *view* can still start on an odd byte, which faults the device rather + # than returning wrong data, so the fast path must decline it. + t, h, w = 2, 64, 96 + n = t * h * w * 3 + g = torch.Generator(device="cuda").manual_seed(0) + base = torch.randint(0, 256, (n + 4,), dtype=torch.uint8, device="cuda", generator=g) + frames = base[byte_offset : byte_offset + n].view(t, h, w, 3) + assert frames.is_contiguous() + assert torch.equal(resize_area_u8(frames, 2), reference.resize_area_u8(frames, 2)) + + def test_axis_table_cache_is_bounded(self): + # The cache key carries caller-supplied source dimensions, so an + # unbounded cache would let a long-lived worker retain a GPU table for + # every resolution it ever served. + resize_module._cubic_tables_x.cache_clear() + clip = _clip(64, 96, t=1) + for dst_w in range(8, 8 + 2 * resize_module._TABLE_CACHE_ENTRIES): + resize_cubic_u8(clip, dst_w, 16) + info = resize_module._cubic_tables_x.cache_info() + assert info.currsize <= resize_module._TABLE_CACHE_ENTRIES + + def test_area_rejects_unsupported_geometry(self): + frames = _clip(64, 64) + with pytest.raises(ValueError, match=r"factor=3, expected one of \(2, 4\)"): + resize_area_u8(frames, 3) + with pytest.raises(ValueError, match="65x64 not divisible by factor=2"): + resize_area_u8(_clip(64, 65), 2) + + +class TestCannyKernel: + @pytest.mark.parametrize("size", [(96, 128), (704, 1280), (91, 161)]) + @pytest.mark.parametrize("thresholds", sorted(set(EDGE_PRESETS.values()))) + def test_matches_reference(self, size, thresholds): + h, w = size + frames = _clip(h, w).permute(3, 0, 1, 2).contiguous() + low, high = thresholds + got = canny_edges(frames, low, high) + assert got.shape == (3, h, w) and got.dtype == torch.uint8 + assert torch.equal(got, reference.canny_edges(frames, low, high)) + + def test_output_is_binary(self): + got = canny_edges(_clip(96, 128).permute(3, 0, 1, 2).contiguous(), 100, 200) + assert set(got.unique().tolist()) <= {0, 255} + + def test_higher_thresholds_give_sparser_edges(self): + frames = _clip(96, 128).permute(3, 0, 1, 2).contiguous() + counts = [ + (canny_edges(frames, lo, hi) > 0).sum().item() + for lo, hi in sorted(set(EDGE_PRESETS.values())) + ] + assert counts == sorted(counts, reverse=True) + + +class TestBilateralKernel: + @pytest.mark.parametrize("size", [(48, 64), (128, 96)]) + @pytest.mark.parametrize("params", [(9, 75.0, 75.0), (31, 150.0, 100.0), (13, 60.0, 40.0)]) + def test_matches_reference(self, size, params): + h, w = size + frames = _clip(h, w) + got = bilateral_filter(frames, *params) + assert got.shape == frames.shape and got.dtype == torch.uint8 + assert torch.equal(got, reference.bilateral_filter(frames, *params)) + + def test_preserves_flat_regions(self): + # Every weight is equal over a constant patch, so the filter is the + # identity there regardless of sigma. + flat = torch.full((2, 32, 32, 3), 100, dtype=torch.uint8, device="cuda") + assert torch.equal(bilateral_filter(flat, 9, 75.0, 75.0), flat) + + +class TestKernelInputValidation: + def test_rejects_cpu_tensors(self): + cpu = torch.zeros(1, 8, 8, 3, dtype=torch.uint8) + with pytest.raises(ValueError, match="requires a CUDA tensor, got device=cpu"): + bilateral_filter(cpu, 9, 75.0, 75.0) + with pytest.raises(ValueError, match="requires a CUDA tensor, got device=cpu"): + resize_linear_u8(cpu, 4, 4) + with pytest.raises(ValueError, match="requires a CUDA tensor, got device=cpu"): + canny_edges(torch.zeros(3, 1, 8, 8, dtype=torch.uint8), 100, 200) + + def test_rejects_non_contiguous(self): + # The kernels address storage densely and never read strides, so a + # strided view used to be accepted and silently produce wrong pixels. + g = torch.Generator(device="cuda").manual_seed(0) + chw = torch.randint(0, 256, (2, 3, 64, 96), dtype=torch.uint8, device="cuda", generator=g) + view = chw.permute(0, 2, 3, 1) # valid [T, H, W, C] shape, not contiguous + assert not view.is_contiguous() + for call in ( + lambda: resize_linear_u8(view, 48, 32), + lambda: resize_cubic_u8(view, 48, 32), + lambda: resize_area_u8(view, 2), + lambda: bilateral_filter(view, 9, 75.0, 75.0), + ): + with pytest.raises(ValueError, match="requires a contiguous tensor"): + call() + + def test_canny_rejects_strided_frame_planes(self): + # canny takes dim 0's stride, but the [T, H, W] block behind it must + # still be dense -- a channel-permuted view is not. + g = torch.Generator(device="cuda").manual_seed(0) + thw = torch.randint(0, 256, (2, 3, 64, 96), dtype=torch.uint8, device="cuda", generator=g) + view = thw.permute(1, 0, 2, 3) + assert not view[0].is_contiguous() + with pytest.raises(ValueError, match="slice along dim 0 must be contiguous"): + canny_edges(view, 100, 200) + + @pytest.mark.parametrize("window", [(0, 8), (8, 24), (24, 32)]) + def test_canny_reads_a_windowed_clip_in_place(self, window): + # Slicing [C, T, H, W] along T leaves dim 0 striding over the *whole* + # clip, so this used to need a .contiguous() copy per window. The result + # must be bit-identical to materializing it. + start, stop = window + frames = _clip(64, 96, t=32).permute(3, 0, 1, 2).contiguous() + view = frames[:, start:stop] + assert not view.is_contiguous() and view.stride(0) == 32 * 64 * 96 + assert torch.equal(canny_edges(view, 100, 200), canny_edges(view.contiguous(), 100, 200)) + + def test_rejects_non_uint8(self): + # Silently casting would be a data copy on the inference path; the + # kernels require the caller to hand over the dtype they expect. + f32 = torch.zeros(1, 8, 8, 3, dtype=torch.float32, device="cuda") + with pytest.raises(TypeError, match="requires uint8 frames, got dtype=torch.float32"): + bilateral_filter(f32, 9, 75.0, 75.0) + with pytest.raises(TypeError, match="requires uint8 frames, got dtype=torch.float32"): + resize_cubic_u8(f32, 4, 4) + + +class TestTransferControlGeneration: + """The Cosmos3 entry points that compose the kernels above.""" + + @pytest.mark.parametrize("preset", sorted(EDGE_PRESETS)) + def test_edge_control_shape_and_broadcast(self, preset): + frames = _clip(64, 96).permute(3, 0, 1, 2).contiguous() + edge = make_edge_control(frames, preset) + assert edge.shape == frames.shape and edge.dtype == torch.uint8 + assert edge.is_cuda and edge.is_contiguous() + # the single-channel edge map is broadcast across RGB + assert torch.equal(edge[0], edge[1]) and torch.equal(edge[0], edge[2]) + + def test_edge_control_uses_every_channel(self): + # R and G step in opposite directions, so the luma is flat across the + # seam: a grayscale-first detector sees nothing, while per-channel + # selection sees a full-scale edge in both. + frames = torch.zeros(3, 1, 32, 32, dtype=torch.uint8, device="cuda") + frames[0, :, :, :16] = 255 + frames[1, :, :, 16:] = 255 + frames[2] = 128 + assert make_edge_control(frames, "medium").any() + + @pytest.mark.parametrize("preset", sorted(BLUR_PRESETS)) + def test_blur_control_shape(self, preset): + frames = _clip(64, 128).permute(3, 0, 1, 2).contiguous() + blurred = make_blur_control(frames, preset) + assert blurred.shape == frames.shape and blurred.dtype == torch.uint8 + assert blurred.is_cuda and blurred.is_contiguous() + + def test_blur_none_preset_is_identity(self): + frames = _clip(64, 128).permute(3, 0, 1, 2).contiguous() + assert torch.equal(make_blur_control(frames, "none"), frames) + + def test_blur_reduces_variance(self): + frames = _clip(64, 128).permute(3, 0, 1, 2).contiguous() + sharp = frames.to(torch.float32) + for preset in ("low", "medium", "high"): + blurred = make_blur_control(frames, preset).to(torch.float32) + assert blurred.var().item() < sharp.var().item() + + @pytest.mark.parametrize("preset", ["medium", "high"]) + def test_generation_is_window_invariant(self, preset, monkeypatch): + # Control generation is windowed to bound preprocessing memory. Frames + # are independent, so the window size is a memory/parallelism knob and + # must not move a single pixel -- if this fails, some kernel grew a + # dependency across the temporal axis. + frames = _clip(64, 128, t=5).permute(3, 0, 1, 2).contiguous() + monkeypatch.setattr(transfer_module, "CONTROL_FRAME_WINDOW", frames.shape[1]) + edge_unwindowed = make_edge_control(frames, preset) + blur_unwindowed = make_blur_control(frames, preset) + + for window in (1, 2, 4): + monkeypatch.setattr(transfer_module, "CONTROL_FRAME_WINDOW", window) + assert torch.equal(make_edge_control(frames, preset), edge_unwindowed) + assert torch.equal(make_blur_control(frames, preset), blur_unwindowed) + + @pytest.mark.parametrize("preset", ["nonsense", ""]) + def test_unknown_presets_raise(self, preset): + frames = _clip(32, 32).permute(3, 0, 1, 2).contiguous() + with pytest.raises(ValueError, match="Unsupported Cosmos3 edge preset"): + make_edge_control(frames, preset) + with pytest.raises(ValueError, match="Unsupported Cosmos3 blur preset"): + make_blur_control(frames, preset) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index ac81efb0a56c..afa08f920bf9 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -987,6 +987,118 @@ def test_t2i_and_video_rejected(self, cosmos3_pipeline): ) +class TestCosmos3TransferRouting: + def test_transfer_rejects_an_image_reference(self): + """`_forward_transfer` takes no image, so a request carrying both used + to have its image silently dropped. The sibling guards already reject + transfer with image output and with audio; this one completes them.""" + from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import resolve_transfer_config + + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.action_gen = False + # __new__ skips __init__, where the real pipeline resolves this. + pipeline.family = QWEN3_RECIPE.name + pipeline.audio_gen = False + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 1.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def generation_default_overrides(self): + return {} + + pipeline.sampling = FakeSampling() + pipeline._forward_transfer = lambda **kwargs: None + # A precomputed control and no video: an existing guard already rejects + # image+video, so this is the shape where the image used to reach + # `_forward_transfer` and be discarded. + cfg = resolve_transfer_config( + {"edge": _V2V_FIXTURE_MP4.read_bytes()}, + SimpleNamespace(num_frames=93, guidance_scale=None), + None, + ) + + with pytest.raises(ValueError, match="cannot be combined with an image reference"): + pipeline.forward( + prompt="bounce", + image="frame.png", + transfer_config=cfg, + height=16, + width=16, + num_frames=5, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + frame_rate=8.0, + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=None, + use_guardrails=False, + ) + + def test_transfer_use_system_prompt_defaults_off(self): + """Reference parity: transfer defaults ``use_system_prompt=False`` even + when a video input is present — V2V's default-True rule must not leak + into the transfer branch (vllm-omni ``_forward_transfer`` defaults False). + An explicit request value is still honored.""" + from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import resolve_transfer_config + + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) + pipeline.action_gen = False + # __new__ skips __init__, where the real pipeline resolves this. + pipeline.family = QWEN3_RECIPE.name + pipeline.audio_gen = False + + class FakeSampling: + is_distilled = False + checkpoint_flow_shift = 1.0 + + def validate_request(self, num_inference_steps, guidance_scale): + return None + + def generation_default_overrides(self): + return {} + + pipeline.sampling = FakeSampling() + captured = {} + + def fake_forward_transfer(**kwargs): + captured.update(kwargs) + return None + + pipeline._forward_transfer = fake_forward_transfer + cfg = resolve_transfer_config( + {"edge": True}, SimpleNamespace(num_frames=93, guidance_scale=None), None + ) + + for explicit, expected in ((None, False), (True, True)): + captured.clear() + pipeline.forward( + prompt="bounce", + video=_V2V_FIXTURE_MP4.read_bytes(), + transfer_config=cfg, + height=16, + width=16, + num_frames=5, + num_inference_steps=1, + guidance_scale=1.0, + seed=1, + max_sequence_length=8, + frame_rate=8.0, + use_duration_template=False, + use_resolution_template=False, + use_system_prompt=explicit, + use_guardrails=False, + ) + assert captured["use_system_prompt"] is expected + + @pytest.mark.integration @pytest.mark.cosmos3_t2i @pytest.mark.high_cuda_memory diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py new file mode 100644 index 000000000000..2102d701867b --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transfer.py @@ -0,0 +1,1243 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""CPU unit tests for Cosmos3 Transfer (control-video conditioning). + +Ported from vllm-omni ``tests/diffusion/models/cosmos3/test_cosmos3_pipeline.py`` +(post-PR #4379) and adapted to TRT-LLM APIs, plus TRT-LLM-specific coverage for +the tensor-direct media decode and the chunk arithmetic. The diffuse-transfer +CFG tests assert the exact combination arithmetic (254/152/104/508) via +deterministic stubs, so any drift in the nested control/text CFG math fails +loudly. +""" + +import os +import pickle +from pathlib import Path +from types import SimpleNamespace + +os.environ["TLLM_DISABLE_MPI"] = "1" +os.environ["TRTLLM_DISABLE_COSMOS3_GUARDRAILS"] = "1" + +import pytest +import torch +import torch.nn as nn + +from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 as pipeline_module +from tensorrt_llm._torch.visual_gen.models.cosmos3 import transfer as transfer_module +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_EXTRA_SPECS, + COSMOS3_GENERATION_DEFAULTS, + VIDEO_RES_SIZE_INFO, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import DISTILLED_GUIDANCE_SCALE +from tensorrt_llm._torch.visual_gen.models.cosmos3.transfer import ( + BILATERAL_D, + BILATERAL_SIGMA_COLOR, + BILATERAL_SIGMA_SPACE, + _scaled_bilateral_params, + decode_media_to_uint8_cthw, + find_closest_target_size, + load_or_compute_control_frames, + pad_temporal_frames, + resolve_transfer_config, + uint8_cthw_to_normalized_5d, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + QWEN3_RECIPE, + TransformerOutput, +) +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer +from tensorrt_llm.media.decoding import VideoStreamInfo +from tensorrt_llm.visual_gen.params import VisualGenParams + +pytestmark = pytest.mark.cosmos3 + +# What a plain Cosmos3 video request advertises, and so what the executor +# merges into every request before infer() sees it. +_ADVERTISED_VIDEO_DEFAULTS = COSMOS3_GENERATION_DEFAULTS[("qwen3", "video")] + + +def _ids(value: int) -> torch.Tensor: + return torch.tensor([[value]], dtype=torch.long) + + +def _mask() -> torch.Tensor: + return torch.ones(1, 1, dtype=torch.long) + + +def _fake_decode_window(num_frames: int): + """Stand in for the NVDEC window decode, returning ``[T, H, W, 3]`` uint8.""" + + def _decode(data, *, first_frame, last_frame, target_h, target_w, device): + return torch.zeros(num_frames, target_h, target_w, 3, dtype=torch.uint8) + + return _decode + + +def _stub_info(height, width, frame_rate=None): + """Stand in for the container-header probe.""" + return lambda _data: VideoStreamInfo(height, width, frame_rate) + + +def _req(**overrides): + """Params as a caller supplied them: what is passed is marked caller intent. + + A real ``VisualGenParams``, not a stand-in, because telling a caller's value + from an executor-merged default is exactly what ``model_fields_set`` + carries — and a ``SimpleNamespace`` cannot express the difference. + """ + return VisualGenParams(**overrides) + + +def _merged_req(**merged): + """Params as the executor hands them to ``infer()``. + + The values are present but marked as pipeline defaults rather than caller + intent, mirroring ``DiffusionExecutor._merge_defaults``. + """ + params = VisualGenParams(**merged) + for field_name in merged: + params.model_fields_set.discard(field_name) + return params + + +class StubScheduler: + def __init__(self, timesteps=None): + self.timesteps = torch.tensor(timesteps or [9, 3], dtype=torch.int64) + self.config = SimpleNamespace( + num_train_timesteps=1000, flow_shift=1.0, use_karras_sigmas=True + ) + self.set_timesteps_calls = [] + self.sigmas_calls = [] + self.step_generators = [] + + def set_timesteps(self, num_steps=None, device=None, sigmas=None): + if sigmas is not None: + # Distilled: the policy installs a fixed sigma list, not a count. + self.sigmas_calls.append(list(sigmas)) + self.timesteps = torch.arange(len(sigmas), 0, -1, dtype=torch.int64) + return + self.set_timesteps_calls.append((num_steps, device)) + self.timesteps = torch.arange(num_steps, 0, -1, dtype=torch.int64) + + def step(self, noise_pred, timestep, latents, return_dict=False, generator=None): + # `generator` is accepted but nothing else is: a new pipeline-side + # argument must fail loudly here, not get silently swallowed. + assert return_dict is False + if generator is not None: + self.step_generators.append(generator) + return (latents + noise_pred,) + + +class StubTransformer(nn.Module): + """Deterministic transformer: returns full(token + 100·has_control). + + Also locks the calling convention: ``timestep`` must be the normalized + value and ``raw_timestep`` the raw scheduler value (the regression we + fixed after the VSA rebase). + """ + + def __init__(self): + super().__init__() + self.device = torch.device("cpu") + self.cached_kv = None + self.cached_freqs_gen = None + self.calls = [] + self.reset_calls = 0 + + def reset_cache(self): + self.reset_calls += 1 + self.cached_kv = None + self.cached_freqs_gen = None + + def forward(self, *, hidden_states, timestep, raw_timestep, text_ids, text_mask, **kwargs): + del text_mask + token = int(text_ids.reshape(-1)[0].item()) if text_ids.numel() else 0 + control_latents = kwargs.get("control_latents") + torch.testing.assert_close(timestep, raw_timestep / self.calls_num_train_timesteps) + self.calls.append({"token": token, "has_control": control_latents is not None}) + if self.cached_kv is None: + marker = torch.tensor([token], dtype=torch.float32) + self.cached_kv = [(marker, marker + 100)] + self.cached_freqs_gen = (marker + 200, marker + 300) + control_bonus = 100 if control_latents is not None else 0 + video = torch.full_like(hidden_states, float(token + control_bonus)) + return TransformerOutput(video=video, image=video) + + calls_num_train_timesteps = 1000 + + +class StubSamplingPolicy: + """Base-checkpoint stand-in for Cosmos3SamplingPolicy. + + Transfer programs the scheduler through the policy, so the stub records + what it was asked for; ``is_distilled`` flips the distilled contract on. + """ + + def __init__(self, is_distilled=False, fixed_sigmas=(1.0, 0.75, 0.5, 0.25)): + self.is_distilled = is_distilled + self.fixed_sigmas = fixed_sigmas + self.set_timesteps_calls = [] + self.step_kwargs_calls = 0 + self.flow_shift_calls = [] + + def set_flow_shift(self, scheduler, target_shift, use_karras_sigmas=None): + """Record the programmed shift and return the scheduler to install. + + Returns a distinct instance, like the real policy: the base scheduler + is kept pristine so shifts never accumulate, which means the caller has + to assign the result. Implemented here so the tests drive the + pipeline's real ``_scheduler_for`` instead of a stand-in -- stubbing + the pipeline method would re-create the hole that let a call to a + deleted helper ship green. + """ + del scheduler + self.flow_shift_calls.append(target_shift) + return StubScheduler() + + def set_timesteps(self, scheduler, num_inference_steps, device=None): + self.set_timesteps_calls.append(num_inference_steps) + if self.is_distilled: + scheduler.set_timesteps(sigmas=list(self.fixed_sigmas), device=device) + else: + scheduler.set_timesteps(num_inference_steps, device=device) + + def scheduler_step_kwargs(self, generator): + self.step_kwargs_calls += 1 + return {"generator": generator} if self.is_distilled else {} + + def generation_default_overrides(self): + # Mirrors the real policy rather than returning {}, so a distilled stub + # still overrides the table the way the checkpoint would. + if not self.is_distilled: + return {} + return { + "num_inference_steps": len(self.fixed_sigmas), + "guidance_scale": DISTILLED_GUIDANCE_SCALE, + } + + def num_steps(self, default): + return len(self.fixed_sigmas) if self.is_distilled else default + + +def _started_timer() -> CudaPhaseTimer: + """A timer in the state `forward()` hands to `_forward_transfer`.""" + timer = CudaPhaseTimer() + timer.mark_pre_start() + return timer + + +def _make_pipeline(sampling=None): + pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + nn.Module.__init__(pipeline) + pipeline.transformer = StubTransformer() + # __new__ skips __init__, which is where the real pipeline resolves this + # from the transformer config; the mode-defaults tables are keyed on it. + pipeline.family = QWEN3_RECIPE.name + pipeline.scheduler = StubScheduler() + pipeline.sampling = sampling or StubSamplingPolicy() + pipeline.safety_checker = None + pipeline.pipeline_config = SimpleNamespace(torch_dtype=torch.float32) + pipeline.vae_scale_factor_temporal = 4 + pipeline._guidance_scale = None + pipeline._num_timesteps = None + # Stub VAE encode: temporal subsample stands in for compression. + pipeline._encode_video_tensor = lambda video: video[:, :, :: pipeline.vae_scale_factor_temporal] + return pipeline + + +# ============================================================================= +# Config resolution (transfer.py) +# ============================================================================= + + +class TestTransferConfig: + def test_resolve_defaults_for_edge(self): + cfg = resolve_transfer_config({"edge": True}, _req()) + assert cfg is not None + assert list(cfg.hints) == ["edge"] + # The per-hint preset applies when the caller omitted guidance_scale, + # matching both references. The executor merges a pipeline default into + # every request, so `model_fields_set` -- not "is the value None" -- is + # what separates caller intent from a merged default. + assert cfg.guidance_scale == 3.0 + assert ( + resolve_transfer_config({"edge": True}, _req(guidance_scale=6.0)).guidance_scale == 6.0 + ) + # An executor-merged default must not read as caller intent, or the + # preset becomes unreachable for every request that goes through infer(). + merged = resolve_transfer_config({"edge": True}, _merged_req(guidance_scale=7.0)) + assert merged.guidance_scale == 3.0 + assert cfg.control_guidance == 1.5 + assert cfg.flow_shift == 10.0 + assert cfg.num_video_frames_per_chunk == 93 + assert cfg.share_vision_temporal_positions is True + + def test_control_directive_is_appended_by_default(self): + """Reference parity: cosmos-framework names the active control modality + in the user prompt unless the caller opts out. The system prompt is + untouched, which keeps the text in the training distribution.""" + cfg = resolve_transfer_config({"edge": True}, _req()) + assert cfg.emphasize_control_in_prompt is True + emphasized = cfg.emphasized_prompt("a robot dancing") + assert emphasized.startswith("a robot dancing") + assert "Follow the edge control video precisely" in emphasized + + def test_control_directive_names_every_active_hint(self): + cfg = resolve_transfer_config({"edge": True, "seg": b"clip"}, _req()) + # Hint order is TRANSFER_HINT_KEYS order, not caller order, so the + # directive text is stable across equivalent requests. + assert "Follow the edge, seg control video precisely" in cfg.emphasized_prompt("x") + + def test_control_directive_can_be_disabled(self): + cfg = resolve_transfer_config({"edge": True, "emphasize_control_in_prompt": False}, _req()) + assert cfg.emphasize_control_in_prompt is False + assert cfg.emphasized_prompt("a robot dancing") == "a robot dancing" + + def test_no_hints_resolves_none(self): + assert resolve_transfer_config({}, _req()) is None + assert resolve_transfer_config({"guidance_scale": 3.0}, _req()) is None + + def test_wsm_fps_preset_default_and_override(self): + """The override is a request field, not an extra param. + + `frame_rate` reaches the pipeline as a declared generation field, so + that is the only spelling a real request can use -- an `extra_params` + copy of it would be rejected as an undeclared key by + `validate_visual_gen_params` before `infer()` ever runs. + """ + assert resolve_transfer_config({"wsm": True}, _req()).fps == 10 + assert resolve_transfer_config({"wsm": True}, _req(frame_rate=24.0)).fps == 24.0 + # The rate the executor merges into every request is not caller intent, + # so it must not defeat the preset. + assert resolve_transfer_config({"wsm": True}, _merged_req(frame_rate=24.0)).fps == 10 + + def test_wsm_clip_presets_survive_merged_request_defaults(self): + """wsm wants 101 frames at 10 fps, but `num_frames` and `frame_rate` + are advertised defaults the executor merges into every request before + `infer()` sees it. A request carrying only those merged values must + still get the preset; values the caller actually chose must win.""" + merged = _merged_req( + num_frames=_ADVERTISED_VIDEO_DEFAULTS["num_frames"], + frame_rate=_ADVERTISED_VIDEO_DEFAULTS["frame_rate"], + ) + cfg = resolve_transfer_config({"wsm": True}, merged) + assert (cfg.num_frames, cfg.fps) == (101, 10) + + chosen = _req(num_frames=200, frame_rate=30.0) + cfg = resolve_transfer_config({"wsm": True}, chosen) + assert (cfg.num_frames, cfg.fps) == (200, 30.0) + + def test_advertised_defaults_are_left_intact(self): + """The preset is recovered inside transfer, not by nulling the model's + published defaults — clients read those to learn the output shape.""" + # Positive, not merely equal to the table they are read from: the + # advertised entry *is* that table, so an identity check would hold + # even if both were zeroed to make the preset win. + assert _ADVERTISED_VIDEO_DEFAULTS["num_frames"] > 0 + assert _ADVERTISED_VIDEO_DEFAULTS["frame_rate"] > 0 + + def test_precomputed_control_bytes_reach_the_decoder(self, monkeypatch): + monkeypatch.setattr( + transfer_module, "decode_video_reference_window", _fake_decode_window(2) + ) + cfg = resolve_transfer_config({"edge": {"control": b"\x00control"}}, _req()) + loaded = load_or_compute_control_frames( + cfg.hints["edge"], + height=8, + width=8, + max_frames=2, + input_frames=None, + device=torch.device("cpu"), + ) + assert tuple(loaded.shape) == (3, 2, 8, 8) and loaded.dtype == torch.uint8 + + +# ============================================================================= +# Media helpers (transfer.py + utils.py) +# ============================================================================= + + +class TestTransferMediaHelpers: + def test_pad_temporal_frames_reflects(self): + # Reference parity: [0, 3, 6] padded to 5 reflects the tail -> [0, 3, 6, 6, 3]. + frames = torch.arange(3 * 3, dtype=torch.uint8).reshape(1, 3, 1, 3) + assert pad_temporal_frames(frames, 5)[0, :, 0, 0].tolist() == [0, 3, 6, 6, 3] + + def test_malformed_hints_are_client_errors(self): + # The worker classifier maps ValueError to a client error (400) and + # anything else to an unclassified server fault (500). A caller's + # malformed hint must not be reported as our failure. + for payload in (123, ["edge.mp4"], {"control": 123}): + with pytest.raises(ValueError): + resolve_transfer_config({"edge": payload}, _req()) + with pytest.raises(ValueError): + decode_media_to_uint8_cthw( + "not-bytes", height=8, width=8, max_frames=1, device=torch.device("cpu") + ) + + def test_bilateral_params_scale_with_resolution(self): + # Tuned at a 720p reference: a 72px longest side is 1/10 of it, so the + # diameter and both sigmas scale down by the same factor. + assert _scaled_bilateral_params(72, 72) == (3, 15.0, 10.0) + assert _scaled_bilateral_params(720, 720) == ( + BILATERAL_D + 1, # 30 is even; diameters are forced odd + float(BILATERAL_SIGMA_COLOR), + float(BILATERAL_SIGMA_SPACE), + ) + # The longest side drives the scale, and the sigmas have a floor of 1. + assert _scaled_bilateral_params(4, 1280) == _scaled_bilateral_params(1280, 4) + assert _scaled_bilateral_params(1, 1) == (1, 1.0, 1.0) + + def test_generated_control_hints_require_input_frames(self): + for key in ("edge", "blur"): + cfg = resolve_transfer_config({key: True}, _req()) + with pytest.raises(ValueError, match="requires either a video input"): + load_or_compute_control_frames( + cfg.hints[key], + height=8, + width=8, + max_frames=1, + input_frames=None, + device=torch.device("cpu"), + ) + + +class TestSourceDerivedDefaults: + """Unset output size and frame rate follow the reference, worker-side. + + Previously only the offline example fitted the aspect, so a served portrait + or square reference was center-cropped into the default landscape bucket. + The probe reads the container header (no GPU, no frame decoded). + """ + + REFERENCE = Path(__file__).parent / "test_data" / "cosmos3_v2v_ref_9f_bframes.mp4" + + def _infer_req(self, _params=None, **extra): + # Executor-merged shape: num_frames/frame_rate carry pipeline defaults, + # height/width are declared None, and nothing reads as caller intent. + params = ( + _params + if _params is not None + else _merged_req( + num_frames=COSMOS3_720P_PARAMS["num_frames"], + frame_rate=COSMOS3_720P_PARAMS["frame_rate"], + max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + seed=0, + ) + ) + params.extra_params = dict(extra) + return SimpleNamespace(params=params, prompt="a prompt") + + def _captured(self, req): + pipeline = _make_pipeline() # rank is 0 while dist is uninitialized + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + pipeline.infer(req) + return captured + + def _size(self, req): + captured = self._captured(req) + return captured["width"], captured["height"] + + def test_square_reference_picks_the_square_bucket(self): + # The checked-in reference is 64x64, so a real header probe end to end. + req = self._infer_req(video=self.REFERENCE.read_bytes()) + assert self._size(req) == VIDEO_RES_SIZE_INFO["720"]["1,1"] + + # source_hw is (height, width); the bucket table is keyed (width, height). + @pytest.mark.parametrize( + "source_hw, bucket", + [ + ((320, 192), "9,16"), # portrait + ((192, 320), "16,9"), # landscape + ((1104, 832), "3,4"), # tall, not 9:16 + ((832, 1104), "4,3"), # wide, not 16:9 + ((600, 600), "1,1"), # square + ], + ) + def test_aspect_selects_the_matching_bucket(self, source_hw, bucket, monkeypatch): + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(*source_hw)) + req = self._infer_req(video=b"stand-in for encoded bytes") + assert self._size(req) == VIDEO_RES_SIZE_INFO["720"][bucket] + + def test_explicit_dimensions_win(self, monkeypatch): + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(320, 192)) + req = self._infer_req(video=b"stand-in") + req.params.height, req.params.width = 704, 1280 # assignment marks them + assert self._size(req) == (1280, 704) + + def test_a_half_specified_size_is_left_alone(self, monkeypatch): + # Overriding the unset half of a stated intent would be worse than + # leaving the request on the mode defaults. + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(320, 192)) + req = self._infer_req(video=b"stand-in") + req.params.width = 1280 + assert self._size(req) == (1280, COSMOS3_720P_PARAMS["height"]) + + def test_no_reference_keeps_the_mode_defaults(self): + req = self._infer_req() + assert self._size(req) == (COSMOS3_720P_PARAMS["width"], COSMOS3_720P_PARAMS["height"]) + + def test_unreadable_reference_falls_back_to_defaults(self): + # A convenience probe must not fail the request; the real decode still + # reports the problem properly. + req = self._infer_req(video=b"not a video container") + assert self._size(req) == (COSMOS3_720P_PARAMS["width"], COSMOS3_720P_PARAMS["height"]) + + def test_transfer_fits_to_a_control_when_there_is_no_video(self, monkeypatch): + # Transfer can run on precomputed controls alone. + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(320, 192)) + req = self._infer_req(edge={"control": b"stand-in"}) + assert self._size(req) == VIDEO_RES_SIZE_INFO["720"]["9,16"] + + # --- frame rate ------------------------------------------------------- + + def test_source_frame_rate_is_adopted_when_unset(self, monkeypatch): + # Emitting an 8 fps source at the merged 24 fps default replays it at + # 3x speed and misreports its duration to the text conditioning. + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(192, 320, 8.0)) + req = self._infer_req(video=b"stand-in") + assert self._captured(req)["frame_rate"] == 8.0 + + def test_explicit_frame_rate_wins(self, monkeypatch): + # The whole point of the caller-intent bit: an explicit 24.0 is the + # same *value* the executor would have merged, but not the same intent. + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(192, 320, 8.0)) + for chosen in (24.0, 30.0): + req = self._infer_req(_params=_req(seed=0, frame_rate=chosen), video=b"stand-in") + assert self._captured(req)["frame_rate"] == chosen + + def test_default_stands_without_a_usable_source_rate(self, monkeypatch): + monkeypatch.setattr(pipeline_module, "video_stream_info", _stub_info(192, 320, None)) + req = self._infer_req(video=b"stand-in") + assert self._captured(req)["frame_rate"] == COSMOS3_720P_PARAMS["frame_rate"] + + def test_default_stands_without_a_reference(self): + assert ( + self._captured(self._infer_req())["frame_rate"] == (COSMOS3_720P_PARAMS["frame_rate"]) + ) + + +class TestTransferPreflightValidation: + """Deterministic client mistakes must 400 at enqueue, not 202 then fail. + + These validators run in the coordinator (``visual_gen/params.py``), so they + must stay in step with ``resolve_transfer_config``: anything they reject has + to be something the worker would have rejected too, only later. + """ + + # Minimal ISO-BMFF header, enough for the container sniff. + MP4 = b"\x00\x00\x00\x18ftypisom\x00\x00\x02\x00isomiso2" + + def _check(self, key, value): + COSMOS3_EXTRA_SPECS[key].validator(value) + + def test_validators_are_picklable(self): + # Specs are pickled to the coordinator in the READY handshake, so a + # closure or lambda here would break serving but pass every other test. + for key in ("edge", "blur", "depth", "seg", "wsm", "control_guidance_interval"): + validator = COSMOS3_EXTRA_SPECS[key].validator + assert pickle.loads(pickle.dumps(validator)) is validator + + @pytest.mark.parametrize("key", ["edge", "blur"]) + def test_generated_hints_accept_auto_compute_and_controls(self, key): + self._check(key, True) + self._check(key, self.MP4) + self._check(key, {"control": self.MP4}) + self._check(key, {}) + + @pytest.mark.parametrize("key", ["depth", "seg", "wsm"]) + def test_precomputed_hints_reject_auto_compute(self, key): + # These have no generator; `true` used to 202 and then fail in the worker. + for value in (True, {}): + with pytest.raises(ValueError, match="no on-the-fly generator"): + self._check(key, value) + self._check(key, self.MP4) # a real control clip is still fine + + def test_presets_are_checked_against_the_tables(self): + self._check("edge", {"preset_edge_threshold": "very_high"}) + self._check("blur", {"preset_blur_strength": "very_high"}) + # Empty/None means "default", exactly as resolve_transfer_config reads it. + self._check("edge", {"preset_edge_threshold": ""}) + with pytest.raises(ValueError, match="unsupported preset_edge_threshold"): + self._check("edge", {"preset_edge_threshold": "sharpish"}) + with pytest.raises(ValueError, match="unsupported preset_blur_strength"): + self._check("blur", {"preset_blur_strength": "soupy"}) + + def test_rejects_paths_and_undecodable_bytes(self): + with pytest.raises(ValueError, match="control_path"): + self._check("edge", {"control_path": "/tmp/control.mp4"}) + with pytest.raises(ValueError, match="not a recognized video container"): + self._check("edge", b"not a video") + # Bare bytes and the object form are held to the same bar. + with pytest.raises(ValueError, match="not a recognized video container"): + self._check("edge", {"control": b"not a video"}) + with pytest.raises(ValueError, match="Omit the key"): + self._check("edge", False) + + def test_interval_must_be_an_ordered_pair(self): + self._check("control_guidance_interval", [0.1, 0.9]) + with pytest.raises(ValueError, match="exactly two values"): + self._check("control_guidance_interval", [0.5]) + with pytest.raises(ValueError, match="ordered as"): + self._check("control_guidance_interval", [0.9, 0.1]) + + def test_frame_counts_are_bounded(self): + self._check("num_video_frames_per_chunk", 93) + self._check("num_conditional_frames", 0) + for key in ("num_video_frames_per_chunk", "max_frames"): + with pytest.raises(ValueError, match="positive frame count"): + self._check(key, 0) + for key in ("num_conditional_frames", "num_first_chunk_conditional_frames"): + with pytest.raises(ValueError, match="non-negative frame count"): + self._check(key, -1) + + +class TestTransferFrameConversions: + """Unit tests for the transfer control-frame conversion helpers (CPU-only).""" + + def test_uint8_cthw_to_normalized_5d_maps_0_255_to_pm1(self): + black = torch.zeros(3, 2, 4, 5, dtype=torch.uint8) + out = uint8_cthw_to_normalized_5d(black, dtype=torch.float32) + assert out.shape == (1, 3, 2, 4, 5) and out.dtype == torch.float32 + assert torch.allclose(out, torch.full_like(out, -1.0)) # 0 -> -1 + white = torch.full((3, 1, 2, 2), 255, dtype=torch.uint8) + assert torch.allclose( + uint8_cthw_to_normalized_5d(white, dtype=torch.float32), + torch.ones(1, 3, 1, 2, 2), # 255 / 127.5 - 1 == 1.0 + ) + + def test_uint8_cthw_to_normalized_5d_rejects_bad_shape(self): + with pytest.raises(ValueError, match="3, T, H, W"): + uint8_cthw_to_normalized_5d( + torch.zeros(2, 4, 4, dtype=torch.uint8), dtype=torch.float32 + ) + + +class TestTransferControlPayloads: + """Controls cross the extra-param boundary as encoded bytes, like ``video``.""" + + def test_hint_accepts_bytes_bare_and_in_an_object(self): + assert resolve_transfer_config({"edge": b"MP4"}, _req()).hints["edge"].control == b"MP4" + cfg = resolve_transfer_config({"depth": {"control": b"MP4"}}, _req()) + assert cfg.hints["depth"].control == b"MP4" + + def test_hint_rejects_a_bare_path(self): + with pytest.raises(ValueError, match="not a path"): + resolve_transfer_config({"edge": "/tmp/control.mp4"}, _req()) + + def test_hint_rejects_control_path(self): + with pytest.raises(ValueError, match="control_path"): + resolve_transfer_config({"edge": {"control_path": "/tmp/control.mp4"}}, _req()) + + def test_hint_rejects_non_bytes_control(self): + with pytest.raises(ValueError, match="encoded MP4/AVI bytes"): + resolve_transfer_config({"edge": {"control": torch.zeros(3, 1, 4, 4)}}, _req()) + + def test_decode_asks_for_the_leading_window_and_returns_cthw(self, monkeypatch): + seen = {} + + def fake_decode(data, *, first_frame, last_frame, target_h, target_w, device): + seen.update(data=data, first=first_frame, last=last_frame, h=target_h, w=target_w) + return torch.zeros(3, target_h, target_w, 3, dtype=torch.uint8) + + monkeypatch.setattr(transfer_module, "decode_video_reference_window", fake_decode) + out = decode_media_to_uint8_cthw( + b"clip", height=8, width=6, max_frames=4, device=torch.device("cpu") + ) + assert tuple(out.shape) == (3, 3, 8, 6) and out.dtype == torch.uint8 + assert seen == {"data": b"clip", "first": 0, "last": 3, "h": 8, "w": 6} + + def test_decode_rejects_unencoded_payloads(self): + with pytest.raises(ValueError, match="encoded MP4/AVI bytes"): + decode_media_to_uint8_cthw( + torch.zeros(3, 1, 4, 4), + height=4, + width=4, + max_frames=1, + device=torch.device("cpu"), + ) + + def test_decode_rejects_a_nonpositive_window(self): + with pytest.raises(ValueError, match="max_frames must be positive, got 0"): + decode_media_to_uint8_cthw( + b"clip", height=4, width=4, max_frames=0, device=torch.device("cpu") + ) + + +# ============================================================================= +# diffuse_transfer — nested control/text CFG arithmetic (ported verbatim) +# ============================================================================= + + +class TestDiffuseTransferCFG: + def _run(self, pipeline, *, timesteps, guidance_scale, control_guidance, **overrides): + latents = torch.zeros(1, 2, 1, 1, 1) + velocity_mask = torch.ones(1, 1, 1, 1, 1) + kwargs = dict( + latents=latents, + timesteps=torch.tensor(timesteps), + cond_ids=_ids(2), + cond_mask=_mask(), + uncond_ids=_ids(1), + uncond_mask=_mask(), + guidance_scale=guidance_scale, + control_guidance=control_guidance, + control_guidance_interval=None, + control_latents=[torch.zeros_like(latents)], + shared_kwargs={ + "video_shape": (1, 1, 1), + "fps": 24.0, + "noisy_frame_mask": velocity_mask, + }, + velocity_mask=velocity_mask, + condition_latents=torch.zeros_like(latents), + generator=torch.Generator().manual_seed(0), + ) + kwargs.update(overrides) + return pipeline.diffuse_transfer(**kwargs), latents + + def test_applies_control_and_text_cfg(self): + pipeline = _make_pipeline() + result, latents = self._run( + pipeline, timesteps=[7], guidance_scale=3.0, control_guidance=1.5 + ) + # cond_full=102, no_control=2, uncond=101: + # control_cond = 2 + 1.5*(102-2) = 152; 101 + 3*(152-101) = 254 + assert [(c["token"], c["has_control"]) for c in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 254.0)) + + def test_skips_idle_cfg_branches(self): + control_only = _make_pipeline() + result, latents = self._run( + control_only, timesteps=[7], guidance_scale=1.0, control_guidance=1.5 + ) + assert [(c["token"], c["has_control"]) for c in control_only.transformer.calls] == [ + (2, True), + (2, False), + ] + torch.testing.assert_close(result, torch.full_like(latents, 152.0)) + + text_only = _make_pipeline() + result, latents = self._run( + text_only, timesteps=[7], guidance_scale=3.0, control_guidance=1.0 + ) + assert [(c["token"], c["has_control"]) for c in text_only.transformer.calls] == [ + (2, True), + (1, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 104.0)) + + def test_interval_switches_branch_counts(self): + pipeline = _make_pipeline() + result, latents = self._run( + pipeline, + timesteps=[900, 500, 100], + guidance_scale=3.0, + control_guidance=1.5, + control_guidance_interval=(400.0, 1000.0), + guidance_interval=(800.0, 1000.0), + ) + # t=900: 3 branches -> +254; t=500: control only -> +152; t=100: single -> +102 + assert [(c["token"], c["has_control"]) for c in pipeline.transformer.calls] == [ + (2, True), + (2, False), + (1, True), + (2, True), + (2, False), + (2, True), + ] + torch.testing.assert_close(result, torch.full_like(latents, 508.0)) + + +# ============================================================================= +# _forward_transfer — chunk arithmetic and multichunk stitching +# ============================================================================= + + +class TestForwardTransferChunks: + def test_get_transfer_num_chunks_arithmetic(self): + chunks = Cosmos3OmniMoTPipeline._get_transfer_num_chunks + assert chunks(93, 93, 1) == (1, 93) + assert chunks(189, 93, 1) == (3, 92) + assert chunks(5, 93, 1) == (1, 93) + with pytest.raises(ValueError, match="num_conditional_frames"): + chunks(189, 93, 93) + + def test_decode_window_is_bounded_by_the_output_length(self, monkeypatch): + """The decoder reserves its retention ring from the requested window, so + asking for ``max_frames`` (5000 by default) would reserve ~14 GB at 720p + before a frame lands. Only what gets generated is decoded.""" + pipeline = _make_pipeline() + windows = [] + + class StopAfterDecode(Exception): + pass + + def recording_decode(data, *, first_frame, last_frame, target_h, target_w, device): + windows.append((first_frame, last_frame)) + raise StopAfterDecode + + monkeypatch.setattr(transfer_module, "decode_video_reference_window", recording_decode) + cfg = resolve_transfer_config({"edge": {"control": b"clip"}}, _req()) + assert cfg.max_frames == 5000 # the default ceiling stays in place + + with pytest.raises(StopAfterDecode): + pipeline._forward_transfer( + prompt="transfer", + negative_prompt="", + height=16, + width=16, + max_frames=cfg.max_frames, + num_inference_steps=1, + max_sequence_length=8, + use_system_prompt=False, + use_duration_template=False, + use_resolution_template=False, + seed=1, + frame_rate=24.0, + num_frames=189, + use_guardrails=False, + timer=_started_timer(), + transfer_config=cfg, + video=b"input-clip", + ) + + assert windows == [(0, 188)] + + def test_multichunk_overlap_path(self, monkeypatch): + pipeline = _make_pipeline() + captured = {"targets": [], "conditional_frames": [], "decode_calls": []} + + tokenized = iter([(_ids(2), _mask()), (_ids(1), _mask())]) + pipeline._tokenize_prompt = lambda *args, **kwargs: next(tokenized) + + original_prepare = pipeline._prepare_transfer_latents + + def recording_prepare(target_norm, current_conditional_frames, generator): + captured["targets"].append(target_norm.detach().clone()) + captured["conditional_frames"].append(current_conditional_frames) + return original_prepare(target_norm, current_conditional_frames, generator) + + pipeline._prepare_transfer_latents = recording_prepare + + decoded_chunks = [ + torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2], dtype=torch.float32), + torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5], dtype=torch.float32), + ] + + def fake_decode(latents): + values = decoded_chunks[len(captured["decode_calls"])] + captured["decode_calls"].append(latents.detach().clone()) + return values.view(1, 1, 5, 1, 1).expand(1, 3, 5, 16, 16).clone() + + pipeline._decode_latents_raw = fake_decode + + # Input video: frame0 black (-1 normalized), frame1+ white (+1); the + # control is an all-black clip. Both arrive as encoded bytes and are + # decoded on the worker, so the decode is what the stub stands in for. + def fake_media_decode(data, *, first_frame, last_frame, target_h, target_w, device): + frames = torch.zeros(8, target_h, target_w, 3, dtype=torch.uint8) + if data == b"input-clip": + frames[1:] = 255 + return frames + + monkeypatch.setattr(transfer_module, "decode_video_reference_window", fake_media_decode) + cfg = resolve_transfer_config( + { + "edge": {"control": b"control-clip"}, + "guidance_scale": 1.0, + "control_guidance": 1.0, + "max_frames": 8, + "num_video_frames_per_chunk": 5, + "num_conditional_frames": 1, + "num_first_chunk_conditional_frames": 2, + }, + _req(num_frames=8, guidance_scale=1.0), + ) + + output = pipeline._forward_transfer( + prompt="transfer", + negative_prompt="", + height=16, + width=16, + max_frames=8, + num_inference_steps=1, + max_sequence_length=8, + use_system_prompt=False, + use_duration_template=False, + use_resolution_template=False, + seed=123, + frame_rate=24.0, + num_frames=8, + use_guardrails=False, + timer=_started_timer(), + transfer_config=cfg, + video=b"input-clip", + ) + + assert captured["conditional_frames"] == [2, 1] + assert len(captured["decode_calls"]) == 2 + # (B, T, H, W, C) uint8 -- what postprocess_video_tensor declares. + assert output.video.shape == (1, 8, 16, 16, 3) + # The stitched sequence, run through the same conversion the pipeline + # applies, so this still asserts the chunk arithmetic rather than a + # hand-computed uint8 table. + expected = torch.tensor([-0.6, -0.5, -0.4, -0.3, -0.2, 0.2, 0.3, 0.4]) + torch.testing.assert_close( + output.video[0, :, 0, 0, 0], + pipeline_module.postprocess_video_tensor( + expected.view(1, 1, 8, 1, 1).expand(1, 3, 8, 16, 16) + )[0, :, 0, 0, 0], + ) + # First chunk target: frame0 = normalized black input, frame1 = white, + # remainder filled by repeating the last conditional frame. + torch.testing.assert_close( + captured["targets"][0][:, :, 0], torch.full((1, 3, 16, 16), -1.0) + ) + torch.testing.assert_close(captured["targets"][0][:, :, 1], torch.full((1, 3, 16, 16), 1.0)) + torch.testing.assert_close( + captured["targets"][0][:, :, 2:], torch.full((1, 3, 3, 16, 16), 1.0) + ) + + +class TestControlLengthMismatch: + """Hint lengths decide the output length and how much control is invented.""" + + @staticmethod + def _frames(count: int) -> torch.Tensor: + return torch.zeros(3, count, 16, 16, dtype=torch.uint8) + + def _warning(self, monkeypatch, per_hint: dict, total_frames: int) -> str: + """Warnings emitted for one (hints, total_frames) pair. + + ``tensorrt_llm.logger`` does not propagate to the root logger, so this + captures through the module's logger rather than pytest's caplog. + """ + pipeline = _make_pipeline() + seen: list[str] = [] + monkeypatch.setattr( + pipeline_module.logger, + "warning", + lambda *args, **kwargs: seen.append(" ".join(str(a) for a in args)), + ) + pipeline._warn_on_control_length_mismatch(per_hint, total_frames) + return " ".join(seen) + + def test_warns_when_a_hint_is_mostly_padding(self, monkeypatch): + text = self._warning(monkeypatch, {"edge": self._frames(200), "seg": self._frames(50)}, 200) + assert "seg: 50" in text + assert "200 frames" in text + # The hint that set the length was not padded, so it is not named. + assert "edge:" not in text + + def test_silent_for_a_short_tail_difference(self, monkeypatch): + """A few frames of ping-pong at the tail is the normal case the + reference pads for; warning on it would train people to ignore this.""" + assert ( + self._warning(monkeypatch, {"edge": self._frames(200), "seg": self._frames(195)}, 200) + == "" + ) + + def test_silent_when_nothing_is_padded(self, monkeypatch): + """A request that pins `num_frames` down to the shortest clip truncates + the long hint rather than padding the short one, so no control is + invented and there is nothing to warn about.""" + assert ( + self._warning(monkeypatch, {"edge": self._frames(200), "seg": self._frames(50)}, 50) + == "" + ) + + def test_output_length_follows_the_longest_hint(self, monkeypatch): + """`edge` sorts before `seg` in TRANSFER_HINT_KEYS, so taking the first + hint's length would truncate the longer seg control to 4 frames and cut + the output in half -- on nothing the request expressed.""" + pipeline = _make_pipeline() + tokenized = iter([(_ids(2), _mask()), (_ids(1), _mask())]) + pipeline._tokenize_prompt = lambda *a, **k: next(tokenized) + pipeline._decode_latents_raw = lambda latents: torch.zeros(1, 3, 5, 16, 16) + + lengths = {b"edge-clip": 4, b"seg-clip": 8} + + def fake_media_decode(data, *, first_frame, last_frame, target_h, target_w, device): + count = min(lengths[data], last_frame + 1) + return torch.zeros(count, target_h, target_w, 3, dtype=torch.uint8) + + monkeypatch.setattr(transfer_module, "decode_video_reference_window", fake_media_decode) + cfg = resolve_transfer_config( + { + "edge": {"control": b"edge-clip"}, + "seg": {"control": b"seg-clip"}, + "max_frames": 8, + "num_video_frames_per_chunk": 5, + "num_conditional_frames": 1, + }, + _req(num_frames=8, guidance_scale=1.0), + ) + output = pipeline._forward_transfer( + prompt="transfer", + negative_prompt="", + height=16, + width=16, + max_frames=8, + num_inference_steps=1, + max_sequence_length=8, + use_system_prompt=False, + use_duration_template=False, + use_resolution_template=False, + seed=1, + frame_rate=24.0, + num_frames=8, + use_guardrails=False, + timer=_started_timer(), + transfer_config=cfg, + video=None, + ) + # (B, T, H, W, C): frames are dim 1. + assert output.video.shape[1] == 8 + + +class TestTransferFrameRateResolution: + """What `_forward_transfer` actually emits at, not what `infer()` hands it. + + `resolve_transfer_config` runs before the pipeline probes the source, and + `_forward_transfer` prefers ``transfer_config.fps`` over its ``frame_rate`` + argument. A test that stubs ``forward()`` and asserts what ``infer()`` + passes cannot see that, and will pass while the value is discarded one + layer down -- which is exactly what happened. + """ + + def _emitted_fps(self, monkeypatch, *, hints, caller, infer_fps): + pipeline = _make_pipeline() + monkeypatch.setattr( + transfer_module, "decode_video_reference_window", _fake_decode_window(5) + ) + tokenized = iter([(_ids(2), _mask()), (_ids(1), _mask())]) + pipeline._tokenize_prompt = lambda *a, **k: next(tokenized) + pipeline._decode_latents_raw = lambda latents: torch.zeros(1, 3, 5, 16, 16) + + params = _merged_req(frame_rate=COSMOS3_720P_PARAMS["frame_rate"], num_frames=5) + for key, value in caller.items(): + setattr(params, key, value) # assignment marks it as caller intent + cfg = resolve_transfer_config({**hints, "num_video_frames_per_chunk": 5}, params) + return pipeline._forward_transfer( + prompt="transfer", + negative_prompt="", + height=16, + width=16, + max_frames=cfg.max_frames, + num_inference_steps=35, + max_sequence_length=8, + use_system_prompt=False, + use_duration_template=False, + use_resolution_template=False, + seed=1, + frame_rate=infer_fps, # what infer() resolved and passed down + num_frames=5, + use_guardrails=False, + timer=_started_timer(), + transfer_config=cfg, + video=None, + ).frame_rate + + def test_source_rate_reaches_the_output(self, monkeypatch): + # The regression: config.fps used to capture the executor-merged 24 and + # shadow the rate the pipeline inferred from the source. + fps = self._emitted_fps( + monkeypatch, hints={"edge": {"control": b"clip"}}, caller={}, infer_fps=8.0 + ) + assert fps == 8.0 + + def test_explicit_request_rate_wins_over_the_source(self, monkeypatch): + fps = self._emitted_fps( + monkeypatch, + hints={"edge": {"control": b"clip"}}, + caller={"frame_rate": 30.0}, + infer_fps=30.0, + ) + assert fps == 30.0 + + def test_a_pinned_frame_count_keeps_the_default_rate(self, monkeypatch): + # `seconds` is converted to num_frames at the default rate before the + # worker sees the media, so adopting the source rate here would change + # the duration the caller asked for. + fps = self._emitted_fps( + monkeypatch, + hints={"edge": {"control": b"clip"}}, + caller={"num_frames": 5}, + infer_fps=COSMOS3_720P_PARAMS["frame_rate"], + ) + assert fps == COSMOS3_720P_PARAMS["frame_rate"] + + def test_wsm_preset_outranks_the_source(self, monkeypatch): + fps = self._emitted_fps( + monkeypatch, hints={"wsm": {"control": b"clip"}}, caller={}, infer_fps=8.0 + ) + assert fps == 10 + + +class TestTransferSamplingAndSafety: + """The transfer branch must not skip what every other mode goes through.""" + + def _run(self, pipeline, monkeypatch, *, use_guardrails=False, cfg_extra=None, tokenize=None): + monkeypatch.setattr( + transfer_module, "decode_video_reference_window", _fake_decode_window(5) + ) + tokenized = iter([(_ids(2), _mask()), (_ids(1), _mask())]) + pipeline._tokenize_prompt = tokenize or (lambda *a, **k: next(tokenized)) + pipeline._decode_latents_raw = lambda latents: torch.zeros(1, 3, 5, 16, 16) + cfg = resolve_transfer_config( + {"edge": {"control": b"clip"}, "num_video_frames_per_chunk": 5, **(cfg_extra or {})}, + _req(num_frames=5, guidance_scale=1.0), + ) + return pipeline._forward_transfer( + prompt="transfer", + negative_prompt="", + height=16, + width=16, + max_frames=cfg.max_frames, + num_inference_steps=35, + max_sequence_length=8, + use_system_prompt=False, + use_duration_template=False, + use_resolution_template=False, + seed=1, + frame_rate=24.0, + num_frames=5, + use_guardrails=use_guardrails, + timer=_started_timer(), + transfer_config=cfg, + video=None, + ) + + def test_control_directive_reaches_the_tokenizer(self, monkeypatch): + """Resolving the directive is not enough -- it has to reach the text + the model actually conditions on, and only the positive prompt.""" + pipeline = _make_pipeline() + seen = [] + tokenized = iter([(_ids(2), _mask()), (_ids(1), _mask())]) + + def capture(prompt, *args, **kwargs): + seen.append(prompt) + return next(tokenized) + + self._run(pipeline, monkeypatch, tokenize=capture) + assert "Follow the edge control video precisely" in seen[0] + assert "Follow the edge" not in seen[1], "directive leaked into the negative prompt" + + def test_transfer_installs_a_shifted_scheduler(self, monkeypatch): + """Transfer owns its scheduler setup, so it must actually perform it. + + ``forward()`` deliberately skips the scheduler rebuild when a transfer + config is present, which makes this the only place the shift is applied + -- and the only thing standing between a request and whatever schedule + the previous request left on the worker. Asserted through the real + ``_scheduler_for``: a rename of that helper breaks here rather than + being papered over by a stub. + """ + sampling = StubSamplingPolicy() + pipeline = _make_pipeline(sampling=sampling) + pipeline.scheduler = StubScheduler() + baseline = pipeline.scheduler + self._run(pipeline, monkeypatch) + assert sampling.flow_shift_calls == [ + transfer_module.TRANSFER_DEFAULTS["edge"]["flow_shift"] + ], "transfer did not program the hint's flow shift" + # Installed, not merely built: a bare call that drops the returned + # scheduler leaves the previous request's schedule in place. + assert pipeline.scheduler is not baseline, ( + "transfer built a scheduler but never installed it" + ) + + def test_schedule_is_programmed_through_the_sampling_policy(self, monkeypatch): + """A distilled checkpoint runs a fixed sigma list, not a step count.""" + sampling = StubSamplingPolicy(is_distilled=True) + pipeline = _make_pipeline(sampling=sampling) + self._run(pipeline, monkeypatch) + assert sampling.set_timesteps_calls, "transfer bypassed the sampling policy" + # Distilled: the policy substitutes its own step count for the request's. + assert pipeline.scheduler.sigmas_calls == [list(sampling.fixed_sigmas)] + + def test_distilled_step_draws_noise_from_the_seeded_generator(self, monkeypatch): + """Unseeded SDE noise diverges the replicated latents across ranks.""" + sampling = StubSamplingPolicy(is_distilled=True) + pipeline = _make_pipeline(sampling=sampling) + self._run(pipeline, monkeypatch) + assert pipeline.scheduler.step_generators, "scheduler.step() got no generator" + assert all(g is not None for g in pipeline.scheduler.step_generators) + + def test_base_checkpoint_passes_no_generator(self, monkeypatch): + pipeline = _make_pipeline(sampling=StubSamplingPolicy(is_distilled=False)) + self._run(pipeline, monkeypatch) + assert pipeline.scheduler.step_generators == [] + + def test_output_is_screened_when_guardrails_are_on(self, monkeypatch): + pipeline = _make_pipeline() + seen = {} + + class Checker: + pass + + pipeline.safety_checker = Checker() + monkeypatch.setattr( + pipeline_module, + "check_video_safety", + lambda video, checker: seen.setdefault("called", True) and video, + ) + self._run(pipeline, monkeypatch, use_guardrails=True) + assert seen.get("called"), "transfer returned generated video unscreened" + + def test_phase_timings_are_populated(self, monkeypatch): + output = self._run(_make_pipeline(), monkeypatch) + assert output.denoise > 0.0, "transfer reported a zero denoise phase" + + +class TestFindClosestTargetSize: + """Maps a source frame onto the aspect-ratio-closest output bucket for a + resolution level, returning ``(width, height)``.""" + + @pytest.mark.parametrize( + "h, w, resolution, expected", + [ + # Exact aspect ratios at the 720 level resolve to their own bucket. + (720, 1280, 720, (1280, 720)), # 16:9 landscape + (1280, 720, 720, (720, 1280)), # 9:16 portrait + (512, 512, 720, (960, 960)), # 1:1 square + (768, 1024, 720, (1104, 832)), # 4:3 landscape + (1024, 768, 720, (832, 1104)), # 3:4 portrait + # Other levels select from that level's own table. + (720, 1280, 480, (832, 480)), + (256, 256, 256, (256, 256)), + ], + ) + def test_maps_to_matching_bucket(self, h, w, resolution, expected): + assert find_closest_target_size(h, w, resolution) == expected + + def test_returns_width_height_order(self): + # A landscape source (w > h) must yield a landscape bucket (target_w > + # target_h); guards against an (h, w) transposition of the return value. + target_w, target_h = find_closest_target_size(720, 1280, 720) + assert (target_w, target_h) == (1280, 720) + assert target_w > target_h + + def test_picks_nearest_when_no_exact_match(self): + # A 2:1 ultra-wide source has no exact bucket; the closest ratio at the + # 720 level is 16:9 (1280x720). + assert find_closest_target_size(500, 1000, 720) == (1280, 720) + + def test_resolution_accepts_int_or_str(self): + assert find_closest_target_size(720, 1280, 720) == find_closest_target_size( + 720, 1280, "720" + ) + + def test_unknown_resolution_raises(self): + with pytest.raises(ValueError, match="Unknown Cosmos3 transfer resolution"): + find_closest_target_size(720, 1280, 1080) 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 ea2a96e6f12a..218d580e27ee 100644 --- a/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py +++ b/tests/unittest/_torch/visual_gen/test_visual_gen_utils.py @@ -33,8 +33,9 @@ class _StubExtraParamSpec: - def __init__(self, default: Any = None) -> None: + def __init__(self, default: Any = None, type: str = "str") -> None: self.default = default + self.type = type class _StubVisualGen: @@ -615,3 +616,56 @@ def test_empty_extras_dict_normalizes_to_none(self): params = self._make_params() _merge_extra_params(params, request_extras=None, extra_param_specs={}) assert params.extra_params is None + + +# ============================================================================= +# Inline binary extra params — base64 in, bytes out +# ============================================================================= + + +class TestInlineMediaDecoding: + """A pipeline that declares a ``bytes`` extra param must receive bytes. + + JSON has no byte type, so an HTTP client can only inline binary as base64. + Decoding happens here rather than in the pipeline, which keeps a + bytes-only contract. Cosmos3 transfer's precomputed controls + (`depth`/`seg`/`wsm`) are unreachable over serving without it. + """ + + def _generator(self): + return _StubVisualGen( + extra_param_specs={ + "video": _StubExtraParamSpec(type="bytes"), + "edge": _StubExtraParamSpec(type="bool_or_bytes_or_dict"), + "resolution": _StubExtraParamSpec(type="str"), + } + ) + + def test_base64_extra_param_reaches_the_pipeline_as_bytes(self): + request = VideoGenerationRequest( + prompt="storm", + extra_params={"video": base64.b64encode(b"\x00mp4").decode()}, + ) + params = parse_visual_gen_params(request, "id-b64", self._generator()) + assert params.extra_params["video"] == b"\x00mp4" + + def test_nested_control_reaches_the_pipeline_as_bytes(self): + request = VideoGenerationRequest( + prompt="storm", + extra_params={"edge": {"control": base64.b64encode(b"\x00ctrl").decode()}}, + ) + params = parse_visual_gen_params(request, "id-nested", 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()) + 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())