diff --git a/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml b/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml index b67ab39e235b..fd08a83432a7 100644 --- a/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml +++ b/examples/visual_gen/configs/cosmos3-nano-1gpu.yaml @@ -13,20 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. -# 1-GPU Cosmos3 (Nano / Super) with FP8 dynamic quantization. +# 1-GPU Cosmos3 (Nano / Super). # Model: nvidia/Cosmos3-Nano or nvidia/Cosmos3-Super # Shared by offline examples (--visual_gen_args) and trtllm-serve. # # Cosmos3 constraints: VANILLA attention only; -# no Attention2D / Ring. Use CFG + Ulysses for multi-GPU (see cosmos3-super-4gpu.yaml). -quant_config: - quant_algo: FP8 - dynamic: true - ignore: ["language_model.*", "vae2llm", "llm2vae", "time_embedder.*"] +# Use CFG + Ulysses for multi-GPU (see cosmos3-super-4gpu.yaml). attention_config: backend: VANILLA parallel_config: cfg_size: 1 ulysses_size: 1 -cuda_graph_config: - enable: false diff --git a/examples/visual_gen/configs/cosmos3-super-4gpu.yaml b/examples/visual_gen/configs/cosmos3-super-4gpu.yaml index 34ddec38ceea..0aa77d45c001 100644 --- a/examples/visual_gen/configs/cosmos3-super-4gpu.yaml +++ b/examples/visual_gen/configs/cosmos3-super-4gpu.yaml @@ -13,16 +13,12 @@ # See the License for the specific language governing permissions and # limitations under the License. -# 4-GPU Cosmos3-Super with FP8 dynamic quantization (CFG + Ulysses + parallel VAE). +# 4-GPU Cosmos3-Super with (CFG + Ulysses + parallel VAE). # Launch with 4 processes, e.g. torchrun --nproc_per_node=4 ... # Model: nvidia/Cosmos3-Super # Shared by offline examples (--visual_gen_args) and trtllm-serve. # # GPU layout: cfg_size=2 (positive | negative) x ulysses_size=2 (sequence split). -quant_config: - quant_algo: FP8 - dynamic: true - ignore: ["language_model.*", "vae2llm", "llm2vae", "time_embedder.*"] attention_config: backend: VANILLA parallel_config: diff --git a/examples/visual_gen/models/cosmos3/README.md b/examples/visual_gen/models/cosmos3/README.md new file mode 100644 index 000000000000..69be21fe4880 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/README.md @@ -0,0 +1,77 @@ +# Cosmos3 Text(+Image)-to-Video(+Audio) generation + +Cosmos3 supports four 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. +- **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 + +Pass the Hub ID or local path via `--model`: + +- [`nvidia/Cosmos3-Nano`](https://huggingface.co/nvidia/Cosmos3-Nano) +- [`nvidia/Cosmos3-Super`](https://huggingface.co/nvidia/Cosmos3-Super) + +## Guardrails + +Guardrails are enabled by default (required by the [NVIDIA Open Model License Agreement](https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license)). Install and authenticate as follows: + +```bash +pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python +``` + +Accept the terms for the guardrail checkpoint at https://huggingface.co/nvidia/Cosmos-1.0-Guardrail and set a valid `HF_TOKEN` (the checkpoint is downloaded automatically on first run). + +To run without guardrails (you are responsible for safe deployment): + +```bash +export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 +``` + +## Deployment configs + +See `examples/visual_gen/configs/`: + +- `cosmos3-nano-1gpu.yaml` — 1 GPU +- `cosmos3-super-4gpu.yaml` — 4 GPU, CFG + Ulysses + parallel VAE + +Example prompts live under `prompts/` (mirroring `cosmos3-internal/inputs/omni`). + +## Usage + +```bash +# T2V: text-to-video +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# I2V/TI2V: image-conditioned video (vision_path is read from the prompt file; +# local path, file://, http(s):// URL, or data: URI are all accepted) +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/i2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# I2V with an explicit conditioning image (overrides the prompt file) +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/i2v.json \ + --image_path https://example.com/frame.jpg \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# T2AV: text-to-video with synchronized audio +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2av.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + +# T2I: text-to-image +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2i.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --output_path output.png + +# Inline prompt (--prompt or a JSON file path) +python cosmos3.py --model nvidia/Cosmos3-Nano \ + --prompt "A cute puppy playing with a ball in a park" \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml +``` diff --git a/examples/visual_gen/models/cosmos3/cosmos3.py b/examples/visual_gen/models/cosmos3/cosmos3.py new file mode 100644 index 000000000000..de9e9e5010ad --- /dev/null +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -0,0 +1,207 @@ +#!/usr/bin/env python3 +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +import argparse +import json +import os +from pathlib import Path +from typing import Any, Dict, Optional + +from tensorrt_llm import VisualGen, VisualGenArgs + +_SCRIPT_DIR = Path(__file__).resolve().parent + + +def _resolve_path(path: str) -> str: + candidate = Path(path) + if candidate.is_file(): + return str(candidate.resolve()) + relative_to_script = _SCRIPT_DIR / path + if relative_to_script.is_file(): + return str(relative_to_script.resolve()) + return path + + +def load_prompt_file(path: str) -> Dict[str, Any]: + """Load a Cosmos3 omni prompt JSON (``prompt``, optional ``vision_path``, etc.).""" + resolved = _resolve_path(path) + with open(resolved, encoding="utf-8") as f: + data = json.load(f) + 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.") + return data + + +def resolve_prompt_and_options( + *, + prompt: Optional[str], + prompt_file: Optional[str], + image_path: Optional[str], + enable_audio: bool, + output_type: str, +) -> tuple[str, Optional[str], bool, str]: + """Merge CLI args with optional prompt-file defaults.""" + prompt_data: Dict[str, Any] = {} + if prompt_file is not None: + prompt_data = load_prompt_file(prompt_file) + + resolved_prompt = prompt + if resolved_prompt is None: + resolved_prompt = prompt_data.get("prompt") + if not resolved_prompt: + raise ValueError("Provide --prompt or --prompt_file with a 'prompt' field.") + + resolved_image = image_path + if resolved_image is None: + resolved_image = prompt_data.get("vision_path") or prompt_data.get("image_path") + + resolved_enable_audio = enable_audio or bool(prompt_data.get("enable_audio", False)) + + resolved_output_type = output_type + model_mode = str(prompt_data.get("model_mode", "")).lower() + if model_mode == "text2image" and output_type == "video": + resolved_output_type = "image" + + return resolved_prompt, resolved_image, resolved_enable_audio, resolved_output_type + + +def main(): + parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video(+Audio) example") + parser.add_argument( + "--model", + type=str, + default="nvidia/Cosmos3-Nano", + help="Model path or HuggingFace Hub ID (nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super)", + ) + parser.add_argument( + "--visual_gen_args", + dest="visual_gen_args", + type=str, + default=None, + help="Path to YAML config (same as trtllm-serve --visual_gen_args)", + ) + parser.add_argument( + "--prompt", + type=str, + default=None, + help="Text prompt for generation (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)", + ) + parser.add_argument( + "--negative_prompt", + type=str, + default="cosmos3_negative_prompt.json", + help="Text prompt or path to JSON file for negative prompt", + ) + parser.add_argument( + "--image_path", + type=str, + default=None, + help="Optional conditioning image path or URL for I2V/TI2V", + ) + parser.add_argument( + "--output_path", + type=str, + default="cosmos3_output.mp4", + help="Path to save the output video", + ) + parser.add_argument( + "--disable_duration_template", + action="store_true", + help="Disable duration metadata template (enabled by default, matching cosmos-framework CLI)", + ) + parser.add_argument( + "--disable_resolution_template", + action="store_true", + help="Disable resolution metadata template (enabled by default, matching cosmos-framework CLI)", + ) + parser.add_argument( + "--use_system_prompt", action="store_true", help="Use system prompt in prompt" + ) + parser.add_argument("--enable_audio", action="store_true", help="Enable audio generation") + parser.add_argument( + "--output_type", type=str, default="video", help="Output type (video, image)" + ) + + # Guardrails + parser.add_argument( + "--disable_guardrails", action="store_true", help="NOT RECOMMENDED: Disable guardrails" + ) + args = parser.parse_args() + + prompt, image_path, enable_audio, output_type = resolve_prompt_and_options( + prompt=args.prompt, + prompt_file=args.prompt_file, + image_path=args.image_path, + enable_audio=args.enable_audio, + output_type=args.output_type, + ) + + # Engine config from shared YAML (optional); model-specific defaults apply otherwise. + extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None + visual_gen = VisualGen(model=args.model, args=extra_args) + + # --- Model-specific: T2V / TI2V request construction --- + # Query per-model defaults (resolution, steps, guidance, seed, etc.). + params = visual_gen.default_params + 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 + + if args.disable_duration_template: + params.extra_params["use_duration_template"] = False + if args.disable_resolution_template: + params.extra_params["use_resolution_template"] = False + params.extra_params["use_system_prompt"] = args.use_system_prompt + params.extra_params["enable_audio"] = enable_audio + params.extra_params["use_guardrails"] = not args.disable_guardrails + params.extra_params["output_type"] = output_type + + 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) + + output = visual_gen.generate( + inputs=prompt, + params=params, + ) + + output.save(args.output_path) + print(f"Saved: {args.output_path}") + print(output.metrics) + + +if __name__ == "__main__": + main() diff --git a/examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json b/examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json new file mode 100644 index 000000000000..44dff2693426 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json @@ -0,0 +1,108 @@ +{ + "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 \u2014 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 \u2014 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.", + "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 \u2014 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 \u2014 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\u00e9 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 \u2014 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 \u2014 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 \u2014 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 \u2014 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/examples/visual_gen/models/cosmos3/prompts/i2v.json b/examples/visual_gen/models/cosmos3/prompts/i2v.json new file mode 100644 index 000000000000..27bc79101a1e --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/i2v.json @@ -0,0 +1,5 @@ +{ + "model_mode": "image2video", + "prompt": "The video opens with a view of a testing environment, characterized by a large wooden table at the center. On this table, two robot arms are positioned at opposite ends, with the left arm closer to the camera and the right arm further away. Between the hands lies a dark wooden shelf with a red spherical object on its top rack, likely serving as a platform or obstacle. In the background, various pieces of equipment, including a tripod, a chair, are visible. A person wearing a blue jacket and black pants stands near the center of the room, observing the experiment, with a static hand position throughout. The floor is tiled with a patterned design, and additional items like a small robot figure and some cables can be seen scattered around the space. As the video progresses, the right robotic hand extends outward, moving from its initial position towards the red spherical object on the shelf. The hand then picks up the object and places it on the lowest rack of the shelf, completing a smooth, deliberate manipulation. The left robotic hand remains stationary throughout the sequence. No new objects appear in the video; all existing elements maintain their positions except for the movement of the right robotic hand. The scene concludes with the right robotic hand returning to its initial position, while the left hand continues to rest on the table. The overall environment remains unchanged, with the focus remaining on the interaction between the robotic hands and the wooden block, highlighting precise control during the demonstration.", + "vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_153.jpg" +} diff --git a/examples/visual_gen/models/cosmos3/prompts/t2av.json b/examples/visual_gen/models/cosmos3/prompts/t2av.json new file mode 100644 index 000000000000..1fde4ab55c31 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/t2av.json @@ -0,0 +1,5 @@ +{ + "model_mode": "text2video", + "prompt": "The video opens with a view of a well-lit indoor space featuring a wooden display case with compartments filled with various fruits, including bananas, apples, pears, oranges, and carambolas. The bananas are neatly arranged in the middle compartment, while apples are in the left and a mix of pears, oranges, and carambolas are in the right. Two robotic arms with grippers are positioned at the bottom of the frame, with the one on the left remaining stationary, partially obscuring the apples. The robotic arm on the right begins its action, extending towards the right side of the display case. It carefully picks up a pear from the fruit section, placing it into a plastic bag in the shopping cart nearby, which has red handles. After securing the pear, the arm retracts back to its original position. The process repeats as the robotic arm picks up an orange and places it in the bag, followed by a carambola. The final frame captures the robotic arm returning to its initial position, leaving the display case and surrounding area unchanged. The video showcases a seamless and efficient automated fruit-picking process, highlighting the precision and efficiency of modern robotics in a retail setting.", + "enable_audio": true +} diff --git a/examples/visual_gen/models/cosmos3/prompts/t2i.json b/examples/visual_gen/models/cosmos3/prompts/t2i.json new file mode 100644 index 000000000000..7454c8449c08 --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/t2i.json @@ -0,0 +1,4 @@ +{ + "model_mode": "text2image", + "prompt": "A medium shot of a modern robotics research laboratory with white walls and a gray floor. A robotic arm with a metallic finish is mounted on a clean white workbench, its gripper positioned above a row of small colored objects. A laptop and neatly arranged tools sit beside the robot. A large monitor on the wall behind displays a software interface. The scene is brightly lit by overhead fluorescent lights." +} diff --git a/examples/visual_gen/models/cosmos3/prompts/t2v.json b/examples/visual_gen/models/cosmos3/prompts/t2v.json new file mode 100644 index 000000000000..727e107cf14c --- /dev/null +++ b/examples/visual_gen/models/cosmos3/prompts/t2v.json @@ -0,0 +1,4 @@ +{ + "model_mode": "text2video", + "prompt": "The video opens with a view of a well-lit indoor space featuring a wooden display case with compartments filled with various fruits, including bananas, apples, pears, oranges, and carambolas. The bananas are neatly arranged in the middle compartment, while apples are in the left and a mix of pears, oranges, and carambolas are in the right. Two robotic arms with grippers are positioned at the bottom of the frame, with the one on the left remaining stationary, partially obscuring the apples. The robotic arm on the right begins its action, extending towards the right side of the display case. It carefully picks up a pear from the fruit section, placing it into a plastic bag in the shopping cart nearby, which has red handles. After securing the pear, the arm retracts back to its original position. The process repeats as the robotic arm picks up an orange and places it in the bag, followed by a carambola. The final frame captures the robotic arm returning to its initial position, leaving the display case and surrounding area unchanged. The video showcases a seamless and efficient automated fruit-picking process, highlighting the precision and efficiency of modern robotics in a retail setting." +} diff --git a/examples/visual_gen/models/cosmos3_ti2v.py b/examples/visual_gen/models/cosmos3_ti2v.py deleted file mode 100644 index b9efb9cdec81..000000000000 --- a/examples/visual_gen/models/cosmos3_ti2v.py +++ /dev/null @@ -1,133 +0,0 @@ -#!/usr/bin/env python3 -# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -r"""Cosmos3 Text(+Image)-to-Video generation. - -Cosmos3 OmniMoT supports text-only (T2V) and image-conditioned (I2V/TI2V) -generation from the same checkpoint. Pass ``--image_path`` to condition on a -reference frame. - -Checkpoints (pass the Hub ID or local path via ``--model``): - -- `nvidia/Cosmos3-Nano `_ -- `nvidia/Cosmos3-Super `_ - -Guardrails are enabled by default (required by the -`NVIDIA Open Model License Agreement -`_). -Install and authenticate as follows:: - - pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python - -Accept the terms for the guardrail checkpoint at -https://huggingface.co/nvidia/Cosmos-1.0-Guardrail and set a valid ``HF_TOKEN`` -(the checkpoint is downloaded automatically on first run). - -To run without guardrails (you are responsible for safe deployment):: - - export TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 - -Deployment configs (``examples/visual_gen/configs/``): - -- ``cosmos3-nano-1gpu.yaml`` — 1 GPU, FP8 dynamic quant -- ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE - -Usage: - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \\ - --prompt "The video opens with a view of a well-lit indoor space featuring a " \\ - "wooden display case with compartments filled with various fruits, " \\ - "including bananas, apples, pears, oranges, and carambolas. " \\ - "The bananas are neatly arranged in the middle compartment, while apples " \\ - "are in the left and a mix of pears, oranges, and carambolas are in the " \\ - "right. " \\ - "Two robotic arms with grippers are positioned at the bottom of the frame, " \\ - "with the one on the left remaining stationary, partially obscuring the " \\ - "apples. " \\ - "The robotic arm on the right begins its action, extending towards the " \\ - "right side of the display case. " \\ - "It carefully picks up a pear from the fruit section, placing it into a " \\ - "plastic bag in the shopping cart nearby, which has red handles. " \\ - "After securing the pear, the arm retracts back to its original position. " \\ - "The process repeats as the robotic arm picks up an orange and places it " \\ - "in the bag, followed by a carambola. " \\ - "The final frame captures the robotic arm returning to its initial " \\ - "position, leaving the display case and surrounding area unchanged. " \\ - "The video showcases a seamless and efficient automated fruit-picking " \\ - "process, highlighting the precision and efficiency of modern robotics " \\ - "in a retail setting." \\ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml -""" - -import argparse - -from tensorrt_llm import VisualGen, VisualGenArgs - - -def main(): - parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video example") - parser.add_argument( - "--model", - type=str, - default="nvidia/Cosmos3-Nano", - help="Model path or HuggingFace Hub ID (nvidia/Cosmos3-Nano, nvidia/Cosmos3-Super)", - ) - parser.add_argument( - "--visual_gen_args", - dest="visual_gen_args", - type=str, - default=None, - help="Path to YAML config (same as trtllm-serve --visual_gen_args)", - ) - parser.add_argument( - "--prompt", - type=str, - required=True, - help="Text prompt for generation", - ) - parser.add_argument( - "--image_path", - type=str, - default=None, - help="Optional conditioning image path for I2V/TI2V", - ) - parser.add_argument( - "--output_path", - type=str, - default="cosmos3_ti2v_output.mp4", - help="Path to save the output video", - ) - args = parser.parse_args() - - # Engine config from shared YAML (optional); model-specific defaults apply otherwise. - extra_args = VisualGenArgs.from_yaml(args.visual_gen_args) if args.visual_gen_args else None - visual_gen = VisualGen(model=args.model, args=extra_args) - - # --- Model-specific: T2V / TI2V request construction --- - # Query per-model defaults (resolution, steps, guidance, seed, etc.). - params = visual_gen.default_params - if args.image_path is not None: - params.image = args.image_path - - output = visual_gen.generate( - inputs=args.prompt, - params=params, - ) - - output.save(args.output_path) - print(f"Saved: {args.output_path}") - - -if __name__ == "__main__": - main() diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index e54e818356c4..f5747544946d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -30,11 +30,24 @@ "width": 1280, "num_inference_steps": 35, "guidance_scale": 6.0, - "max_sequence_length": 1024, + "max_sequence_length": 4096, "num_frames": 189, "frame_rate": 24.0, } +# Text-to-image (``output_type="image"``) defaults. Applied by the pipeline when +# the corresponding request field still carries the merged video default, since +# the executor merges a single ``default_generation_params`` dict (the video +# params above) into the request before ``infer()`` runs. +COSMOS3_T2I_PARAMS = { + "height": 1024, + "width": 1024, + "num_inference_steps": 50, + "guidance_scale": 7.0, + "flow_shift": 3.0, + "guidance_interval": (400.0, 1000.0), +} + COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", @@ -56,4 +69,14 @@ default=True, description="Whether to use the guardrails.", ), + "enable_audio": ExtraParamSchema( + type="bool", + default=False, + description="Whether to enable audio generation.", + ), + "output_type": ExtraParamSchema( + type="Literal['video', 'image']", + default="video", + description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", + ), } diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py new file mode 100644 index 000000000000..1ec9f64b57ee --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py @@ -0,0 +1,242 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +import math +from typing import Any, Dict + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch.nn import Parameter +from torch.nn.utils import spectral_norm, weight_norm + +# --------------------------------------------------------------------------- +# Activations +# --------------------------------------------------------------------------- + + +class SnakeBeta(nn.Module): + """ + A modified Snake function which uses separate parameters for the magnitude of the periodic components + Shape: + - Input: (B, C, T) + - Output: (B, C, T), same shape as the input + Parameters: + - alpha - trainable parameter that controls frequency + - beta - trainable parameter that controls magnitude + References: + - This activation function is a modified version + based on this paper by Liu Ziyin, Tilman Hartwig, Masahito Ueda: + https://arxiv.org/abs/2006.08195 + Examples: + >>> a1 = snakebeta(256) + >>> x = torch.randn(256) + >>> x = a1(x) + """ + + def __init__( + self, + in_features: int, + alpha: float = 1.0, + alpha_trainable: bool = True, + alpha_logscale: bool = True, + ) -> None: + super().__init__() + self.in_features = in_features + + self.alpha_logscale = alpha_logscale + param_shape = (1, in_features, 1) + if self.alpha_logscale: + self.alpha = Parameter(torch.zeros(param_shape) * alpha) + self.beta = Parameter(torch.zeros(param_shape) * alpha) + else: + self.alpha = Parameter(torch.ones(param_shape) * alpha) + self.beta = Parameter(torch.ones(param_shape) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + def forward(self, x: torch.Tensor) -> torch.Tensor: + # Keep compatibility with checkpoints storing Snake params as either [C] or [1, C, 1]. + alpha = self.alpha if self.alpha.ndim == 3 else self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta if self.beta.ndim == 3 else self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + + return x + (1.0 / (beta + 1e-9)) * pow(torch.sin(x * alpha), 2) + + +# --------------------------------------------------------------------------- +# WN wrappers +# --------------------------------------------------------------------------- + + +def WNConv1d(*args: Any, **kwargs: Any) -> nn.Conv1d: + """Weight-normalized 1D convolution.""" + return weight_norm(nn.Conv1d(*args, **kwargs)) + + +def WNConvTranspose1d(*args: Any, **kwargs: Any) -> nn.ConvTranspose1d: + """Weight-normalized 1D transpose convolution.""" + return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) + + +# --------------------------------------------------------------------------- +# EnCodec-style conv helpers (SConv1d / SConvTranspose1d) +# --------------------------------------------------------------------------- + +CONV_NORMALIZATIONS = frozenset( + ["none", "weight_norm", "spectral_norm", "time_layer_norm", "layer_norm", "time_group_norm"] +) + + +def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Module: + assert norm in CONV_NORMALIZATIONS + if norm == "weight_norm": + return weight_norm(module) + elif norm == "spectral_norm": + return spectral_norm(module) + return module + + +class ConvLayerNorm(nn.Module): + """LayerNorm over the channel dim of a ``[N, C, T]`` conv output. + + ``nn.LayerNorm`` normalizes the trailing dimension, so it cannot be applied + directly to ``[N, C, T]`` tensors. This wrapper moves the channel axis last, + normalizes, then restores the original layout. + """ + + def __init__(self, num_channels: int, **norm_kwargs: Any) -> None: + super().__init__() + self.norm = nn.LayerNorm(num_channels, **norm_kwargs) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.norm(x.transpose(1, 2)).transpose(1, 2) + + +def get_norm_module( + module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs +) -> nn.Module: + assert norm in CONV_NORMALIZATIONS + if norm in ("layer_norm", "time_layer_norm"): + assert isinstance(module, nn.modules.conv._ConvNd) + return ConvLayerNorm(module.out_channels, **norm_kwargs) + elif norm == "time_group_norm": + if causal: + raise ValueError("GroupNorm doesn't support causal evaluation.") + assert isinstance(module, nn.modules.conv._ConvNd) + return nn.GroupNorm(1, module.out_channels, **norm_kwargs) + return nn.Identity() + + +def pad1d( + x: torch.Tensor, paddings: tuple, mode: str = "constant", value: float = 0.0 +) -> torch.Tensor: + """Tiny wrapper around F.pad that handles reflect padding on short inputs.""" + length = x.shape[-1] + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + if mode == "reflect": + max_pad = max(padding_left, padding_right) + extra_pad = 0 + if length <= max_pad: + extra_pad = max_pad - length + 1 + x = F.pad(x, (0, extra_pad)) + padded = F.pad(x, paddings, mode, value) + end = padded.shape[-1] - extra_pad + return padded[..., :end] + return F.pad(x, paddings, mode, value) + + +def unpad1d(x: torch.Tensor, paddings: tuple) -> torch.Tensor: + """Remove padding from x. Only for 1D.""" + padding_left, padding_right = paddings + assert padding_left >= 0 and padding_right >= 0, (padding_left, padding_right) + assert (padding_left + padding_right) <= x.shape[-1] + end = x.shape[-1] - padding_right + return x[..., padding_left:end] + + +class NormConvTranspose1d(nn.Module): + """ConvTranspose1d with optional weight_norm / spectral_norm.""" + + def __init__( + self, + *args, + causal: bool = False, + norm: str = "none", + norm_kwargs: Dict[str, Any] = {}, + **kwargs, + ): + super().__init__() + self.convtr = apply_parametrization_norm(nn.ConvTranspose1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.convtr, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.convtr(x) + x = self.norm(x) + return x + + +class SConvTranspose1d(nn.Module): + """ConvTranspose1d with builtin asymmetric/causal padding and normalization.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + causal: bool = False, + norm: str = "none", + trim_right_ratio: float = 1.0, + norm_kwargs: Dict[str, Any] = {}, + ): + super().__init__() + self.convtr = NormConvTranspose1d( + in_channels, + out_channels, + kernel_size, + stride, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + self.causal = causal + self.trim_right_ratio = trim_right_ratio + assert self.causal or self.trim_right_ratio == 1.0, ( + "`trim_right_ratio` != 1.0 only makes sense for causal convolutions" + ) + assert 0.0 <= self.trim_right_ratio <= 1.0 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + kernel_size = self.convtr.convtr.kernel_size[0] + stride = self.convtr.convtr.stride[0] + padding_total = kernel_size - stride + + y = self.convtr(x) + + if self.causal: + padding_right = math.ceil(padding_total * self.trim_right_ratio) + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + else: + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + y = unpad1d(y, (padding_left, padding_right)) + return y 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 2de422dcefb5..1ab9acd30f62 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json import math import os import time @@ -30,25 +31,25 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger -from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS +from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS from .guardrails import check_video_safety, download_guardrail_checkpoint +from .sound_tokenizer import LatentAutoEncoderV2 from .transformer_cosmos3 import Cosmos3VFMTransformer -COSMOS3_DEFAULT_NEGATIVE_PROMPT = ( - "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " - "over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, " - "underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, " - "low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, " - "unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of " - "poor quality." -) +COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" COSMOS3_DEFAULT_SYSTEM_PROMPT = ( "You are a helpful assistant who will generate videos from a given prompt." ) +COSMOS3_T2I_SYSTEM_PROMPT = ( + "You are a helpful assistant who will generate images from a given prompt." +) 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" @@ -64,11 +65,27 @@ ) class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): + primary_pretrained_config = pipeline_config.primary_pretrained_config + self.audio_gen = False + self.action_gen = False + if getattr( + primary_pretrained_config, + "audio_gen", + getattr(primary_pretrained_config, "sound_gen", False), + ): + logger.info("Initializing Cosmos3OmniMoTPipeline with audio generation.") + self.audio_gen = True + + if getattr(primary_pretrained_config, "action_gen", False): + logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") + self.action_gen = True + super().__init__(pipeline_config) def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") - self.transformer = Cosmos3VFMTransformer(self.pipeline_config.model_configs["transformer"]) + model_config = self.pipeline_config.model_configs["transformer"] + self.transformer = Cosmos3VFMTransformer(model_config) def load_weights(self, weights: dict) -> None: if self.transformer is not None and hasattr(self.transformer, "load_weights"): @@ -81,6 +98,18 @@ def load_standard_components( ) -> None: skip_components = skip_components or [] + if self.audio_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: + logger.info("Loading audio tokenizer...") + self.audio_tokenizer = ( + LatentAutoEncoderV2.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.SOUND_TOKENIZER, + ) + .to(device) + .to(self.dtype) + .eval() + ) + if PipelineComponent.TOKENIZER not in skip_components: logger.info("Loading tokenizer...") self.tokenizer = Qwen2Tokenizer.from_pretrained( @@ -114,6 +143,18 @@ def load_standard_components( checkpoint_dir, subfolder=PipelineComponent.SCHEDULER, ) + # Snapshot the checkpoint scheduler config so the scheduler can be + # rebuilt at request time when a mode-specific ``flow_shift`` is + # needed (T2I uses shift=3.0; T2V/I2V keep the checkpoint default). + self._base_scheduler_config = self.scheduler.config + self._engine_init_flow_shift = float( + getattr(self.scheduler.config, "flow_shift", 1.0) or 1.0 + ) + self._current_flow_shift = self._engine_init_flow_shift + if self.audio_gen: + # Separate instance so video and audio scheduler states don't collide + # (UniPC mutates internal correction buffers on every .step() call). + self.audio_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) # 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" @@ -144,6 +185,23 @@ def load_standard_components( self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + def _set_flow_shift(self, target_shift: float) -> None: + """Rebuild the UniPC scheduler with ``flow_shift=target_shift`` if needed. + + T2I uses ``flow_shift=3.0`` while T2V/I2V use the checkpoint default. + ``self._current_flow_shift`` is tracked explicitly so a prior T2I rebuild + does not leak into a subsequent video request. + """ + if not hasattr(self, "_base_scheduler_config"): + return + target = float(target_shift) + if target == float(self._current_flow_shift): + return + self.scheduler = UniPCMultistepScheduler.from_config( + self._base_scheduler_config, flow_shift=target + ) + self._current_flow_shift = target + @property def default_warmup_resolutions(self): return [(720, 1280)] @@ -174,28 +232,74 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], use_guardrails=False, image=None, + enable_audio=False, ) + @staticmethod + def _resolve_t2i_default(merged_value, video_default, t2i_default): + """Pick the T2I default when the field still carries the merged video default. + + The executor merges a single ``default_generation_params`` dict (the + video params) into the request before ``infer()``, so an unspecified + field arrives equal to its video default. For T2I we substitute the + T2I default in that case while honoring any explicit user override. + """ + return t2i_default if merged_value == video_default else merged_value + 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" + + height = req.params.height + width = req.params.width + num_inference_steps = req.params.num_inference_steps + guidance_scale = req.params.guidance_scale + if is_t2i: + height = self._resolve_t2i_default( + height, COSMOS3_720P_PARAMS["height"], COSMOS3_T2I_PARAMS["height"] + ) + width = self._resolve_t2i_default( + width, COSMOS3_720P_PARAMS["width"], COSMOS3_T2I_PARAMS["width"] + ) + num_inference_steps = self._resolve_t2i_default( + num_inference_steps, + COSMOS3_720P_PARAMS["num_inference_steps"], + COSMOS3_T2I_PARAMS["num_inference_steps"], + ) + guidance_scale = self._resolve_t2i_default( + guidance_scale, + COSMOS3_720P_PARAMS["guidance_scale"], + COSMOS3_T2I_PARAMS["guidance_scale"], + ) + return self.forward( prompt=req.prompt, negative_prompt=req.params.negative_prompt, image=req.params.image, - height=req.params.height, - width=req.params.width, + height=height, + width=width, num_frames=req.params.num_frames, - num_inference_steps=req.params.num_inference_steps, - guidance_scale=req.params.guidance_scale, + num_inference_steps=num_inference_steps, + guidance_scale=guidance_scale, seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, frame_rate=req.params.frame_rate, - use_duration_template=req.params.extra_params.get("use_duration_template", True), - use_resolution_template=req.params.extra_params.get("use_resolution_template", True), - use_system_prompt=req.params.extra_params.get("use_system_prompt", False), - use_guardrails=req.params.extra_params.get("use_guardrails", True), + use_duration_template=extra_params.get( + "use_duration_template", + COSMOS3_EXTRA_SPECS["use_duration_template"].default, + ), + use_resolution_template=extra_params.get( + "use_resolution_template", + COSMOS3_EXTRA_SPECS["use_resolution_template"].default, + ), + use_system_prompt=extra_params.get("use_system_prompt", False), + use_guardrails=extra_params.get("use_guardrails", True), + enable_audio=extra_params.get("enable_audio", False), + output_type=output_type, ) - def _format_prompt_with_template( + def _apply_metadata_templates( self, prompt: str, *, @@ -203,22 +307,73 @@ def _format_prompt_with_template( width: int, num_frames: int, frame_rate: float, - use_duration_template: bool = True, - use_resolution_template: bool = True, + duration_template: Optional[str] = COSMOS3_DURATION_TEMPLATE, + resolution_template: Optional[str] = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + force_duration_template: bool = False, ) -> str: - prompt = prompt.strip() + """Append duration and resolution metadata to a plain-text prompt. - if use_duration_template and num_frames > 1: + ``duration_template`` / ``resolution_template`` of ``None`` disables that + template. JSON prompts are handled by ``_format_prompt_with_metadata``. + """ + 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): duration = num_frames / frame_rate - dur_text = COSMOS3_DURATION_TEMPLATE.format(duration=duration, fps=frame_rate) - prompt = prompt.rstrip(".") + ". " + dur_text - - prompt = prompt.strip() - if use_resolution_template: - res_text = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE.format(height=height, width=width) - prompt = prompt.rstrip(".") + ". " + res_text - - return prompt + parts.append(duration_template.format(duration=duration, fps=frame_rate).rstrip(".")) + if resolution_template is not None: + parts.append(resolution_template.format(height=height, width=width).rstrip(".")) + if not parts: + return "" + return ". ".join(parts) + "." + + def _format_prompt_with_metadata( + self, + prompt: str, + *, + height: int, + width: int, + num_frames: int, + frame_rate: float, + duration_template: Optional[str], + resolution_template: Optional[str], + force_duration_template: bool = False, + ) -> str: + """Apply cosmos-framework-style metadata to plain text or JSON prompts.""" + stripped = prompt.strip() + if stripped.startswith("{"): + try: + data = json.loads(stripped) + except json.JSONDecodeError: + data = None + else: + if isinstance(data, dict): + 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 + ) + 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) + + return self._apply_metadata_templates( + prompt, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + duration_template=duration_template, + resolution_template=resolution_template, + force_duration_template=force_duration_template, + ) def _resize_and_center_crop_image( self, image: PIL.Image.Image, height: int, width: int @@ -237,14 +392,18 @@ def _resize_and_center_crop_image( @nvtx_range("_tokenize_prompt", color="blue") def _tokenize_prompt( - self, text: str, max_sequence_length: int, use_system_prompt: bool = False + self, + text: str, + max_sequence_length: int, + use_system_prompt: bool = False, + system_prompt: Optional[str] = None, ): """Tokenize a prompt using the Qwen2 chat template. Returns (input_ids, attention_mask) as [1, S] tensors on device. """ conversations = ( - [{"role": "system", "content": COSMOS3_DEFAULT_SYSTEM_PROMPT}] + [{"role": "system", "content": system_prompt or COSMOS3_DEFAULT_SYSTEM_PROMPT}] if use_system_prompt else [] ) @@ -434,6 +593,21 @@ def _decode_latents(self, latents): video = postprocess_video_tensor(video) return video + # ========================================================================= + # Audio generation + # ========================================================================= + + def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: + """Decode audio latent tokens back to waveform. + + Args: + latent: Audio latent tensor of shape (B, C, T). + + Returns: + Waveform tensor of shape (B, audio_channels, N_samples). + """ + return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -457,6 +631,8 @@ def forward( use_resolution_template: bool = COSMOS3_EXTRA_SPECS["use_resolution_template"].default, use_system_prompt: bool = COSMOS3_EXTRA_SPECS["use_system_prompt"].default, use_guardrails: bool = COSMOS3_EXTRA_SPECS["use_guardrails"].default, + enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, + output_type: str = COSMOS3_EXTRA_SPECS["output_type"].default, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -464,6 +640,25 @@ def forward( use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + # Text-to-image mode: same checkpoint/forward path as T2V, but a single + # latent frame, image-flavored prompt templates, flow_shift=3.0, a CFG + # guidance interval, and an image (rather than video) output. + is_t2i = str(output_type).lower() == "image" + guidance_interval = None + if is_t2i: + if image is not None: + raise ValueError( + "Cosmos3 text-to-image (output_type='image') does not accept an image input." + ) + num_frames = 1 + enable_audio = False + guidance_interval = COSMOS3_T2I_PARAMS["guidance_interval"] + self._set_flow_shift(COSMOS3_T2I_PARAMS["flow_shift"]) + else: + # Restore the checkpoint flow_shift in case a prior T2I request + # rebuilt the scheduler with shift=3.0. + self._set_flow_shift(getattr(self, "_engine_init_flow_shift", 1.0)) + if isinstance(prompt, str): prompt = [prompt] batch_size = len(prompt) @@ -481,7 +676,7 @@ def forward( ) # Text guardrail — check both positive and user-supplied negative prompts. - # None negative_prompt means the hardcoded default will be used (safe); skip it. + # None negative_prompt means the empty default will be used (safe); skip it. text_blocked = torch.zeros((), device=self.device, dtype=torch.int32) if self.rank == 0 and use_guardrails and self.safety_checker is not None: prompts_to_check = list(prompt) @@ -506,25 +701,40 @@ def forward( if negative_prompt is None: negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT - negative_prompt = self._format_prompt_with_template( + # Positive prompt: forward duration/resolution templates. T2I has no + # duration concept (single image) and uses the image-flavored + # resolution template. + use_duration_template = use_duration_template and not is_t2i + dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None + if use_resolution_template: + res_tmpl = ( + COSMOS3_IMAGE_RESOLUTION_TEMPLATE if is_t2i else COSMOS3_DEFAULT_RESOLUTION_TEMPLATE + ) + else: + res_tmpl = None + + # 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( negative_prompt, height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, - use_duration_template=use_duration_template, - use_resolution_template=use_resolution_template, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + force_duration_template=False, ) prompt = [ - self._format_prompt_with_template( + self._format_prompt_with_metadata( p, height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, - use_duration_template=use_duration_template, - use_resolution_template=use_resolution_template, + duration_template=dur_tmpl, + resolution_template=res_tmpl, ) for p in prompt ] @@ -534,15 +744,18 @@ def forward( # 1. Tokenize prompts (no separate text encoder — transformer embeds internally) logger.info("Tokenizing prompts...") - cond_ids, cond_mask = self._tokenize_prompt(prompt, max_sequence_length, use_system_prompt) + system_prompt = COSMOS3_T2I_SYSTEM_PROMPT if is_t2i else COSMOS3_DEFAULT_SYSTEM_PROMPT + cond_ids, cond_mask = self._tokenize_prompt( + prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt + ) uncond_ids, uncond_mask = self._tokenize_prompt( - negative_prompt, max_sequence_length, use_system_prompt + negative_prompt, max_sequence_length, use_system_prompt, system_prompt=system_prompt ) # 2. Prepare latents if image is not None: if isinstance(image, str): - image = PIL.Image.open(image).convert("RGB") + image = load_image(image, format="pil") if isinstance(image, PIL.Image.Image): image = image.convert("RGB") @@ -570,6 +783,23 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) + # 3b. Audio noise init — latent length matches diffusers Cosmos3OmniPipeline.prepare_latents. + do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") + audio_latents = None + if do_audio: + audio_cfg = self.audio_tokenizer.model_config + n_audio_samples = int(num_frames / frame_rate * audio_cfg["sampling_rate"]) + hop_size = math.prod(audio_cfg["dec_strides"]) + T_audio = (n_audio_samples + hop_size - 1) // hop_size + audio_latents = randn_tensor( + (1, self.transformer.audio_dim, T_audio), + generator=generator, + device=self.device, + dtype=latents.dtype, + ) + # Audio uses the same scheduler type/config as video. + self.audio_scheduler.set_timesteps(num_inference_steps, device=self.device) + # 4. Build forward_fn for the denoise loop def forward_fn( latent_input, @@ -584,7 +814,9 @@ def forward_fn( Since Cosmos3 embeds text internally, we pass token IDs via extra_tensors rather than through encoder_hidden_states. """ - noise_pred = self.transformer( + current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None + + result = self.transformer( hidden_states=latent_input, timestep=timestep, attention_timestep=timestep / self.scheduler.config.num_train_timesteps, @@ -593,10 +825,18 @@ def forward_fn( video_shape=video_shape, fps=frame_rate, noisy_frame_mask=velocity_mask, + audio_latents=current_audio, ) + + video_noise_pred = result.video + audio_noise_pred = result.audio + if velocity_mask is not None: - noise_pred = noise_pred * velocity_mask - return noise_pred + video_noise_pred = video_noise_pred * velocity_mask + + if audio_noise_pred is not None: + return video_noise_pred, {"audio": audio_noise_pred} + return video_noise_pred # 5. Build CFG tensors — text_ids and text_mask need to be split for CFG # BasePipeline.denoise batches [uncond, cond] when guidance_scale > 1 @@ -610,7 +850,8 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - latents = self.denoise( + extra_streams = {"audio": (audio_latents, self.audio_scheduler)} if do_audio else None + denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, prompt_embeds=cond_ids, # placeholder — actual conditioning via extra_cfg_tensors @@ -618,10 +859,20 @@ def forward_fn( guidance_scale=guidance_scale, forward_fn=forward_fn, extra_cfg_tensors=extra_cfg_tensors, + extra_streams=extra_streams, + guidance_interval=guidance_interval, ) + + if extra_streams is not None: + latents, extra_latents = denoise_result + audio_latents = extra_latents.get("audio") + else: + latents = denoise_result + audio_latents = None + timer.mark_post_start() - # 7. Decode + # 7. Decode video logger.info("Decoding video...") decode_start = time.time() @@ -631,7 +882,13 @@ def forward_fn( video = self.decode_latents(latents, self._decode_latents) - # Video guardrails + # 7b. Decode audio + waveform = None + if do_audio and audio_latents is not None: + logger.info("Decoding audio...") + waveform = self.decode_audio(audio_latents) # [B, audio_channels, N_samples] + + # Video guardrail if self.rank == 0: logger.info(f"Video decoded in {time.time() - decode_start:.2f}s") logger.info(f"Total pipeline time: {time.time() - pipeline_start:.2f}s") @@ -640,4 +897,19 @@ def forward_fn( video = check_video_safety(video, self.safety_checker) timer.mark_end() - return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate)) + + if is_t2i: + # Collapse the single decoded frame [B, T=1, H, W, C] -> [B, H, W, C]. + image = video[:, 0] if video is not None else None + return timer.fill(PipelineOutput(image=image)) + + return timer.fill( + PipelineOutput( + video=video, + frame_rate=frame_rate, + audio=waveform, + audio_sample_rate=self.audio_tokenizer.model_config["sampling_rate"] + if waveform is not None + else None, + ) + ) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py new file mode 100644 index 000000000000..6dc3ac898150 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -0,0 +1,593 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-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. + +import json +import math +import os +from typing import Any, Dict, Literal, Optional + +import torch +from torch import Tensor, nn +from torch.nn.utils import remove_weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +from tensorrt_llm.logger import logger + +from .modules import SConvTranspose1d, SnakeBeta, WNConv1d, WNConvTranspose1d + + +def _resolve_activation_name( + model_config: Dict[str, Any], use_snake: bool +) -> Literal["elu", "snakebeta", "none"]: + if not use_snake: + return "elu" + activation = model_config.get("activation", "snakebeta") + if activation in ("snake", "snakebeta"): + return "snakebeta" + if activation == "none": + return "none" + raise ValueError(f"Unknown activation {activation}") + + +def _resolve_decoder_out_channels(model_config: Dict[str, Any]) -> int: + if "dec_out_channels" in model_config: + return model_config["dec_out_channels"] + out_channels = model_config["input_channels"] + if model_config.get("stereo", False): + out_channels *= 2 + return out_channels + + +def _extract_decoder_state_dict( + state_dict: Dict[str, Tensor], +) -> Dict[str, Tensor]: + """Return checkpoint weights keyed for ``LatentAutoEncoderV2.decoder``.""" + prefixed = {key: value for key, value in state_dict.items() if key.startswith("decoder.")} + if prefixed: + return prefixed + + # Legacy checkpoints may omit the decoder. prefix. + return {f"decoder.{key}": value for key, value in state_dict.items()} + + +def get_activation( + activation: Literal["elu", "snake", "snakebeta", "none"], + antialias: bool = False, + channels: Optional[int] = None, + use_cuda_kernel: bool = False, + snake_logscale: bool = True, +) -> nn.Module: + """ + Get activation module by name. + + Args: + activation: Activation type ('elu', 'snakebeta', or 'none') + antialias: Whether to wrap with anti-aliasing + channels: Number of channels (required for snake activation) + use_cuda_kernel: Whether to use CUDA kernel (not supported) + snake_logscale: Whether SnakeBeta uses log-scaled parameters + + Returns: + Activation module + """ + if activation == "elu": + act = nn.ELU() + elif activation in ("snake", "snakebeta"): + if channels is None: + raise ValueError("channels is required for snake activation") + act = SnakeBeta(channels, alpha_logscale=snake_logscale) + elif activation == "none": + act = nn.Identity() + else: + raise ValueError(f"Unknown activation {activation}") + + if use_cuda_kernel: + raise NotImplementedError("CUDA kernel activation not supported") + + if antialias: + raise NotImplementedError("antialias activation not supported") + + return act + + +class ResidualUnit(nn.Module): + """ + Residual unit with dilated convolutions. + Used in OobleckDecoderBlock. + + Args: + in_channels: Number of input channels + out_channels: Number of output channels + dilation: Dilation rate + kernel_size: Convolution kernel size (default: 7) + use_snake: Whether to use Snake activation (default: False) + antialias_activation: Whether to use anti-aliasing (default: False) + causal: Whether to use causal convolutions (default: False) + padding_mode: Padding mode for convolutions (default: 'zeros') + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + dilation: int, + kernel_size: int = 7, + use_snake: bool = False, + antialias_activation: bool = False, + causal: bool = False, + padding_mode: str = "zeros", + activation: Literal["elu", "snakebeta", "none"] = "elu", + snake_logscale: bool = True, + use_cuda_kernel: bool = False, + ) -> None: + super().__init__() + + self.dilation = dilation + self.causal = causal + self.kernel_size = kernel_size + + if causal: + self.padding = dilation * (kernel_size - 1) + else: + self.padding = (dilation * (kernel_size - 1)) // 2 + + self.padding_mode = padding_mode + activation_name = activation if use_snake else "elu" + + self.snake1 = get_activation( + activation_name, + antialias=antialias_activation, + channels=out_channels, + snake_logscale=snake_logscale, + use_cuda_kernel=use_cuda_kernel, + ) + self.conv1 = WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + dilation=dilation, + padding=self.padding, + padding_mode=self.padding_mode, + ) + self.snake2 = get_activation( + activation_name, + antialias=antialias_activation, + channels=out_channels, + snake_logscale=snake_logscale, + use_cuda_kernel=use_cuda_kernel, + ) + self.conv2 = WNConv1d( + in_channels=out_channels, out_channels=out_channels, kernel_size=1, padding=0 + ) + + def forward(self, x: Tensor) -> Tensor: + """ + Forward pass. + + Args: + x: Input tensor of shape (B, C, T) + + Returns: + Output tensor of shape (B, C, T) + """ + output_tensor = self.conv1(self.snake1(x)) + output_tensor = self.conv2(self.snake2(output_tensor)) + + if self.causal: + output_tensor = output_tensor[:, :, : -self.padding] + res = x[:, :, : -self.padding] + return res + output_tensor + + res = x + padding = (res.shape[-1] - output_tensor.shape[-1]) // 2 + if padding > 0: + res = res[..., padding:-padding] + return res + output_tensor + + +class OobleckDecoderBlock(nn.Module): + """ + Oobleck decoder block with upsampling and residual units. + + Args: + in_channels: Number of input channels + out_channels: Number of output channels + stride: Upsampling stride + use_snake: Whether to use Snake activation (default: False) + antialias_activation: Whether to use anti-aliasing (default: False) + use_nearest_upsample: Whether to use nearest neighbor upsampling (default: False) + causal: Whether to use causal convolutions (default: False) + padding_mode: Padding mode for convolutions (default: 'zeros') + """ + + def __init__( + self, + in_channels: int, + out_channels: int, + stride: int, + use_snake: bool = False, + antialias_activation: bool = False, + use_nearest_upsample: bool = False, + causal: bool = False, + padding_mode: str = "zeros", + activation: Literal["elu", "snakebeta", "none"] = "elu", + snake_logscale: bool = True, + use_cuda_kernel: bool = False, + ) -> None: + super().__init__() + + self.causal = causal + activation_name = activation if use_snake else "elu" + + self.snake1 = get_activation( + activation_name, + antialias=antialias_activation, + channels=in_channels, + snake_logscale=snake_logscale, + use_cuda_kernel=use_cuda_kernel, + ) + self.conv_t1 = self._create_upsample_layer( + in_channels, out_channels, stride, use_nearest_upsample, causal, padding_mode + ) + res_unit_kwargs = { + "use_snake": use_snake, + "causal": causal, + "padding_mode": padding_mode, + "activation": activation, + "snake_logscale": snake_logscale, + "use_cuda_kernel": use_cuda_kernel, + "antialias_activation": antialias_activation, + } + self.res_unit1 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=1, + **res_unit_kwargs, + ) + self.res_unit2 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=3, + **res_unit_kwargs, + ) + self.res_unit3 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=9, + **res_unit_kwargs, + ) + + def _create_upsample_layer( + self, + in_channels: int, + out_channels: int, + stride: int, + use_nearest_upsample: bool, + causal: bool, + padding_mode: str, + ) -> nn.Module: + """ + Create upsampling layer based on configuration. + + Note: padding_mode parameter is not used in this function. + """ + + if ( + causal + ): # use EnCodec's SConvTransposed1d for convenience. padding_mode is reflect by default + assert not use_nearest_upsample, ( + "use_nearest_upsample is not implemented for causal mode!" + ) + upsample_layer = SConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + causal=True, + norm="weight_norm", + ) + else: + if use_nearest_upsample: + upsample_layer = nn.Sequential( + nn.Upsample(scale_factor=stride, mode="nearest"), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=1, + bias=False, + padding="same", + ), + ) + else: + # WNConvTranspose1d only supports zeros padding mode so it's hardcoded + upsample_layer = WNConvTranspose1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + output_padding=stride % 2, + padding_mode="zeros", + ) + + return upsample_layer + + def forward(self, x: Tensor) -> Tensor: + """ + Forward pass. + + Args: + x: Input tensor of shape (B, C, T) + + Returns: + Output tensor of shape (B, C, T_upsampled) + """ + x = self.snake1(x) + x = self.conv_t1(x) + x = self.res_unit1(x) + x = self.res_unit2(x) + return self.res_unit3(x) + + def remove_weight_norm(self) -> None: + """Remove weight normalization from all layers.""" + for layer in [self.conv_t1, self.res_unit1, self.res_unit2, self.res_unit3]: + try: + remove_weight_norm(layer) + except (ValueError, AttributeError): + pass + + +class TrimPadding(nn.Module): + """ + Used for causal convolution support of a conv layer wrapped with nn.Sequential + """ + + def __init__(self, padding: int) -> None: + super().__init__() + self.padding = padding + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return x[:, :, : -self.padding] + + +class OobleckDecoder(nn.Module): + """ + Oobleck Decoder for audio synthesis. + + Decodes latent representations into audio waveforms using + upsampling blocks with optional Snake activation and anti-aliasing. + """ + + def __init__( + self: "OobleckDecoder", + model_config: Dict[str, Any], + ) -> None: + super().__init__() + + self.model_config = model_config + + latent_dim = model_config["vocoder_input_dim"] + out_channels = _resolve_decoder_out_channels(model_config) + + channels = model_config["dec_dim"] + c_mults = model_config["dec_c_mults"] + strides = model_config["dec_strides"] + use_snake = model_config["dec_use_snake"] + use_nearest_upsample = model_config["dec_use_nearest_upsample"] + antialias_activation = model_config["dec_anti_aliasing"] + causal = model_config["causal"] + final_tanh = model_config["dec_use_tanh_at_final"] + padding_mode = model_config["padding_mode"] + snake_logscale = model_config.get("snake_logscale", True) + use_cuda_kernel = model_config.get("use_cuda_kernel", False) + activation = _resolve_activation_name(model_config, use_snake) + block_kwargs = { + "use_snake": use_snake, + "antialias_activation": antialias_activation, + "use_nearest_upsample": use_nearest_upsample, + "causal": causal, + "padding_mode": padding_mode, + "activation": activation, + "snake_logscale": snake_logscale, + "use_cuda_kernel": use_cuda_kernel, + } + + c_mults = [1, *c_mults] + + self.depth = len(c_mults) + + self.first_padding = 6 if causal else 3 + self.conv1 = WNConv1d( + in_channels=latent_dim, + out_channels=c_mults[-1] * channels, + kernel_size=7, + padding=self.first_padding, + padding_mode=padding_mode, + ) + self.conv1_trim = TrimPadding(self.first_padding) if causal else nn.Identity() + + blocks = [] + for i in range(self.depth - 1, 0, -1): + blocks += [ + OobleckDecoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i - 1] * channels, + stride=strides[i - 1], + **block_kwargs, + ) + ] + self.block = nn.ModuleList(blocks) + + self.final_padding = 6 if causal else 3 + self.snake1 = get_activation( + activation, + antialias=antialias_activation, + channels=c_mults[0] * channels, + snake_logscale=snake_logscale, + use_cuda_kernel=use_cuda_kernel, + ) + self.conv2 = WNConv1d( + in_channels=c_mults[0] * channels, + out_channels=out_channels, + kernel_size=7, + padding=self.final_padding, + padding_mode=padding_mode, + bias=False, + ) + self.conv2_trim = TrimPadding(self.final_padding) if causal else nn.Identity() + self.final_activation = nn.Tanh() if final_tanh else nn.Identity() + + def forward(self: "OobleckDecoder", x: torch.Tensor) -> torch.Tensor: + causal = self.model_config.get("causal", False) + if causal: + x = self.conv1(x) + x = self.conv1_trim(x) + for block in self.block: + x = block(x) + x = self.snake1(x) + x = self.conv2(x) + x = self.conv2_trim(x) + return self.final_activation(x) + + x = self.conv1(x) + for block in self.block: + x = block(x) + x = self.snake1(x) + x = self.conv2(x) + if not isinstance(self.final_activation, nn.Identity): + x = self.final_activation(x) + return x + + def remove_weight_norm(self: "OobleckDecoder") -> None: + for module in self.modules(): + if hasattr( + module, "parametrizations" + ): # for new WN implementation using parameterizations + remove_parametrizations(module, "weight") + elif hasattr(module, "weight"): + try: + remove_weight_norm(module) + except ValueError: + pass + + +class LatentAutoEncoderV2(nn.Module): + """ + Decoder-only autoencoder_v2 wrapper for Cosmos3 sound generation. + + Checkpoints store weights under the ``decoder.*`` prefix, e.g. + ``decoder.block.0.conv_t1.weight_g`` and ``decoder.conv1.bias``. + """ + + def __init__(self, model_config: Dict[str, Any]) -> None: + super().__init__() + self.model_config = model_config + self.stereo = model_config.get("stereo", False) + + if model_config.get("encoder_only", False): + raise NotImplementedError("Encoder-only mode not supported") + + dec_type = model_config.get("dec_type", "oobleck") + if dec_type != "oobleck": + raise NotImplementedError( + f"Decoder type '{dec_type}' not supported. Only 'oobleck' is supported." + ) + + self.decoder = OobleckDecoder(model_config) + self.latent_mean = model_config.get("latent_mean", None) + self.latent_std = model_config.get("latent_std", None) + + @classmethod + def from_pretrained( + cls, + checkpoint_dir: str, + subfolder: Optional[str] = None, + dtype: torch.dtype = torch.bfloat16, + device: Optional[torch.device] = None, + **kwargs: Any, + ) -> "LatentAutoEncoderV2": + if subfolder is not None: + checkpoint_dir = os.path.join(checkpoint_dir, subfolder) + + with open(os.path.join(checkpoint_dir, "config.json"), "r") as f: + config = json.load(f) + + model = cls(config) + state_dict: Optional[Dict[str, Any]] = None + + for name in ["diffusion_pytorch_model.safetensors"]: + path = os.path.join(checkpoint_dir, name) + if os.path.exists(path): + from safetensors.torch import load_file + + state_dict = load_file(path, device="cpu") + break + + if state_dict is None: + raise FileNotFoundError( + f"No weight file found in '{checkpoint_dir}'. " + "Expected diffusion_pytorch_model.safetensors." + ) + + decoder_state = _extract_decoder_state_dict(state_dict) + if not decoder_state: + raise FileNotFoundError( + f"No decoder weights found in '{checkpoint_dir}'. " + "Expected keys prefixed with 'decoder.'." + ) + + missing, unexpected = model.load_state_dict(decoder_state, strict=False) + decoder_missing = [key for key in missing if key.startswith("decoder.")] + if decoder_missing: + raise RuntimeError( + f"Failed to load sound tokenizer decoder weights. Missing keys: {decoder_missing}" + ) + if unexpected: + logger.warning(f"Unexpected keys when loading sound tokenizer: {unexpected}") + + # Must remove weight norm AFTER loading the weight_g / weight_v parameters + model.remove_weight_norm() + + model.eval() + for param in model.parameters(): + param.requires_grad = False + + if dtype is not None: + model = model.to(dtype=dtype) + if device is not None: + model = model.to(device=device) + + return model + + def decode(self: "LatentAutoEncoderV2", latent: torch.Tensor) -> torch.Tensor: + """ + Decode latent tokens back to waveform. + + Args: + latent: Latent tensor [B, latent_ch, T_latent]. + + Returns: + Reconstructed waveform [B, audio_channels, T_samples]. + """ + if self.latent_mean is not None and self.latent_std is not None: + latent = latent * self.latent_std + self.latent_mean + + return self.decoder(latent).clamp(-1.0, 1.0) + + def remove_weight_norm(self: "LatentAutoEncoderV2") -> None: + """Remove weight normalization from all components.""" + if self.decoder is not None: + self.decoder.remove_weight_norm() 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 96a103e39d94..91f86ed978c6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -14,6 +14,7 @@ # limitations under the License. import math +from dataclasses import dataclass from typing import Optional, Tuple import torch @@ -58,6 +59,23 @@ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: return output +@dataclass +class TransformerOutput: + """Velocity predictions from Cosmos3VFMTransformer.forward().""" + + video: torch.Tensor + """[B, C, T, H, W] video (or image when T=1) velocity prediction.""" + + image: torch.Tensor + """[B, C, 1, H, W] alias of video for image generation (same tensor).""" + + audio: Optional[torch.Tensor] = None + """[B, audio_dim, T_audio] audio velocity prediction, or None.""" + + action: Optional[torch.Tensor] = None + """[B, T_action, action_dim] action velocity prediction, or None.""" + + def compute_mrope_position_ids_text( num_tokens: int, temporal_offset: int, @@ -333,6 +351,7 @@ def forward( freqs_cos: torch.Tensor, freqs_sin: torch.Tensor, timestep=None, + real_text_lens: Optional[list[int]] = None, ) -> torch.Tensor: """ Args: @@ -356,16 +375,33 @@ def forward( q, k = self.apply_qk_norm(q, k) q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) - k_all = torch.cat([k_und, k], dim=1).contiguous() - v_all = torch.cat([v_und, v], dim=1).contiguous() - - out = self._attn_impl( - q, - k_all, - v_all, - attention_mask=PredefinedAttentionMask.FULL, - timestep=timestep, - ) + if real_text_lens is not None and batch_size > 1: + outs = [] + for b in range(batch_size): + Lb = int(real_text_lens[b]) + k_all_b = torch.cat([k_und[b : b + 1, :Lb], k[b : b + 1]], dim=1) + v_all_b = torch.cat([v_und[b : b + 1, :Lb], v[b : b + 1]], dim=1) + outs.append( + self._attn_impl( + q[b : b + 1], + k_all_b, + v_all_b, + attention_mask=PredefinedAttentionMask.FULL, + timestep=timestep, + ) + ) + out = torch.cat(outs, dim=0) + else: + k_all = torch.cat([k_und, k], dim=1).contiguous() + v_all = torch.cat([v_und, v], dim=1).contiguous() + + out = self._attn_impl( + q, + k_all, + v_all, + attention_mask=PredefinedAttentionMask.FULL, + timestep=timestep, + ) return self.to_out[0](out) @@ -485,6 +521,7 @@ def forward( v_und: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor], timestep=None, + real_text_lens: Optional[list[int]] = None, ) -> torch.Tensor: residual = hidden_states hidden_states = self.input_layernorm(hidden_states) @@ -497,6 +534,7 @@ def forward( freqs_cos=cos, freqs_sin=sin, timestep=timestep, + real_text_lens=real_text_lens, ) hidden_states = residual + hidden_states @@ -665,6 +703,8 @@ class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) pretrained_config = model_config.pretrained_config + self.audio_gen = getattr(pretrained_config, "sound_gen", False) + self.action_gen = getattr(pretrained_config, "action_gen", False) self.hidden_size = pretrained_config.hidden_size self.num_hidden_layers = pretrained_config.num_hidden_layers @@ -684,6 +724,13 @@ def __init__(self, model_config: DiffusionModelConfig): self.num_kv_heads = pretrained_config.num_key_value_heads self.enable_fps_modulation = pretrained_config.enable_fps_modulation + if self.audio_gen: + self.audio_dim = pretrained_config.sound_dim + self.audio_latent_fps = pretrained_config.sound_latent_fps + self.temporal_compression_factor_audio = ( + pretrained_config.temporal_compression_factor_sound + ) + if pretrained_config.position_embedding_type != "unified_3d_mrope": raise ValueError( f"Position embedding type {pretrained_config.position_embedding_type} not supported" @@ -722,6 +769,12 @@ def __init__(self, model_config: DiffusionModelConfig): self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) + if self.audio_gen: + # Projections for audio modality (mirrors cosmos3-internal Cosmos3VFMNetwork) + self.audio2llm = nn.Linear(self.audio_dim, self.hidden_size) + self.llm2audio = nn.Linear(self.hidden_size, self.audio_dim) + self.audio_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + # try timestep embedder in float32 if acc loss self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=torch.bfloat16) @@ -858,6 +911,59 @@ def _compute_rope_freqs( freqs_gen = (cos_gen.unsqueeze(2), sin_gen.unsqueeze(2)) return freqs_und, freqs_gen + # ------------------------------------------------------------------------- + # Audio helpers + # ------------------------------------------------------------------------- + + def _compute_audio_rope_freqs( + self, + T_audio: int, + text_mask: torch.Tensor, + fps_audio: float, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute mRoPE cos/sin for audio tokens. + + Audio tokens use a 1×1 spatial grid (H=W=1) aligned with the vision + temporal axis at the audio latent rate. This mirrors the cosmos3-internal + ``sequence_packing.py`` treatment where audio mRoPE uses + ``get_3d_mrope_ids_vae_tokens(grid_h=1, grid_w=1, tcf=1)``. + """ + B = text_mask.shape[0] + text_lengths = text_mask.sum(dim=1).long() + + audio_pos_list = [] + for b in range(B): + real_len = int(text_lengths[b].item()) + _, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) + # Audio tokens share the vision temporal space; use modality margin offset. + s_pos, _ = compute_mrope_position_ids_vision( + T_audio, + 1, # grid_h + 1, # grid_w + temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, + fps=fps_audio, + base_fps=self.base_fps, + temporal_compression_factor=1, # audio latent is already at audio_latent_fps + enable_fps_modulation=self.enable_fps_modulation, + ) + audio_pos_list.append(s_pos) + + audio_pos_ids = torch.stack(audio_pos_list, dim=1).to(device) # [3, B, T_audio] + rotary_emb = self.language_model.rotary_emb + _dummy = torch.tensor([], dtype=dtype, device=device) + cos_a, sin_a = rotary_emb(_dummy, position_ids=audio_pos_ids) + return cos_a.unsqueeze(2), sin_a.unsqueeze(2) # [B, T_audio, 1, head_dim] + + def pack_audio_latents(self, audio_latents: torch.Tensor) -> torch.Tensor: + """[B, audio_dim, T_audio] → [B, T_audio, audio_dim].""" + return audio_latents.permute(0, 2, 1) + + def unpack_audio_latents(self, hidden_audio: torch.Tensor) -> torch.Tensor: + """[B, T_audio, audio_dim] → [B, audio_dim, T_audio].""" + return hidden_audio.permute(0, 2, 1) + def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None @@ -872,8 +978,9 @@ def forward( video_shape: Optional[Tuple[int, int, int]] = None, fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, + audio_latents: Optional[torch.Tensor] = None, **kwargs, - ) -> torch.Tensor: + ) -> "TransformerOutput": """ Forward pass for parallel denoising. @@ -891,14 +998,21 @@ def forward( timestep embedding, predict velocity) and 0=conditioned (clean context, skip timestep embedding). None means all frames noisy (T2V mode). + audio_latents: Optional [B, audio_dim, T_audio] noisy audio latents. + When provided, audio tokens are appended to the generation + sequence and an audio velocity is returned alongside the video + velocity. Requires ``audio_gen=True`` in the pretrained config. Returns: - [B, C, T, H, W] velocity prediction + TransformerOutput with video (and image alias) always set. + audio is set to the predicted audio velocity when audio_latents is + provided; otherwise None. action is always None for now. """ del kwargs # Kept for diffusers API compatibility. T, H, W = video_shape Hp, Wp, _, _ = self._pad_to_patch_size(H, W) max_real_len = text_mask.sum(dim=1).max().item() + real_text_lens = text_mask.sum(dim=1).tolist() hidden_gen = self.vae2llm(self.patchify(hidden_states, T, H, W)) @@ -957,9 +1071,35 @@ def forward( else: self.cached_kv = cached_kv_full + # --- Audio token injection ------------------------------------------------- + T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp + T_audio = 0 + if audio_latents is not None and self.audio_gen: + T_audio = audio_latents.shape[2] + hidden_audio = self.pack_audio_latents(audio_latents).to(hidden_gen.dtype) + hidden_audio = self.audio2llm(hidden_audio) + self.audio_modality_embed + hidden_audio = hidden_audio + time_embed.unsqueeze(1) + cos_a, sin_a = self._compute_audio_rope_freqs( + T_audio, + text_mask, + float(self.audio_latent_fps), + hidden_states.device, + hidden_gen.dtype, + ) + # [B, T_vid+T_audio, hidden_size] + hidden_gen = torch.cat([hidden_gen, hidden_audio], dim=1) + cos_v, sin_v = self.cached_freqs_gen + freqs_gen_combined = ( + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), + ) + else: + freqs_gen_combined = self.cached_freqs_gen + # -------------------------------------------------------------------------- + S_gen = hidden_gen.shape[1] hidden_gen = self.sharder.shard(hidden_gen, dim=1, pad_to_multiple=True) - cos, sin = self.cached_freqs_gen + cos, sin = freqs_gen_combined cos = self.sharder.shard(cos, dim=1, pad_to_multiple=True) sin = self.sharder.shard(sin, dim=1, pad_to_multiple=True) freqs_gen = (cos, sin) @@ -969,18 +1109,40 @@ def forward( if not self.sharder.is_active: k_und = k_und[:, :max_real_len] v_und = v_und[:, :max_real_len] - hidden_gen = layer( - hidden_gen, - k_und, - v_und, - freqs_gen, - timestep=attention_timestep, - ) + hidden_gen = layer( + hidden_gen, + k_und, + v_und, + freqs_gen, + timestep=attention_timestep, + real_text_lens=real_text_lens, + ) + else: + hidden_gen = layer( + hidden_gen, + k_und, + v_und, + freqs_gen, + timestep=attention_timestep, + ) hidden_gen = self.sharder.gather(hidden_gen, dim=1, unpad_to=S_gen) hidden_gen = self.norm_moe_gen(hidden_gen) - return self.unpatchify(self.llm2vae(hidden_gen), T, H, W) + + # --- Decode video velocity ------------------------------------------------ + video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) + + # --- Decode audio velocity (if requested) --------------------------------- + audio_vel = None + if T_audio > 0 and audio_latents is not None and self.audio_gen: + # hidden_gen[:, T_vid_tokens:] → [B, T_audio, hidden_size] + # → llm2audio → [B, T_audio, audio_dim] → unpack → [B, audio_dim, T_audio] + audio_vel = self.unpack_audio_latents( + self.llm2audio(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_audio]) + ) + + return TransformerOutput(video=video_vel, image=video_vel, audio=audio_vel) def load_weights(self, weights: dict) -> None: """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. @@ -994,8 +1156,6 @@ def load_weights(self, weights: dict) -> None: "lm_head.", "action_modality_embed", "action_proj_", - "audio_modality_embed", - "audio_proj_", ) for key, value in weights.items(): @@ -1004,26 +1164,38 @@ def load_weights(self, weights: dict) -> None: if k.startswith(skip_prefixes): continue - if k.startswith(("vae2llm.", "llm2vae.")): - remapped[k] = value - 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.". + if k.startswith("model."): + k = k[len("model.") :] if k.startswith("proj_in."): remapped[k.replace("proj_in.", "vae2llm.", 1)] = value continue + if k.startswith("proj_out."): remapped[k.replace("proj_out.", "llm2vae.", 1)] = value continue + if k.startswith("audio_proj_in."): + remapped[k.replace("audio_proj_in.", "audio2llm.", 1)] = value + continue + + if k.startswith("audio_proj_out."): + remapped[k.replace("audio_proj_out.", "llm2audio.", 1)] = value + continue + + if k.startswith("audio_modality_embed"): + remapped[k] = value + continue + if k.startswith("time_embedder.linear"): k = k.replace("time_embedder.linear_1.", "time_embedder.mlp.linear_1.") k = k.replace("time_embedder.linear_2.", "time_embedder.mlp.linear_2.") remapped[k] = value continue - if k.startswith("model."): - k = k[len("model.") :] - # embed_tokens and norm → language_model.* if k.startswith("embed_tokens.") or k.startswith("norm."): remapped[f"language_model.{k}"] = value @@ -1145,6 +1317,11 @@ def post_load_weights(self) -> None: self.vae2llm.to(target_dtype) self.llm2vae.to(target_dtype) + if self.audio_gen: + self.audio2llm.to(target_dtype) + self.llm2audio.to(target_dtype) + self.audio_modality_embed.data = self.audio_modality_embed.data.to(target_dtype) + for _, module in self.named_modules(): if isinstance(module, Linear) or isinstance(module, Qwen3VLTextRMSNorm): module.post_load_weights() diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index 9f1d0452e26a..13cbe77bdf12 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -216,9 +216,7 @@ def world_size(self): @property def dtype(self): - if hasattr(self, "transformer"): - return next(self.transformer.parameters()).dtype - return torch.float32 + return self.pipeline_config.torch_dtype @property def device(self): @@ -731,6 +729,26 @@ def _rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): noise_pred_rescaled = noise_cfg * (std_text / std_cfg) return guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg + @staticmethod + def _resolve_step_guidance_scale( + t: torch.Tensor, + guidance_scale: float, + guidance_interval: Optional[Tuple[float, float]] = None, + guidance_scale_2: Optional[float] = None, + boundary_timestep: Optional[float] = None, + ) -> float: + """Per-step CFG scale, including two-stage and guidance-interval gating.""" + current = guidance_scale + t_scalar = t.item() if t.dim() == 0 else t[0].item() + if guidance_scale_2 is not None and boundary_timestep is not None: + if t_scalar < boundary_timestep: + current = guidance_scale_2 + if guidance_interval is not None: + interval_lo, interval_hi = guidance_interval + if not (interval_lo <= t_scalar <= interval_hi): + current = 1.0 + return current + def _setup_cfg_config( self, guidance_scale, prompt_embeds, neg_prompt_embeds, extra_cfg_tensors=None ): @@ -889,9 +907,10 @@ def _denoise_step_standard( guidance_scale, guidance_rescale, local_extras, + do_cfg: bool = False, ): """Execute single denoising step without CFG parallel.""" - if guidance_scale > 1.0: + if do_cfg: latent_input = torch.cat([latents] * 2) # Duplicate extra stream latents for CFG extra_stream_input = { @@ -924,7 +943,7 @@ def _denoise_step_standard( t_transformer = time.time() - t_start c_start = time.time() - if guidance_scale > 1.0: + if do_cfg: noise_pred_uncond, noise_pred_text = noise_pred.chunk(2) noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_text - noise_pred_uncond) @@ -989,6 +1008,7 @@ def denoise( extra_streams: Optional[Dict[str, Tuple[torch.Tensor, Any]]] = None, guidance_scale_2: Optional[float] = None, boundary_timestep: Optional[float] = None, + guidance_interval: Optional[Tuple[float, float]] = None, post_step_fn: Optional[Callable] = None, ): """Execute denoising loop with optional CFG parallel and TeaCache support. @@ -1019,6 +1039,9 @@ def denoise( to guidance_scale_2 when timestep < boundary_timestep. boundary_timestep: Optional timestep boundary for two-stage denoising. Switches guidance scale when crossing this threshold. + guidance_interval: Optional ``(lo, hi)`` scheduler-timestep range in which CFG + is active. Outside the interval the effective scale is 1.0 + (conditional prediction only); both branches still run. post_step_fn: Optional callable applied to latents after each scheduler step. Signature: post_step_fn(latents) -> latents Use for constraints that must hold throughout denoising. @@ -1057,6 +1080,7 @@ def denoise( cfg_config = self._setup_cfg_config( guidance_scale, prompt_embeds, neg_prompt_embeds, extra_cfg_tensors ) + do_cfg = guidance_scale > 1.0 do_cfg_parallel = cfg_config["enabled"] prompt_embeds = cfg_config["prompt_embeds"] local_extras = cfg_config["local_extras"] @@ -1090,12 +1114,13 @@ def denoise( step_start = time.time() - # Two-stage denoising: switch guidance scale at boundary - current_guidance_scale = guidance_scale - if guidance_scale_2 is not None and boundary_timestep is not None: - t_scalar = t.item() if t.dim() == 0 else t[0].item() - if t_scalar < boundary_timestep: - current_guidance_scale = guidance_scale_2 + current_guidance_scale = self._resolve_step_guidance_scale( + t, + guidance_scale, + guidance_interval, + guidance_scale_2, + boundary_timestep, + ) # Denoise with nvtx_range(f"denoise_step {i}"): @@ -1134,6 +1159,7 @@ def denoise( current_guidance_scale, guidance_rescale, local_extras, + do_cfg=do_cfg, ) # Scheduler step for all streams diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 4300a2943e80..8fe7753ce7f7 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -56,6 +56,7 @@ class PipelineComponent(str, Enum): SCHEDULER = "scheduler" IMAGE_ENCODER = "image_encoder" IMAGE_PROCESSOR = "image_processor" + SOUND_TOKENIZER = "sound_tokenizer" @dataclass diff --git a/tests/integration/defs/examples/visual_gen/test_visual_gen.py b/tests/integration/defs/examples/visual_gen/test_visual_gen.py index 42825e570ea6..b2d351b2aa29 100644 --- a/tests/integration/defs/examples/visual_gen/test_visual_gen.py +++ b/tests/integration/defs/examples/visual_gen/test_visual_gen.py @@ -1647,7 +1647,7 @@ def test_qwen_image_example(_visual_gen_deps, llm_root, llm_venv): def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): - """Run examples/visual_gen/models/cosmos3_ti2v.py with FP8 config end-to-end. + """Run examples/visual_gen/models/cosmos3/cosmos3.py with FP8 config end-to-end. Validates that the Cosmos3-Nano example script and ``configs/cosmos3-nano-1gpu.yaml`` work together as documented. Uses the local Cosmos3-Nano checkpoint and @@ -1660,7 +1660,9 @@ def test_cosmos3_example(_visual_gen_deps, llm_root, llm_venv): os.makedirs(out_dir, exist_ok=True) output_path = os.path.join(out_dir, "cosmos3_output.mp4") - script_path = os.path.join(llm_root, "examples", "visual_gen", "models", "cosmos3_ti2v.py") + script_path = os.path.join( + llm_root, "examples", "visual_gen", "models", "cosmos3", "cosmos3.py" + ) config_path = os.path.join( llm_root, "examples", "visual_gen", "configs", "cosmos3-nano-1gpu.yaml" ) 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 15ba3cb5e8d9..38e8675c3f0f 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 @@ -110,6 +110,24 @@ _TIMESTEP = 500.0 _FPS = 24.0 +# Audio (sound) modality: audio tokens are appended to the gen sequence, so the +# combined seq becomes video_tokens (8) + T_audio. Keep T_audio even so the +# combined length stays divisible by Ulysses=2 (the sharder also pads, but this +# keeps the parity comparison free of padding artifacts). +_AUDIO_DIM = 16 +_T_AUDIO = 4 +_SOUND_LATENT_FPS = 24.0 + +# Same architecture as _COSMOS3_TEST_CONFIG, with the audio modality enabled. +# The transformer reads audio attributes via the legacy ``sound_*`` keys. +_COSMOS3_AUDIO_CONFIG = dict( + **_COSMOS3_TEST_CONFIG, + sound_gen=True, + sound_dim=_AUDIO_DIM, + sound_latent_fps=_SOUND_LATENT_FPS, + temporal_compression_factor_sound=1, +) + SEED_WEIGHTS = 123 SEED_INPUT = 456 SEED_COND_TEXT = 42 @@ -374,7 +392,35 @@ def _forward(model: Cosmos3VFMTransformer, device: torch.device, text_seed: int) text_mask=text_mask, video_shape=video_shape, fps=_FPS, + ).video + + +def _forward_with_audio( + model: Cosmos3VFMTransformer, device: torch.device, text_seed: int +) -> Tuple[torch.Tensor, torch.Tensor]: + """Forward an audio-enabled model; returns (video_velocity, audio_velocity).""" + channels = _COSMOS3_TEST_CONFIG["latent_channel"] + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + device, channels=channels, text_seed=text_seed + ) + # Deterministic audio noise, independent of the (seed-controlled) video/text + # inputs, so ref and parallel models see identical audio_latents. + torch.manual_seed(SEED_INPUT + 1) + audio_latents = ( + torch.randn(hs.shape[0], _AUDIO_DIM, _T_AUDIO, device=device, dtype=hs.dtype) * 0.1 + ) + model.reset_cache() + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=_FPS, + audio_latents=audio_latents, ) + return out.video, out.audio def _build_ref_and_parallel( @@ -385,8 +431,9 @@ def _build_ref_and_parallel( attn2d_row_size: int = 1, attn2d_col_size: int = 1, backend: str = "VANILLA", - pretrained_dict: dict = _COSMOS3_TEST_CONFIG, + pretrained_dict: dict = None, ) -> Tuple[Cosmos3VFMTransformer, Cosmos3VFMTransformer, VisualGenMapping, torch.device]: + pretrained_dict = pretrained_dict if pretrained_dict is not None else _COSMOS3_TEST_CONFIG device = torch.device(f"cuda:{dist.get_rank() % torch.cuda.device_count()}") torch.manual_seed(SEED_WEIGHTS) @@ -481,6 +528,37 @@ def _logic_cosmos3_ulysses_vs_single_gpu(rank, world_size): ) +def _logic_cosmos3_ulysses_audio_vs_single_gpu(rank, world_size): + ref_model, ulysses_model, _, device = _build_ref_and_parallel( + ulysses_size=world_size, pretrained_dict=_COSMOS3_AUDIO_CONFIG + ) + text_seed = _cfg_text_seed(rank, tp_size=1, ulysses_size=world_size, cfg_size=1) + + ref_video, ref_audio = _forward_with_audio(ref_model, device, text_seed) + ulysses_video, ulysses_audio = _forward_with_audio(ulysses_model, device, text_seed) + + if rank == 0: + vdiff = (ulysses_video.float() - ref_video.float()).abs() + adiff = (ulysses_audio.float() - ref_audio.float()).abs() + print( + f"[ulysses={world_size}+audio] " + f"video max_abs_diff={vdiff.max().item():.6e}, " + f"audio max_abs_diff={adiff.max().item():.6e}", + flush=True, + ) + + _assert_parity( + ulysses_video, + ref_video, + msg=f"Rank {rank}: Ulysses+audio VIDEO differs from single-GPU reference", + ) + _assert_parity( + ulysses_audio, + ref_audio, + msg=f"Rank {rank}: Ulysses+audio AUDIO differs from single-GPU reference", + ) + + def _logic_cosmos3_tp_ulysses_vs_single_gpu(rank, world_size): tp_size = 2 ulysses_size = 2 @@ -624,6 +702,12 @@ 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_ulysses2_audio_vs_single_gpu(self): + """Ulysses parity with the audio modality on: video + audio tokens are + sharded together across the sequence dimension.""" + self._skip_if_unavailable() + run_test_in_distributed(world_size=2, test_fn=_logic_cosmos3_ulysses_audio_vs_single_gpu) + @pytest.mark.gpu4 def test_tp2_ulysses2_vs_single_gpu(self): self._skip_if_unavailable() diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 61bf5dfd1f90..c0329b71a060 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -4,13 +4,20 @@ """Smoke tests for Cosmos3OmniMoTPipeline. Loads Cosmos3-Nano when available, runs end-to-end generation, and asserts -valid uint8 video outputs. No diffusers reference comparison. +valid uint8 video/image outputs and float32 audio when enabled. No diffusers +reference comparison. Run all pipeline smoke tests: pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s -m cosmos3 -Run single mode: - pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s -m cosmos3_i2v +Run T2I only: + pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s -m cosmos3_t2i + +Run audio only: + pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s -m cosmos3_audio + +Run prompt metadata unit tests (no GPU): + pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -k FormatPromptWithMetadata Override checkpoint: DIFFUSION_MODEL_PATH_COSMOS3=/path/to/Cosmos3-Nano \\ @@ -18,6 +25,7 @@ """ import gc +import json import os from pathlib import Path @@ -28,7 +36,13 @@ import pytest import torch -from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import Cosmos3OmniMoTPipeline +from tensorrt_llm._torch.visual_gen.models.cosmos3.defaults import COSMOS3_T2I_PARAMS +from tensorrt_llm._torch.visual_gen.models.cosmos3.pipeline_cosmos3 import ( + COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + COSMOS3_DURATION_TEMPLATE, + COSMOS3_IMAGE_RESOLUTION_TEMPLATE, + Cosmos3OmniMoTPipeline, +) from tensorrt_llm._torch.visual_gen.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -79,6 +93,12 @@ def _checkpoint(env_var: str, default_name: str) -> str: GUIDANCE_SCALE = 6.0 FRAME_RATE = 24.0 +# T2I smoke resolution — smaller than the 1024 default to keep CI memory down; +# ``output_type="image"`` still exercises flow_shift and guidance_interval. +T2I_HEIGHT = 512 +T2I_WIDTH = 512 +T2I_GUIDANCE_SCALE = COSMOS3_T2I_PARAMS["guidance_scale"] + COSMOS3_FP8_QUANT_CONFIG = { "quant_algo": "FP8", "dynamic": True, @@ -103,15 +123,24 @@ def _load_pipeline(checkpoint_path: str, **visual_gen_kwargs): return PipelineLoader(args).load(skip_warmup=True) -def _run_forward(pipeline, *, image=None, num_frames=NUM_FRAMES, **extra): +def _run_forward( + pipeline, + *, + image=None, + num_frames=NUM_FRAMES, + height=HEIGHT, + width=WIDTH, + guidance_scale=GUIDANCE_SCALE, + **extra, +): return pipeline.forward( prompt=PROMPT, image=image, - height=HEIGHT, - width=WIDTH, + height=height, + width=width, num_frames=num_frames, num_inference_steps=NUM_STEPS, - guidance_scale=GUIDANCE_SCALE, + guidance_scale=guidance_scale, seed=SEED, frame_rate=FRAME_RATE, use_guardrails=False, @@ -141,6 +170,51 @@ def _assert_valid_video( assert vf.min() >= 0 and vf.max() <= 255 +def _assert_valid_image( + image: torch.Tensor, + *, + height: int = T2I_HEIGHT, + width: int = T2I_WIDTH, +): + """PipelineOutput.image is (B, H, W, C) uint8 per output.py.""" + assert image is not None + assert image.dtype == torch.uint8 + assert image.dim() == 4, f"Expected (B,H,W,C), got {image.shape}" + batch, h, w, c = image.shape + assert batch == 1 + assert h == height and w == width + assert c == 3 + img = image.float() + assert not torch.isnan(img).any() + assert not torch.isinf(img).any() + assert img.min() >= 0 and img.max() <= 255 + + +def _assert_valid_audio( + audio: torch.Tensor, + audio_sample_rate: int, +): + """PipelineOutput.audio is (B, C, T) float32.""" + assert audio is not None + assert audio_sample_rate is not None and audio_sample_rate > 0 + assert audio.dtype == torch.float32 + assert audio.dim() == 3, f"Expected (B,C,T), got {audio.shape}" + batch, channels, samples = audio.shape + assert batch == 1 + assert channels >= 1 + assert samples > 0 + af = audio.float() + assert not torch.isnan(af).any() + assert not torch.isinf(af).any() + + +def _require_audio_pipeline(pipeline) -> None: + if not getattr(pipeline, "audio_gen", False): + pytest.skip("Checkpoint does not enable audio generation") + if not hasattr(pipeline, "audio_tokenizer"): + pytest.skip("Audio tokenizer was not loaded for this pipeline") + + def _make_test_image() -> PIL.Image.Image: image_path = os.environ.get("COSMOS3_TEST_IMAGE") if image_path and os.path.exists(image_path): @@ -148,6 +222,157 @@ def _make_test_image() -> PIL.Image.Image: return PIL.Image.new("RGB", (WIDTH, HEIGHT), color=(64, 128, 192)) +@pytest.fixture +def cosmos3_format_pipeline(): + """Minimal pipeline for prompt formatting helpers (no checkpoint).""" + return Cosmos3OmniMoTPipeline.__new__(Cosmos3OmniMoTPipeline) + + +def _format_prompt_with_metadata( + pipeline, + prompt: str, + *, + height: int = HEIGHT, + width: int = WIDTH, + num_frames: int = 189, + frame_rate: float = FRAME_RATE, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + force_duration_template: bool = False, +) -> str: + return pipeline._format_prompt_with_metadata( + prompt, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + duration_template=duration_template, + resolution_template=resolution_template, + force_duration_template=force_duration_template, + ) + + +class TestFormatPromptWithMetadataPlainText: + def test_appends_duration_and_resolution(self, cosmos3_format_pipeline): + result = _format_prompt_with_metadata(cosmos3_format_pipeline, "A cat on a beach") + assert result.startswith("A cat on a beach.") + assert "7.9 seconds long" in result + assert "720x1280" in result + + def test_matches_apply_metadata_templates(self, cosmos3_format_pipeline): + prompt = "Mountain lake at sunrise" + via_format = _format_prompt_with_metadata(cosmos3_format_pipeline, prompt) + via_apply = cosmos3_format_pipeline._apply_metadata_templates( + prompt, + height=HEIGHT, + width=WIDTH, + num_frames=189, + frame_rate=FRAME_RATE, + duration_template=COSMOS3_DURATION_TEMPLATE, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + assert via_format == via_apply + + def test_templates_disabled_returns_prompt_only(self, cosmos3_format_pipeline): + result = _format_prompt_with_metadata( + cosmos3_format_pipeline, + "Plain prompt", + duration_template=None, + resolution_template=None, + ) + assert result == "Plain prompt." + + def test_empty_prompt_with_templates(self, cosmos3_format_pipeline): + result = _format_prompt_with_metadata(cosmos3_format_pipeline, "") + assert "7.9 seconds long" in result + assert "720x1280" in result + + def test_invalid_json_prefix_falls_back_to_append(self, cosmos3_format_pipeline): + result = _format_prompt_with_metadata(cosmos3_format_pipeline, "{not valid json") + assert result.startswith("{not valid json.") + assert "720x1280" in result + + def test_json_array_falls_back_to_append(self, cosmos3_format_pipeline): + result = _format_prompt_with_metadata(cosmos3_format_pipeline, '["a", "b"]') + assert result.startswith('["a", "b"].') + assert "720x1280" in result + + +class TestFormatPromptWithMetadataJson: + def test_injects_metadata_fields(self, cosmos3_format_pipeline): + prompt = json.dumps({"prompt": "A foundry pour", "subjects": []}) + result = _format_prompt_with_metadata(cosmos3_format_pipeline, prompt) + 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" + + def test_overwrites_existing_metadata_fields(self, cosmos3_format_pipeline): + prompt = json.dumps( + { + "prompt": "test", + "duration": "5s", + "fps": 30, + "resolution": {"W": 640, "H": 480}, + "aspect_ratio": "3,4", + } + ) + 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" + + def test_single_frame_skips_duration_by_default(self, cosmos3_format_pipeline): + prompt = json.dumps({"prompt": "still life"}) + 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 data["resolution"] == {"W": 1280, "H": 720} + + def test_single_frame_duration_when_forced(self, cosmos3_format_pipeline): + prompt = json.dumps({"prompt": "still life"}) + data = json.loads( + _format_prompt_with_metadata( + cosmos3_format_pipeline, + prompt, + num_frames=1, + force_duration_template=True, + ) + ) + assert data["duration"] == "0.0s" + + def test_non_integer_fps_preserved(self, cosmos3_format_pipeline): + prompt = json.dumps({"prompt": "test"}) + data = json.loads( + _format_prompt_with_metadata(cosmos3_format_pipeline, prompt, frame_rate=23.976) + ) + assert data["fps"] == 23.976 + + def test_resolution_only_when_duration_template_disabled(self, cosmos3_format_pipeline): + prompt = json.dumps({"prompt": "test"}) + data = json.loads( + _format_prompt_with_metadata( + cosmos3_format_pipeline, + prompt, + duration_template=None, + resolution_template=COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, + ) + ) + assert "duration" not in data + assert "fps" not in data + assert data["resolution"] == {"W": 1280, "H": 720} + + @pytest.fixture(scope="class") def cosmos3_pipeline(): checkpoint = _require_checkpoint() @@ -158,20 +383,6 @@ def cosmos3_pipeline(): torch.cuda.empty_cache() -@pytest.mark.integration -class TestCosmos3PipelineLoad: - def test_load_pipeline(self): - checkpoint = _require_checkpoint() - pipeline = _load_pipeline(checkpoint) - try: - assert isinstance(pipeline, Cosmos3OmniMoTPipeline) - assert pipeline.transformer is not None - finally: - del pipeline - gc.collect() - torch.cuda.empty_cache() - - @pytest.mark.integration @pytest.mark.cosmos3_t2v @pytest.mark.high_cuda_memory @@ -198,9 +409,28 @@ def test_i2v_smoke(self, cosmos3_pipeline): @pytest.mark.high_cuda_memory class TestCosmos3T2I: def test_t2i_smoke(self, cosmos3_pipeline): - result = _run_forward(cosmos3_pipeline, image=None, num_frames=1) - _assert_valid_video(result.video, num_frames=1) + result = _run_forward( + cosmos3_pipeline, + image=None, + output_type="image", + height=T2I_HEIGHT, + width=T2I_WIDTH, + guidance_scale=T2I_GUIDANCE_SCALE, + ) + assert result.video is None + _assert_valid_image(result.image, height=T2I_HEIGHT, width=T2I_WIDTH) + + +@pytest.mark.integration +@pytest.mark.cosmos3_audio +@pytest.mark.high_cuda_memory +class TestCosmos3Audio: + def test_audio_smoke(self, cosmos3_pipeline): + _require_audio_pipeline(cosmos3_pipeline) + result = _run_forward(cosmos3_pipeline, enable_audio=True) + _assert_valid_video(result.video, num_frames=NUM_FRAMES) assert result.frame_rate == FRAME_RATE + _assert_valid_audio(result.audio, result.audio_sample_rate) @pytest.mark.integration @@ -236,9 +466,8 @@ def test_template_variants( @pytest.mark.cosmos3_t2v @pytest.mark.high_cuda_memory class TestCosmos3NegativePrompt: - @pytest.mark.parametrize("negative_prompt", [None, ""], ids=["default", "empty"]) - def test_negative_prompt(self, cosmos3_pipeline, negative_prompt): - result = _run_forward(cosmos3_pipeline, negative_prompt=negative_prompt) + def test_default_negative_prompt(self, cosmos3_pipeline): + result = _run_forward(cosmos3_pipeline, negative_prompt=None) _assert_valid_video(result.video, num_frames=NUM_FRAMES) diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py index 4349d0184ca2..5d1100b394b1 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_transformer.py @@ -109,6 +109,30 @@ def _load_model_config(checkpoint_dir: str) -> DiffusionModelConfig: return DiffusionPipelineConfig.from_pretrained(checkpoint_dir, args=args).primary_model_config +def _enable_audio( + model_config: DiffusionModelConfig, + *, + audio_dim: int = 16, + audio_latent_fps: float = 24.0, + temporal_compression_factor: int = 1, +) -> DiffusionModelConfig: + """Pin the audio (sound) modality on with small, test-friendly dimensions. + + The Cosmos3 checkpoint already enables sound by default; this overrides the + audio dims so random-weight builds stay light and assertions can rely on a + known ``audio_dim``. The transformer reads audio attributes via ``sound_*`` + fallbacks (see ``Cosmos3VFMTransformer.__init__``), so we set those legacy + keys. ``pretrained_config`` is a ``SimpleNamespace``, so attributes can be + set freely. + """ + cfg = model_config.pretrained_config + cfg.sound_gen = True + cfg.sound_dim = audio_dim + cfg.sound_latent_fps = audio_latent_fps + cfg.temporal_compression_factor_sound = temporal_compression_factor + return model_config + + def _init_all_weights(model: torch.nn.Module, std: float = 0.02) -> None: with torch.no_grad(): for name, param in model.named_parameters(): @@ -197,7 +221,7 @@ def test_sanity_forward(self, cosmos3_model_config): text_mask=text_mask, video_shape=video_shape, ) - _assert_finite_output(out, hs.shape) + _assert_finite_output(out.video, hs.shape) @pytest.mark.high_cuda_memory def test_reset_cache(self, cosmos3_model_config): @@ -221,8 +245,8 @@ def test_reset_cache(self, cosmos3_model_config): text_mask=text_mask, video_shape=video_shape, ) - _assert_finite_output(out1, hs.shape) - _assert_finite_output(out2, hs.shape) + _assert_finite_output(out1.video, hs.shape) + _assert_finite_output(out2.video, hs.shape) @pytest.mark.high_cuda_memory def test_sanity_forward_i2v_mask(self, cosmos3_model_config): @@ -243,7 +267,126 @@ def test_sanity_forward_i2v_mask(self, cosmos3_model_config): video_shape=video_shape, noisy_frame_mask=noisy_frame_mask, ) - _assert_finite_output(out, hs.shape) + _assert_finite_output(out.video, hs.shape) + + +@pytest.mark.integration +class TestCosmos3Audio: + """Audio (sound) modality — Nano architecture, random weights, audio_gen on. + + Loads the Nano transformer config and flips on the audio modality so the + audio projection heads and sound-token injection path are exercised without + needing an audio-capable checkpoint. + """ + + AUDIO_DIM = 16 + T_AUDIO = 8 + + @pytest.fixture(autouse=True) + def _require_cuda(self): + if not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + @pytest.fixture + def audio_model_config(self): + # Function-scoped + freshly loaded so we never mutate a config shared + # with the video-only test classes. + checkpoint_dir = _require_checkpoint() + model_config = _load_model_config(checkpoint_dir) + return _enable_audio(model_config, audio_dim=self.AUDIO_DIM) + + @pytest.fixture + def cosmos3_model_config_noaudio(self): + # The Cosmos3 checkpoint enables sound by default, so explicitly disable + # it to exercise the video-only construction path. + checkpoint_dir = _require_checkpoint() + model_config = _load_model_config(checkpoint_dir) + model_config.pretrained_config.sound_gen = False + return model_config + + def test_audio_model_structure(self, audio_model_config): + model = Cosmos3VFMTransformer(model_config=audio_model_config) + assert model.audio_gen is True + assert model.audio_dim == self.AUDIO_DIM + assert hasattr(model, "audio2llm") + assert hasattr(model, "llm2audio") + assert hasattr(model, "audio_modality_embed") + # audio2llm: audio_dim -> hidden_size, llm2audio: hidden_size -> audio_dim + assert model.audio2llm.in_features == self.AUDIO_DIM + assert model.audio2llm.out_features == model.hidden_size + assert model.llm2audio.in_features == model.hidden_size + assert model.llm2audio.out_features == self.AUDIO_DIM + assert model.audio_modality_embed.shape == (model.hidden_size,) + + def test_video_only_model_has_no_audio_heads(self, cosmos3_model_config_noaudio): + model = Cosmos3VFMTransformer(model_config=cosmos3_model_config_noaudio) + assert model.audio_gen is False + assert not hasattr(model, "audio2llm") + assert not hasattr(model, "llm2audio") + + @pytest.mark.high_cuda_memory + def test_forward_with_audio(self, audio_model_config): + cfg = audio_model_config.pretrained_config + model = _build_random_weight_model(audio_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + audio_latents = torch.randn(1, model.audio_dim, self.T_AUDIO, device=DEVICE, dtype=DTYPE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + audio_latents=audio_latents, + ) + # Video velocity is unchanged in shape; audio velocity mirrors the input. + _assert_finite_output(out.video, hs.shape) + assert out.audio is not None + _assert_finite_output(out.audio, torch.Size([1, model.audio_dim, self.T_AUDIO])) + + @pytest.mark.high_cuda_memory + def test_forward_without_audio_latents_returns_none(self, audio_model_config): + """An audio-capable model still returns audio=None when no audio is passed.""" + cfg = audio_model_config.pretrained_config + model = _build_random_weight_model(audio_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel + ) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + ) + _assert_finite_output(out.video, hs.shape) + assert out.audio is None + + @pytest.mark.high_cuda_memory + def test_forward_with_audio_multiframe(self, audio_model_config): + """Audio injection works alongside a multi-frame video sequence.""" + cfg = audio_model_config.pretrained_config + model = _build_random_weight_model(audio_model_config) + hs, ts, text_ids, text_mask, video_shape = _cosmos3_inputs( + DEVICE, channels=cfg.latent_channel, t=3 + ) + audio_latents = torch.randn(1, model.audio_dim, self.T_AUDIO, device=DEVICE, dtype=DTYPE) + with torch.inference_mode(): + out = model( + hidden_states=hs, + timestep=ts, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=24.0, + audio_latents=audio_latents, + ) + _assert_finite_output(out.video, hs.shape) + _assert_finite_output(out.audio, torch.Size([1, model.audio_dim, self.T_AUDIO])) @pytest.mark.integration @@ -281,7 +424,7 @@ def test_load_weights_and_forward(self, cosmos3_transformer): text_mask=text_mask, video_shape=video_shape, ) - _assert_finite_output(out, hs.shape) + _assert_finite_output(out.video, hs.shape) @pytest.mark.parametrize("quant_algo", ["FP8"]) def test_load_fp8_quantization(self, quant_algo: str): @@ -310,7 +453,7 @@ def test_load_fp8_quantization(self, quant_algo: str): text_mask=text_mask, video_shape=video_shape, ) - _assert_finite_output(out, hs.shape) + _assert_finite_output(out.video, hs.shape) finally: del pipeline gc.collect()