diff --git a/docs/source/models/supported-models.md b/docs/source/models/supported-models.md index 9c09a53d1c9f..fd77fd29eb72 100644 --- a/docs/source/models/supported-models.md +++ b/docs/source/models/supported-models.md @@ -206,6 +206,7 @@ For full documentation, see the [Visual Generation](./visual-generation.md) page | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | | `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) | ### Feature Matrix diff --git a/docs/source/models/visual-generation.md b/docs/source/models/visual-generation.md index 70821ea81ef1..da537acff32a 100644 --- a/docs/source/models/visual-generation.md +++ b/docs/source/models/visual-generation.md @@ -44,6 +44,7 @@ TensorRT-LLM **VisualGen** provides a unified inference stack for diffusion mode | `nvidia/Cosmos3-Super` | Text-to-Image, Text-to-Video, Image-to-Video | | `nvidia/Cosmos3-Super-Text2Image-4Step` | Text-to-Image (DMD2-distilled, fixed 4-step schedule) | | `nvidia/Cosmos3-Super-Image2Video-4Step` | Image-to-Video (DMD2-distilled, fixed 4-step schedule) | +| `nvidia/Cosmos3-Edge` | Text-to-Image, Text-to-Video, Image-to-Video (Nemotron-dense backbone, 480p-native) | Models are auto-detected from the checkpoint directory. Diffusers-format models are detected via `model_index.json`; LTX-2 monolithic safetensors checkpoints are detected via embedded metadata. The `AutoPipeline` registry selects the appropriate pipeline class automatically. diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md index 4ea3f64b4d0b..5b6981bb5647 100644 --- a/examples/visual_gen/models/cosmos3/README.md +++ b/examples/visual_gen/models/cosmos3/README.md @@ -5,7 +5,7 @@ Cosmos3 supports the following generation modes from a single checkpoint: - **T2V** — text-to-video (`prompts/t2v.json`). - **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)). +- **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. - **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 @@ -16,6 +16,7 @@ Pass the Hub ID or local path via `--model`: - [`nvidia/Cosmos3-Super`](https://huggingface.co/nvidia/Cosmos3-Super) - [`nvidia/Cosmos3-Super-Text2Image-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Text2Image-4Step) — DMD2-distilled text-to-image: fixed 4-step schedule with classifier-free guidance baked into the weights. Steps/guidance are read from the checkpoint; conflicting request values are rejected. Use with `configs/cosmos3-t2i-1gpu.yaml`. - [`nvidia/Cosmos3-Super-Image2Video-4Step`](https://huggingface.co/nvidia/Cosmos3-Super-Image2Video-4Step) — DMD2-distilled image-to-video: same fixed 4-step, guidance-baked-in contract. The default omni video shape (720p × 189 frames) is the deployed shape, so no dedicated config is needed. This checkpoint declares `default_use_system_prompt: true` in its `model_index.json`, which the pipeline applies automatically (override with `--use_system_prompt` / `--no-use_system_prompt`). +- [`nvidia/Cosmos3-Edge`](https://huggingface.co/nvidia/Cosmos3-Edge) — 4B Nemotron-dense backbone supporting **T2I / T2V / I2V only**: no audio tower, and the checkpoint's action weights are not supported by this pipeline yet. 480p-native defaults (832×480 × 121 frames, 50 UniPC steps on the checkpoint-declared native flow schedule with shift 3.0, guidance 5.0; T2I defaults to 640×640), so no dedicated config is needed. The model card validates 256p/480p, 50–150 frames, and 12–30 FPS; requests outside that envelope run with an advisory log. ## Guardrails @@ -48,6 +49,25 @@ See `examples/visual_gen/configs/`: Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`). +### Prompt inputs + +`--prompt` and `--negative_prompt` each accept **either literal text or a path to a +prompt file**, chosen by whether the value names an existing file. `--prompt_file` +and `--negative_prompt_file` accept a path only and fail if the file is missing, so +use those when a silent fallback to literal text would be a bug (scripts, CI). + +A prompt file may hold any of three shapes: + +| Shape | Example | Notes | +|---|---|---| +| Omni prompt object | `prompts/t2v.json` | `prompt` plus optional `model_mode`, `vision_path`, `enable_audio`, which supply defaults for the matching flags | +| Structured caption | a checkpoint's `assets/example_i2v_prompt.json` | the object *is* the caption; carries no options | +| Plain text | any `.txt` | used verbatim | + +Structured captions are what the model cards ship and what the checkpoints were +tuned on; they give noticeably cleaner output than a one-line summary. +`--negative_prompt` defaults to `cosmos3_negative_prompt.json` in this directory. + ## Usage ```bash @@ -104,7 +124,17 @@ python cosmos3.py --model nvidia/Cosmos3-Super-Image2Video-4Step \ --image_path https://example.com/frame.jpg \ --output_path output.mp4 -# Inline prompt (--prompt or a JSON file path) +# 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 +# hf download nvidia/Cosmos3-Edge --local-dir Cosmos3-Edge +python cosmos3.py --model nvidia/Cosmos3-Edge \ + --prompt Cosmos3-Edge/assets/example_i2v_prompt.json \ + --negative_prompt Cosmos3-Edge/assets/negative_prompt.json \ + --image_path Cosmos3-Edge/assets/example_i2v_input.jpg \ + --output_path output.mp4 + +# Inline prompt 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 31aa0cb3f936..57acd701506f 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -31,6 +31,9 @@ _SCRIPT_DIR = Path(__file__).resolve().parent +DEFAULT_PROMPT_FILE = "prompts/t2v.json" +DEFAULT_NEGATIVE_PROMPT_FILE = "cosmos3_negative_prompt.json" + def _resolve_path(path: str) -> str: candidate = Path(path) @@ -42,18 +45,76 @@ def _resolve_path(path: str) -> str: return path -def load_prompt_file(path: str) -> Dict[str, Any]: - """Load a Cosmos3 omni prompt JSON (``prompt``, optional ``vision_path``, etc.).""" +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)) + + +def _read_prompt_payload(path: str) -> Any: + """Read a prompt file, decoding it as JSON when it parses and as text otherwise.""" resolved = _resolve_path(path) + if not os.path.isfile(resolved): + raise ValueError(f"Prompt file {path!r} does not exist (resolved to {resolved!r}).") with open(resolved, encoding="utf-8") as f: - data = json.load(f) + raw = f.read() + try: + return json.loads(raw) + except json.JSONDecodeError: + return raw.strip() + + +def load_prompt_file(path: str) -> Dict[str, Any]: + """Load a Cosmos3 prompt file. + + Three shapes are accepted: an omni prompt object (``prompt`` plus optional + ``vision_path`` / ``model_mode`` / ``enable_audio``), a structured caption + object such as the ``assets/*_prompt.json`` files shipped with a checkpoint, + or plain text. The latter two carry no options, so they yield ``prompt`` only. + """ + data = _read_prompt_payload(path) + if isinstance(data, str): + if not data: + raise ValueError(f"Prompt file {path!r} is empty.") + return {"prompt": data} if not isinstance(data, dict): - raise ValueError(f"Prompt file must be a JSON object, got {type(data)!r}.") - if not data.get("prompt"): - raise ValueError(f"Prompt file {resolved!r} is missing a non-empty 'prompt' field.") + raise ValueError( + f"Prompt file {path!r} must hold a JSON object or text, got {type(data).__name__}." + ) + if "prompt" not in data: + if not data: + raise ValueError(f"Prompt file {path!r} is an empty JSON object.") + return {"prompt": json.dumps(data)} + if not data["prompt"]: + raise ValueError(f"Prompt file {path!r} is missing a non-empty 'prompt' field.") return data +def load_negative_prompt_file(path: str) -> str: + """Load a negative prompt file (structured JSON object or plain text).""" + data = _read_prompt_payload(path) + if isinstance(data, dict): + return json.dumps(data) + if isinstance(data, str): + return data + raise ValueError( + f"Negative prompt file {path!r} must hold a JSON object or text, got {type(data).__name__}." + ) + + +def resolve_negative_prompt( + *, + negative_prompt: Optional[str], + negative_prompt_file: Optional[str], +) -> str: + """Pick the negative prompt: ``--negative_prompt``, then the file, then the default.""" + if negative_prompt is not None: + # --negative_prompt takes either literal text or a path to a prompt file. + if _is_prompt_file(negative_prompt): + return load_negative_prompt_file(negative_prompt) + return negative_prompt + return load_negative_prompt_file(negative_prompt_file or DEFAULT_NEGATIVE_PROMPT_FILE) + + def resolve_prompt_and_options( *, prompt: Optional[str], @@ -67,7 +128,15 @@ def resolve_prompt_and_options( if prompt_file is not None: prompt_data = load_prompt_file(prompt_file) - resolved_prompt = prompt + inline_prompt: Optional[str] = None + if prompt is not None: + # --prompt takes either literal text or a path to a prompt file. + if _is_prompt_file(prompt): + prompt_data = {**prompt_data, **load_prompt_file(prompt)} + else: + inline_prompt = prompt + + resolved_prompt = inline_prompt if resolved_prompt is None: resolved_prompt = prompt_data.get("prompt") if not resolved_prompt: @@ -93,7 +162,8 @@ def main(): "--model", type=str, default="nvidia/Cosmos3-Nano", - help="Model path or HuggingFace Hub ID (nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super)", + help="Model path or HuggingFace Hub ID " + "(nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super, nvidia/Cosmos3-Edge)", ) parser.add_argument( "--visual_gen_args", @@ -106,19 +176,27 @@ def main(): "--prompt", type=str, default=None, - help="Text prompt for generation (overrides --prompt_file when both are set)", + help="Prompt text, or a path to a prompt file (overrides --prompt_file when both are set)", ) parser.add_argument( "--prompt_file", type=str, - default="prompts/t2v.json", - help="Path to a JSON prompt file (default: prompts/t2v.json)", + default=DEFAULT_PROMPT_FILE, + help=f"Path to a prompt file; must exist (default: {DEFAULT_PROMPT_FILE})", ) parser.add_argument( "--negative_prompt", type=str, - default="cosmos3_negative_prompt.json", - help="Text prompt or path to JSON file for negative prompt", + default=None, + help="Negative prompt text, or a path to a negative prompt file " + f"(overrides --negative_prompt_file; default: {DEFAULT_NEGATIVE_PROMPT_FILE})", + ) + parser.add_argument( + "--negative_prompt_file", + type=str, + default=None, + help=f"Path to a negative prompt file; must exist " + f"(default: {DEFAULT_NEGATIVE_PROMPT_FILE})", ) parser.add_argument( "--image_path", @@ -187,15 +265,10 @@ def main(): if image_path is not None: params.image = image_path - negative_prompt_path = _resolve_path(args.negative_prompt) - if args.negative_prompt is not None: - if os.path.isfile(negative_prompt_path) and negative_prompt_path.endswith(".json"): - with open(negative_prompt_path, encoding="utf-8") as f: - negative_prompt = json.load(f) - else: - negative_prompt = args.negative_prompt - else: - negative_prompt = None + negative_prompt = resolve_negative_prompt( + negative_prompt=args.negative_prompt, + negative_prompt_file=args.negative_prompt_file, + ) if args.disable_duration_template: params.extra_params["use_duration_template"] = False @@ -210,12 +283,7 @@ def main(): if args.video_path is not None: params.extra_params["video"] = Path(args.video_path).read_bytes() - if negative_prompt is None: - params.negative_prompt = None - elif isinstance(negative_prompt, str): - params.negative_prompt = negative_prompt - else: - params.negative_prompt = json.dumps(negative_prompt) + params.negative_prompt = negative_prompt output = visual_gen.generate( inputs=prompt, diff --git a/tensorrt_llm/_torch/visual_gen/executor.py b/tensorrt_llm/_torch/visual_gen/executor.py index 166951919e4f..25edec04570a 100644 --- a/tensorrt_llm/_torch/visual_gen/executor.py +++ b/tensorrt_llm/_torch/visual_gen/executor.py @@ -418,6 +418,12 @@ def _merge_defaults(self, req: DiffusionRequest): ): continue setattr(params, field_name, default_value) + # Marks it as a pipeline default rather than caller intent, so + # request-dependent defaults stay re-resolvable; assigning the + # field re-marks it. + # Assumes model_fields_set is the live __pydantic_fields_set__, not a + # copy; TestDefaultMarksThroughRealPath fails loudly if that changes. + params.model_fields_set.discard(field_name) # Extra param defaults — fill all declared keys so infer() can use direct access specs = self.pipeline.extra_param_specs diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index af1e5f9272c2..359907801466 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -111,15 +111,67 @@ def _validate_video_reference(video) -> None: "guidance_interval": (400.0, 1000.0), } -# Fields merged by the executor into every request. Mode-dependent values -# remain None until infer() selects the request mode; key membership also -# declares these fields supported during request validation. -COSMOS3_PIPELINE_DEFAULTS = { - **COSMOS3_720P_PARAMS, - "height": None, - "width": None, - "num_inference_steps": None, - "guidance_scale": None, +# Edge (Nemotron-dense backbone) is 480p-native. Video values follow the +# model-card I2V command (T2V mirrors it — the model card documents I2V only); +# ``flow_shift`` rides the checkpoint-declared native flow schedule. T2I values +# are the cosmos-framework t2i mode defaults at Edge's native resolution +# (480p at 1:1 aspect), with full-range CFG. +COSMOS3_EDGE_VIDEO_PARAMS = { + "height": 480, + "width": 832, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "max_sequence_length": 4096, + "num_frames": 121, + "frame_rate": 24.0, + "flow_shift": 3.0, +} + +COSMOS3_EDGE_T2I_PARAMS = { + "height": 640, + "width": 640, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "flow_shift": 3.0, + "guidance_interval": None, +} + +# Model-card validated envelope for Edge; advisory only (the reference +# runtime accepts a wider range), surfaced as a log line per request. +COSMOS3_EDGE_ENVELOPE = { + "num_frames": (50, 150), + "frame_rate": (12.0, 30.0), + "max_sequence_length": 4096, + "resolutions": frozenset( + { + (640, 640), + (544, 736), + (736, 544), + (480, 832), + (832, 480), + (256, 256), + (256, 320), + (320, 256), + (192, 320), + (320, 192), + } + ), +} + +# (family, mode) → generation defaults. Family is the architecture recipe +# name resolved from the transformer config; mode is the request's output +# type — never inferred from the checkpoint name (a task-specialized +# checkpoint can still be asked to run any mode). +COSMOS3_GENERATION_DEFAULTS: Dict = { + ("qwen3", "video"): COSMOS3_720P_PARAMS, + ("qwen3", "image"): COSMOS3_T2I_PARAMS, + ("nemotron_dense", "video"): COSMOS3_EDGE_VIDEO_PARAMS, + ("nemotron_dense", "image"): COSMOS3_EDGE_T2I_PARAMS, +} + +# Families without an entry get no envelope advisory. +COSMOS3_ENVELOPES: Dict = { + "nemotron_dense": COSMOS3_EDGE_ENVELOPE, } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/negative_prompt.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/negative_prompt.py new file mode 100644 index 000000000000..b4cd5c3ae389 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/negative_prompt.py @@ -0,0 +1,136 @@ +# 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. + +# ruff: noqa: E501 +"""Cosmos3's default negative prompt for video modes. + +Verbatim copy of the reference's ``inference/defaults/neg_prompts.json``, which every +video mode (``text2video``, ``image2video``, ``video2video``, ``audio_image2video``) +points at via ``negative_prompt_file``. Image modes declare none. + +Kept as a Python literal rather than a data file so it ships with the package in every +install mode without a ``package_data`` entry. Key order is significant: the serialized +form must match the reference byte for byte, and ``json.dumps`` preserves insertion +order. +""" + +COSMOS3_VIDEO_NEGATIVE_PROMPT = { + "subjects": [ + { + "description": "Blurry, poorly defined subjects with inconsistent shapes and unrealistic proportions.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0, + }, + { + "description": "Extremely low-quality subjects with visible rendering artifacts, broken mesh geometry, and completely unrealistic proportions throughout.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0, + }, + { + "description": "Poorly generated subjects exhibiting all hallmarks of failed neural rendering — flickering edges, inconsistent depth, and uncanny spatial relationships.", + "appearance_details": "Distorted features, visible compression artifacts, muddy textures lacking fine detail, color bleeding between elements, and unnatural skin tones or surface textures that appear artificial or computer-generated.", + "relationship": "Subjects appear disconnected from the environment, floating or improperly grounded in the scene without proper occlusion or spatial coherence.", + "location": "Subjects are poorly placed within the frame, appearing at awkward positions that violate basic compositional rules.", + "relative_size": "Inconsistent scale relationships between subjects and the environment, with objects appearing too large or too small relative to their surroundings.", + "orientation": "Unnatural orientations that defy physics and spatial logic.", + "pose": "Stiff, mannequin-like poses with unnatural joint angles and impossible limb positions that look computer-generated.", + "action": "Incoherent motion with visible frame-to-frame discontinuities. Movement appears as a slideshow rather than smooth animation. Limbs and appendages pop between positions without interpolation.", + "state_changes": "Visual state transitions are abrupt and jarring. Colors shift without motivation. Surface textures flicker between different materials randomly. Outlines shimmer and vibrate.", + "clothing": "Clothing appears painted on with no sense of material weight or drape. Fabric textures are flat and repeat visibly.", + "expression": "Frozen, uncanny valley expressions or expressions that change abruptly without natural transition.", + "gender": "", + "age": "", + "skin_tone_and_texture": "Waxy, plastic-looking skin with visible artifacts and inconsistent texture resolution across the frame.", + "facial_features": "Asymmetric facial features, extra fingers or limbs, teeth that appear blurry or malformed.", + "number_of_subjects": 0, + "number_of_arms": 0, + "number_of_legs": 0, + }, + ], + "background_setting": "A poorly rendered, flat background with visible seams, repeated textures, and inconsistent depth cues. The environment lacks volumetric depth and appears as a painted backdrop rather than a three-dimensional space. Vegetation looks like flat cutouts with no volumetric depth. The background appears to have been composited from multiple source materials at different resolutions, creating visible seams and edge artifacts where elements meet. Textures swim and shift across surfaces in a way that breaks the illusion of solidity — patterns drift laterally rather than staying anchored to the geometry they belong to. Background elements flicker in and out of existence between frames, particularly at the edges of the field of view. The rendering resolution is visibly lower for distant elements, creating a jarring transition between near and far objects. Cloud textures repeat obviously in the sky with visible tiling. Water surfaces lack proper reflection and refraction, appearing as flat animated textures. Fog and atmospheric effects pop in and out rather than smoothly transitioning. Trees and vegetation exhibit obvious LOD (level-of-detail) switching. Building facades have inconsistent window spacing and pattern repetition. The overall scene feels like a poorly assembled collage of individually rendered elements rather than a coherent whole.", # codespell:ignore + "lighting": { + "conditions": "Harsh, flat lighting with no natural variation. The scene appears uniformly lit as if by a single overhead fluorescent light, removing all sense of depth and atmosphere.", + "direction": "Inconsistent light sources — shadows point in multiple contradictory directions, breaking physical plausibility.", + "shadows": "Hard-edged, unrealistic shadows that pop in and out of existence between frames. Some objects cast no shadows while others have impossibly dark ones that don't animate smoothly with the object's motion. Shadow edges exhibit visible staircase aliasing artifacts. Shadow maps appear to have been rendered at extremely low resolution, creating blocky patterns. Self-shadowing on characters shows visible peter-panning artifacts where shadows detach from their source. Contact shadows between objects and the ground appear and disappear as objects move slightly. Shadow color is pure black with no ambient contribution, creating an unnaturally harsh contrast that flattens the image. Multiple shadow cascades have visible boundaries where resolution changes. The shadow rendering appears to be temporally unstable — even static objects have shadows that shimmer and crawl frame to frame, breaking the illusion of a stable light source.", + "illumination_effect": "No bounce light, no ambient occlusion, no subtle color interactions between surfaces. The scene looks like a poorly lit 3D render from the early 2000s.", + }, + "aesthetics": { + "composition": "Cluttered, poorly framed composition with no clear focal point. Important elements are cut off by the frame edges. The rule of thirds is completely ignored, leading to an unbalanced and visually unpleasant arrangement.", + "color_scheme": "Oversaturated, garish colors that clash violently. Color banding is visible in gradient areas. The overall palette feels artificial and digitally processed rather than natural.", + "mood_atmosphere": "Unsettling, uncanny atmosphere that fails to evoke any intended emotional response. The scene feels lifeless and sterile despite attempting to portray dynamic action.", + "patterns": "Visible tiling artifacts in textures, moiré patterns, and aliasing on edges.", + }, + "cinematography": { + "camera_motion": "Extremely shaky, unstable camera with visible rolling shutter artifacts. The motion is jerky and discontinuous, causing motion sickness and making the scene impossible to follow.", + "framing": "Poorly framed shots that cut off important elements and include unnecessary empty space.", + "camera_angle": "Awkward, disorienting camera angles that provide no useful spatial information about the scene. The camera path exhibits visible mathematical artifacts suggesting simple interpolation between keyframes rather than natural camera operation. Camera motion is completely disconnected from the scene content — panning away from action, dollying during dialogue, and shaking during still moments. The camera appears to pass through solid objects occasionally. Zoom is applied digitally rather than optically, revealing progressively worse resolution. Camera motion exhibits non-physical acceleration profiles — instant starts and stops rather than smooth ease-in/ease-out. Rolling shutter simulation is applied inconsistently, present in some frames but not others. The camera occasionally exhibits impossible motion like teleporting between positions. Virtual camera stabilization creates an uncanny floating sensation disconnected from any physical camera rig.", + "depth_of_field": "Uniform focus throughout, creating a flat, documentary-like appearance with no cinematic depth separation.", + "focus": "Soft, out-of-focus imagery with visible chromatic aberration and lens distortion that was not corrected in post-processing.", + "lens_focal_length": "Inappropriate focal length causing barrel distortion and unnatural perspective compression.", + }, + "style_medium": "Low quality compressed digital video with visible encoding artifacts", + "artistic_style": "Amateur, unpolished with inconsistent visual style", + "context": "A poorly produced video with numerous technical and artistic flaws that detract from any intended narrative or visual impact.", + "actions": [ + { + "time": "0:00-0:08", + "description": "Subjects attempt to move but their motion is jerky, temporally inconsistent, and physically implausible. Background elements flicker and shift between frames.", + } + ], + "text_and_signage_elements": [], + "segments": [ + { + "segment_index": 0, + "time_range": "0:00-0:08", + "description": "A single continuous shot suffering from severe temporal inconsistencies — subjects that morph and deform between frames, backgrounds that shift and wobble, and rendering quality that fluctuates visibly over time. Motion blur is applied incorrectly, smearing in directions that don't match actual movement. Frame-to-frame coherence breaks down with individual pixels changing color randomly in flat areas. Texture detail level fluctuates between frames as if the rendering budget varied shot to shot. Color grading drifts over the duration with no creative motivation. Noise patterns change between frames in ways that draw attention rather than being invisible. Overall visual quality degrades progressively from start to finish.", + "key_changes": "No meaningful progression or narrative development. Visual quality degrades over time.", + "camera": "Unstable, poorly controlled camera work with visible mathematical interpolation artifacts.", + } + ], + "transitions": [], + "temporal_caption": "The scene opens at 0.0 seconds with a poorly rendered establishing shot that immediately reveals low production quality. At 1.0 seconds, subjects begin to move but their motion is jerky and inconsistent, with limbs bending at unnatural angles and objects clipping through each other. From 2.0 to 4.0 seconds, the camera shakes violently while the scene exhibits visible compression artifacts, color banding in the sky, and flickering in the shadows. Between 4.0 and 6.0 seconds, temporal coherence breaks down as elements appear and disappear between frames, textures swim and morph unnaturally, and the lighting shifts abruptly without physical cause. In the final 2 seconds, the overall visual quality deteriorates further with increasing noise, blur, and a general loss of spatial coherence that makes the scene nearly unwatchable. Additionally, the frame rate appears inconsistent with visible judder and stuttering throughout. Color temperature shifts randomly between warm and cool tones with no motivation. The encode quality degrades in complex regions showing macro-blocking and mosquito noise around moving edges. Temporal noise patterns are spatially correlated, creating swimming artifacts on flat surfaces.", + "audio_description": "", + "physical_realism": "No adherence to physical laws. Objects defy gravity, pass through solid surfaces, and change mass and momentum without cause. Fluid dynamics, cloth simulation, and rigid body physics are all fundamentally broken. Furthermore, conservation of energy is violated as objects gain or lose kinetic energy spontaneously. Elastic collisions produce inelastic results and vice versa. Surface friction is inconsistent — objects slide on rough surfaces while sticking to smooth ones. Air resistance appears to affect only some objects while others move through the atmosphere unimpeded.", +} 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 077a6bbc3b7e..dbe25e0eec86 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -13,18 +13,19 @@ # See the License for the specific language governing permissions and # limitations under the License. +import functools import json import math import os import time -from typing import Iterable, List, Optional, Union +from typing import Any, Iterable, List, Optional, Union import PIL.Image import torch from diffusers import AutoencoderKLWan from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor -from transformers import Qwen2Tokenizer +from transformers import AutoTokenizer from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline @@ -40,19 +41,46 @@ from tensorrt_llm.media.decoding import decode_video_reference_window from .defaults import ( - COSMOS3_720P_PARAMS, + COSMOS3_ENVELOPES, COSMOS3_EXTRA_SPECS, - COSMOS3_PIPELINE_DEFAULTS, - COSMOS3_T2I_PARAMS, + COSMOS3_GENERATION_DEFAULTS, _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 .sound_tokenizer import LatentAutoEncoderV2 -from .transformer_cosmos3 import Cosmos3VFMTransformer +from .transformer_cosmos3 import NEMOTRON_DENSE_RECIPE, Cosmos3VFMTransformer, resolve_arch_recipe +# Image modes declare no negative prompt in the reference +# while every video mode points at ``neg_prompts.json``. COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" + + +@functools.lru_cache(maxsize=1) +def default_video_negative_prompt() -> str: + """The reference's default negative prompt for video modes. + + Serialized the way the reference loads it -- ``json.dumps(json.loads(...))`` -- + so the text reaching the tokenizer is byte-identical. + """ + return json.dumps(COSMOS3_VIDEO_NEGATIVE_PROMPT) + + +def default_negative_prompt(output_type: str) -> str: + """Default negative prompt for a request, keyed on output kind not request mode. + + The reference wires its negative prompt into every video mode and none of the + image ones, so anything producing an image defaults to empty. + """ + return ( + COSMOS3_DEFAULT_NEGATIVE_PROMPT + if output_type == "image" + else default_video_negative_prompt() + ) + + # NOTE: Intentional typo in "give" instead of "given" to match training setup. COSMOS3_DEFAULT_SYSTEM_PROMPT = ( "You are a helpful assistant who will generate videos from a give prompt." @@ -60,12 +88,88 @@ COSMOS3_T2I_SYSTEM_PROMPT = ( "You are a helpful assistant who will generate images from a give prompt." ) +COSMOS3_V2V_FLOW_SHIFT = 10.0 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." TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" +# ``W,H`` bucket names the reference builds requests from. A request there starts +# as a (resolution, bucket) pair and the bucket string is carried into the prompt +# verbatim; we only ever see the resolved height/width, so map back to the nearest +# bucket. Emitting the exact reduced ratio instead would put a string the model +# never saw in training into the caption (832x480 reduces to "15,26", not "16,9"). +COSMOS3_ASPECT_RATIO_BUCKETS = ("1,1", "4,3", "3,4", "16,9", "9,16") + + +def _aspect_ratio_bucket(height: int, width: int) -> str: + """Nearest reference aspect-ratio bucket for a resolved frame size.""" + if height <= 0 or width <= 0: + raise ValueError( + f"Cosmos3 aspect ratio needs positive dimensions, got height={height}, width={width}." + ) + ratio = width / height + return min( + COSMOS3_ASPECT_RATIO_BUCKETS, + key=lambda bucket: abs( + math.log(ratio / (int(bucket.split(",")[0]) / int(bucket.split(",")[1]))) + ), + ) + + +def _validate_sampling_recipe(family: str, use_native_flow_schedule: bool, sampling) -> None: + """Family, model_index schedule flag, and scheduler recipe must form a + known-supported combination — the pieces come from three different + checkpoint files and a mismatch samples the wrong trajectory silently. + """ + if family == NEMOTRON_DENSE_RECIPE.name: + if sampling.is_distilled: + raise ValueError( + "Distilled (fixed-sigma FlowMatchEuler) sampling is not supported for " + "the Edge (nemotron_dense) family; no such checkpoint exists." + ) + if not use_native_flow_schedule: + raise ValueError( + "Edge (nemotron_dense) checkpoints must declare " + "use_native_flow_schedule: true in model_index.json. Without it the " + "checkpoint's karras scheduler config would sample the wrong " + "trajectory; a missing flag means a broken or stale conversion." + ) + elif use_native_flow_schedule: + raise ValueError( + "use_native_flow_schedule is only supported for the Edge " + f"(nemotron_dense) family, but this checkpoint's family is {family!r}." + ) + + +def _assert_anchor_matches(image_latent: torch.Tensor, latents: torch.Tensor) -> None: + """The I2V conditioning frame must be writable into the denoised latents as-is. + + Both derive from the pipeline dtype/device, so a mismatch means an upstream + change broke that. Slice assignment would hide it behind a per-step cast + rather than fail, which is why this is checked and never coerced. + """ + if image_latent.dtype != latents.dtype or image_latent.device != latents.device: + raise RuntimeError( + "Cosmos3 I2V conditioning latent must match the denoised latents: got " + f"conditioning {image_latent.dtype} on {image_latent.device}, expected " + f"{latents.dtype} on {latents.device}." + ) + + +def _validate_temporal_compression(transformer, vae_scale_factor_temporal: int) -> None: + """A config-declared temporal compression factor must match the VAE.""" + if ( + getattr(transformer, "temporal_compression_factor_declared", False) + and transformer.temporal_compression_factor != vae_scale_factor_temporal + ): + raise ValueError( + f"Transformer config declares temporal_compression_factor=" + f"{transformer.temporal_compression_factor}, but the VAE reports " + f"scale_factor_temporal={vae_scale_factor_temporal}." + ) + def _condition_pixel_frame_count( condition_video_latent_indexes: Iterable[int], @@ -101,6 +205,7 @@ def _load_reference_image(path: str): "nvidia/Cosmos3-Super-Image2Video-4Step", "nvidia/Cosmos3-Super-Text2Image", "nvidia/Cosmos3-Super-Text2Image-4Step", + "nvidia/Cosmos3-Edge", ], doc="Cosmos3 Omnimodal world models.", ) @@ -108,6 +213,9 @@ class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): primary_pretrained_config = pipeline_config.primary_pretrained_config self.audio_gen = False + # Checkpoint fact vs runtime capability: the checkpoint may ship + # action weights, but action generation is not implemented here. + self.has_action_weights = False self.action_gen = False # Pre-load placeholder; load_standard_components derives the real # policy from the checkpoint's scheduler via from_scheduler(). @@ -116,6 +224,10 @@ def __init__(self, pipeline_config): # omitted value survives the executor's default merge and reaches # forward() as "unset". self.default_use_system_prompt = False + self.use_native_flow_schedule = False + self.family = resolve_arch_recipe(primary_pretrained_config).name + # Schedulers are identified by their resolved (flow_shift, karras) pair + self._scheduler_cache: dict = {} if getattr( primary_pretrained_config, "audio_gen", @@ -125,11 +237,80 @@ def __init__(self, pipeline_config): self.audio_gen = True if getattr(primary_pretrained_config, "action_gen", False): - logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") - self.action_gen = True + logger.info( + "Checkpoint declares action weights; action generation is not supported " + "by this pipeline (weights are skipped)." + ) + self.has_action_weights = True super().__init__(pipeline_config) + def _mode_params(self, mode: str) -> dict: + """Generation default table for this checkpoint family and request mode.""" + return COSMOS3_GENERATION_DEFAULTS[(self.family, mode)] + + def _resolve_generation_params(self, mode: str, **values) -> dict: + """Fill None values: sampling-policy overrides win, then the mode + table, then the video table (for fields the image table omits).""" + mode_params = self._mode_params(mode) + video_params = self._mode_params("video") + sampling_overrides = self.sampling.generation_default_overrides() + resolved = {} + for key, value in values.items(): + if value is None: + if key in sampling_overrides: + value = sampling_overrides[key] + else: + value = mode_params.get(key, video_params.get(key)) + resolved[key] = value + return resolved + + def _log_envelope_advisory( + self, + *, + is_t2i: bool, + height: int, + width: int, + num_frames: int, + frame_rate: float, + max_sequence_length: int, + ) -> None: + """One advisory line for requests outside the model-card envelope. + + The envelope is documented support, not enforced validation: the + reference runtime accepts a wider range, so out-of-envelope requests + run — they just carry no quality claim. Families without a declared + envelope get no advisory. + + One line per request, not one per rank: every rank runs this code on a + TP/Ulysses worker. + """ + if self.rank != 0: + return + env = COSMOS3_ENVELOPES.get(self.family) + if env is None: + return + outside = [] + if (height, width) not in env["resolutions"]: + outside.append(f"{width}x{height} resolution") + if not is_t2i: + lo, hi = env["num_frames"] + if not lo <= num_frames <= hi: + outside.append(f"num_frames={num_frames} (validated: {lo}-{hi})") + lo, hi = env["frame_rate"] + if not lo <= frame_rate <= hi: + outside.append(f"frame_rate={frame_rate} (validated: {lo}-{hi})") + if max_sequence_length > env["max_sequence_length"]: + outside.append( + f"max_sequence_length={max_sequence_length} (validated: " + f"<= {env['max_sequence_length']})" + ) + if outside: + logger.warning( + "Request is outside the model-card validated envelope " + f"({'; '.join(outside)}); generation proceeds but quality may degrade." + ) + def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") model_config = self.pipeline_config.model_configs["transformer"] @@ -157,6 +338,9 @@ def load_standard_components( self.default_use_system_prompt = bool( model_index.get("default_use_system_prompt", self.default_use_system_prompt) ) + self.use_native_flow_schedule = bool( + model_index.get("use_native_flow_schedule", self.use_native_flow_schedule) + ) if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading audio tokenizer...") @@ -172,7 +356,7 @@ def load_standard_components( if PipelineComponent.TOKENIZER not in skip_components: logger.info("Loading tokenizer...") - self.tokenizer = Qwen2Tokenizer.from_pretrained( + self.tokenizer = AutoTokenizer.from_pretrained( checkpoint_dir, subfolder="text_tokenizer", ) @@ -195,6 +379,7 @@ def load_standard_components( self.vae_scale_factor_spatial = getattr( self.vae.config, "scale_factor_spatial", self.vae_scale_factor_spatial ) + _validate_temporal_compression(self.transformer, self.vae_scale_factor_temporal) self.transformer.temporal_compression_factor = self.vae_scale_factor_temporal if PipelineComponent.SCHEDULER not in skip_components: @@ -203,11 +388,18 @@ def load_standard_components( # checkpoints, FlowMatchEuler (fixed stochastic schedule) for # distilled ones. The policy holds the derived immutable facts. self.scheduler = load_scheduler(checkpoint_dir) - self.sampling = Cosmos3SamplingPolicy.from_scheduler(self.scheduler) + self.sampling = Cosmos3SamplingPolicy.from_scheduler( + self.scheduler, native_flow_schedule=self.use_native_flow_schedule + ) + _validate_sampling_recipe(self.family, self.use_native_flow_schedule, self.sampling) + # Each stream's variants derive from that stream's own instance, so + # keep both as untouched bases for _scheduler_for(). + self._base_scheduler = self.scheduler if self.audio_gen: # Separate instance so video and audio scheduler states don't # collide (schedulers mutate internal state on every .step()). self.audio_scheduler = type(self.scheduler).from_config(self.scheduler.config) + self._base_audio_scheduler = self.audio_scheduler # Re-check the env var in case it was changed after initialization like in unit tests. guardrails_disabled = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -240,11 +432,12 @@ def load_standard_components( @property def default_warmup_resolutions(self): - return [(720, 1280)] + video = self._mode_params("video") + return [(video["height"], video["width"])] @property def default_warmup_num_frames(self): - return [189] + return [self._mode_params("video")["num_frames"]] @property def default_warmup_steps(self): @@ -253,7 +446,17 @@ def default_warmup_steps(self): @property def default_generation_params(self): - return {**COSMOS3_PIPELINE_DEFAULTS, **self.sampling.generation_default_overrides()} + """Fields merged by the executor into every request. + + These are the video-mode values — what an unmodified request runs. + A request that selects another mode re-resolves them in ``infer()``, + which tells a merged default from a caller-supplied value via + ``model_fields_set``. Key membership also declares these fields + supported during request validation. ``flow_shift`` is + pipeline-internal, not a request field. + """ + defaults = {k: v for k, v in self._mode_params("video").items() if k != "flow_shift"} + return {**defaults, **self.sampling.generation_default_overrides()} def classify_request_failure(self, exc: BaseException) -> Optional[str]: """Cosmos3 rejects unusable request content with ``ValueError`` and @@ -276,7 +479,7 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N defaults = self.default_generation_params guidance_scale = defaults["guidance_scale"] if guidance_scale is None: - guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] + guidance_scale = self._mode_params("video")["guidance_scale"] with torch.no_grad(): self.forward( prompt="warmup", @@ -293,38 +496,121 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N enable_audio=False, ) - def _apply_flow_shift( - self, target_shift: Optional[float], *, use_karras_sigmas: Optional[bool] = None - ) -> None: - """Rebuild both stream schedulers for the requested sampling knobs. - - Video and audio denoise in lockstep in one loop, so a mode that - rebuilds only the video scheduler leaves audio on the checkpoint's - sigmas and the two streams step on different schedules. + def _scheduler_for( + self, + target_shift: Optional[float], + use_karras_sigmas: Optional[bool] = None, + *, + stream: str = "video", + ) -> Any: + """The scheduler for one resolved sampling configuration, built once. + + A scheduler's identity is its resolved ``(flow_shift, karras)`` pair, + not the request mode: modes that resolve to the same pair share an + instance, and a pair is never rebuilt once seen. Streams get separate + instances at the same configuration because schedulers mutate internal + state on every ``.step()`` — video and audio denoise in lockstep, so + they must share the knobs but not the object. """ - self.scheduler = self.sampling.set_flow_shift( - self.scheduler, target_shift, use_karras_sigmas=use_karras_sigmas - ) - if getattr(self, "audio_scheduler", None) is not None: - self.audio_scheduler = self.sampling.set_flow_shift( - self.audio_scheduler, target_shift, use_karras_sigmas=use_karras_sigmas + if stream == "audio": + base = getattr(self, "_base_audio_scheduler", None) or self.audio_scheduler + else: + base = getattr(self, "_base_scheduler", None) or self.scheduler + # Only configurations this checkpoint can resolve to on its own are + # memoized. ``flow_shift`` is a caller-supplied float with no bounded + # domain, so caching every value seen would let a client grow the cache + # for the worker's lifetime; a one-off value builds a scheduler that is + # discarded with the request instead. + if target_shift is not None and target_shift not in self._cacheable_flow_shifts(): + return self.sampling.set_flow_shift( + base, target_shift, use_karras_sigmas=use_karras_sigmas + ) + + cache = getattr(self, "_scheduler_cache", None) + if cache is None: + cache = self._scheduler_cache = {} + key = (target_shift, use_karras_sigmas, stream) + if key not in cache: + cache[key] = self.sampling.set_flow_shift( + base, target_shift, use_karras_sigmas=use_karras_sigmas ) + return cache[key] + + def _cacheable_flow_shifts(self) -> frozenset: + """Flow shifts reachable without a caller override, for this checkpoint. + + 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. + """ + cached = getattr(self, "_cacheable_flow_shifts_cache", None) + if cached is not None: + return cached + shifts = {COSMOS3_V2V_FLOW_SHIFT} + checkpoint_shift = getattr(self.sampling, "checkpoint_flow_shift", None) + if checkpoint_shift is not None: + shifts.add(float(checkpoint_shift)) + for mode in ("video", "image"): + table = COSMOS3_GENERATION_DEFAULTS.get((self.family, mode)) or {} + mode_shift = table.get("flow_shift") + if mode_shift is not None: + shifts.add(float(mode_shift)) + cached = self._cacheable_flow_shifts_cache = frozenset(shifts) + return cached + + def _release_scheduler_solver_state(self) -> None: + """Drop the multistep solver's retained model outputs after a request. + + UniPC keeps ``solver_order`` previous outputs, which are latent-sized -- + tens of MB of device memory that a cached scheduler would otherwise pin + until its next use. ``set_timesteps`` resets them at the start of every + request anyway, so this only shortens how long they are held; it frees + references rather than allocating, and generation is strictly serial, so + nothing else can be mid-loop on these instances. + + The live schedulers are covered too: a caller-supplied ``flow_shift`` + outside the cacheable set builds a one-off instance that never enters the + cache, and it stays reachable here until the next request replaces it. + """ + cached_schedulers = self._scheduler_cache.values() + live_schedulers = (getattr(self, "scheduler", None), getattr(self, "audio_scheduler", None)) + for scheduler in (*cached_schedulers, *live_schedulers): + if scheduler is None: + continue + order = getattr(getattr(scheduler, "config", None), "solver_order", None) + if order is None: + continue + if getattr(scheduler, "model_outputs", None) is not None: + scheduler.model_outputs = [None] * order + if getattr(scheduler, "timestep_list", None) is not None: + scheduler.timestep_list = [None] * order 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" - # None = unset; resolve by mode exactly once. Non-None values pass through. - mode_params = COSMOS3_T2I_PARAMS if is_t2i else COSMOS3_720P_PARAMS - - def resolved(value, field_name): - return value if value is not None else mode_params[field_name] - - height = resolved(req.params.height, "height") - width = resolved(req.params.width, "width") - num_inference_steps = resolved(req.params.num_inference_steps, "num_inference_steps") - guidance_scale = resolved(req.params.guidance_scale, "guidance_scale") + # Caller-assigned values win. Anything still carrying a pipeline + # default — unset, or merged by the executor from the video table — + # resolves against this request's own mode exactly once. + specified = req.params.model_fields_set + + def as_given(field_name): + value = getattr(req.params, field_name) + return value if field_name in specified else None + + resolved = self._resolve_generation_params( + "image" if is_t2i else "video", + height=as_given("height"), + width=as_given("width"), + num_inference_steps=as_given("num_inference_steps"), + guidance_scale=as_given("guidance_scale"), + ) + height = resolved["height"] + 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( @@ -370,16 +656,22 @@ def _apply_metadata_templates( resolution_template: Optional[str] = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, force_duration_template: bool = False, ) -> str: - """Append duration and resolution metadata to a plain-text prompt. + """Append duration and resolution metadata as sentences. ``duration_template`` / ``resolution_template`` of ``None`` disables that - template. JSON prompts are handled by ``_format_prompt_with_metadata``. + template. A JSON positive prompt instead gets the metadata injected as + object fields by ``_format_prompt_with_metadata``; negative prompts always + come here, matching the reference, so a JSON negative prompt keeps its + serialized form and gains the sentences after it. """ parts: List[str] = [] head = prompt.rstrip(".").strip() if head: parts.append(head) if duration_template is not None and (num_frames > 1 or force_duration_template): + # Fractional on purpose: the reference's text path keeps the exact value + # and lets the template's own precision render it, unlike its JSON path, + # which truncates (cosmos-framework _format_prompt_with_template). duration = num_frames / frame_rate parts.append(duration_template.format(duration=duration, fps=frame_rate).rstrip(".")) if resolution_template is not None: @@ -412,16 +704,20 @@ def _format_prompt_with_metadata( if duration_template is not None and ( num_frames > 1 or force_duration_template ): - duration = num_frames / frame_rate - data["duration"] = f"{duration:.1f}s" - data["fps"] = ( - int(frame_rate) if frame_rate == int(frame_rate) else frame_rate - ) + # Truncated, not rounded, and integer-valued even though the + # text template above stays fractional: both mirror the + # reference (cosmos-framework _format_json_prompt_with_template). + data["duration"] = f"{int(num_frames / frame_rate)}s" + data["fps"] = float(frame_rate) + else: + # A still carries no duration: drop whatever the caller's + # JSON declared rather than leaving it stale. + data.pop("duration", None) + data.pop("fps", None) if resolution_template is not None: - data["resolution"] = {"W": width, "H": height} - divisor = math.gcd(height, width) - data["aspect_ratio"] = f"{height // divisor},{width // divisor}" - return json.dumps(data, ensure_ascii=False) + data["resolution"] = {"H": int(height), "W": int(width)} + data["aspect_ratio"] = _aspect_ratio_bucket(height, width) + return json.dumps(data) return self._apply_metadata_templates( prompt, @@ -617,6 +913,7 @@ def _conditioning_anchor_post_step(self, image_latent: Optional[torch.Tensor]): def post_step_fn(latents: torch.Tensor) -> torch.Tensor: # In-place: writes one latent frame, no full-tensor copies. + _assert_anchor_matches(image_latent, latents) latents[:, :, 0:1] = image_latent return latents @@ -816,11 +1113,10 @@ def forward( """Run one generation. ``infer()`` is the resolved entry point. Production requests arrive through ``infer()`` with fully resolved - values; the signature defaults are the base-checkpoint *video* table - values for direct internal callers. ``forward()`` cannot tell a - signature default from an explicit argument, so on distilled - checkpoints (which fix steps/guidance and reject anything else) direct - callers must pass checkpoint-valid sampling values. + values; unset (None) numeric parameters resolve here from the same + per-variant mode tables, so direct internal callers get + checkpoint-appropriate values (including the fixed distilled + steps/guidance). ``use_system_prompt=None`` means "unset": V2V always uses the system prompt, and every other mode takes the checkpoint-declared default, so @@ -841,19 +1137,35 @@ def forward( raise ValueError(f"output_type must be 'video' or 'image', got {output_type!r}.") is_t2i = output_type == "image" + mode_params = self._mode_params(output_type) + resolved = self._resolve_generation_params( + output_type, + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, + max_sequence_length=max_sequence_length, + frame_rate=frame_rate, + ) + height = resolved["height"] + width = resolved["width"] + num_frames = resolved["num_frames"] + num_inference_steps = resolved["num_inference_steps"] + guidance_scale = resolved["guidance_scale"] + max_sequence_length = resolved["max_sequence_length"] + frame_rate = resolved["frame_rate"] + self.sampling.validate_request(num_inference_steps, guidance_scale) - # A distilled checkpoint's step count and guidance are checkpoint facts, - # not mode defaults, so they resolve before the mode tables below. - # Requests through infer() already carry them via - # default_generation_params; a direct forward() call would otherwise - # fall through to the base tables and run CFG at 6.0 against weights - # with guidance baked in. - checkpoint_defaults = self.sampling.generation_default_overrides() - if num_inference_steps is None: - num_inference_steps = checkpoint_defaults.get("num_inference_steps") - if guidance_scale is None: - guidance_scale = checkpoint_defaults.get("guidance_scale") + self._log_envelope_advisory( + is_t2i=is_t2i, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + max_sequence_length=max_sequence_length, + ) if image is not None and video is not None: raise ValueError( @@ -870,6 +1182,7 @@ def forward( use_system_prompt = is_v2v or self.default_use_system_prompt else: use_system_prompt = bool(use_system_prompt) + guidance_interval = None if is_t2i: if image is not None: @@ -880,39 +1193,24 @@ def forward( # T2I force-disables audio instead of rejecting it, so an image # request never trips the audio-weight presence check below. enable_audio = False - height = height or COSMOS3_T2I_PARAMS["height"] - width = width or COSMOS3_T2I_PARAMS["width"] - num_inference_steps = num_inference_steps or COSMOS3_T2I_PARAMS["num_inference_steps"] - if guidance_scale is None: - guidance_scale = COSMOS3_T2I_PARAMS["guidance_scale"] - guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] - self._apply_flow_shift( - flow_shift if flow_shift is not None else COSMOS3_T2I_PARAMS["flow_shift"], - ) + guidance_interval = mode_params["guidance_interval"] + + # Flow shift is a mode table fact unless the request overrides it, and + # V2V additionally wants the uniform sigma grid. Both streams take the + # same knobs so video and audio never step on different schedules. + mode_shift = mode_params.get("flow_shift") + if mode_shift is None: + mode_shift = self.sampling.checkpoint_flow_shift + if is_v2v: + # V2V wants a stronger shift and the uniform sigma schedule. + target_shift = COSMOS3_V2V_FLOW_SHIFT if flow_shift is None else flow_shift + target_karras = False else: - height = height or COSMOS3_720P_PARAMS["height"] - width = width or COSMOS3_720P_PARAMS["width"] - num_frames = num_frames or COSMOS3_720P_PARAMS["num_frames"] - num_inference_steps = num_inference_steps or COSMOS3_720P_PARAMS["num_inference_steps"] - if guidance_scale is None: - guidance_scale = COSMOS3_720P_PARAMS["guidance_scale"] - if is_v2v: - # V2V wants a stronger shift and the uniform sigma schedule. - self._apply_flow_shift( - flow_shift if flow_shift is not None else 10.0, - use_karras_sigmas=False, - ) - else: - # Restore the checkpoint sampling knobs in case a prior T2I or - # V2V request rebuilt the scheduler with mode-specific values. - self._apply_flow_shift( - flow_shift if flow_shift is not None else self.sampling.checkpoint_flow_shift, - use_karras_sigmas=None, - ) - - max_sequence_length = max_sequence_length or COSMOS3_720P_PARAMS["max_sequence_length"] - if frame_rate is None: - frame_rate = COSMOS3_720P_PARAMS["frame_rate"] + 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 self.rank == 0: logger.info( @@ -971,7 +1269,7 @@ def forward( generator = torch.Generator(device=self.device).manual_seed(seed) if negative_prompt is None: - negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT + negative_prompt = default_negative_prompt(output_type) # Positive prompt: forward duration/resolution templates. T2I has no # duration concept (single image) and uses the image-flavored @@ -987,7 +1285,11 @@ def forward( # Negative prompt: mirror positive metadata (cosmos-framework CLI default # when ``negative_prompt_keep_metadata`` promotes mode to ``same``). - negative_prompt = self._format_prompt_with_metadata( + # Always the plain-text templates, never the JSON field injection the + # positive branch uses: the reference appends these sentences to the + # negative prompt whether or not it is a JSON object, so a JSON negative + # prompt ends up as the serialized object followed by the sentences. + negative_prompt = self._apply_metadata_templates( negative_prompt, height=height, width=width, @@ -1243,6 +1545,8 @@ def post_step_fn(step_latents): latents = denoise_result audio_latents = None + self._release_scheduler_solver_state() + timer.mark_post_start() # 7. Decode video @@ -1251,6 +1555,7 @@ def post_step_fn(step_latents): if image_latent is not None: # In-place: the loop output is consumed only by the decode below. + _assert_anchor_matches(image_latent, latents) latents[:, :, 0:1] = image_latent video = self.decode_latents(latents, self._decode_latents) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py index 59e36e203ae6..cfdbc3531356 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sampling.py @@ -18,7 +18,8 @@ ``scheduler/scheduler_config.json``: * ``UniPCMultistepScheduler`` without fixed sigmas — base checkpoints: - request tables drive steps/guidance; T2I rebuilds with ``flow_shift=3.0``. + request tables drive steps/guidance, and the flow shift comes from the mode + table unless the request overrides it. * ``FlowMatchEulerDiscreteScheduler`` with ``stochastic_sampling`` enabled and a nonempty ``fixed_step_sampler_config.t_list`` — distilled checkpoints: the step count is locked to the schedule, classifier-free @@ -36,6 +37,7 @@ from dataclasses import dataclass from typing import Any, Optional +import numpy as np from diffusers import FlowMatchEulerDiscreteScheduler, UniPCMultistepScheduler from tensorrt_llm.logger import logger @@ -105,9 +107,14 @@ class Cosmos3SamplingPolicy: fixed_sigmas: "tuple[float, ...] | None" = None # Checkpoint scheduler config, kept for flow-shift rebuilds (UniPC only). unipc_base_config: Optional[Any] = None + # Checkpoint-declared (model_index): UniPC runs on explicit linear flow + # sigmas with a runtime shift instead of the config's karras grid. + native_flow_schedule: bool = False @classmethod - def from_scheduler(cls, scheduler: Any) -> "Cosmos3SamplingPolicy": + def from_scheduler( + cls, scheduler: Any, native_flow_schedule: bool = False + ) -> "Cosmos3SamplingPolicy": """Derive the policy from a loaded scheduler's config. Valid recipes: UniPC without fixed sigmas (base) and stochastic @@ -129,7 +136,11 @@ def from_scheduler(cls, scheduler: Any) -> "Cosmos3SamplingPolicy": ) if is_unipc and fixed_sigmas is None: - return cls(fixed_sigmas=None, unipc_base_config=scheduler.config) + return cls( + fixed_sigmas=None, + unipc_base_config=scheduler.config, + native_flow_schedule=native_flow_schedule, + ) if is_flow_match and fixed_sigmas is not None: if not _config_get(scheduler.config, "stochastic_sampling", False): @@ -167,7 +178,7 @@ def is_distilled(self) -> bool: def generation_default_overrides(self) -> dict: """Checkpoint-mandated overrides of the table generation defaults. - Merged over ``COSMOS3_720P_PARAMS`` by the pipeline's + Merged over the checkpoint family's video table by the pipeline's ``default_generation_params``, so executor-merged requests arrive carrying the checkpoint's true defaults. """ @@ -207,6 +218,13 @@ def set_timesteps(self, scheduler: Any, num_inference_steps: int, device: Any) - """Program a scheduler for one generation: fixed sigmas or a step count.""" if self.is_distilled: scheduler.set_timesteps(sigmas=list(self.fixed_sigmas), device=device) + elif self.native_flow_schedule: + # The PyTorch-backend base grid: linear flow sigmas over + # (1 - 1/T, 0]. UniPC applies its flow_shift to provided sigmas; + # a numpy array is required (a list breaks diffusers 0.39). + num_train = int(_config_get(scheduler.config, "num_train_timesteps", 1000)) + sigmas = np.linspace(1.0 - 1.0 / num_train, 0.0, num_inference_steps + 1)[:-1] + scheduler.set_timesteps(num_inference_steps, device=device, sigmas=sigmas) else: scheduler.set_timesteps(num_inference_steps, device=device) @@ -244,9 +262,10 @@ def set_flow_shift( Current values are read from the supplied scheduler's own config, so no tracking state exists to diverge. ``None`` means "whatever the checkpoint shipped" for either knob: V2V passes - ``use_karras_sigmas=False`` to force the uniform sigma schedule. - Structural no-op for distilled checkpoints (no UniPC base config) and - when neither knob is requested. + ``use_karras_sigmas=False`` to force the uniform sigma schedule, and a + checkpoint on the native flow schedule needs the same for the same + reason. Structural no-op for distilled checkpoints (no UniPC base + config) and when neither knob is requested. """ if self.unipc_base_config is None: return scheduler @@ -257,6 +276,11 @@ def set_flow_shift( target_shift = self.checkpoint_flow_shift if target_shift is None else float(target_shift) current_karras = bool(_config_get(scheduler.config, "use_karras_sigmas", False)) base_karras = bool(_config_get(self.unipc_base_config, "use_karras_sigmas", False)) + # The native flow schedule is defined on explicit linear sigmas, which + # UniPC's karras branch discards; treat it as an implicit request for + # the uniform grid rather than a separate code path. + if use_karras_sigmas is None and self.native_flow_schedule: + use_karras_sigmas = False target_karras = base_karras if use_karras_sigmas is None else bool(use_karras_sigmas) if current_shift == target_shift and current_karras == target_karras: 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 1a85cf2b5784..64deb4a3f618 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -25,7 +25,9 @@ from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask from tensorrt_llm._torch.modules.embedding import Embedding from tensorrt_llm._torch.modules.gated_mlp import GatedMLP -from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.modules.linear import Linear, WeightMode +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.utils import relu2 from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig from tensorrt_llm._torch.visual_gen.models.modeling import BaseDiffusionModel from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode @@ -39,6 +41,7 @@ PRETRAINED_CONFIG_COMPAT_DEFAULTS = { "position_embedding_type": "unified_3d_mrope", "max_position_embeddings": 262144, + "temporal_compression_factor": 4, "temporal_compression_factor_sound": 1, } @@ -56,6 +59,138 @@ def apply_pretrained_config_compat_defaults( return pretrained_config +COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" + + +def resolve_rope_axes_dim(pretrained_config) -> list: + """MRoPE axes: the explicit top-level ``rope_axes_dim`` wins (diffusers + precedence); the legacy ``rope_scaling["mrope_section"]`` is the fallback. + Both-declared-but-contradictory, wrong length, or a sum that does not + cover half the head dim are config errors. + """ + top_level = getattr(pretrained_config, "rope_axes_dim", None) + rope_scaling = getattr(pretrained_config, "rope_scaling", None) or {} + nested = rope_scaling.get("mrope_section") + + if top_level is not None and nested is not None and list(top_level) != list(nested): + raise ValueError( + f"Cosmos3 config declares contradictory MRoPE axes: rope_axes_dim=" + f"{list(top_level)} vs rope_scaling.mrope_section={list(nested)}." + ) + axes = top_level if top_level is not None else nested + if axes is None: + raise ValueError( + "Cosmos3 config declares neither rope_axes_dim nor rope_scaling.mrope_section." + ) + axes = list(axes) + half_head_dim = pretrained_config.head_dim // 2 + if len(axes) != 3 or sum(axes) != half_head_dim: + raise ValueError( + f"Cosmos3 MRoPE axes {axes} must have 3 entries summing to " + f"head_dim/2 = {half_head_dim}." + ) + return axes + + +@dataclass(frozen=True) +class Cosmos3ArchRecipe: + """Complete architecture combination selected by ``backbone_type``. + + Module construction derives from the recipe, never from individual config + flags, so a mixed config cannot half-apply. + """ + + name: str + gated_mlp: bool + und_qk_norm: bool + use_und_k_norm_for_gen: bool + nemotron_norms: bool + + +QWEN3_RECIPE = Cosmos3ArchRecipe( + name="qwen3", + gated_mlp=True, + und_qk_norm=True, + use_und_k_norm_for_gen=False, + nemotron_norms=False, +) + +NEMOTRON_DENSE_RECIPE = Cosmos3ArchRecipe( + name="nemotron_dense", + gated_mlp=False, + und_qk_norm=False, + use_und_k_norm_for_gen=True, + nemotron_norms=True, +) + + +def resolve_arch_recipe(pretrained_config) -> Cosmos3ArchRecipe: + """Select and validate the architecture recipe declared by the config.""" + backbone_type = getattr(pretrained_config, "backbone_type", None) + if backbone_type is None: + recipe = QWEN3_RECIPE + expected_flags = { + "hidden_act": (None, "silu"), + "qk_norm_for_text": (None, True), + "use_und_k_norm_for_gen": (None, False), + } + elif backbone_type == COSMOS3_EDGE_BACKBONE_TYPE: + recipe = NEMOTRON_DENSE_RECIPE + expected_flags = { + "hidden_act": ("relu2",), + "qk_norm_for_text": (False,), + "use_und_k_norm_for_gen": (True,), + "sound_gen": (False,), + "attention_bias": (False,), + "rms_norm_eps": (1e-5,), + } + else: + raise ValueError( + f"Unsupported Cosmos3 transformer backbone_type={backbone_type!r}; " + f"supported: absent (Qwen3 family) or {COSMOS3_EDGE_BACKBONE_TYPE!r}." + ) + + for key, allowed in expected_flags.items(): + actual = getattr(pretrained_config, key, None) + if actual not in allowed: + raise ValueError( + f"Cosmos3 config contradicts the {recipe.name!r} recipe: " + f"{key}={actual!r}, expected one of {allowed}." + ) + + # Latent-geometry invariants are validated for the Edge recipe only: every + # published Edge checkpoint declares 48/2, so a different value means a + # wrong or corrupt config. The qwen3 recipe is left unchecked because + # long-standing reduced-dimension test fixtures build it with tiny latent + # channels; real qwen3 checkpoints are protected by weight-shape checks. + if recipe is NEMOTRON_DENSE_RECIPE: + invariants = {"latent_channel": 48, "latent_patch_size": 2} + for key, expected in invariants.items(): + actual = getattr(pretrained_config, key, None) + if actual != expected: + raise ValueError( + f"Unsupported Cosmos3 Edge transformer config: {key}={actual!r}, " + f"expected {expected}." + ) + + # A declared patch_latent_dim must agree with the latent geometry it is + # derived from (the transformer recomputes it and would silently ignore an + # inconsistent declaration). + declared_patch_dim = getattr(pretrained_config, "patch_latent_dim", None) + latent_channel = getattr(pretrained_config, "latent_channel", None) + latent_patch_size = getattr(pretrained_config, "latent_patch_size", None) + if None not in (declared_patch_dim, latent_channel, latent_patch_size): + expected_patch_dim = (latent_patch_size**2) * latent_channel + if declared_patch_dim != expected_patch_dim: + raise ValueError( + f"Inconsistent Cosmos3 transformer config: patch_latent_dim=" + f"{declared_patch_dim}, but latent_patch_size**2 * latent_channel = " + f"{expected_patch_dim}." + ) + + return recipe + + class Qwen3VLTextRMSNorm(nn.Module): def __init__( self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.bfloat16 @@ -80,6 +215,19 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return output +class NemotronRMSNorm(Qwen3VLTextRMSNorm): + """RMSNorm with the weight multiply in float32 before downcast. + + ``F.rms_norm`` is bit-exact to this flavor for bf16 inputs, unlike the + parent's post-downcast bf16 weight multiply. + """ + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return F.rms_norm( + hidden_states, (hidden_states.shape[-1],), self.weight, self.variance_epsilon + ) + + @dataclass class TransformerOutput: """Velocity predictions from Cosmos3VFMTransformer.forward().""" @@ -253,6 +401,7 @@ def __init__( num_key_value_heads: int, head_dim: int, model_config: DiffusionModelConfig, + recipe: Cosmos3ArchRecipe, layer_idx: int = 0, module_name: Optional[str] = None, ): @@ -270,14 +419,28 @@ def __init__( module_name=module_name, enable_sequence_parallel=False, ) - self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) - self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + # Attention Q/K norms run the fp32-weight-multiply flavor in both + # recipes (this path has always been F.rms_norm); only the layernorms + # differ per recipe. + eps = model_config.pretrained_config.rms_norm_eps + if recipe.und_qk_norm: + self.norm_q = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + self.norm_k = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + else: + self.norm_q = None + self.norm_k = None + if recipe.use_und_k_norm_for_gen: + self.k_norm_und_for_gen = NemotronRMSNorm( + hidden_size=head_dim, eps=eps, dtype=torch.bfloat16 + ) + else: + self.k_norm_und_for_gen = None def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Per-head RMSNorm on 4D tensors [B, S, H, D].""" - q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) - k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) - return q, k + if self.norm_q is None: + return q, k + return self.norm_q(q), self.norm_k(k) def forward_with_kv( self, @@ -294,9 +457,20 @@ def forward_with_kv( k = k.view(batch_size, seq_len, self.local_num_key_value_heads, self.head_dim) v = v.view(batch_size, seq_len, self.local_num_key_value_heads, self.head_dim) + # The gen tower consumes a separately normed view of the raw und keys; + # the und self-attention below must not see that norm. + k_for_gen = None + if self.k_norm_und_for_gen is not None: + k_for_gen = self.k_norm_und_for_gen(k) + q, k = self.apply_qk_norm(q, k) q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + if k_for_gen is not None: + _, k_for_gen = qwen3_apply_rotary_pos_emb(q, k_for_gen, freqs_cos, freqs_sin) + else: + k_for_gen = k + out = self._attn_impl( q, k, @@ -305,7 +479,7 @@ def forward_with_kv( timestep=timestep, ) - return self.to_out[0](out), k, v + return self.to_out[0](out), k_for_gen, v def forward(self): raise NotImplementedError( @@ -331,6 +505,7 @@ def __init__( num_key_value_heads: int, head_dim: int, model_config: DiffusionModelConfig, + recipe: Cosmos3ArchRecipe, layer_idx: int = 0, module_name: Optional[str] = None, ): @@ -355,14 +530,15 @@ def __init__( ) model_config.attention.backend = original_backend - self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) - self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + # Same flavor note as Cosmos3CausalAttention: attention Q/K norms are + # fp32-weight-multiply in both recipes. + eps = model_config.pretrained_config.rms_norm_eps + self.norm_q = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) + self.norm_k = NemotronRMSNorm(hidden_size=head_dim, eps=eps, dtype=torch.bfloat16) def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: """Per-head RMSNorm on 4D tensors [B, S, H, D].""" - q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) - k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) - return q, k + return self.norm_q(q), self.norm_k(k) def forward( self, @@ -427,14 +603,46 @@ def forward( return self.to_out[0](out) +def _build_cosmos3_mlp( + model_config: DiffusionModelConfig, recipe: Cosmos3ArchRecipe, layer_idx: int +) -> nn.Module: + hidden_size = model_config.pretrained_config.hidden_size + intermediate_size = model_config.pretrained_config.intermediate_size + if recipe.gated_mlp: + return GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + reduce_output=model_config.mapping.tp_size > 1, + ) + return MLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + activation=relu2, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + reduce_output=model_config.mapping.tp_size > 1, + ) + + +def _layer_norm_cls(recipe: Cosmos3ArchRecipe): + return NemotronRMSNorm if recipe.nemotron_norms else Qwen3VLTextRMSNorm + + class Cosmos3UndDecoderLayer(nn.Module): """Understanding pathway decoder layer: causal self-attention + MLP.""" - def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + def __init__( + self, model_config: DiffusionModelConfig, layer_idx: int, recipe: Cosmos3ArchRecipe + ): super().__init__() self.layer_idx = layer_idx hidden_size = model_config.pretrained_config.hidden_size - intermediate_size = model_config.pretrained_config.intermediate_size self.self_attn = Cosmos3CausalAttention( hidden_size=hidden_size, @@ -442,28 +650,22 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): num_key_value_heads=model_config.pretrained_config.num_key_value_heads, head_dim=model_config.pretrained_config.head_dim, model_config=model_config, + recipe=recipe, layer_idx=layer_idx, module_name=f"layers.{layer_idx}.self_attn", ) - self.input_layernorm = Qwen3VLTextRMSNorm( + norm_cls = _layer_norm_cls(recipe) + self.input_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.post_attention_layernorm = Qwen3VLTextRMSNorm( + self.post_attention_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.mlp = GatedMLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - bias=False, - dtype=torch.bfloat16, - config=model_config, - layer_idx=layer_idx, - reduce_output=model_config.mapping.tp_size > 1, - ) + self.mlp = _build_cosmos3_mlp(model_config, recipe, layer_idx) def forward( self, @@ -500,11 +702,12 @@ def forward( class Cosmos3GenDecoderLayer(nn.Module): """Generation pathway decoder layer: cross-attention (to UND K/V) + MLP.""" - def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + def __init__( + self, model_config: DiffusionModelConfig, layer_idx: int, recipe: Cosmos3ArchRecipe + ): super().__init__() self.layer_idx = layer_idx hidden_size = model_config.pretrained_config.hidden_size - intermediate_size = model_config.pretrained_config.intermediate_size self.cross_attention = Cosmos3CrossAttention( hidden_size=hidden_size, @@ -512,28 +715,22 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): num_key_value_heads=model_config.pretrained_config.num_key_value_heads, head_dim=model_config.pretrained_config.head_dim, model_config=model_config, + recipe=recipe, layer_idx=layer_idx, module_name=f"layers.{layer_idx}.cross_attention", ) - self.input_layernorm = Qwen3VLTextRMSNorm( + norm_cls = _layer_norm_cls(recipe) + self.input_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.post_attention_layernorm = Qwen3VLTextRMSNorm( + self.post_attention_layernorm = norm_cls( hidden_size=hidden_size, eps=model_config.pretrained_config.rms_norm_eps, dtype=torch.bfloat16, ) - self.mlp = GatedMLP( - hidden_size=hidden_size, - intermediate_size=intermediate_size, - bias=False, - dtype=torch.bfloat16, - config=model_config, - layer_idx=layer_idx, - reduce_output=model_config.mapping.tp_size > 1, - ) + self.mlp = _build_cosmos3_mlp(model_config, recipe, layer_idx) def forward( self, @@ -614,11 +811,16 @@ def _compute_default_rope_parameters( class Qwen3VLTextRotaryEmbedding(nn.Module): def __init__(self, model_config: DiffusionModelConfig): super().__init__() - self.rope_type = model_config.pretrained_config.rope_scaling["rope_type"] + # Edge checkpoints omit rope_type from rope_scaling, and a checkpoint + # declaring the axes via top-level ``rope_axes_dim`` may omit the block + # entirely -- ``resolve_rope_axes_dim`` below supports that, so read it + # with the same tolerance rather than failing before reaching it. + rope_scaling = getattr(model_config.pretrained_config, "rope_scaling", None) or {} + self.rope_type = rope_scaling.get("rope_type", "default") self.max_seq_len_cached = model_config.pretrained_config.max_position_embeddings self.original_max_seq_len = model_config.pretrained_config.max_position_embeddings - self.mrope_section = model_config.pretrained_config.rope_scaling["mrope_section"] + self.mrope_section = resolve_rope_axes_dim(model_config.pretrained_config) inv_freq, self.attention_scaling = _compute_default_rope_parameters(model_config) self.register_buffer("inv_freq", inv_freq, persistent=False) @@ -675,7 +877,7 @@ class Cosmos3LanguageModel(nn.Module): computed once and reused across all sampling steps. """ - def __init__(self, model_config: DiffusionModelConfig): + def __init__(self, model_config: DiffusionModelConfig, recipe: Cosmos3ArchRecipe): super().__init__() hidden_size = model_config.pretrained_config.hidden_size num_hidden_layers = model_config.pretrained_config.num_hidden_layers @@ -688,7 +890,10 @@ def __init__(self, model_config: DiffusionModelConfig): ) self.rotary_emb = Qwen3VLTextRotaryEmbedding(model_config) self.layers = nn.ModuleList( - [Cosmos3UndDecoderLayer(model_config, layer_idx=i) for i in range(num_hidden_layers)] + [ + Cosmos3UndDecoderLayer(model_config, layer_idx=i, recipe=recipe) + for i in range(num_hidden_layers) + ] ) def forward( @@ -723,9 +928,14 @@ def forward( class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) + self.temporal_compression_factor_declared = ( + getattr(model_config.pretrained_config, "temporal_compression_factor", None) is not None + ) pretrained_config = apply_pretrained_config_compat_defaults(model_config.pretrained_config) + self.recipe = resolve_arch_recipe(pretrained_config) self.audio_gen = getattr(pretrained_config, "sound_gen", False) - self.action_gen = getattr(pretrained_config, "action_gen", False) + # Config fact only: the transformer never constructs action modules. + self.has_action_weights = getattr(pretrained_config, "action_gen", False) self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -735,8 +945,9 @@ def __init__(self, model_config: DiffusionModelConfig): self.timestep_scale = pretrained_config.timestep_scale self.base_fps = pretrained_config.base_fps - # Comes from VAE. Updated after VAE is loaded. - self.temporal_compression_factor = 4 + # Config-declared (compat default 4); cross-checked against the VAE + # after component loading. + self.temporal_compression_factor = pretrained_config.temporal_compression_factor self.unified_3d_mrope_temporal_modality_margin = ( pretrained_config.unified_3d_mrope_temporal_modality_margin @@ -785,7 +996,7 @@ def __init__(self, model_config: DiffusionModelConfig): "Ring parallelism is not supported for Cosmos3 cross-attention." ) - self.language_model = Cosmos3LanguageModel(model_config) + self.language_model = Cosmos3LanguageModel(model_config, self.recipe) self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) @@ -801,12 +1012,12 @@ def __init__(self, model_config: DiffusionModelConfig): self.gen_layers = nn.ModuleList( [ - Cosmos3GenDecoderLayer(model_config, layer_idx=i) + Cosmos3GenDecoderLayer(model_config, layer_idx=i, recipe=self.recipe) for i in range(self.num_hidden_layers) ] ) - self.norm_moe_gen = Qwen3VLTextRMSNorm( + self.norm_moe_gen = _layer_norm_cls(self.recipe)( hidden_size=self.hidden_size, eps=pretrained_config.rms_norm_eps, ) @@ -1182,19 +1393,22 @@ def load_weights(self, weights: dict) -> None: "action_modality_embed", "action_proj_", ) + skipped_keys = [] for key, value in weights.items(): k = key - if k.startswith(skip_prefixes): - continue - - # Normalize a leading "model." prefix up front so every remap below - # matches whether or not the checkpoint namespaces top-level tensors - # (e.g. "model.audio_proj_in.weight") under "model.". + # Normalize a leading "model." prefix up front so every skip and + # remap below matches whether or not the checkpoint namespaces + # top-level tensors (e.g. "model.audio_proj_in.weight") under + # "model.". if k.startswith("model."): k = k[len("model.") :] + if k.startswith(skip_prefixes): + skipped_keys.append(k) + continue + if k.startswith("proj_in."): remapped[k.replace("proj_in.", "vae2llm.", 1)] = value continue @@ -1221,11 +1435,17 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue - # embed_tokens and norm → language_model.* - if k.startswith("embed_tokens.") or k.startswith("norm."): + # embed_tokens → language_model.* + if k.startswith("embed_tokens."): remapped[f"language_model.{k}"] = value continue + # Und final norm: normalizes the und hidden state consumed only by + # lm_head; the generation path uses per-layer und K/V exclusively. + if k.startswith("norm."): + skipped_keys.append(k) + continue + # norm_moe_gen stays at top level if k.startswith("norm_moe_gen."): remapped[k] = value @@ -1252,6 +1472,7 @@ def load_weights(self, weights: dict) -> None: "self_attn.to_out.": f"{und_lp}.self_attn.to_out.0.", "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + "self_attn.k_norm_und_for_gen.": f"{und_lp}.self_attn.k_norm_und_for_gen.", } # --- GEN attention → gen_layers.{i}.cross_attention.* --- @@ -1303,35 +1524,78 @@ def load_weights(self, weights: dict) -> None: } loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + # Coverage is default-fail in both directions: every parameter of every + # constructed module must receive a checkpoint tensor, and every mapped + # checkpoint tensor must land on a parameter. The remap-stage skip list + # above is the only source of intentional omissions. + missing = [] + consumed = set() + for param_name, param in self._parameters.items(): - if param is not None and param_name in remapped: + if param is None: + continue + if param_name in remapped: param.data.copy_(remapped[param_name].to(param.dtype)) + consumed.add(param_name) + else: + missing.append(param_name) + + def _linear_source_prefixes(module: Linear, name: str) -> list: + weights_config = getattr(module, "weights_loading_config", None) + weight_mode = getattr(weights_config, "weight_mode", None) + if weight_mode in (WeightMode.FUSED_QKV_LINEAR, WeightMode.FUSED_GATE_UP_LINEAR): + parent = name.rsplit(".", 1)[0] + for suffix, sources in loader.params_map.items(): + if name == suffix or name.endswith("." + suffix): + return [f"{parent}.{source}." for source in sources] + return [f"{name}."] - loaded_linear = 0 - loaded_other = 0 - skipped_modules = [] for name, module in self.named_modules(): - if len(module._parameters) == 0: + if not name or len(module._parameters) == 0: continue if isinstance(module, Linear): weight_dicts = loader.get_linear_weights(module, name, remapped) - if weight_dicts: + # Fused sources yield one dict each; an empty dict means that + # source is absent, so a partially covered fusion also fails. + if weight_dicts and all(weight_dicts): loader.load_linear_weights(module, name, weight_dicts) - loaded_linear += 1 + prefixes = _linear_source_prefixes(module, name) + consumed.update(k for k in remapped if any(k.startswith(p) for p in prefixes)) else: - skipped_modules.append(f"{name}(Linear)") + missing.append(f"{name}(Linear)") else: module_weights = loader.filter_weights(name, remapped) - if module_weights: - loaded_other += 1 - else: - has_params = any(p is not None for p in module._parameters.values()) - if has_params and name: - skipped_modules.append(f"{name}({type(module).__name__})") for param_name, param in module._parameters.items(): - if param is not None and param_name in module_weights: + if param is None: + continue + if param_name in module_weights: param.data.copy_(module_weights[param_name].to(param.dtype)) + consumed.add(f"{name}.{param_name}") + else: + missing.append(f"{name}.{param_name}") + + if missing: + preview = ", ".join(missing[:10]) + suffix = " ..." if len(missing) > 10 else "" + raise ValueError( + f"Cosmos3 checkpoint is missing weights for {len(missing)} constructed " + f"parameter(s)/module(s): {preview}{suffix}" + ) + unconsumed = [k for k in remapped if k not in consumed] + if unconsumed: + preview = ", ".join(unconsumed[:10]) + suffix = " ..." if len(unconsumed) > 10 else "" + logger.warning( + f"{len(unconsumed)} mapped checkpoint tensor(s) matched no constructed " + f"parameter: {preview}{suffix}" + ) + if skipped_keys: + logger.info( + f"Skipped {len(skipped_keys)} intentionally unused checkpoint tensors " + f"(lm_head / und final norm / action heads): " + f"{', '.join(sorted(set(k.split('.')[0] for k in skipped_keys)))}" + ) def post_load_weights(self) -> None: """Post-load processing: dtype conversion and Linear finalization.""" diff --git a/tensorrt_llm/visual_gen/visual_gen.py b/tensorrt_llm/visual_gen/visual_gen.py index d597bc88b238..4e0edff9f9a0 100644 --- a/tensorrt_llm/visual_gen/visual_gen.py +++ b/tensorrt_llm/visual_gen/visual_gen.py @@ -293,16 +293,11 @@ def extra_param_specs(self) -> Dict[str, "ExtraParamSchema"]: @property def default_params(self) -> "VisualGenParams": - """Returns a ``VisualGenParams`` with the loaded pipeline's defaults. + """Returns a ``VisualGenParams`` with all defaults resolved for the loaded pipeline. Universal fields (height, width, etc.) are filled from the - pipeline's defaults. Pipelines with mode-dependent defaults - (e.g. Cosmos3, where text-to-image and video requests use - different resolutions) leave such fields as ``None``; they are - resolved per request from the output mode, so ``None`` here - means "the mode's default", not "unset". All declared - ``extra_params`` keys are included with their defaults - (``None`` for params without one). + pipeline's defaults. All declared ``extra_params`` keys are + included with their defaults (``None`` for params without one). Use this to inspect what the model will use, then modify and pass to ``generate()``:: @@ -321,7 +316,13 @@ def default_params(self) -> "VisualGenParams": if extra: kwargs["extra_params"] = extra - return VisualGenParams(**kwargs) + params = VisualGenParams(**kwargs) + # These came from the pipeline, not the caller. Un-marking them keeps + # request-dependent defaults resolvable after a round trip through + # this object; assigning any of them re-marks it automatically. + for field_name in self.executor.default_generation_params: + params.model_fields_set.discard(field_name) + return params @set_api_status("prototype") def generate( diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json new file mode 100644 index 000000000000..00f4e3ececa9 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_i2v_lpips_golden_video.json @@ -0,0 +1,27 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "video": "cosmos3_edge_i2v_lpips_golden_video.mp4", + "prompt": "The orange sphere slowly rises while the camera pans right across the scene", + "conditioning_image": "deterministic 832x480 image drawn by _write_cosmos3_edge_conditioning_image in test_visual_gen.py", + "height": 480, + "width": 832, + "num_frames": 29, + "num_inference_steps": 10, + "guidance_scale": 5.0, + "frame_rate": 24.0, + "lpips_threshold": 0.13, + "measured_lpips_at_creation": 0.0778, + "threshold_rationale": "0.0778 measured cross-stack at 10 steps (I2V accumulates cross-stack drift faster than T2V: 0.1105 at the deployed 50 steps), plus ~0.04 cross-host headroom. The failure signal is far away: a wrong-seed run against this golden measures LPIPS 0.858. The deployed 50-step I2V shape is exercised by test_cosmos3_edge_i2v_example.", + "notes": "Per-step masked-velocity parity vs diffusers is 0.8-1.5 percent rel (noisy frames); diffusers zeroes the conditioned frame's velocity while TRT-LLM masks it in the pipeline - equivalent for the scheduler." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json new file mode 100644 index 000000000000..0a5331b8bbc1 --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2i_lpips_golden.json @@ -0,0 +1,24 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "image": "cosmos3_edge_t2i_lpips_golden.png", + "prompt": "A ceramic teapot pouring steaming tea into a cup, morning window light", + "height": 640, + "width": 640, + "num_frames": 1, + "num_inference_steps": 50, + "guidance_scale": 4.0, + "lpips_threshold": 0.05, + "measured_lpips_at_creation": 0.0056, + "threshold_rationale": "0.0056 measured cross-stack on B200; 0.05 matches the FLUX/QwenImage image-gate convention." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json new file mode 100644 index 000000000000..904374e075ef --- /dev/null +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/cosmos3_edge_t2v_lpips_golden_video.json @@ -0,0 +1,25 @@ +{ + "source": "diffusers Cosmos3OmniPipeline on diffusers main (reference implementation, not a TRT-LLM self-golden)", + "diffusers_reference": "huggingface/diffusers#14181 'Cosmos3 edge support' + #14246 'Fix Cosmos3 Edge generator K normalization'", + "diffusers_version": "0.40.0.dev0", + "diffusers_commit": "2919c50968389232c527bdab1a3af69cef01ed07", + "scheduler_override": "UniPCMultistepScheduler.from_config(checkpoint config, use_karras_sigmas=False, flow_shift=3.0). With the checkpoint's use_native_flow_schedule=true this reproduces the cosmos-framework PyTorch backend schedule (fm_solvers_unipc @ 117c7d2) to fp32-ulp: timesteps bit-identical, full synthetic step() trajectories agree to <=1.6e-7 rel (see TestNativeFlowSchedule fixtures). Stock diffusers is NOT used as-is because its karras branch swallows the native flow sigmas.", + "prompt_text_matching": "The golden run passed pre-formatted cond AND uncond texts with add_duration_template=False and add_resolution_template=False. Both texts were produced by TRT-LLM's _format_prompt_with_metadata (keep-metadata negative-prompt semantics, matching cosmos-framework's CLI default rather than diffusers' inverse templates), so both stacks tokenize identical sequences in both CFG branches.", + "model": "Cosmos3-Edge", + "seed": 42, + "generator": "torch.Generator(device='cuda').manual_seed(42); initial latents match TRT-LLM's randn_tensor draw bit-for-bit (same shape/dtype/generator semantics)", + "use_system_prompt": false, + "torch_dtype": "bfloat16", + "lpips_net": "alex", + "video": "cosmos3_edge_t2v_lpips_golden_video.mp4", + "prompt": "A red ball rolls across a wooden floor, casting a soft shadow.", + "height": 480, + "width": 832, + "num_frames": 29, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "frame_rate": 24.0, + "lpips_threshold": 0.1, + "measured_lpips_at_creation": 0.0447, + "threshold_rationale": "0.0447 measured cross-stack (TRT-LLM VANILLA attention vs diffusers main) on B200 with matched noise and matched CFG texts, plus headroom for the ~0.04 cross-host kernel drift documented in _preserve_lpips_candidate_on_failure." +} diff --git a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip index 0b4fe6b6f6ad..14baa5a55a91 100644 --- a/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip +++ b/tests/integration/defs/examples/visual_gen/golden/visual_gen_lpips/visual_gen_lpips_golden_media.zip @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:358b9578166226f6ae065bf5ffd186f1f4001e66867eb701e469bc76daf9a4fb -size 33672147 +oid sha256:c37594d4b002605335b2bdf1041dcec572222ea980f9c78133a4189c35559f28 +size 34525771 diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py index f2677ab6e464..bb72a4656f9f 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen_cosmos3.py @@ -163,6 +163,9 @@ def _run_cosmos3_lpips_pipeline(num_frames, video=None): with torch.no_grad(): result = pipeline.forward( prompt=COSMOS3_LPIPS_PROMPT, + # The goldens were generated against an empty uncond branch, + # so pin it rather than inheriting the video-mode default. + negative_prompt="", seed=COSMOS3_LPIPS_SEED, height=COSMOS3_LPIPS_HEIGHT, width=COSMOS3_LPIPS_WIDTH, @@ -262,6 +265,9 @@ def _generate_cosmos3_feature_image(case, output_path): _assert_feature_quantization_installed(pipeline, case.features) result = pipeline.forward( prompt=COSMOS3_LPIPS_PROMPT, + # The goldens were generated against an empty uncond branch, + # so pin it rather than inheriting the video-mode default. + negative_prompt="", seed=COSMOS3_LPIPS_SEED, height=COSMOS3_LPIPS_HEIGHT, width=COSMOS3_LPIPS_WIDTH, @@ -566,6 +572,9 @@ def _run_cosmos3_i2v_4step_lpips_pipeline(image_path): with torch.no_grad(): result = pipeline.forward( prompt=COSMOS3_I2V_4STEP_LPIPS_PROMPT, + # The goldens were generated against an empty uncond branch, + # so pin it rather than inheriting the video-mode default. + negative_prompt="", seed=COSMOS3_LPIPS_SEED, image=image_path, height=COSMOS3_LPIPS_HEIGHT, @@ -636,3 +645,247 @@ def test_cosmos3_i2v_4step_lpips_against_golden(_visual_gen_deps, request, tmp_p "cosmos3_i2v_4step_lpips_golden_video.mp4", ) _assert_lpips_below_threshold(score, COSMOS3_I2V_4STEP_LPIPS_THRESHOLD) + + +def _write_cosmos3_edge_conditioning_image(path): + """Deterministic 832x480 conditioning image (Edge's native 480p 16:9).""" + from PIL import Image, ImageDraw + + image = Image.new("RGB", (832, 480)) + draw = ImageDraw.Draw(image) + for y in range(480): + draw.line([(0, y), (832, y)], fill=(30, 60 + y // 6, 140)) + draw.ellipse([320, 130, 520, 330], fill=(230, 120, 40), outline=(255, 255, 255), width=4) + draw.rectangle([60, 330, 260, 450], fill=(40, 160, 90)) + image.save(path) + + +def test_cosmos3_edge_i2v_example(_visual_gen_deps, llm_root, llm_venv): + """Run the Edge checkpoint through the recommended invocation. + + Validates the documented deployment for ``Cosmos3-Edge``: the example + script with a conditioning image and no config override (the Edge + defaults — 480p x 121 frames, 50 UniPC steps on the native flow schedule, + guidance 5.0, shift 3.0 — are the deployed shape). The run must produce a + video. + """ + model_path = _lpips_model_path("Cosmos3-Edge") + _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + + out_dir = os.path.join( + llm_venv.get_working_directory(), "visual_gen_output", "cosmos3_edge_i2v_example" + ) + os.makedirs(out_dir, exist_ok=True) + image_path = os.path.join(out_dir, "conditioning.png") + _write_cosmos3_edge_conditioning_image(image_path) + output_path = os.path.join(out_dir, "cosmos3_edge_i2v_output.mp4") + if os.path.exists(output_path): + os.remove(output_path) + + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) + assert os.path.isfile(script_path), f"Example script not found: {script_path}" + + venv_check_call( + llm_venv, + [ + script_path, + "--model", + model_path, + "--prompt", + "The orange sphere slowly rises while the camera pans right across the scene", + "--image_path", + image_path, + "--output_path", + output_path, + ], + env={"TRTLLM_DISABLE_COSMOS3_GUARDRAILS": "1"}, + ) + assert os.path.isfile(output_path), f"Example did not produce output at {output_path}" + assert os.path.getsize(output_path) > 0, f"Example produced an empty video at {output_path}" + + +# Edge LPIPS gates compare against diffusers-main reference goldens with the +# scheduler patched to the cosmos-framework native flow schedule; full +# provenance in golden/visual_gen_lpips/cosmos3_edge_*.json. The I2V gate runs +# 10 steps (cross-stack drift accumulates per step; the deployed 50-step shape +# is covered by test_cosmos3_edge_i2v_example). +COSMOS3_EDGE_LPIPS_SEED = 42 +COSMOS3_EDGE_LPIPS_FRAME_RATE = 24.0 +COSMOS3_EDGE_LPIPS_NUM_FRAMES = 29 +COSMOS3_EDGE_T2V_LPIPS_PROMPT = "A red ball rolls across a wooden floor, casting a soft shadow." +COSMOS3_EDGE_T2V_LPIPS_STEPS = 50 +COSMOS3_EDGE_T2V_LPIPS_THRESHOLD = 0.1 +COSMOS3_EDGE_I2V_LPIPS_PROMPT = ( + "The orange sphere slowly rises while the camera pans right across the scene" +) +COSMOS3_EDGE_I2V_LPIPS_STEPS = 10 +COSMOS3_EDGE_I2V_LPIPS_THRESHOLD = 0.13 +COSMOS3_EDGE_T2I_LPIPS_PROMPT = ( + "A ceramic teapot pouring steaming tea into a cup, morning window light" +) +COSMOS3_EDGE_T2I_LPIPS_STEPS = 50 +COSMOS3_EDGE_T2I_LPIPS_THRESHOLD = 0.05 + + +def _run_cosmos3_edge_lpips_pipeline(**forward_kwargs): + """Run the Cosmos3-Edge pipeline and return the PipelineOutput. + + VANILLA attention, compile-off; guardrails disabled for the run. + """ + guardrails_env_key = "TRTLLM_DISABLE_COSMOS3_GUARDRAILS" + previous_guardrails_env = os.environ.get(guardrails_env_key) + os.environ[guardrails_env_key] = "1" + try: + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader + from tensorrt_llm.visual_gen.args import ( + AttentionConfig, + CompilationConfig, + TorchCompileConfig, + VisualGenArgs, + ) + + model_path = _lpips_model_path("Cosmos3-Edge") + _skip_if_missing(model_path, "Cosmos3-Edge checkpoint", is_dir=True) + _disable_inductor_compile_worker_quiesce() + args = VisualGenArgs( + model=model_path, + compilation_config=CompilationConfig(skip_warmup=True), + torch_compile_config=TorchCompileConfig(enable=False), + attention_config=AttentionConfig(backend="VANILLA"), + ) + pipeline = PipelineLoader(args).load(skip_warmup=True) + try: + # The goldens were generated against an empty uncond branch, so pin it + # here rather than inheriting the video-mode default negative prompt. + forward_kwargs.setdefault("negative_prompt", "") + with torch.no_grad(): + result = pipeline.forward( + seed=COSMOS3_EDGE_LPIPS_SEED, + use_guardrails=False, + **forward_kwargs, + ) + if result is not None: + if result.video is not None: + result.video = result.video.detach().cpu() + if result.image is not None: + result.image = result.image.detach().cpu() + return result + finally: + del pipeline + _cleanup_cuda() + finally: + if previous_guardrails_env is None: + os.environ.pop(guardrails_env_key, None) + else: + os.environ[guardrails_env_key] = previous_guardrails_env + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_t2v_lpips_against_golden(_visual_gen_deps, request, tmp_path): + generated_path = tmp_path / "cosmos3_edge_t2v_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_t2v_lpips_golden_video.mp4", "Cosmos3-Edge T2V LPIPS golden video" + ) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_T2V_LPIPS_PROMPT, + height=480, + width=832, + num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, + num_inference_steps=COSMOS3_EDGE_T2V_LPIPS_STEPS, + guidance_scale=5.0, + frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, + ) + assert result is not None and result.video is not None, "Edge T2V produced no video" + _save_lpips_video_mp4(result.video, generated_path, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_t2v", + "video", + COSMOS3_EDGE_T2V_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + COSMOS3_EDGE_T2V_LPIPS_THRESHOLD, + generated_path, + "cosmos3_edge_t2v_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_T2V_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_i2v_lpips_against_golden(_visual_gen_deps, request, tmp_path): + generated_path = tmp_path / "cosmos3_edge_i2v_generated.mp4" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_i2v_lpips_golden_video.mp4", "Cosmos3-Edge I2V LPIPS golden video" + ) + image_path = tmp_path / "cosmos3_edge_i2v_conditioning.png" + _write_cosmos3_edge_conditioning_image(str(image_path)) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_I2V_LPIPS_PROMPT, + image=str(image_path), + height=480, + width=832, + num_frames=COSMOS3_EDGE_LPIPS_NUM_FRAMES, + num_inference_steps=COSMOS3_EDGE_I2V_LPIPS_STEPS, + guidance_scale=5.0, + frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE, + ) + assert result is not None and result.video is not None, "Edge I2V produced no video" + _save_lpips_video_mp4(result.video, generated_path, frame_rate=COSMOS3_EDGE_LPIPS_FRAME_RATE) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_i2v", + "video", + COSMOS3_EDGE_I2V_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + COSMOS3_EDGE_I2V_LPIPS_THRESHOLD, + generated_path, + "cosmos3_edge_i2v_lpips_golden_video.mp4", + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_I2V_LPIPS_THRESHOLD) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_cosmos3_edge_t2i_lpips_against_golden(request, tmp_path): + from tensorrt_llm.media.encoding import save_image + + generated_path = tmp_path / "cosmos3_edge_t2i_generated.png" + golden_path = _golden_media_path( + tmp_path, "cosmos3_edge_t2i_lpips_golden.png", "Cosmos3-Edge T2I LPIPS golden image" + ) + result = _run_cosmos3_edge_lpips_pipeline( + prompt=COSMOS3_EDGE_T2I_LPIPS_PROMPT, + height=640, + width=640, + num_inference_steps=COSMOS3_EDGE_T2I_LPIPS_STEPS, + guidance_scale=4.0, + output_type="image", + ) + assert result is not None and result.image is not None, "Edge T2I produced no image" + save_image(result.image[0], str(generated_path)) + score = _run_lpips_eval( + tmp_path, + "cosmos3_edge_t2i", + "image", + COSMOS3_EDGE_T2I_LPIPS_PROMPT, + golden_path, + generated_path, + ) + _preserve_lpips_candidate_on_failure( + request, + score, + COSMOS3_EDGE_T2I_LPIPS_THRESHOLD, + generated_path, + "cosmos3_edge_t2i_lpips_golden.png", + ) + _assert_lpips_below_threshold(score, COSMOS3_EDGE_T2I_LPIPS_THRESHOLD) diff --git a/tests/integration/test_lists/test-db/l0_b200.yml b/tests/integration/test_lists/test-db/l0_b200.yml index 406b01f25807..c256fae08247 100644 --- a/tests/integration/test_lists/test-db/l0_b200.yml +++ b/tests/integration/test_lists/test-db/l0_b200.yml @@ -246,6 +246,8 @@ l0_b200: - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_example - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_layered_example - examples/visual_gen/test_visual_gen_qwen_image.py::test_qwen_image_edit_example + - unittest/_torch/visual_gen/test_cosmos3_edge.py + - unittest/_torch/visual_gen/test_cosmos3_example_prompts.py # ------------- Host perf module regression tests (6 representative scenarios) --------------- - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_gen_only_bs8] - perf/host_perf/test_module_scheduler.py::test_scheduler_production[production_mixed_32gen_4ctx] @@ -274,6 +276,7 @@ l0_b200: - unittest/disaggregated/test_openai_server_info.py - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_t2i_4step_example TIMEOUT (30) - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_i2v_4step_example TIMEOUT (45) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_edge_i2v_example TIMEOUT (30) - condition: ranges: system_gpu_count: @@ -370,6 +373,9 @@ l0_b200: - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_t2v_lpips_against_golden TIMEOUT (15) - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_nano_v2v_lpips_against_golden TIMEOUT (10) - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_i2v_4step_lpips_against_golden TIMEOUT (20) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_edge_t2v_lpips_against_golden TIMEOUT (20) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_edge_i2v_lpips_against_golden TIMEOUT (15) + - examples/visual_gen/test_visual_gen_cosmos3.py::test_cosmos3_edge_t2i_lpips_against_golden TIMEOUT (15) - examples/visual_gen/test_visual_gen_wan.py::test_fastwan_lpips_against_golden TIMEOUT (10) - visual_gen/test_visual_gen_benchmark.py::test_offline_benchmark - visual_gen/test_visual_gen_benchmark.py::test_online_benchmark[openai-videos] diff --git a/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py b/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py new file mode 100644 index 000000000000..3a6b1bf0ef7e --- /dev/null +++ b/tests/unittest/_torch/visual_gen/cosmos3_edge_diffusers_parity.py @@ -0,0 +1,131 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +"""Per-step Edge transformer parity: TRT-LLM vs diffusers main. + +Not a pytest module — run as a subprocess by test_cosmos3_edge.py, because +diffusers main (>= 0.40, first with the Edge classes) cannot be imported into +a process that already imported the pinned diffusers. +Usage: DIFFUSERS_MAIN_PATH=/path/to/diffusers python cosmos3_edge_diffusers_parity.py + +Runs the diffusers Cosmos3OmniPipeline for 2 steps (guidance off, fixed +latents), captures each transformer call's (latent, timestep, velocity), +then replays the same inputs through the TRT-LLM Edge transformer and +compares velocities. +""" + +import os +import sys + +sys.path.insert(0, os.path.join(os.environ["DIFFUSERS_MAIN_PATH"], "src")) +os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + +CKPT = sys.argv[1] +DEV = "cuda" +H, W, FRAMES, STEPS = 192, 320, 9, 2 +PROMPT = "A red cube sits on a wooden table." + + +def main(): + import diffusers + import torch + + expected_root = os.path.realpath(os.environ["DIFFUSERS_MAIN_PATH"]) + assert os.path.realpath(diffusers.__file__).startswith(expected_root), diffusers.__file__ + print("diffusers:", diffusers.__version__) + from diffusers import Cosmos3OmniPipeline + + pipe = Cosmos3OmniPipeline.from_pretrained(CKPT, torch_dtype=torch.bfloat16).to(DEV) + + records = [] + orig_forward = pipe.transformer.forward + + def spy(**kwargs): + out = orig_forward(**kwargs) + velocity = out[0][0] if isinstance(out, tuple) else out.sample[0] + records.append( + { + "latent": kwargs["vision_tokens"][0].detach().clone(), + "t": kwargs["vision_timesteps"].detach().reshape(-1)[0].clone(), + "vel": velocity.detach().clone(), + } + ) + return out + + pipe.transformer.forward = spy + + latent_shape = (1, 48, (FRAMES - 1) // 4 + 1, H // 16, W // 16) + init_latents = torch.randn( + latent_shape, generator=torch.Generator().manual_seed(0), dtype=torch.float32 + ).to(DEV, torch.bfloat16) + + cond_ids, _ = pipe.tokenize_prompt( + PROMPT, negative_prompt=None, num_frames=FRAMES, height=H, width=W, fps=24.0 + ) + + pipe( + prompt=PROMPT, + num_frames=FRAMES, + height=H, + width=W, + fps=24.0, + num_inference_steps=STEPS, + guidance_scale=1.0, + generator=torch.Generator().manual_seed(0), + latents=init_latents.clone(), + output_type="latent", + ) + print(f"captured {len(records)} transformer calls") + + del pipe.transformer + del pipe + torch.cuda.empty_cache() + + from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader + from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + + args = VisualGenArgs(model=CKPT, torch_compile_config=TorchCompileConfig(enable=False)) + trt_pipe = PipelineLoader(args).load( + skip_warmup=True, + skip_components=[ + PipelineComponent.VAE, + PipelineComponent.SCHEDULER, + PipelineComponent.TOKENIZER, + PipelineComponent.SOUND_TOKENIZER, + ], + ) + transformer = trt_pipe.transformer + + text_ids = torch.tensor([cond_ids], dtype=torch.long, device=DEV) + text_mask = torch.ones_like(text_ids) + latent_t, latent_h, latent_w = latent_shape[2:] + + transformer.reset_cache() + for step, rec in enumerate(records): + latent = rec["latent"].to(DEV, torch.bfloat16) + if latent.dim() == 4: + latent = latent.unsqueeze(0) + t = rec["t"].float().reshape(1).to(DEV) + with torch.inference_mode(): + out = transformer( + hidden_states=latent, + timestep=t / 1000.0, + raw_timestep=t, + text_ids=text_ids, + text_mask=text_mask, + video_shape=(latent_t, latent_h, latent_w), + fps=24.0, + ) + ours = out.video[0].float().cpu() + ref = rec["vel"].float().cpu() + if ref.dim() == 5: + ref = ref[0] + diff = (ours - ref).abs() + rel = diff.max() / ref.abs().max() + print( + f"step {step}: t={t.item():7.2f} max|d|={diff.max():.5f} " + f"mean|d|={diff.mean():.6f} rel={rel:.5f} |ref|max={ref.abs().max():.3f}" + ) + + +if __name__ == "__main__": + main() diff --git a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py index e5b770f5dfa8..afc39db47aa4 100644 --- a/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py +++ b/tests/unittest/_torch/visual_gen/multi_gpu/test_cosmos3_transformer_parallel.py @@ -34,6 +34,7 @@ ) from tensorrt_llm._torch.visual_gen.mapping import VisualGenMapping from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + COSMOS3_EDGE_BACKBONE_TYPE, Cosmos3VFMTransformer, ) @@ -47,6 +48,9 @@ MODULES_AVAILABLE = True except ImportError: MODULES_AVAILABLE = False + # Module-level configs below reference this; every test skips in this + # branch, but the definitions still have to import cleanly. + COSMOS3_EDGE_BACKBONE_TYPE = "cosmos3_edge_nemotron_dense" # Attention2D (attn2d) wraps the compute backend in Attention2DAttention, which # requires (a) an LSE-capable inner backend — only FA4, VANILLA does not support @@ -78,7 +82,7 @@ num_attention_heads=8, num_key_value_heads=4, head_dim=64, - rope_scaling={"rope_type": "default", "mrope_section": [16, 12, 12]}, + rope_scaling={"rope_type": "default", "mrope_section": [12, 10, 10]}, rms_norm_eps=1e-6, vocab_size=1024, rope_theta=1_000_000.0, @@ -91,13 +95,37 @@ # attn2d needs an LSE-capable backend (FA4); FA4's CUTE kernels run with head_dim=128. # Same architecture as _COSMOS3_TEST_CONFIG otherwise (heads still divisible by -# Ulysses=2 for the attn2d+ulysses case). mrope_section is unchanged: the interleave -# slices clip to head_dim//2 (=64 here), so [16, 12, 12] stays valid. +# Ulysses=2 for the attn2d+ulysses case). mrope_section must sum to head_dim//2 +# (=64 here), so this config uses the real checkpoint sections. _COSMOS3_FA4_CONFIG = dict( _COSMOS3_TEST_CONFIG, head_dim=128, hidden_size=8 * 128, intermediate_size=1024, + rope_scaling={"rope_type": "default", "mrope_section": [24, 20, 20]}, +) + +# Edge (Nemotron-dense) recipe. The point of covering it here is the MLP: the +# Qwen recipe builds a GatedMLP whose gate_proj/up_proj the loader fuses into +# gate_up_proj, while Edge builds a plain relu² MLP with no fusion — a +# different column/row sharding path that the configs above never reach. The +# latent-geometry values are the recipe's validated invariants and cannot +# shrink; only the backbone dimensions do. +_COSMOS3_EDGE_CONFIG = dict( + _COSMOS3_TEST_CONFIG, + backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, + hidden_act="relu2", + use_und_k_norm_for_gen=True, + attention_bias=False, + sound_gen=False, + sound_dim=None, + latent_channel=48, + latent_patch_size=2, + patch_latent_dim=192, + rope_axes_dim=[12, 10, 10], + qk_norm_for_text=False, + rms_norm_eps=1e-5, + temporal_compression_factor=4, ) # Video: [B, C, T, H, W]. patch_size=2 → seq_len = T * (H/2) * (W/2). @@ -379,8 +407,16 @@ def _cosmos3_inputs( return hidden_states, timestep, text_ids, text_mask, (_LATENT_T, _LATENT_H, _LATENT_W) -def _forward(model: Cosmos3VFMTransformer, device: torch.device, text_seed: int) -> torch.Tensor: - channels = _COSMOS3_TEST_CONFIG["latent_channel"] +def _forward( + model: Cosmos3VFMTransformer, + device: torch.device, + text_seed: int, + pretrained_dict: dict = None, +) -> torch.Tensor: + # Latent channel count is a recipe invariant, not a free knob: the Edge + # recipe validates 48 where the Qwen test config uses 4. + config = pretrained_dict if pretrained_dict is not None else _COSMOS3_TEST_CONFIG + channels = config["latent_channel"] hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( device, channels=channels, text_seed=text_seed ) @@ -509,6 +545,36 @@ def _logic_cosmos3_tp_vs_single_gpu(rank, world_size): _assert_parity(tp_out, ref_out, msg=f"Rank {rank}: TP output differs from single-GPU reference") +def _logic_cosmos3_edge_tp_vs_single_gpu(rank, world_size): + ref_model, tp_model, _, device = _build_ref_and_parallel( + tp_size=world_size, pretrained_dict=_COSMOS3_EDGE_CONFIG + ) + text_seed = _cfg_text_seed(rank, tp_size=world_size, ulysses_size=1, cfg_size=1) + + ref_out = _forward(ref_model, device, text_seed, _COSMOS3_EDGE_CONFIG) + tp_out = _forward(tp_model, device, text_seed, _COSMOS3_EDGE_CONFIG) + + _assert_parity( + tp_out, ref_out, msg=f"Rank {rank}: Edge TP output differs from single-GPU reference" + ) + + +def _logic_cosmos3_edge_ulysses_vs_single_gpu(rank, world_size): + ref_model, ulysses_model, _, device = _build_ref_and_parallel( + ulysses_size=world_size, pretrained_dict=_COSMOS3_EDGE_CONFIG + ) + text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) + + ref_out = _forward(ref_model, device, text_seed, _COSMOS3_EDGE_CONFIG) + ulysses_out = _forward(ulysses_model, device, text_seed, _COSMOS3_EDGE_CONFIG) + + _assert_parity( + ulysses_out, + ref_out, + msg=f"Rank {rank}: Edge Ulysses output differs from single-GPU reference", + ) + + def _logic_cosmos3_ulysses_vs_single_gpu(rank, world_size): ref_model, ulysses_model, _, device = _build_ref_and_parallel(ulysses_size=world_size) text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) @@ -705,6 +771,18 @@ def test_ulysses2_vs_single_gpu(self): self._skip_if_unavailable() run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_vs_single_gpu) + def test_edge_tp2_vs_single_gpu(self): + """Edge's non-gated relu² MLP shards without the gate_up fusion the + Qwen recipe uses, so column/row splitting takes a different path.""" + self._skip_if_unavailable() + run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_edge_tp_vs_single_gpu) + + def test_edge_ulysses2_vs_single_gpu(self): + """Edge under sequence sharding: no und Q/K norm, and the reasoner's + keys are normed only where the generator consumes them.""" + self._skip_if_unavailable() + run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_edge_ulysses_vs_single_gpu) + def test_ulysses2_audio_vs_single_gpu(self): """Ulysses parity with the audio modality on: video + audio tokens are sharded together across the sequence dimension.""" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py index e949c0183be5..02280c2f4419 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_distilled.py @@ -6,8 +6,10 @@ defaults, mode resolution, and the guidance-1.0 denoise-loop contract.""" import json +import pickle from pathlib import Path from types import SimpleNamespace +from unittest.mock import MagicMock, patch import pytest import torch @@ -23,8 +25,10 @@ Cosmos3SamplingPolicy, load_scheduler, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import QWEN3_RECIPE from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY, AutoPipeline from tensorrt_llm._torch.visual_gen.profiler import VisualGenProfiler +from tensorrt_llm.visual_gen.params import VisualGenParams pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] @@ -83,35 +87,52 @@ def _bare_pipeline(**attrs) -> Cosmos3OmniMoTPipeline: defaults = dict( audio_gen=False, action_gen=False, + has_action_weights=False, sampling=Cosmos3SamplingPolicy(), default_use_system_prompt=False, + family=QWEN3_RECIPE.name, + use_native_flow_schedule=False, ) defaults.update(attrs) + defaults.setdefault("_scheduler_cache", {}) + defaults.setdefault("_base_scheduler", defaults.get("scheduler")) for key, value in defaults.items(): setattr(pipeline, key, value) return pipeline def _fake_request(output_type: str = "video", **param_overrides) -> SimpleNamespace: - """A DiffusionRequest look-alike with executor-merged (None = unset) params.""" - params = SimpleNamespace( - height=None, - width=None, - num_inference_steps=None, - guidance_scale=None, + """A DiffusionRequest look-alike carrying executor-merged params. + + Mirrors ``_merge_defaults``: mode-independent defaults are written in and + then un-marked in ``model_fields_set``, so only ``param_overrides`` count + as caller-supplied. Use ``_merged_request`` for values that arrived from + the pipeline's default table rather than from the caller. + """ + params = VisualGenParams( num_frames=COSMOS3_720P_PARAMS["num_frames"], max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], frame_rate=COSMOS3_720P_PARAMS["frame_rate"], seed=0, - negative_prompt=None, - image=None, extra_params={"output_type": output_type}, ) + # ``seed`` is materialized by generate(), not merged, so it stays marked. + for key in ("num_frames", "max_sequence_length", "frame_rate"): + params.model_fields_set.discard(key) for key, value in param_overrides.items(): setattr(params, key, value) return SimpleNamespace(prompt="x", params=params) +def _merged_request(output_type: str = "video", **merged) -> SimpleNamespace: + """Like ``_fake_request``, but the values arrive as executor-merged + pipeline defaults (present, yet not marked caller-supplied).""" + req = _fake_request(output_type, **merged) + for key in merged: + req.params.model_fields_set.discard(key) + return req + + class TestSchedulerLoading: def test_flow_match_declared(self, tmp_path): _write_scheduler_config(tmp_path, DISTILLED_SCHEDULER_CONFIG) @@ -350,15 +371,19 @@ def test_distilled_defaults_report_checkpoint_truth(self): params = _bare_pipeline(sampling=_distilled_policy()).default_generation_params assert params["num_inference_steps"] == 4 assert params["guidance_scale"] == DISTILLED_GUIDANCE_SCALE - assert params["height"] is None # mode-dependent, resolved in infer() + assert params["height"] == COSMOS3_720P_PARAMS["height"] assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] - def test_base_defaults_leave_mode_dependent_fields_unset(self): + def test_base_defaults_report_the_video_table(self): + """The declared defaults are concrete video-mode values, so + ``VisualGen.default_params`` can show what an unmodified request runs. + A request in another mode re-resolves them in ``infer()``.""" params = _bare_pipeline().default_generation_params for field in ("height", "width", "num_inference_steps", "guidance_scale"): - assert params[field] is None + assert params[field] == COSMOS3_720P_PARAMS[field] assert params["num_frames"] == COSMOS3_720P_PARAMS["num_frames"] assert params["max_sequence_length"] == COSMOS3_720P_PARAMS["max_sequence_length"] + assert "flow_shift" not in params class TestInferModeResolution: @@ -396,6 +421,114 @@ def test_distilled_merged_defaults_pass_through(self): assert got["guidance_scale"] == DISTILLED_GUIDANCE_SCALE assert got["height"] == COSMOS3_T2I_PARAMS["height"] + def test_merged_video_defaults_reresolve_for_t2i(self): + """The regression this contract exists for: values merged from the + video table must not be mistaken for caller intent when the request + selects text-to-image.""" + req = _merged_request( + "image", + height=COSMOS3_720P_PARAMS["height"], + width=COSMOS3_720P_PARAMS["width"], + num_inference_steps=COSMOS3_720P_PARAMS["num_inference_steps"], + guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + ) + got = self._captured_forward_kwargs(_bare_pipeline(), req) + for field in ("height", "width", "num_inference_steps", "guidance_scale"): + assert got[field] == COSMOS3_T2I_PARAMS[field] + + def test_distilled_merged_steps_survive_reresolution(self): + """Distilled steps/guidance are checkpoint facts, not mode defaults: + re-resolution must keep them even though they are unmarked.""" + req = _merged_request("image", num_inference_steps=4, guidance_scale=1.0) + got = self._captured_forward_kwargs(_bare_pipeline(sampling=_distilled_policy()), req) + assert got["num_inference_steps"] == 4 + assert got["guidance_scale"] == DISTILLED_GUIDANCE_SCALE + + def test_caller_values_survive_a_mode_switch(self): + """An explicitly assigned value outranks the mode table.""" + req = _merged_request("image", height=COSMOS3_720P_PARAMS["height"]) + req.params.width = 777 # caller assignment marks it + got = self._captured_forward_kwargs(_bare_pipeline(), req) + assert got["width"] == 777 + assert got["height"] == COSMOS3_T2I_PARAMS["height"] + + +class TestDefaultMarksThroughRealPath: + """The unmarking contract, exercised through the production functions. + + The fixtures above unmark fields by hand, so they would stay green if the + ``model_fields_set`` bookkeeping were dropped from ``default_params`` or + ``_merge_defaults``. These drive the real code instead. + """ + + MODE_FIELDS = ("height", "width", "num_inference_steps", "guidance_scale") + + def _visual_gen(self, pipeline): + from tensorrt_llm.visual_gen import VisualGen + + with patch.object(VisualGen, "__init__", lambda self, *a, **kw: None): + vg = VisualGen.__new__(VisualGen) + vg.executor = MagicMock() + vg.executor.default_generation_params = pipeline.default_generation_params + vg.executor.extra_param_specs = pipeline.extra_param_specs + return vg + + def _merge(self, pipeline, params): + from tensorrt_llm._torch.visual_gen.executor import DiffusionExecutor, DiffusionRequest + + executor = MagicMock() + executor.pipeline = pipeline # real object: no stray truthy getattr results + req = DiffusionRequest(request_id=0, prompt=["x"], params=params) + DiffusionExecutor._merge_defaults(executor, req) + return req + + def test_default_params_are_concrete_yet_unmarked(self): + params = self._visual_gen(_bare_pipeline()).default_params + for field in self.MODE_FIELDS: + assert getattr(params, field) == COSMOS3_720P_PARAMS[field] + assert field not in params.model_fields_set + + def test_merge_defaults_fills_without_marking(self): + pipeline = _bare_pipeline() + req = self._merge(pipeline, VisualGenParams()) + for field in self.MODE_FIELDS: + assert getattr(req.params, field) == COSMOS3_720P_PARAMS[field] + assert field not in req.params.model_fields_set + + def test_assignment_marks_the_field(self): + req = self._merge(_bare_pipeline(), VisualGenParams()) + assert "height" not in req.params.model_fields_set + req.params.height = 256 + assert "height" in req.params.model_fields_set + + def test_deep_copy_and_pickle_preserve_marks(self): + params = self._visual_gen(_bare_pipeline()).default_params + params.width = 640 # one marked, the rest not + + for label, clone in ( + ("model_copy", params.model_copy(deep=True)), + ("pickle", pickle.loads(pickle.dumps(params))), # noqa: S301 - test fixture + ): + assert "width" in clone.model_fields_set, label + assert "height" not in clone.model_fields_set, label + assert clone.height == COSMOS3_720P_PARAMS["height"], label + + def test_full_chain_mode_switch_resolves_to_t2i(self): + """default_params -> switch mode -> deep copy -> pickle -> merge -> infer.""" + pipeline = _bare_pipeline() + params = self._visual_gen(pipeline).default_params + params.extra_params["output_type"] = "image" + + params = pickle.loads(pickle.dumps(params.model_copy(deep=True))) # noqa: S301 + req = self._merge(pipeline, params) + + captured = {} + pipeline.forward = lambda **kwargs: captured.update(kwargs) + pipeline.infer(req) + + for field in self.MODE_FIELDS: + assert captured[field] == COSMOS3_T2I_PARAMS[field] + class TestPipelineSchedulerLoading: def test_distilled_checkpoint_loads_flow_match(self, tmp_path): @@ -639,6 +772,28 @@ def test_anchor_writes_only_frame_zero_in_place(self): assert torch.all(latents[:, :, 0:1] == self.CLEAN) assert torch.equal(latents[:, :, 1:], untouched) + def test_anchor_rejects_dtype_mismatch(self): + """Slice assignment would silently cast; the anchor must refuse instead, + so a dtype drift surfaces as an error rather than a per-step conversion.""" + pipeline = _bare_pipeline(sampling=_distilled_policy()) + post_step_fn = pipeline._conditioning_anchor_post_step( + self._clean_frame().to(torch.float16) + ) + latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.float32) + + with pytest.raises(RuntimeError, match="must match the denoised latents"): + post_step_fn(latents) + + def test_anchor_accepts_matching_dtype(self): + pipeline = _bare_pipeline(sampling=_distilled_policy()) + post_step_fn = pipeline._conditioning_anchor_post_step( + self._clean_frame().to(torch.bfloat16) + ) + latents = torch.zeros(1, 4, 3, 2, 2, dtype=torch.bfloat16) + + post_step_fn(latents) + assert torch.all(latents[:, :, 0:1] == self.CLEAN) + def _run_denoise(self, with_anchor: bool): """Run the real BasePipeline.denoise loop with a perturbing scheduler, recording what the transformer receives at every step.""" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py new file mode 100644 index 000000000000..ffe3777774d4 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_edge.py @@ -0,0 +1,1478 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Cosmos3-Edge (Nemotron-dense backbone) tests. + +Unit tests run on a reduced config mirroring the Edge checkpoint's exact key +set: recipe validation, Nemotron norm semantics, the generator-only und +K-norm, native flow schedule parity against cosmos-framework, strict weight +loading, and per-family defaults. Checkpoint-gated tests cover the real +checkpoint (tokenizer, recipe/scheduler wiring, load + forward). + +Override checkpoint: + DIFFUSION_MODEL_PATH_COSMOS3_EDGE=/path/to/Cosmos3-Edge \\ + pytest tests/unittest/_torch/visual_gen/test_cosmos3_edge.py -v +""" + +import gc +import os +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from diffusers import UniPCMultistepScheduler + +from tensorrt_llm._torch.modules.mlp import MLP +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.models.cosmos3 import pipeline_cosmos3 as pipeline_module +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import ( + COSMOS3_720P_PARAMS, + COSMOS3_EDGE_T2I_PARAMS, + COSMOS3_EDGE_VIDEO_PARAMS, + COSMOS3_GENERATION_DEFAULTS, +) +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.sampling import Cosmos3SamplingPolicy +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + COSMOS3_EDGE_BACKBONE_TYPE, + NEMOTRON_DENSE_RECIPE, + QWEN3_RECIPE, + Cosmos3VFMTransformer, + NemotronRMSNorm, + Qwen3VLTextRMSNorm, + resolve_arch_recipe, +) +from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineComponent, PipelineLoader +from tensorrt_llm._torch.visual_gen.pipeline_registry import PIPELINE_REGISTRY +from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs + +pytestmark = [pytest.mark.cosmos3, pytest.mark.usefixtures("disable_cosmos3_guardrails")] + +DEVICE = "cuda" + + +@pytest.fixture(autouse=True) +def _cleanup_gpu(): + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + yield + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + +def _require_edge_checkpoint() -> str: + """Resolve the Edge checkpoint lazily so unit tests collect and run on + machines without model storage; only checkpoint-gated tests skip.""" + path = os.environ.get("DIFFUSION_MODEL_PATH_COSMOS3_EDGE") + if not path: + root = Path(os.environ.get("LLM_MODELS_ROOT", "/home/scratch.trt_llm_data_ci/llm-models/")) + if not root.exists(): + root = Path("/scratch/trt_llm_data/llm-models/") + path = str(root / "Cosmos3-Edge") + if not os.path.isdir(path): + pytest.skip(f"Checkpoint not found: {path}") + return path + + +def _reduced_edge_config() -> SimpleNamespace: + # Key set mirrors the Edge checkpoint's transformer/config.json verbatim + # (including the missing rope_type in rope_scaling); only sizes shrink. + return SimpleNamespace( + action_dim=8, + action_gen=True, + attention_bias=False, + attention_dropout=0.0, + backbone_type=COSMOS3_EDGE_BACKBONE_TYPE, + base_fps=24, + enable_fps_modulation=True, + head_dim=8, + hidden_act="relu2", + hidden_size=32, + intermediate_size=64, + # Latent geometry stays at the real invariant values (validated for + # the Edge recipe); only backbone dimensions shrink. + latent_channel=48, + latent_patch_size=2, + num_attention_heads=4, + num_embodiment_domains=32, + num_hidden_layers=2, + num_key_value_heads=2, + patch_latent_dim=192, + qk_norm_for_text=False, + rms_norm_eps=1e-5, + rope_axes_dim=[2, 1, 1], + rope_scaling={"mrope_section": [2, 1, 1]}, + rope_theta=100000000, + sound_dim=None, + sound_gen=False, + temporal_compression_factor=4, + timestep_scale=0.001, + unified_3d_mrope_reset_spatial_ids=True, + unified_3d_mrope_temporal_modality_margin=15000, + use_und_k_norm_for_gen=True, + vocab_size=64, + ) + + +def _reduced_edge_model_config() -> DiffusionModelConfig: + model_config = DiffusionModelConfig(pretrained_config=_reduced_edge_config()) + model_config.attention.backend = "VANILLA" + return model_config + + +def _reduced_qwen3_config(**overrides) -> SimpleNamespace: + cfg = _reduced_edge_config() + cfg.backbone_type = None + cfg.hidden_act = "silu" + cfg.qk_norm_for_text = True + cfg.use_und_k_norm_for_gen = False + cfg.rms_norm_eps = 1e-6 + cfg.rope_theta = 5000000 + cfg.rope_scaling = {"mrope_section": [2, 1, 1], "rope_type": "default"} + for key, value in overrides.items(): + setattr(cfg, key, value) + return cfg + + +def _reduced_qwen3_model_config(**overrides) -> DiffusionModelConfig: + model_config = DiffusionModelConfig(pretrained_config=_reduced_qwen3_config(**overrides)) + model_config.attention.backend = "VANILLA" + return model_config + + +def _synthetic_state_dict(cfg: SimpleNamespace) -> dict: + """A complete synthetic checkpoint for either recipe (diffusers key layout).""" + h, d = cfg.hidden_size, cfg.head_dim + q_dim = cfg.num_attention_heads * d + kv_dim = cfg.num_key_value_heads * d + gated = cfg.hidden_act == "silu" + sd = { + "embed_tokens.weight": torch.randn(cfg.vocab_size, h), + "lm_head.weight": torch.randn(cfg.vocab_size, h), + "norm.weight": torch.ones(h), + "norm_moe_gen.weight": torch.ones(h), + "proj_in.weight": torch.randn(h, cfg.patch_latent_dim), + "proj_in.bias": torch.randn(h), + "proj_out.weight": torch.randn(cfg.patch_latent_dim, h), + "proj_out.bias": torch.randn(cfg.patch_latent_dim), + "time_embedder.linear_1.weight": torch.randn(h, 256), + "time_embedder.linear_1.bias": torch.randn(h), + "time_embedder.linear_2.weight": torch.randn(h, h), + "time_embedder.linear_2.bias": torch.randn(h), + } + if getattr(cfg, "sound_gen", False): + sd.update( + { + "audio_proj_in.weight": torch.randn(h, cfg.sound_dim), + "audio_proj_in.bias": torch.randn(h), + "audio_proj_out.weight": torch.randn(cfg.sound_dim, h), + "audio_proj_out.bias": torch.randn(cfg.sound_dim), + "audio_modality_embed": torch.randn(h), + } + ) + for i in range(cfg.num_hidden_layers): + p = f"layers.{i}" + sd.update( + { + f"{p}.self_attn.to_q.weight": torch.randn(q_dim, h), + f"{p}.self_attn.to_k.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_v.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_out.weight": torch.randn(h, q_dim), + f"{p}.self_attn.add_q_proj.weight": torch.randn(q_dim, h), + f"{p}.self_attn.add_k_proj.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.add_v_proj.weight": torch.randn(kv_dim, h), + f"{p}.self_attn.to_add_out.weight": torch.randn(h, q_dim), + f"{p}.self_attn.norm_added_q.weight": torch.ones(d), + f"{p}.self_attn.norm_added_k.weight": torch.ones(d), + f"{p}.input_layernorm.weight": torch.ones(h), + f"{p}.input_layernorm_moe_gen.weight": torch.ones(h), + f"{p}.post_attention_layernorm.weight": torch.ones(h), + f"{p}.post_attention_layernorm_moe_gen.weight": torch.ones(h), + f"{p}.mlp.up_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp.down_proj.weight": torch.randn(h, cfg.intermediate_size), + f"{p}.mlp_moe_gen.up_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp_moe_gen.down_proj.weight": torch.randn(h, cfg.intermediate_size), + } + ) + if gated: + sd.update( + { + f"{p}.self_attn.norm_q.weight": torch.ones(d), + f"{p}.self_attn.norm_k.weight": torch.ones(d), + f"{p}.mlp.gate_proj.weight": torch.randn(cfg.intermediate_size, h), + f"{p}.mlp_moe_gen.gate_proj.weight": torch.randn(cfg.intermediate_size, h), + } + ) + else: + sd[f"{p}.self_attn.k_norm_und_for_gen.weight"] = torch.ones(d) + return sd + + +def _edge_state_dict(cfg: SimpleNamespace) -> dict: + return _synthetic_state_dict(cfg) + + +class TestArchRecipe: + def test_edge_config_resolves_nemotron_recipe(self): + recipe = resolve_arch_recipe(_reduced_edge_config()) + assert recipe is NEMOTRON_DENSE_RECIPE + + def test_qwen3_resolves_without_backbone_type(self): + cfg = SimpleNamespace(hidden_act="silu", qk_norm_for_text=True) + assert resolve_arch_recipe(cfg) is QWEN3_RECIPE + + def test_unknown_backbone_raises(self): + cfg = _reduced_edge_config() + cfg.backbone_type = "cosmos4_hybrid" + with pytest.raises(ValueError, match="cosmos4_hybrid"): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize( + "key,value", + [ + ("hidden_act", "silu"), + ("qk_norm_for_text", True), + ("use_und_k_norm_for_gen", False), + ("sound_gen", True), + ("sound_gen", None), + ("attention_bias", True), + ("rms_norm_eps", 1e-6), + ], + ) + def test_edge_flag_contradiction_raises(self, key, value): + cfg = _reduced_edge_config() + setattr(cfg, key, value) + with pytest.raises(ValueError, match=key): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize("make_config", [_reduced_edge_config, lambda: _reduced_qwen3_config()]) + def test_inconsistent_patch_latent_dim_raises(self, make_config): + cfg = make_config() + cfg.patch_latent_dim = cfg.patch_latent_dim + 1 + with pytest.raises(ValueError, match="patch_latent_dim"): + resolve_arch_recipe(cfg) + + +class TestRopeAxesResolution: + """The explicit top-level rope_axes_dim wins over the legacy + rope_scaling.mrope_section (diffusers precedence); disagreement and + head_dim mismatches are config errors.""" + + def _resolve(self, cfg): + from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import ( + resolve_rope_axes_dim, + ) + + return resolve_rope_axes_dim(cfg) + + def test_top_level_only(self): + cfg = _reduced_edge_config() + cfg.rope_scaling = {} + assert self._resolve(cfg) == [2, 1, 1] + + def test_nested_only(self): + cfg = _reduced_edge_config() + del cfg.rope_axes_dim + assert self._resolve(cfg) == [2, 1, 1] + + def test_contradiction_raises(self): + cfg = _reduced_edge_config() + cfg.rope_axes_dim = [1, 2, 1] + with pytest.raises(ValueError, match="contradictory"): + self._resolve(cfg) + + @pytest.mark.parametrize("axes", [[2, 2, 1], [2, 2], [4]]) + def test_head_dim_mismatch_raises(self, axes): + cfg = _reduced_edge_config() + cfg.rope_axes_dim = axes + cfg.rope_scaling = {} + with pytest.raises(ValueError, match="head_dim"): + self._resolve(cfg) + + def test_neither_declared_raises(self): + cfg = _reduced_edge_config() + del cfg.rope_axes_dim + cfg.rope_scaling = {} + with pytest.raises(ValueError, match="neither"): + self._resolve(cfg) + + def test_transformer_builds_from_top_level_only(self): + cfg = _reduced_edge_config() + cfg.rope_scaling = {} + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + model = Cosmos3VFMTransformer(model_config) + assert model.language_model.rotary_emb.mrope_section == [2, 1, 1] + + def test_qwen3_flag_contradiction_raises(self): + cfg = SimpleNamespace(hidden_act="relu2") + with pytest.raises(ValueError, match="hidden_act"): + resolve_arch_recipe(cfg) + + @pytest.mark.parametrize("key,value", [("latent_channel", 4), ("latent_patch_size", 4)]) + def test_edge_latent_geometry_invariants(self, key, value): + cfg = _reduced_edge_config() + setattr(cfg, key, value) + with pytest.raises(ValueError, match=key): + resolve_arch_recipe(cfg) + + +class TestNemotronNormSemantics: + """Pins the one intentional numerics fork: fp32 weight multiply, then + downcast — vs the Qwen flavor's bf16 multiply after downcast.""" + + def test_matches_fp32_weight_multiply_bit_exact(self): + torch.manual_seed(0) + x = (torch.randn(256, 8) * 3).bfloat16() + weight = (torch.randn(8) * 2 + 1.5).bfloat16() + + nemotron = NemotronRMSNorm(hidden_size=8, eps=1e-5) + qwen = Qwen3VLTextRMSNorm(hidden_size=8, eps=1e-5) + with torch.no_grad(): + nemotron.weight.copy_(weight) + qwen.weight.copy_(weight) + + xf = x.float() + normed = xf * torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + 1e-5) + reference = (weight.float() * normed).to(torch.bfloat16) + + assert torch.equal(nemotron(x), reference) + assert not torch.equal(qwen(x), reference) + + +class TestEdgeTransformerStructure: + @pytest.fixture(scope="class") + def model(self): + return Cosmos3VFMTransformer(_reduced_edge_model_config()) + + def test_recipe_and_flags(self, model): + assert model.recipe is NEMOTRON_DENSE_RECIPE + assert model.audio_gen is False + assert model.has_action_weights is True + assert model.temporal_compression_factor == 4 + assert model.temporal_compression_factor_declared is True + + def test_und_attention_norms(self, model): + attn = model.language_model.layers[0].self_attn + assert attn.norm_q is None + assert attn.norm_k is None + assert isinstance(attn.k_norm_und_for_gen, NemotronRMSNorm) + + def test_nemotron_norms_everywhere(self, model): + und = model.language_model.layers[0] + gen = model.gen_layers[0] + for norm in ( + und.input_layernorm, + und.post_attention_layernorm, + und.self_attn.k_norm_und_for_gen, + gen.input_layernorm, + gen.post_attention_layernorm, + gen.cross_attention.norm_q, + gen.cross_attention.norm_k, + model.norm_moe_gen, + ): + assert isinstance(norm, NemotronRMSNorm) + assert norm.variance_epsilon == 1e-5 + + def test_relu2_mlp_no_gate(self, model): + for layer in (model.language_model.layers[0], model.gen_layers[0]): + assert isinstance(layer.mlp, MLP) + assert not hasattr(layer.mlp, "gate_proj") + + def test_rope_defaults_without_rope_type(self, model): + assert model.language_model.rotary_emb.rope_type == "default" + + +class TestGeneratorOnlyKNorm: + """Mirror of diffusers' Edge regression: perturbing k_norm_und_for_gen + must not change the und causal attention output, only the gen-facing K.""" + + def test_k_norm_touches_only_gen_facing_keys(self): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + torch.manual_seed(0) + model = Cosmos3VFMTransformer(_reduced_edge_model_config()).to(DEVICE).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(0, 0.02) + model.post_load_weights() + + attn = model.language_model.layers[0].self_attn + hidden = torch.randn(1, 4, 32, dtype=torch.bfloat16, device=DEVICE) + cos = torch.ones(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) + sin = torch.zeros(1, 4, 1, 8, dtype=torch.bfloat16, device=DEVICE) + + with torch.inference_mode(): + out_before, k_gen_before, v_before = attn.forward_with_kv(hidden, cos, sin) + # The rope above is identity (cos=1, sin=0), so the cached gen K + # must be exactly the Nemotron-normed raw K. + _, k_raw, _ = attn.get_qkv(hidden) + k_raw = k_raw.view(1, 4, attn.local_num_key_value_heads, attn.head_dim) + assert torch.equal(k_gen_before, attn.k_norm_und_for_gen(k_raw)) + + attn.k_norm_und_for_gen.weight.fill_(2.0) + out_after, k_gen_after, v_after = attn.forward_with_kv(hidden, cos, sin) + + assert torch.equal(out_before, out_after) + assert torch.equal(v_before, v_after) + assert not torch.equal(k_gen_before, k_gen_after) + + def test_nano_qk_norm_stays_byte_identical(self): + """Regression pin for the qwen3 recipe: apply_qk_norm must keep + today's exact ``F.rms_norm`` semantics (a future refactor routing it + through ``Qwen3VLTextRMSNorm.forward`` would change Nano numerics).""" + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + import torch.nn.functional as F + + torch.manual_seed(0) + model = Cosmos3VFMTransformer(_reduced_qwen3_model_config()).to(DEVICE).eval() + with torch.no_grad(): + for param in model.parameters(): + param.normal_(0, 0.02) + attn = model.language_model.layers[0].self_attn + + q = torch.randn(1, 4, 4, 8, dtype=torch.bfloat16, device=DEVICE) + k = torch.randn(1, 4, 2, 8, dtype=torch.bfloat16, device=DEVICE) + q_normed, k_normed = attn.apply_qk_norm(q, k) + assert torch.equal( + q_normed, F.rms_norm(q, (8,), attn.norm_q.weight, attn.norm_q.variance_epsilon) + ) + assert torch.equal( + k_normed, F.rms_norm(k, (8,), attn.norm_k.weight, attn.norm_k.variance_epsilon) + ) + + +# The relevant subset of the Edge checkpoint's scheduler config (values +# verbatim; identical to Nano's). +EDGE_UNIPC_CONFIG = { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "use_karras_sigmas": True, + "sigma_max": 200.0, + "sigma_min": 0.147, + "solver_order": 2, + "solver_type": "bh2", + "final_sigmas_type": "zero", + "timestep_spacing": "linspace", + "lower_order_final": True, +} + +# Recorded from cosmos-framework 117c7d2 (`fm_solvers_unipc.py` +# FlowUniPCMultistepScheduler, num_train_timesteps=1000, shift=1.0, +# use_dynamic_shifting=False; set_timesteps(steps, shift=shift)). Each +# trajectory runs the synthetic velocity v = 0.05*x + 0.3*sin(t/1000) - 0.1 in +# float64 from x0 = linspace(-1, 1, 8).reshape(1, 2, 2, 2) through every +# step(). +COSMOS_FRAMEWORK_FIXTURES = { + (3.0, 10): { + "timesteps": [999, 963, 922, 874, 817, 749, 666, 562, 428, 249], + "sigmas": [ + 0.99966645, + 0.96394110, + 0.92272168, + 0.87463522, + 0.81780970, + 0.74962479, + 0.66629612, + 0.56214833, + 0.42826521, + 0.24979164, + 0.0, + ], + "final": [ + -0.995152214121, + -0.723385865647, + -0.451619517172, + -0.179853168698, + 0.091913179777, + 0.363679528251, + 0.635445876726, + 0.907212225200, + ], + }, + (10.0, 7): { + "timesteps": [999, 983, 961, 930, 882, 799, 624], + "sigmas": [ + 0.99989992, + 0.98349357, + 0.96140891, + 0.93008101, + 0.88217115, + 0.79977584, + 0.62472641, + 0.0, + ], + "final": [ + -1.040962681673, + -0.769320359241, + -0.497678036809, + -0.226035714377, + 0.045606608054, + 0.317248930486, + 0.588891252918, + 0.860533575350, + ], + }, + (5.0, 13): { + "timesteps": [999, 983, 964, 943, 918, 888, 853, 810, 757, 689, 599, 475, 293], + "sigmas": [ + 0.99979985, + 0.98339677, + 0.96469206, + 0.94316465, + 0.91812354, + 0.88863194, + 0.85338765, + 0.81052577, + 0.75727713, + 0.68934584, + 0.59968787, + 0.47589558, + 0.29389268, + 0.0, + ], + "final": [ + -0.998983254680, + -0.727228149256, + -0.455473043832, + -0.183717938407, + 0.088037167017, + 0.359792272441, + 0.631547377865, + 0.903302483290, + ], + }, +} + + +class TestNativeFlowSchedule: + def _native_policy_and_scheduler(self, shift: float): + scheduler = UniPCMultistepScheduler.from_config(EDGE_UNIPC_CONFIG) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler, native_flow_schedule=True) + return policy, policy.set_flow_shift(scheduler, shift) + + def test_set_flow_shift_disables_karras(self): + policy, scheduler = self._native_policy_and_scheduler(3.0) + assert float(scheduler.config.flow_shift) == 3.0 + assert scheduler.config.use_karras_sigmas is False + # Already matching → same instance. + assert policy.set_flow_shift(scheduler, 3.0) is scheduler + + def test_karras_rebuilt_even_at_checkpoint_shift(self): + """flow_shift 1.0 equals the checkpoint value, but the native flow + schedule still requires the karras grid off.""" + policy, scheduler = self._native_policy_and_scheduler(1.0) + assert scheduler.config.use_karras_sigmas is False + + def test_non_native_keeps_checkpoint_config(self): + scheduler = UniPCMultistepScheduler.from_config(EDGE_UNIPC_CONFIG) + policy = Cosmos3SamplingPolicy.from_scheduler(scheduler, native_flow_schedule=False) + assert policy.set_flow_shift(scheduler, 1.0) is scheduler + assert scheduler.config.use_karras_sigmas is True + + @pytest.mark.parametrize("shift,steps", sorted(COSMOS_FRAMEWORK_FIXTURES)) + def test_matches_cosmos_framework_reference(self, shift, steps): + fixture = COSMOS_FRAMEWORK_FIXTURES[(shift, steps)] + policy, scheduler = self._native_policy_and_scheduler(shift) + policy.set_timesteps(scheduler, num_inference_steps=steps, device="cpu") + + assert scheduler.timesteps.tolist() == fixture["timesteps"] + np.testing.assert_allclose(scheduler.sigmas.tolist(), fixture["sigmas"], atol=1e-6) + + x = torch.linspace(-1.0, 1.0, 8, dtype=torch.float64).reshape(1, 2, 2, 2) + for t in scheduler.timesteps: + v = 0.05 * x + 0.3 * torch.sin(t.double() / 1000.0) - 0.1 + x = scheduler.step(v, t, x, return_dict=False)[0] + np.testing.assert_allclose(x.flatten().tolist(), fixture["final"], atol=1e-9) + + +class TestStrictLoading: + def _model(self) -> Cosmos3VFMTransformer: + return Cosmos3VFMTransformer(_reduced_edge_model_config()) + + def test_full_checkpoint_loads(self): + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd["layers.0.self_attn.k_norm_und_for_gen.weight"] = torch.full((8,), 0.5) + model = self._model() + model.load_weights(sd) + loaded = model.language_model.layers[0].self_attn.k_norm_und_for_gen.weight + assert torch.equal(loaded.cpu().float(), torch.full((8,), 0.5)) + + @pytest.mark.parametrize( + "missing_key", + [ + "layers.0.self_attn.k_norm_und_for_gen.weight", + "layers.0.mlp.up_proj.weight", + "layers.1.mlp_moe_gen.down_proj.weight", + "layers.0.input_layernorm_moe_gen.weight", + "embed_tokens.weight", + ], + ) + def test_missing_weight_raises(self, missing_key): + sd = _edge_state_dict(_reduced_edge_config()) + del sd[missing_key] + with pytest.raises(ValueError, match="missing weights"): + self._model().load_weights(sd) + + def test_partial_fused_qkv_raises(self): + sd = _edge_state_dict(_reduced_edge_config()) + del sd["layers.0.self_attn.add_v_proj.weight"] + with pytest.raises(ValueError, match="missing weights"): + self._model().load_weights(sd) + + def test_intentional_skips_are_logged_with_names(self, monkeypatch): + """The skip log must name the skipped tensor families (dynamic + content, not just the logger's static category text).""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + infos = [] + monkeypatch.setattr(tf_module.logger, "info", infos.append) + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd.update( + { + "action_modality_embed": torch.randn(cfg.hidden_size), + "action_proj_in.fc.weight": torch.randn(4, 8), + "action_proj_in.bias.weight": torch.randn(4, 8), + "action_proj_out.fc.weight": torch.randn(4, 8), + "action_proj_out.bias.weight": torch.randn(4, 8), + } + ) + model = self._model() + model.load_weights(sd) + + param_names = {name for name, _ in model.named_parameters()} + assert not any("lm_head" in name for name in param_names) + assert "language_model.norm.weight" not in param_names + skip_logs = [m for m in infos if "intentionally unused" in m] + assert len(skip_logs) == 1 + # The dynamic name list follows the final colon; the static category + # text also mentions lm_head/norm, so assert the parsed set exactly. + skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} + assert skipped_families == { + "action_modality_embed", + "action_proj_in", + "action_proj_out", + "lm_head", + "norm", + } + + def test_model_prefixed_skip_keys_are_intentional(self, monkeypatch): + """Checkpoints that namespace top-level tensors under "model." must + have their skip keys recognized (prefix normalization runs before the + skip check), not warned about as unknown.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + warnings = [] + infos = [] + monkeypatch.setattr(tf_module.logger, "warning", warnings.append) + monkeypatch.setattr(tf_module.logger, "info", infos.append) + cfg = _reduced_edge_config() + sd = _edge_state_dict(cfg) + sd["model.lm_head.weight"] = torch.randn(cfg.vocab_size, cfg.hidden_size) + sd["model.action_modality_embed"] = torch.randn(cfg.hidden_size) + self._model().load_weights(sd) + + assert not any("unknown checkpoint key" in m for m in warnings) + skip_logs = [m for m in infos if "intentionally unused" in m] + assert len(skip_logs) == 1 + skipped_families = {name.strip() for name in skip_logs[0].rsplit(": ", 1)[1].split(",")} + assert {"lm_head", "action_modality_embed"} <= skipped_families + + def test_unconsumed_mapped_tensor_warns(self, monkeypatch): + """A checkpoint tensor that remaps to a module the recipe didn't + construct must warn, not vanish: Edge has no und norm_q, and a + sound_gen=false model constructs no audio projections.""" + import tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 as tf_module + + warnings = [] + monkeypatch.setattr(tf_module.logger, "warning", warnings.append) + sd = _edge_state_dict(_reduced_edge_config()) + sd["layers.0.self_attn.norm_q.weight"] = torch.ones(8) + sd["audio_proj_in.weight"] = torch.randn(32, 4) + self._model().load_weights(sd) + + matched = [m for m in warnings if "matched no constructed parameter" in m] + assert len(matched) == 1 + assert "norm_q" in matched[0] + assert "audio2llm" in matched[0] + + def test_missing_root_parameter_raises(self): + """Nano-family audio checkpoints carry a root parameter + (audio_modality_embed); its absence must fail, not stay random.""" + cfg = _reduced_qwen3_config( + sound_gen=True, sound_dim=4, sound_latent_fps=25, temporal_compression_factor_sound=1 + ) + sd = _synthetic_state_dict(cfg) + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + + Cosmos3VFMTransformer(model_config).load_weights(dict(sd)) + + del sd["audio_modality_embed"] + model_config = DiffusionModelConfig( + pretrained_config=_reduced_qwen3_config( + sound_gen=True, + sound_dim=4, + sound_latent_fps=25, + temporal_compression_factor_sound=1, + ) + ) + model_config.attention.backend = "VANILLA" + with pytest.raises(ValueError, match="audio_modality_embed"): + Cosmos3VFMTransformer(model_config).load_weights(sd) + + def test_qwen3_synthetic_checkpoint_loads(self): + """Nano-loading regression: the gated recipe (gate/up/down + und QK + norms, no k_norm_und_for_gen) loads cleanly through the same strict + coverage path.""" + cfg = _reduced_qwen3_config() + model_config = DiffusionModelConfig(pretrained_config=cfg) + model_config.attention.backend = "VANILLA" + model = Cosmos3VFMTransformer(model_config) + model.load_weights(_synthetic_state_dict(cfg)) + attn = model.language_model.layers[0].self_attn + assert attn.norm_q is not None + assert attn.k_norm_und_for_gen is None + + +def _bare_pipeline(family: str) -> Cosmos3OmniMoTPipeline: + pipeline = object.__new__(Cosmos3OmniMoTPipeline) + pipeline.family = family + pipeline.sampling = Cosmos3SamplingPolicy() + pipeline.audio_gen = False + pipeline.action_gen = False + pipeline.has_action_weights = family == NEMOTRON_DENSE_RECIPE.name + pipeline.use_native_flow_schedule = family == NEMOTRON_DENSE_RECIPE.name + return pipeline + + +class TestEdgeDefaults: + def test_generation_defaults_matrix(self): + edge_video = COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "video")] + assert edge_video is COSMOS3_EDGE_VIDEO_PARAMS + assert (edge_video["height"], edge_video["width"]) == (480, 832) + assert edge_video["num_frames"] == 121 + assert edge_video["num_inference_steps"] == 50 + assert edge_video["guidance_scale"] == 5.0 + assert edge_video["flow_shift"] == 3.0 + + edge_t2i = COSMOS3_GENERATION_DEFAULTS[(NEMOTRON_DENSE_RECIPE.name, "image")] + assert edge_t2i is COSMOS3_EDGE_T2I_PARAMS + assert (edge_t2i["height"], edge_t2i["width"]) == (640, 640) + assert edge_t2i["guidance_scale"] == 4.0 + assert edge_t2i["guidance_interval"] is None + + assert COSMOS3_GENERATION_DEFAULTS[(QWEN3_RECIPE.name, "video")] is COSMOS3_720P_PARAMS + + def test_executor_defaults_are_edge_shaped(self): + params = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name).default_generation_params + assert params["num_frames"] == 121 + assert params["max_sequence_length"] == 4096 + assert "flow_shift" not in params + # Concrete Edge video-mode values, not Nano/Super's 720p: this is what + # ``VisualGen.default_params`` reports for an Edge checkpoint. + assert (params["height"], params["width"]) == (480, 832) + assert params["num_inference_steps"] == 50 + assert params["guidance_scale"] == 5.0 + + def test_warmup_shape_per_family(self): + edge = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + assert edge.default_warmup_resolutions == [(480, 832)] + assert edge.default_warmup_num_frames == [121] + + qwen3 = _bare_pipeline(QWEN3_RECIPE.name) + assert qwen3.default_warmup_resolutions == [(720, 1280)] + assert qwen3.default_warmup_num_frames == [189] + + def test_hf_id_registered(self): + assert "nvidia/Cosmos3-Edge" in PIPELINE_REGISTRY["Cosmos3OmniMoTPipeline"].hf_ids + + def test_none_params_resolve_from_edge_tables(self): + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + resolved = pipeline._resolve_generation_params( + "video", + height=None, + width=None, + num_frames=None, + num_inference_steps=None, + guidance_scale=None, + max_sequence_length=None, + frame_rate=None, + ) + assert resolved == { + "height": 480, + "width": 832, + "num_frames": 121, + "num_inference_steps": 50, + "guidance_scale": 5.0, + "max_sequence_length": 4096, + "frame_rate": 24.0, + } + # Image mode falls back to the video table for fields the image table + # omits, and explicit values always win. + image = pipeline._resolve_generation_params( + "image", height=None, max_sequence_length=None, num_inference_steps=8 + ) + assert image == {"height": 640, "max_sequence_length": 4096, "num_inference_steps": 8} + + def test_sampling_overrides_beat_tables(self): + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + pipeline.sampling = Cosmos3SamplingPolicy(fixed_sigmas=(1.0, 0.5)) + resolved = pipeline._resolve_generation_params( + "video", num_inference_steps=None, guidance_scale=None, height=None + ) + assert resolved["num_inference_steps"] == 2 + assert resolved["guidance_scale"] == 1.0 + assert resolved["height"] == 480 + + def test_warmup_runs_at_family_guidance(self, monkeypatch): + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + captured = {} + + def fake_forward(**kwargs): + captured.update(kwargs) + + monkeypatch.setattr(pipeline, "forward", fake_forward, raising=False) + pipeline._run_warmup(height=480, width=832, num_frames=121, steps=2) + assert captured["guidance_scale"] == 5.0 + assert captured["max_sequence_length"] == 4096 + + def test_declared_temporal_factor_must_match_vae(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + _validate_temporal_compression, + ) + + declared = SimpleNamespace( + temporal_compression_factor=4, temporal_compression_factor_declared=True + ) + _validate_temporal_compression(declared, 4) + with pytest.raises(ValueError, match="temporal_compression_factor"): + _validate_temporal_compression(declared, 8) + # Nano-style configs don't declare it; the VAE value simply wins. + undeclared = SimpleNamespace( + temporal_compression_factor=4, temporal_compression_factor_declared=False + ) + _validate_temporal_compression(undeclared, 8) + + +class TestSamplingRecipeMatrix: + """Family + model_index schedule flag + scheduler recipe must be validated + together: the three facts come from different checkpoint files, and a + mismatch (e.g. a stale conversion missing use_native_flow_schedule) would + otherwise sample the wrong trajectory silently.""" + + BASE = Cosmos3SamplingPolicy() + DISTILLED = Cosmos3SamplingPolicy(fixed_sigmas=(1.0, 0.5)) + + @pytest.mark.parametrize( + "family,native,sampling,error", + [ + (QWEN3_RECIPE.name, False, BASE, None), + (QWEN3_RECIPE.name, False, DISTILLED, None), + (NEMOTRON_DENSE_RECIPE.name, True, BASE, None), + (NEMOTRON_DENSE_RECIPE.name, False, BASE, "use_native_flow_schedule"), + (QWEN3_RECIPE.name, True, BASE, "use_native_flow_schedule"), + (NEMOTRON_DENSE_RECIPE.name, True, DISTILLED, "istilled"), + (NEMOTRON_DENSE_RECIPE.name, False, DISTILLED, "istilled"), + ], + ) + def test_startup_matrix(self, family, native, sampling, error): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + _validate_sampling_recipe, + ) + + if error is None: + _validate_sampling_recipe(family, native, sampling) + else: + with pytest.raises(ValueError, match=error): + _validate_sampling_recipe(family, native, sampling) + + def test_component_loading_rejects_edge_without_native_flag(self, tmp_path): + """End to end through load_standard_components: a UniPC Edge + checkpoint whose model_index omits the flag must fail at load.""" + import json + + scheduler_dir = tmp_path / "scheduler" + scheduler_dir.mkdir() + (scheduler_dir / "scheduler_config.json").write_text( + json.dumps( + { + "_class_name": "UniPCMultistepScheduler", + "num_train_timesteps": 1000, + "flow_shift": 1.0, + "prediction_type": "flow_prediction", + "use_flow_sigmas": True, + "use_karras_sigmas": True, + "solver_order": 2, + } + ) + ) + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "Cosmos3OmniPipeline"}) + ) + + def fresh_pipeline(): + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + pipeline.default_use_system_prompt = False + # Mirror __init__: the flag starts False and only the checkpoint's + # model_index may turn it on. + pipeline.use_native_flow_schedule = False + return pipeline + + skip = ["text_tokenizer", "tokenizer", "vae", "sound_tokenizer"] + with pytest.raises(ValueError, match="use_native_flow_schedule"): + fresh_pipeline().load_standard_components(str(tmp_path), torch.device("cpu"), skip) + + (tmp_path / "model_index.json").write_text( + json.dumps({"_class_name": "Cosmos3OmniPipeline", "use_native_flow_schedule": True}) + ) + pipeline = fresh_pipeline() + pipeline.load_standard_components(str(tmp_path), torch.device("cpu"), skip) + assert pipeline.sampling.native_flow_schedule is True + scheduler = pipeline._scheduler_for(3.0) + assert scheduler.config.use_karras_sigmas is False + assert float(scheduler.config.flow_shift) == 3.0 + + +class TestSchedulerCacheBounds: + """``flow_shift`` is a caller-supplied float, so the cache must not track it. + + Only shifts this checkpoint can resolve to on its own are memoized; a one-off + value gets a scheduler that is discarded with the request, so a client cannot + grow the cache for the worker's lifetime. + """ + + def _pipeline( + self, family: str = "nemotron_dense", checkpoint_shift: float = 3.0 + ) -> "pipeline_module.Cosmos3OmniMoTPipeline": + pipeline = object.__new__(pipeline_module.Cosmos3OmniMoTPipeline) + pipeline.family = family + pipeline.sampling = SimpleNamespace( + checkpoint_flow_shift=checkpoint_shift, + set_flow_shift=lambda base, shift, **kw: SimpleNamespace(shift=shift, **kw), + ) + pipeline._base_scheduler = SimpleNamespace() + pipeline._base_audio_scheduler = SimpleNamespace() + pipeline._scheduler_cache = {} + pipeline._cacheable_flow_shifts_cache = None + return pipeline + + def test_every_cacheable_shift_source_contributes(self, monkeypatch) -> None: + """Each source must be read independently. + + Distinct values throughout, so dropping any one source -- or hard-coding + today's Edge values, where checkpoint and both mode tables all say 3.0 -- + fails instead of passing by coincidence. + """ + monkeypatch.setitem( + pipeline_module.COSMOS3_GENERATION_DEFAULTS, + ("nemotron_dense", "video"), + {"flow_shift": 4.5}, + ) + monkeypatch.setitem( + pipeline_module.COSMOS3_GENERATION_DEFAULTS, + ("nemotron_dense", "image"), + {"flow_shift": 6.25}, + ) + pipeline = self._pipeline(checkpoint_shift=2.75) + assert pipeline._cacheable_flow_shifts() == frozenset( + {2.75, 4.5, 6.25, pipeline_module.COSMOS3_V2V_FLOW_SHIFT} + ) + + def test_cacheable_set_for_the_shipped_edge_tables(self) -> None: + pipeline = self._pipeline() + # Edge declares 3.0 in both mode tables; V2V contributes its stronger shift. + assert pipeline._cacheable_flow_shifts() == frozenset( + {3.0, pipeline_module.COSMOS3_V2V_FLOW_SHIFT} + ) + + def test_arbitrary_shifts_never_enter_the_cache(self) -> None: + pipeline = self._pipeline() + for i in range(200): + pipeline._scheduler_for(7.0 + i * 1e-3) + assert pipeline._scheduler_cache == {} + + def test_declared_shifts_are_memoized_and_shared(self) -> None: + pipeline = self._pipeline() + first = pipeline._scheduler_for(3.0) + assert pipeline._scheduler_for(3.0) is first, "a declared shift must be reused" + assert len(pipeline._scheduler_cache) == 1 + # Streams key separately: they share knobs but must not share the object. + assert pipeline._scheduler_for(3.0, stream="audio") is not first + + def test_release_drops_retained_solver_state(self) -> None: + pipeline = self._pipeline() + scheduler = SimpleNamespace( + config=SimpleNamespace(solver_order=2), + model_outputs=[torch.zeros(4), torch.zeros(4)], + timestep_list=[torch.zeros(1), torch.zeros(1)], + ) + pipeline._scheduler_cache = {("k",): scheduler} + pipeline._release_scheduler_solver_state() + assert scheduler.model_outputs == [None, None] + assert scheduler.timestep_list == [None, None] + + def test_release_drops_state_from_uncached_live_schedulers(self) -> None: + """The caller-override case: a one-off scheduler never enters the cache. + + Walking only ``_scheduler_cache`` leaves its latent-sized ``model_outputs`` + pinned until the next request replaces the attribute. + """ + + def _one_off(): + return SimpleNamespace( + config=SimpleNamespace(solver_order=2), + model_outputs=[torch.zeros(4), torch.zeros(4)], + timestep_list=[torch.zeros(1), torch.zeros(1)], + ) + + pipeline = self._pipeline() + pipeline._scheduler_for(7.5) # non-cacheable shift: memoizes nothing + assert pipeline._scheduler_cache == {} + pipeline.scheduler = _one_off() + pipeline.audio_scheduler = _one_off() + + pipeline._release_scheduler_solver_state() + + for label, scheduler in ( + ("video", pipeline.scheduler), + ("audio", pipeline.audio_scheduler), + ): + assert scheduler.model_outputs == [None, None], label + assert scheduler.timestep_list == [None, None], label + + def test_release_tolerates_absent_schedulers(self) -> None: + """Edge has no audio tower, so ``audio_scheduler`` may never be assigned.""" + pipeline = self._pipeline() + pipeline._release_scheduler_solver_state() + + +class TestEnvelopeAdvisory: + def _warnings(self, monkeypatch): + records = [] + monkeypatch.setattr(pipeline_module.logger, "warning", records.append) + return records + + def _advise(self, pipeline, **overrides): + kwargs = dict( + is_t2i=False, + height=480, + width=832, + num_frames=121, + frame_rate=24.0, + max_sequence_length=4096, + ) + kwargs.update(overrides) + pipeline._log_envelope_advisory(**kwargs) + + def test_in_envelope_is_silent(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(NEMOTRON_DENSE_RECIPE.name)) + assert records == [] + + def test_out_of_envelope_logs_once(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(NEMOTRON_DENSE_RECIPE.name), num_frames=25) + assert len(records) == 1 + assert "num_frames=25" in records[0] + + def test_only_rank_zero_advises(self, monkeypatch): + """Every rank of a TP/Ulysses worker runs this; only one should speak. + + ``rank`` is a read-only property derived from ``torch.distributed``, so a + non-zero rank is simulated by shadowing it on the class. + """ + records = self._warnings(monkeypatch) + pipeline = _bare_pipeline(NEMOTRON_DENSE_RECIPE.name) + monkeypatch.setattr(type(pipeline), "rank", property(lambda self: 1)) + self._advise(pipeline, num_frames=25) # out of envelope, so rank 0 would warn + assert records == [] + + def test_family_without_envelope_never_logs(self, monkeypatch): + records = self._warnings(monkeypatch) + self._advise(_bare_pipeline(QWEN3_RECIPE.name), num_frames=25, height=13, width=17) + assert records == [] + + +class TestDiffusersParity: + """Per-step velocity parity against diffusers main (first release with the + Edge classes). Runs in a subprocess because diffusers main cannot be + imported next to the pinned diffusers; gated on DIFFUSERS_MAIN_PATH.""" + + def test_per_step_velocity_parity(self): + import re + import subprocess + import sys + + diffusers_main = os.environ.get("DIFFUSERS_MAIN_PATH") + if not diffusers_main: + pytest.skip("Set DIFFUSERS_MAIN_PATH to a diffusers checkout with Edge support") + checkpoint = _require_edge_checkpoint() + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + script = Path(__file__).parent / "cosmos3_edge_diffusers_parity.py" + result = subprocess.run( + [sys.executable, str(script), checkpoint], + env={**os.environ, "DIFFUSERS_MAIN_PATH": diffusers_main}, + capture_output=True, + text=True, + timeout=900, + ) + assert result.returncode == 0, result.stderr[-2000:] + rels = [float(m) for m in re.findall(r"rel=([0-9.]+)", result.stdout)] + assert len(rels) == 2, result.stdout + # bf16 accumulation band across 28 layers with differing attention + # backends; observed 0.007 / 0.016 on B200. + assert all(rel < 0.05 for rel in rels), result.stdout + + +# Recorded from diffusers main 2919c5096 (`Cosmos3OmniPipeline.tokenize_prompt` +# on the real Edge checkpoint): prompt "A red cube sits on a wooden table.", +# num_frames=93, height=480, width=832, fps=24.0, system prompt off +# (checkpoint default), duration + resolution templates on. Cond branch only: +# the negative branch intentionally mirrors cosmos-framework's keep-metadata +# templates, not diffusers' inverse templates. +DIFFUSERS_COND_TOKEN_GOLDEN = [ + 1010, + 10, + 25708, + 1010, + 11, + 1010, + 10, + 3263, + 1010, + 1065, + 4804, + 50061, + 53048, + 1408, + 1261, + 32656, + 4234, + 1046, + 1531, + 7476, + 1395, + 1032, + 1051, + 1046, + 1057, + 12900, + 2730, + 1321, + 1395, + 1307, + 1032, + 1050, + 1052, + 1439, + 8148, + 1046, + 2409, + 7476, + 1395, + 1307, + 1032, + 1052, + 1056, + 1048, + 1120, + 1056, + 1051, + 1050, + 9617, + 1046, + 11, + 1010, + 10, + 1503, + 19464, + 1010, + 12, + 1010, + 11, + 20, +] + +# Uncond (CFG) branch pin: TRT-LLM deliberately mirrors cosmos-framework's +# keep-metadata negative-prompt semantics (same duration/resolution templates +# as the positive branch), which diverges from diffusers' inverse templates — +# so this is a self-golden recorded from this code path on the real Edge +# tokenizer (empty negative prompt, num_frames=93, height=480, width=832, +# fps=24.0, system prompt off). Not yet cross-checked against +# cosmos-framework's own tokenization. +UNCOND_TOKEN_GOLDEN = [ + 1010, + 10, + 25708, + 1010, + 11, + 1010, + 10, + 3263, + 1010, + 1784, + 7476, + 1395, + 1032, + 1051, + 1046, + 1057, + 12900, + 2730, + 1321, + 1395, + 1307, + 1032, + 1050, + 1052, + 1439, + 8148, + 1046, + 2409, + 7476, + 1395, + 1307, + 1032, + 1052, + 1056, + 1048, + 1120, + 1056, + 1051, + 1050, + 9617, + 1046, + 11, + 1010, + 10, + 1503, + 19464, + 1010, + 12, + 1010, + 11, + 20, +] + + +class TestEdgeCheckpoint: + """Gated on the real Cosmos3-Edge checkpoint.""" + + def test_tokenizer_specials_and_chat_template(self): + from transformers import AutoTokenizer + + checkpoint = _require_edge_checkpoint() + tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="text_tokenizer") + assert tokenizer.eos_token_id == 11 + assert tokenizer.pad_token_id == 11 + assert tokenizer.convert_tokens_to_ids("<|vision_start|>") == 20 + ids = tokenizer.apply_chat_template( + [{"role": "user", "content": "hello"}], + tokenize=True, + add_generation_prompt=True, + return_dict=False, + ) + assert isinstance(ids, list) and len(ids) > 0 + + def test_tokenization_matches_diffusers_golden(self): + from transformers import AutoTokenizer + + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DURATION_TEMPLATE, + ) + + checkpoint = _require_edge_checkpoint() + + class _CpuPipeline(Cosmos3OmniMoTPipeline): + device = property(lambda self: torch.device("cpu")) + + pipeline = object.__new__(_CpuPipeline) + pipeline.tokenizer = AutoTokenizer.from_pretrained(checkpoint, subfolder="text_tokenizer") + text = pipeline._format_prompt_with_metadata( + "A red cube sits on a wooden table.", + height=480, + width=832, + num_frames=93, + frame_rate=24.0, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + ids, mask = pipeline._tokenize_prompt(text, 128, use_system_prompt=False) + golden_length = len(DIFFUSERS_COND_TOKEN_GOLDEN) + assert ids[0, :golden_length].tolist() == DIFFUSERS_COND_TOKEN_GOLDEN + assert int(mask.sum()) == golden_length + + uncond_text = pipeline._format_prompt_with_metadata( + "", + height=480, + width=832, + num_frames=93, + frame_rate=24.0, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + uncond_ids, uncond_mask = pipeline._tokenize_prompt( + uncond_text, 128, use_system_prompt=False + ) + uncond_length = len(UNCOND_TOKEN_GOLDEN) + assert uncond_ids[0, :uncond_length].tolist() == UNCOND_TOKEN_GOLDEN + assert int(uncond_mask.sum()) == uncond_length + + def test_model_index_detection(self): + from tensorrt_llm._torch.visual_gen.pipeline_registry import AutoPipeline + + checkpoint = _require_edge_checkpoint() + assert AutoPipeline._detect_from_checkpoint(checkpoint) == "Cosmos3OmniMoTPipeline" + + @pytest.fixture(scope="class") + def edge_pipeline(self): + checkpoint = _require_edge_checkpoint() + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + args = VisualGenArgs( + model=checkpoint, + torch_compile_config=TorchCompileConfig(enable=False), + ) + pipeline = PipelineLoader(args).load( + skip_warmup=True, + skip_components=[PipelineComponent.SOUND_TOKENIZER], + ) + yield pipeline + del pipeline + gc.collect() + torch.cuda.empty_cache() + + def test_recipe_and_scheduler_wiring(self, edge_pipeline): + assert edge_pipeline.family == NEMOTRON_DENSE_RECIPE.name + assert edge_pipeline.transformer.recipe is NEMOTRON_DENSE_RECIPE + assert edge_pipeline.use_native_flow_schedule is True + assert edge_pipeline.sampling.native_flow_schedule is True + # Both Edge tables declare shift 3.0 on the native flow schedule, so + # the two modes resolve to the same (shift, karras) pair and share one + # cached scheduler rather than building two identical ones. + schedulers = { + mode: edge_pipeline._scheduler_for(edge_pipeline._mode_params(mode)["flow_shift"]) + for mode in ("video", "image") + } + for scheduler in schedulers.values(): + assert float(scheduler.config.flow_shift) == 3.0 + assert scheduler.config.use_karras_sigmas is False + assert schedulers["video"] is schedulers["image"] + + def test_load_weights_and_forward(self, edge_pipeline): + transformer = edge_pipeline.transformer + channels = transformer.latent_channel_size + latents = torch.randn(1, channels, 1, 16, 16, dtype=torch.bfloat16, device=DEVICE) + timestep = torch.tensor([999.0], device=DEVICE) + text_ids = torch.randint(0, 1000, (1, 32), device=DEVICE) + text_mask = torch.ones(1, 32, dtype=torch.long, device=DEVICE) + + transformer.reset_cache() + with torch.inference_mode(): + out = transformer( + hidden_states=latents, + timestep=timestep / 1000.0, + raw_timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + video_shape=(1, 16, 16), + fps=24.0, + ) + assert out.video.shape == latents.shape + assert torch.isfinite(out.video.float()).all() + + def test_direct_forward_resolves_edge_defaults(self, edge_pipeline, monkeypatch): + """A direct forward() call with unset numerics must reach the + generation path with Edge-table values (denoise and decode stubbed).""" + captured = {} + + def fake_denoise(**kwargs): + captured["latents"] = kwargs["latents"] + captured["guidance_scale"] = kwargs["guidance_scale"] + captured["scheduler"] = kwargs["scheduler"] + return kwargs["latents"] + + def fake_decode(latents, decode_fn, **kwargs): + captured["decoded"] = True + return torch.zeros(1, 2, 8, 8, 3, dtype=torch.uint8) + + monkeypatch.setattr(edge_pipeline, "denoise", fake_denoise, raising=False) + monkeypatch.setattr(edge_pipeline, "decode_latents", fake_decode, raising=False) + + edge_pipeline.forward(prompt="warmup-shaped request", seed=0, use_guardrails=False) + + assert captured["guidance_scale"] == 5.0 + # Edge video defaults: 121 frames -> 31 latent frames, 480x832 -> 30x52. + assert tuple(captured["latents"].shape) == (1, 48, 31, 30, 52) + assert captured["scheduler"] is edge_pipeline._scheduler_for(3.0) + assert captured["scheduler"].num_inference_steps == 50 + assert captured["decoded"] is True + + def test_t2v_sanity_generation(self, edge_pipeline): + out = edge_pipeline.forward( + prompt="A red ball rolls across a wooden floor.", + seed=0, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=5.0, + use_guardrails=False, + ) + video = out.video + assert video is not None + assert tuple(video.shape) == (1, 9, 192, 320, 3) + assert video.float().std() > 1.0, "generated video is (near-)constant" + + def test_i2v_sanity_generation(self, edge_pipeline): + from PIL import Image + + image = Image.new("RGB", (320, 192)) + for x in range(320): + for y in range(0, 192, 4): + image.putpixel((x, y), (x % 256, (2 * y) % 256, 120)) + out = edge_pipeline.forward( + prompt="The scene slowly brightens.", + seed=0, + image=image, + height=192, + width=320, + num_frames=9, + num_inference_steps=2, + guidance_scale=5.0, + use_guardrails=False, + ) + video = out.video + assert video is not None + assert tuple(video.shape) == (1, 9, 192, 320, 3) + assert video.float().std() > 1.0 + + def test_t2i_sanity_generation(self, edge_pipeline): + out = edge_pipeline.forward( + prompt="A ceramic teapot on a table.", + seed=0, + height=256, + width=256, + num_inference_steps=2, + guidance_scale=4.0, + use_guardrails=False, + output_type="image", + ) + image = out.image + assert image is not None + assert tuple(image.shape) == (1, 256, 256, 3) + assert image.float().std() > 1.0 diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py new file mode 100644 index 000000000000..301c5d435750 --- /dev/null +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_example_prompts.py @@ -0,0 +1,220 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +"""Prompt resolution in the Cosmos3 example CLI. + +``--prompt``/``--negative_prompt`` take literal text or a file path; the +``*_file`` variants take a path only. Passing a checkpoint's structured caption +file to ``--prompt`` used to silently generate from the path string itself. +""" + +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +_PROJECT_ROOT = Path(__file__).resolve().parents[4] +_EXAMPLE_DIR = _PROJECT_ROOT / "examples" / "visual_gen" / "models" / "cosmos3" + + +def _load_example_module(): + spec = importlib.util.spec_from_file_location( + "cosmos3_example_cli", _EXAMPLE_DIR / "cosmos3.py" + ) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +cosmos3 = _load_example_module() + + +def _resolve(prompt=None, prompt_file=None, **kwargs): + kwargs.setdefault("image_path", None) + kwargs.setdefault("enable_audio", False) + kwargs.setdefault("output_type", "video") + return cosmos3.resolve_prompt_and_options(prompt=prompt, prompt_file=prompt_file, **kwargs) + + +class TestPromptAcceptsTextOrPath: + def test_literal_text_is_used_verbatim(self): + prompt, _, _, _ = _resolve(prompt="A cute puppy playing with a ball") + assert prompt == "A cute puppy playing with a ball" + + def test_path_shaped_text_stays_literal_when_no_such_file(self): + """A non-existent path is a prompt, not an error -- prompts may contain slashes.""" + prompt, _, _, _ = _resolve(prompt="assets/example_i2v_prompt.json") + assert prompt == "assets/example_i2v_prompt.json" + + def test_structured_caption_file_becomes_the_prompt(self, tmp_path): + """A checkpoint's assets/*_prompt.json has no 'prompt' key; the object *is* the caption.""" + caption = {"subjects": [{"description": "A car"}], "lighting": "overcast"} + path = tmp_path / "example_i2v_prompt.json" + path.write_text(json.dumps(caption), encoding="utf-8") + + prompt, _, _, _ = _resolve(prompt=str(path)) + + assert json.loads(prompt) == caption + + def test_omni_prompt_file_supplies_options(self, tmp_path): + path = tmp_path / "i2v.json" + path.write_text( + json.dumps( + { + "model_mode": "text2image", + "prompt": "a lighthouse", + "vision_path": "frame.jpg", + "enable_audio": True, + } + ), + encoding="utf-8", + ) + + prompt, image, enable_audio, output_type = _resolve(prompt=str(path)) + + assert prompt == "a lighthouse" + assert image == "frame.jpg" + assert enable_audio is True + assert output_type == "image" + + def test_plain_text_file(self, tmp_path): + path = tmp_path / "prompt.txt" + path.write_text(" the camera pans right \n", encoding="utf-8") + + prompt, _, _, _ = _resolve(prompt=str(path)) + + assert prompt == "the camera pans right" + + def test_prompt_file_path_overrides_prompt_file_flag(self, tmp_path): + override = tmp_path / "override.json" + override.write_text(json.dumps({"prompt": "from --prompt"}), encoding="utf-8") + base = tmp_path / "base.json" + base.write_text(json.dumps({"prompt": "from --prompt_file"}), encoding="utf-8") + + prompt, _, _, _ = _resolve(prompt=str(override), prompt_file=str(base)) + + assert prompt == "from --prompt" + + def test_explicit_image_path_wins_over_prompt_file_vision_path(self, tmp_path): + path = tmp_path / "i2v.json" + path.write_text( + json.dumps({"prompt": "a lighthouse", "vision_path": "from_file.jpg"}), + encoding="utf-8", + ) + + _, image, _, _ = _resolve(prompt=str(path), image_path="from_cli.jpg") + + assert image == "from_cli.jpg" + + +class TestPromptFileIsStrict: + def test_missing_file_raises(self): + with pytest.raises(ValueError, match="does not exist"): + cosmos3.load_prompt_file("no/such/prompt.json") + + def test_literal_text_raises(self): + with pytest.raises(ValueError, match="does not exist"): + cosmos3.load_prompt_file("The camera slowly pans right across the scene") + + def test_empty_prompt_field_raises(self, tmp_path): + path = tmp_path / "empty.json" + path.write_text(json.dumps({"prompt": ""}), encoding="utf-8") + with pytest.raises(ValueError, match="non-empty 'prompt' field"): + cosmos3.load_prompt_file(str(path)) + + def test_empty_object_raises(self, tmp_path): + path = tmp_path / "empty.json" + path.write_text("{}", encoding="utf-8") + with pytest.raises(ValueError, match="empty JSON object"): + cosmos3.load_prompt_file(str(path)) + + def test_json_array_raises(self, tmp_path): + path = tmp_path / "list.json" + path.write_text("[1, 2]", encoding="utf-8") + with pytest.raises(ValueError, match="JSON object or text"): + cosmos3.load_prompt_file(str(path)) + + def test_no_prompt_source_raises(self): + with pytest.raises(ValueError, match="Provide --prompt or --prompt_file"): + _resolve(prompt=None, prompt_file=None) + + +class TestNegativePromptResolution: + def _resolve(self, negative_prompt=None, negative_prompt_file=None): + return cosmos3.resolve_negative_prompt( + negative_prompt=negative_prompt, negative_prompt_file=negative_prompt_file + ) + + def test_literal_text_is_used_verbatim(self): + assert self._resolve(negative_prompt="blurry, low quality") == "blurry, low quality" + + def test_empty_string_disables_the_default(self): + assert self._resolve(negative_prompt="") == "" + + def test_path_loads_the_file(self, tmp_path): + path = tmp_path / "negative_prompt.json" + path.write_text(json.dumps({"subjects": ["blurry"]}), encoding="utf-8") + + assert json.loads(self._resolve(negative_prompt=str(path))) == {"subjects": ["blurry"]} + + def test_negative_prompt_overrides_negative_prompt_file(self, tmp_path): + path = tmp_path / "negative_prompt.json" + path.write_text(json.dumps({"subjects": ["from file"]}), encoding="utf-8") + + assert self._resolve(negative_prompt="from flag", negative_prompt_file=str(path)) == ( + "from flag" + ) + + def test_negative_prompt_file_is_used_when_no_flag(self, tmp_path): + path = tmp_path / "negative_prompt.json" + path.write_text(json.dumps({"subjects": ["from file"]}), encoding="utf-8") + + assert json.loads(self._resolve(negative_prompt_file=str(path))) == { + "subjects": ["from file"] + } + + def test_falls_back_to_bundled_default(self): + assert self._resolve() == cosmos3.load_negative_prompt_file( + cosmos3.DEFAULT_NEGATIVE_PROMPT_FILE + ) + + def test_missing_negative_prompt_file_raises(self): + with pytest.raises(ValueError, match="does not exist"): + self._resolve(negative_prompt_file="no/such/negative.json") + + +class TestNegativePromptFile: + def test_structured_object_is_serialized(self, tmp_path): + payload = {"subjects": [{"description": "Blurry"}]} + path = tmp_path / "negative_prompt.json" + path.write_text(json.dumps(payload), encoding="utf-8") + + assert json.loads(cosmos3.load_negative_prompt_file(str(path))) == payload + + def test_plain_text_file(self, tmp_path): + path = tmp_path / "negative.txt" + path.write_text("blurry, low quality\n", encoding="utf-8") + + assert cosmos3.load_negative_prompt_file(str(path)) == "blurry, low quality" + + def test_missing_file_raises(self): + with pytest.raises(ValueError, match="does not exist"): + cosmos3.load_negative_prompt_file("no/such/negative.json") + + +class TestShippedPromptFiles: + """The files this README tells users to pass must actually load.""" + + @pytest.mark.parametrize("name", ["t2v", "t2i", "i2v", "v2v", "t2av"]) + def test_bundled_prompt_files_load(self, name): + data = cosmos3.load_prompt_file(f"prompts/{name}.json") + assert data["prompt"] + + def test_default_prompt_file(self): + assert cosmos3.load_prompt_file(cosmos3.DEFAULT_PROMPT_FILE)["prompt"] + + def test_default_negative_prompt_file(self): + assert cosmos3.load_negative_prompt_file(cosmos3.DEFAULT_NEGATIVE_PROMPT_FILE) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 714c828cf1d5..ff2285b71fc1 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -53,6 +53,7 @@ _load_reference_image, _normalize_condition_video_latent_indexes, ) +from tensorrt_llm._torch.visual_gen.models.cosmos3.transformer_cosmos3 import QWEN3_RECIPE from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -407,10 +408,13 @@ def test_injects_metadata_fields(self, cosmos3_format_pipeline): data = json.loads(result) assert data["prompt"] == "A foundry pour" assert data["subjects"] == [] - assert data["duration"] == "7.9s" - assert data["fps"] == 24 - assert data["resolution"] == {"W": 1280, "H": 720} - assert data["aspect_ratio"] == "9,16" + # Reference semantics: integer-truncated seconds, float fps, H before W, + # and the aspect-ratio *bucket* rather than the exact reduced ratio. + assert data["duration"] == "7s" + assert data["fps"] == 24.0 + assert data["resolution"] == {"H": 720, "W": 1280} + assert data["aspect_ratio"] == "16,9" + assert '"resolution": {"H": 720, "W": 1280}' in result def test_overwrites_existing_metadata_fields(self, cosmos3_format_pipeline): prompt = json.dumps( @@ -423,10 +427,10 @@ def test_overwrites_existing_metadata_fields(self, cosmos3_format_pipeline): } ) data = json.loads(_format_prompt_with_metadata(cosmos3_format_pipeline, prompt)) - assert data["duration"] == "7.9s" - assert data["fps"] == 24 - assert data["resolution"] == {"W": 1280, "H": 720} - assert data["aspect_ratio"] == "9,16" + assert data["duration"] == "7s" + assert data["fps"] == 24.0 + assert data["resolution"] == {"H": 720, "W": 1280} + assert data["aspect_ratio"] == "16,9" def test_single_frame_skips_duration_by_default(self, cosmos3_format_pipeline): prompt = json.dumps({"prompt": "still life"}) @@ -451,7 +455,54 @@ def test_single_frame_duration_when_forced(self, cosmos3_format_pipeline): force_duration_template=True, ) ) - assert data["duration"] == "0.0s" + assert data["duration"] == "0s" + + def test_still_drops_stale_duration_and_fps(self, cosmos3_format_pipeline): + """A caller's JSON may already declare a duration; a still must not keep it.""" + prompt = json.dumps({"prompt": "still life", "duration": "7s", "fps": 24.0}) + data = json.loads( + _format_prompt_with_metadata( + cosmos3_format_pipeline, + prompt, + num_frames=1, + resolution_template=COSMOS3_IMAGE_RESOLUTION_TEMPLATE, + ) + ) + assert "duration" not in data + assert "fps" not in data + + def test_non_ascii_is_escaped(self, cosmos3_format_pipeline): + """The reference serializes with the json default (``ensure_ascii=True``).""" + result = _format_prompt_with_metadata( + cosmos3_format_pipeline, json.dumps({"prompt": "moiré — artifacts"}) + ) + assert "\\u00e9" in result and "\\u2014" in result + assert "é" not in result and "—" not in result + + @pytest.mark.parametrize( + "height,width,bucket", + [ + (480, 832, "16,9"), + (832, 480, "9,16"), + (640, 640, "1,1"), + (544, 736, "4,3"), + (736, 544, "3,4"), + (720, 1280, "16,9"), + (1024, 1024, "1,1"), + ], + ) + def test_aspect_ratio_maps_to_reference_bucket( + self, cosmos3_format_pipeline, height, width, bucket + ): + data = json.loads( + _format_prompt_with_metadata( + cosmos3_format_pipeline, + json.dumps({"prompt": "test"}), + height=height, + width=width, + ) + ) + assert data["aspect_ratio"] == bucket def test_non_integer_fps_preserved(self, cosmos3_format_pipeline): prompt = json.dumps({"prompt": "test"}) @@ -475,6 +526,132 @@ def test_resolution_only_when_duration_template_disabled(self, cosmos3_format_pi assert data["resolution"] == {"W": 1280, "H": 720} +class TestNegativePromptMetadata: + """The negative prompt takes the sentence-append path even when it is JSON. + + cosmos-framework applies its plain-text formatter to the negative prompt + unconditionally and reserves JSON field injection for the positive prompt, so + a JSON negative prompt must keep its serialized form and gain the sentences + after it -- not grow ``duration``/``fps``/``resolution`` keys inside it. + """ + + NEGATIVE = json.dumps({"subjects": [{"description": "Blurry, poorly defined subjects."}]}) + + def _negative(self, pipeline, **kwargs): + """Format a negative prompt the way ``forward`` does.""" + return pipeline._apply_metadata_templates( + self.NEGATIVE, + height=HEIGHT, + width=WIDTH, + num_frames=189, + frame_rate=FRAME_RATE, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + **kwargs, + ) + + def test_json_negative_keeps_object_and_appends_sentences(self, cosmos3_format_pipeline): + result = self._negative(cosmos3_format_pipeline) + assert result.startswith(self.NEGATIVE.rstrip(".")) + assert result.endswith("This video is of 720x1280 resolution.") + assert "7.9 seconds long" in result + + def test_json_negative_gains_no_injected_fields(self, cosmos3_format_pipeline): + result = self._negative(cosmos3_format_pipeline) + # The metadata must live outside the object, so the result stops being + # parseable JSON and the object itself is untouched. + with pytest.raises(json.JSONDecodeError): + json.loads(result) + for field in ("duration", "fps", "resolution", "aspect_ratio"): + assert f'"{field}"' not in result + + def test_matches_reference_sentence_append(self, cosmos3_format_pipeline): + """Byte-for-byte against cosmos-framework's ``_format_prompt_with_template``.""" + expected = ( + self.NEGATIVE.strip().rstrip(".") + + ". " + + COSMOS3_DURATION_TEMPLATE.format(duration=189 / FRAME_RATE, fps=FRAME_RATE) + ) + expected = ( + expected.strip().rstrip(".") + + ". " + + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE.format(height=HEIGHT, width=WIDTH) + ) + assert self._negative(cosmos3_format_pipeline) == expected.lstrip(".").strip() + + def test_positive_json_still_injects_fields(self, cosmos3_format_pipeline): + """The positive branch keeps field injection -- the two paths differ by design.""" + data = json.loads(_format_prompt_with_metadata(cosmos3_format_pipeline, self.NEGATIVE)) + assert data["resolution"] == {"W": 1280, "H": 720} + assert self._negative(cosmos3_format_pipeline) != _format_prompt_with_metadata( + cosmos3_format_pipeline, self.NEGATIVE + ) + + +class TestDefaultNegativePrompt: + """Video modes inherit the reference's default negative prompt; image modes do not. + + cosmos-framework wires ``negative_prompt_file: neg_prompts.json`` into + ``defaults/{text2video,image2video,video2video,audio_image2video}`` and leaves it + unset for ``text2image``/``image2image``. + """ + + def test_video_default_serializes_like_the_reference(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.negative_prompt import ( + COSMOS3_VIDEO_NEGATIVE_PROMPT, + ) + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + default_video_negative_prompt, + ) + + # The reference loads it as json.dumps(json.loads(...)); ensure_ascii and key + # order both matter, so round-tripping must be a no-op. + text = default_video_negative_prompt() + assert text == json.dumps(COSMOS3_VIDEO_NEGATIVE_PROMPT) + assert json.dumps(json.loads(text)) == text + assert "\\u2014" in text, "non-ASCII must be escaped, as the reference emits it" + + def test_video_default_is_a_json_object_with_expected_shape(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + default_video_negative_prompt, + ) + + data = json.loads(default_video_negative_prompt()) + assert isinstance(data, dict) + for field in ("subjects", "background_setting", "cinematography"): + assert field in data + + def test_image_default_is_empty(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_DEFAULT_NEGATIVE_PROMPT, + ) + + assert COSMOS3_DEFAULT_NEGATIVE_PROMPT == "" + + def test_default_is_cached(self): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + default_video_negative_prompt, + ) + + assert default_video_negative_prompt() is default_video_negative_prompt() + + @pytest.mark.parametrize( + "output_type,expects_video_default", + [("video", True), ("image", False)], + ) + def test_resolution_is_keyed_on_output_kind(self, output_type, expects_video_default): + from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + default_negative_prompt, + default_video_negative_prompt, + ) + + resolved = default_negative_prompt(output_type) + if expects_video_default: + assert resolved == default_video_negative_prompt() + else: + assert resolved == "" + + @pytest.fixture(scope="class") def cosmos3_pipeline(): checkpoint = _require_checkpoint() @@ -666,6 +843,9 @@ def test_v2v_flow_shift_override_request_path(self): pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) pipeline.audio_gen = False + # Bypassing __init__ means the generation-default family has to be set + # here; forward() reads its per-mode table through it. + pipeline.family = QWEN3_RECIPE.name calls = [] token_calls = [] @@ -732,6 +912,9 @@ def test_v2v_rebuilds_the_audio_scheduler_too(self): pipeline = Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) pipeline.transformer = SimpleNamespace(device=torch.device("cpu")) pipeline.audio_gen = True + # Bypassing __init__ means the generation-default family has to be set + # here; forward() reads its per-mode table through it. + pipeline.family = QWEN3_RECIPE.name rebuilt = [] class StopAfterTokenize(Exception): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 9293c371f459..6041300cdf6e 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -477,6 +477,10 @@ def test_load_fp8_quantization(self, quant_algo: str): # --- CPU-only coverage: checkpoint config schema compatibility --- +# Distinguishes "attribute absent" from "attribute present and None". +_OMITTED = object() + + class TestConfigCompatDefaults: """Newer diffusers conversions omit fields older ones carried explicitly.""" @@ -506,6 +510,34 @@ def test_idempotent(self): apply_pretrained_config_compat_defaults(config) assert vars(config) == snapshot + @pytest.mark.parametrize("rope_scaling", [_OMITTED, None, {}], ids=["omitted", "none", "empty"]) + def test_rope_type_tolerates_missing_rope_scaling(self, rope_scaling: object) -> None: + """``rope_axes_dim`` alone is a supported shape, so reading ``rope_type`` + must not fail before ``resolve_rope_axes_dim`` gets to honour it. + + ``omitted`` leaves the attribute off entirely, which is the case the + ``getattr(..., None)`` guard exists for; ``none``/``empty`` only reach the + ``or {}`` half. + """ + from tensorrt_llm._torch.visual_gen.models.cosmos3 import transformer_cosmos3 as tf + + config = SimpleNamespace( + hidden_size=64, + head_dim=16, + rope_axes_dim=[4, 2, 2], + rope_theta=10000.0, + max_position_embeddings=128, + ) + if rope_scaling is not _OMITTED: + config.rope_scaling = rope_scaling + assert hasattr(config, "rope_scaling") is (rope_scaling is not _OMITTED) + apply_pretrained_config_compat_defaults(config) + # Construct for real: the point is that __init__ reaches the resolver + # instead of raising AttributeError on the missing block. + embedding = tf.Qwen3VLTextRotaryEmbedding(SimpleNamespace(pretrained_config=config)) + assert embedding.rope_type == "default" + assert embedding.mrope_section == [4, 2, 2] + class TestI2V4StepConfigShape: """The Image2Video-4Step conversion drops the audio/action towers @@ -557,7 +589,7 @@ def test_constructs_without_audio_or_action_towers(self): model = Cosmos3VFMTransformer(model_config) assert model.audio_gen is False - assert model.action_gen is False + assert model.has_action_weights is False assert not hasattr(model, "audio2llm") assert not hasattr(model, "audio_modality_embed") assert model.base_fps == 16