From 6e616ebbc35acbc956e1d04713289cfbf9ae730e Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 14 May 2026 21:09:37 -0700 Subject: [PATCH 01/20] add sound tokenizer Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/modules.py | 471 +++++++++ .../models/cosmos3/pipeline_cosmos3.py | 51 + .../models/cosmos3/sound_tokenizer.py | 955 ++++++++++++++++++ .../models/cosmos3/transformer_cosmos3.py | 9 + 4 files changed, 1486 insertions(+) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py 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..d95c916cb849 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py @@ -0,0 +1,471 @@ +import math +import warnings +from typing import Any, Dict, Optional, Tuple + +import torch +import torch.nn as nn +import torch.nn.functional as F +from torch import Tensor +from torch.cuda import amp +from torch.nn import Parameter +from torch.nn.utils import spectral_norm, weight_norm + +# --------------------------------------------------------------------------- +# VAE Bottleneck +# --------------------------------------------------------------------------- + + +class VAEBottleneck(nn.Module): + """ + Variational Autoencoder (VAE) bottleneck. + + Applies VAE reparameterization trick during encoding. + """ + + def __init__(self) -> None: + super().__init__() + + def sample(self, mean: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + stdev = nn.functional.softplus(scale) + 1e-4 + var = stdev * stdev + logvar = torch.log(var) + latents = torch.randn_like(mean) * stdev + mean + + kl = (mean * mean + var - logvar - 1).sum(1).mean() + + return latents, kl + + def encode( + self, x: torch.Tensor, return_info: bool = False + ) -> Tuple[torch.Tensor, Dict[str, Any]]: + """ + Encode input through VAE bottleneck. + + Args: + x: Input tensor with shape [B, C*2, T] where C*2 contains + concatenated mean and scale parameters + return_info: Whether to return additional info dict + + Returns: + Sampled latents (and optionally info dict with KL divergence) + """ + info = {} + + mean, scale = x.chunk(2, dim=1) + x, kl = self.sample(mean, scale) + + info["kl"] = kl + + if return_info: + return x, info + else: + return x + + def decode( + self, x: torch.Tensor, return_info: bool = False + ) -> Tuple[torch.Tensor, Dict[str, Any]]: + """ + Decode from latents (identity operation for VAE). + + Args: + x: Latent tensor + return_info: Whether to return additional info dict + + Returns: + Latents (and optionally empty info dict) + """ + info = {} + if return_info: + return x, info + else: + return x + + +# --------------------------------------------------------------------------- +# 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 + if self.alpha_logscale: + self.alpha = Parameter(torch.zeros(in_features) * alpha) + self.beta = Parameter(torch.zeros(in_features) * alpha) + else: + self.alpha = Parameter(torch.ones(in_features) * alpha) + self.beta = Parameter(torch.ones(in_features) * alpha) + + self.alpha.requires_grad = alpha_trainable + self.beta.requires_grad = alpha_trainable + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # [B, C, T] + beta = 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) + + +# --------------------------------------------------------------------------- +# LayerNorm (fp32-safe) +# --------------------------------------------------------------------------- + + +class LayerNorm(nn.Module): + """LayerNorm with optional bias. Forces fp32 to avoid numerical issues.""" + + def __init__(self, size: int, eps: float = 1e-5, use_bias: bool = False) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(size)) + self.bias = nn.Parameter(torch.zeros(size)) if use_bias else None + self.eps = eps + + def forward(self, tensor: Tensor) -> Tensor: + dtype = tensor.dtype + with amp.autocast(enabled=True, dtype=torch.float32): + tensor = F.layer_norm(tensor, self.weight.shape, self.weight, self.bias, self.eps) + return tensor.to(dtype) + + +# --------------------------------------------------------------------------- +# ConvNeXt helpers +# --------------------------------------------------------------------------- + + +def zero_module(module: nn.Module) -> nn.Module: + """Zero out all parameters of a module (identity-friendly init).""" + for p in module.parameters(): + p.detach().zero_() + return module + + +def may_mask(x: Tensor, mask: Optional[Tensor] = None) -> Tensor: + """Apply optional mask tensor to activations.""" + if mask is not None: + x = x * mask + return x + + +# --------------------------------------------------------------------------- +# 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)) + + +# --------------------------------------------------------------------------- +# ConvNeXt block +# --------------------------------------------------------------------------- + + +class ConvNeXtBlock(nn.Module): + """ + ConvNeXt 1D Block adapted from https://github.com/charactr-platform/vocos + which is adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. + Supports causal and non-causal mode. + + Args: + dim (int): Number of input channels. + intermediate_dim (int): Dimensionality of the intermediate layer. + identity_init (bool): If True, initializes the 1x1 conv in residual paths to zero (identity-friendly). + use_snake (bool): If True, uses SnakeBeta activation; otherwise, GELU. + causal (bool): If True, applies causal padding; otherwise, applies symmetric padding for non-causal. + """ + + def __init__( + self, + dim: int, + intermediate_dim: int, + identity_init: bool = False, + use_snake: bool = False, + causal: bool = False, + ): + super().__init__() + self.causal = causal + + if causal: + self.dwconv = nn.Sequential( + nn.ConstantPad1d((6, 0), 0), + nn.Conv1d(dim, dim, kernel_size=7, groups=dim), + ) + else: + self.dwconv = nn.Sequential( + nn.ConstantPad1d((3, 3), 0), + nn.Conv1d(dim, dim, kernel_size=7, groups=dim), + ) + + self.norm = LayerNorm(dim) + self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) + self.act = SnakeBeta(intermediate_dim) if use_snake else nn.GELU() + + if identity_init: + self.pwconv2 = zero_module(nn.Conv1d(intermediate_dim, dim, 1)) + else: + self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1) + + def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor: + residual = x + x = self.dwconv(may_mask(x, mask)) + x = self.norm(x.permute(0, 2, 1)).permute(0, 2, 1) + x = self.pwconv1(x) + x = self.act(x) + x = self.pwconv2(x) + x = residual + x + return may_mask(x, mask) + + def remove_weight_norm(self) -> None: + """No weight norm is applied in ConvNeXtBlock.""" + pass + + +# --------------------------------------------------------------------------- +# 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 + + +def get_norm_module( + module: nn.Module, causal: bool = False, norm: str = "none", **norm_kwargs +) -> nn.Module: + assert norm in CONV_NORMALIZATIONS + if norm == "layer_norm": + assert isinstance(module, nn.modules.conv._ConvNd) + return nn.LayerNorm(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 get_extra_padding_for_conv1d( + x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 +) -> int: + length = x.shape[-1] + n_frames = (length - kernel_size + padding_total) / stride + 1 + ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) + return ideal_length - length + + +def pad1d(x: torch.Tensor, paddings: tuple, mode: str = "zero", 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 NormConv1d(nn.Module): + """Conv1d 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.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) + self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) + self.norm_type = norm + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.conv(x) + x = self.norm(x) + return x + + +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 SConv1d(nn.Module): + """Conv1d with builtin asymmetric/causal padding and normalization.""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: int, + stride: int = 1, + dilation: int = 1, + groups: int = 1, + bias: bool = True, + causal: bool = False, + norm: str = "none", + norm_kwargs: Dict[str, Any] = {}, + pad_mode: str = "reflect", + ): + super().__init__() + if stride > 1 and dilation > 1: + warnings.warn( + "SConv1d has been initialized with stride > 1 and dilation > 1" + f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." + ) + self.conv = NormConv1d( + in_channels, + out_channels, + kernel_size, + stride, + dilation=dilation, + groups=groups, + bias=bias, + causal=causal, + norm=norm, + norm_kwargs=norm_kwargs, + ) + self.causal = causal + self.pad_mode = pad_mode + + def forward(self, x: torch.Tensor) -> torch.Tensor: + kernel_size = self.conv.conv.kernel_size[0] + stride = self.conv.conv.stride[0] + dilation = self.conv.conv.dilation[0] + kernel_size = (kernel_size - 1) * dilation + 1 + padding_total = kernel_size - stride + extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) + if self.causal: + x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) + else: + padding_right = padding_total // 2 + padding_left = padding_total - padding_right + x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) + return self.conv(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..f6ed5cf7ea2a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -35,6 +35,7 @@ from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS from .guardrails import check_video_safety, download_guardrail_checkpoint from .transformer_cosmos3 import Cosmos3VFMTransformer +from .sound_tokenizer import LatentAutoEncoderV2 COSMOS3_DEFAULT_NEGATIVE_PROMPT = ( "The video captures a series of frames showing ugly scenes, static with no motion, motion blur, " @@ -66,6 +67,16 @@ class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): super().__init__(pipeline_config) + self.sound_gen = False + self.action_gen = False + if model_config.pretrained_config.sound_gen: + logger.info("Initializing Cosmos3OmniMoTPipeline with sound generation.") + self.sound_gen = True + + if model_config.pretrained_config.action_gen: + logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") + self.action_gen = True + def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") self.transformer = Cosmos3VFMTransformer(self.pipeline_config.model_configs["transformer"]) @@ -80,6 +91,13 @@ def load_standard_components( self, checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = [] ) -> None: skip_components = skip_components or [] + + if self.sound_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: + logger.info("Loading sound tokenizer...") + self.sound_tokenizer = LatentAutoEncoderV2.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.SOUND_TOKENIZER, + ) if PipelineComponent.TOKENIZER not in skip_components: logger.info("Loading tokenizer...") @@ -434,6 +452,39 @@ def _decode_latents(self, latents): video = postprocess_video_tensor(video) return video + # ========================================================================= + # Sound generation + # ========================================================================= + + def encode_sound(self, waveform: torch.Tensor) -> torch.Tensor: + """Encode audio waveform into latent tokens. + + Args: + waveform: Audio tensor of shape (C, N). A batch dim is added/removed + internally since AVAE expects (B, C, N). + Mono audio is duplicated to stereo if the tokenizer expects 2 channels. + """ + # Ensure correct number of channels (AVAE typically expects stereo) + expected_channels = self.sound_tokenizer.audio_channels + if waveform.shape[0] == 1 and expected_channels == 2: + waveform = waveform.repeat(2, 1) # mono → stereo + elif waveform.shape[0] > expected_channels: + waveform = waveform[:expected_channels] + # AVAE expects (B, C, N) + latent = self.sound_tokenizer.encode(waveform.unsqueeze(0)) # [1,sound_channels,T_sound] + return latent.squeeze(0) # [sound_channels,T_sound] + + def decode_sound(self, latent: torch.Tensor) -> torch.Tensor: + """Decode sound latent tokens back to waveform. + + Args: + latent: Sound latent tensor of shape (C, T). A batch dim is added/removed + internally since AVAE expects (B, C, T). + """ + # AVAE expects (B, C, T) + waveform = self.sound_tokenizer.decode(latent.unsqueeze(0)) # [1,audio_channels,N_samples] + return waveform.squeeze(0) # [audio_channels,N_samples] + # ========================================================================= # Forward (main generation entry point) # ========================================================================= 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..b3c97b9c5814 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -0,0 +1,955 @@ +# 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 functools import partial +from typing import Any, Callable, Dict, Literal, Optional + +import torch +import torch.nn.functional as F +from torch import Tensor, nn +from torch.nn.utils import remove_weight_norm +from torch.nn.utils.parametrize import remove_parametrizations + +from .modules import ( + ConvNeXtBlock, + SConv1d, + SConvTranspose1d, + SnakeBeta, + VAEBottleneck, + WNConv1d, + WNConvTranspose1d, +) + + +def get_activation( + activation: Literal["elu", "snake", "none"], + antialias: bool = False, + channels: Optional[int] = None, + use_cuda_kernel: bool = False, +) -> nn.Module: + """ + Get activation module by name. + + Args: + activation: Activation type ('elu', 'snake', 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) + + Returns: + Activation module + """ + if activation == "elu": + act = nn.ELU() + elif activation == "snake": + act = SnakeBeta(channels) + elif activation == "none": + act = nn.Identity() + else: + raise ValueError(f"Unknown activation {activation}") + + 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", + ) -> 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 + + self.layers = nn.Sequential( + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ), + WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=kernel_size, + dilation=dilation, + padding=self.padding, + padding_mode=self.padding_mode, + ), + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ), + 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) + """ + res = x + + # apply conv layers + x = self.layers(x) + + if self.causal: + # Trim right padding to get the causal output + x = x[:, :, : -self.padding] + + return x + res + + +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", + ) -> None: + super().__init__() + + self.causal = causal + + self.layers = nn.Sequential( + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=in_channels, + ), + self._create_upsample_layer( + in_channels, out_channels, stride, use_nearest_upsample, causal, padding_mode + ), + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=1, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, + ), + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=3, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, + ), + ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=9, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, + ), + ) + + 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) + """ + return self.layers(x) + + def remove_weight_norm(self) -> None: + """Remove weight normalization from all layers.""" + + for layer in self.layers: + try: + remove_weight_norm(layer) + except (ValueError, AttributeError): + # Layer doesn't have weight norm or is not a module with weight norm + 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 = model_config["input_channels"] + if model_config.get("stereo", False): + out_channels *= 2 + + 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"] + + c_mults = [1, *c_mults] + + self.depth = len(c_mults) + + # Padding for the first convolution layer + self.first_padding = 6 if causal else 3 + first_conv = WNConv1d( + in_channels=latent_dim, + out_channels=c_mults[-1] * channels, + kernel_size=7, + padding=self.first_padding, + padding_mode=padding_mode, + ) + + if causal: + first_conv = nn.Sequential(first_conv, TrimPadding(self.first_padding)) + + layers = [first_conv] + + for i in range(self.depth - 1, 0, -1): + layers += [ + OobleckDecoderBlock( + in_channels=c_mults[i] * channels, + out_channels=c_mults[i - 1] * channels, + stride=strides[i - 1], + use_snake=use_snake, + antialias_activation=antialias_activation, + use_nearest_upsample=use_nearest_upsample, + causal=causal, + padding_mode=padding_mode, + ) + ] + + # Padding for the final convolution layer + self.final_padding = 6 if causal else 3 + final_conv = WNConv1d( + in_channels=c_mults[0] * channels, + out_channels=out_channels, + kernel_size=7, + padding=self.final_padding, + padding_mode=padding_mode, + bias=False, + ) + + if causal: + final_conv = nn.Sequential(final_conv, TrimPadding(self.final_padding)) + + layers += [ + get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=c_mults[0] * channels, + ), + final_conv, + nn.Tanh() if final_tanh else nn.Identity(), + ] + + self.layers = nn.Sequential(*layers) + + def forward(self: "OobleckDecoder", x: torch.Tensor) -> torch.Tensor: + x = self.layers(x) + return x + + def remove_weight_norm(self: "OobleckDecoder") -> None: + print("INFO: Removing all weight norm from OobleckDecoder") + for module in self.modules(): + if hasattr( + module, "parametrizations" + ): # for new WN implementation using parameterizations + try: + remove_parametrizations(module, "weight") + except ValueError: + msg = ( + f"[WARNING] No weight norm found in {module} with parameterizations. " + "You can ignore this if you know that this module does not apply weight norm." + ) + print(msg) + elif hasattr(module, "weight"): + try: + remove_weight_norm(module) + except ValueError: + pass + + +class SpectrogramConvNeXtEncoder(nn.Module): + """ + Spectrogram Encoder with ConvNeXtBlocks + + This encoder processes input waveforms by converting them into spectrograms + (magnitude and phase concatenated along the channel dimension) and encodes them + using a sequence of ConvNeXtBlocks and downsampling layers. + + Args (mapped from h): + in_channels (int): Number of input audio channels (1 for mono, 2 for stereo). + channels (int): Base number of channels for the encoder. + latent_dim (int): Dimensionality of the final latent representation. + c_mults (List[int]): Channel multipliers at each depth of the encoder. + strides (List[int]): Downsampling strides for each depth. + num_blocks (int): Number of ConvNeXtBlocks to stack per depth. + identity_init (bool): Whether to initialize the 1x1 convs in residual paths as zeros. + n_fft (int): Number of FFT points for spectrogram computation. + hop_length (int): Hop length for the STFT. + use_snake (bool): Whether to use Snake activation in ConvNeXtBlocks. + causal (bool): If True, uses causal convolutions. + padding_mode (str): Padding mode for convolutions (default: 'zeros'). + + Inputs: + x (torch.Tensor): Input waveform tensor of shape `[batch, in_channels, time]`. + + Outputs: + torch.Tensor: Encoded representation of shape `[batch, time_out, latent_dim]`. + + Forward Pass: + - Converts waveform input into spectrograms (concatenates magnitude and phase). + - Processes the spectrogram through stacked ConvNeXtBlocks and downsampling layers. + - Outputs the final latent representation of specified dimensionality. + + Example: + encoder = SpectrogramConvNeXtEncoder( + in_channels=2, channels=256, latent_dim=128, c_mults=[1, 2, 4], strides=[4, 4, 8] + ) + waveform = torch.randn(8, 2, 65536) # [batch, channels, time] + encoded = encoder(waveform) # Output: [8, time_out, 128] + + NOTE: output is in [B, T, C] to be consistent with other encoders + """ + + def __init__(self, model_config: Dict[str, Any]) -> None: + super().__init__() + self.model_config = model_config + + self.in_channels = model_config["input_channels"] + if model_config.get("stereo", False): + self.in_channels *= 2 + + # if "enc_latent_dim" is found in v2 config, set it as latent_dim + if "enc_latent_dim" in model_config: + self.latent_dim = model_config["enc_latent_dim"] + else: + # if not found, fallback to v1 logic + self.latent_dim = model_config["vocoder_input_dim"] + if model_config["model_type"] == "vae": + self.latent_dim *= 2 + + self.channels = model_config["enc_dim"] + + self.c_mults = model_config["enc_c_mults"] + self.strides = model_config["enc_strides"] + self.num_blocks = model_config["enc_num_blocks"] + self.identity_init = model_config["enc_identity_init"] + self.causal = model_config["causal"] + self.padding_mode = model_config["padding_mode"] + + self.use_snake = model_config["enc_use_snake"] + + # Basic checks + assert len(self.c_mults) == len(self.strides), ( + f"The length of c_mults and strides must match. Got {len(self.c_mults)} vs {len(self.strides)}." + ) + + # Spectrogram function + self.n_fft = model_config["enc_n_fft"] + self.hop_length = model_config["enc_hop_length"] + self.spectrogram_fn = partial( + self.spectrogram, + n_fft=self.n_fft, + hop_length=self.hop_length, + win_length=self.n_fft, + window_fn=torch.hann_window, + ) + + # --------------------------------------------------------------------- + # 1) Initial projection (similar to the first_conv in OobleckEncoder), + # but here we typically use a 1x1 conv for a "spectrogram style" input. + # --------------------------------------------------------------------- + layers = [] + layers.append( + WNConv1d( + (self.n_fft + 2) * self.in_channels, + self.c_mults[0] * self.channels, + kernel_size=1, + bias=False, + ) + ) + + # --------------------------------------------------------------------- + # 2) Stages: For each i in range(len(c_mults)): + # - Stack num_blocks of ConvNeXtBlock + # - Downsample via stride convolution + # --------------------------------------------------------------------- + for i in range(len(self.c_mults)): + dim_in = self.c_mults[i] * self.channels + # Determine output dimension for the block + if i < len(self.c_mults) - 1: # If not the last block + dim_out = self.c_mults[i + 1] * self.channels + else: # For the last block, dim_out is c_mults[-1] * channels + dim_out = self.c_mults[-1] * self.channels + ds_rate = self.strides[i] + + # (a) Repeated ConvNeXtBlocks + for _ in range(self.num_blocks): + layers.append( + ConvNeXtBlock( + dim=dim_in, + intermediate_dim=dim_in * 4, + identity_init=self.identity_init, + use_snake=self.use_snake, + causal=self.causal, + ) + ) + + # (b) Downsampling convolution + layers.append( + self._create_downsample_layer( + dim_in, dim_out, ds_rate, self.causal, self.padding_mode + ) + ) + + # --------------------------------------------------------------------- + # 3) Final projection from the last channel dimension to latent_dim. + # --------------------------------------------------------------------- + layers.append( + WNConv1d(self.c_mults[-1] * self.channels, self.latent_dim, kernel_size=1, bias=False) + ) + + self.layers = nn.Sequential(*layers) + + def spectrogram( + self: "SpectrogramConvNeXtEncoder", + wav: Tensor, + n_fft: int, + hop_length: int, + win_length: int, + window_fn: Callable[[int], torch.Tensor] = torch.hann_window, + ) -> Tensor: + """ + wav: [batch_size?, time_steps], where batch_size? is an optional batch dimension + """ + pad_size_l = (n_fft - hop_length) // 2 + pad_size_r = (n_fft - hop_length) - pad_size_l + with torch.autocast(device_type=wav.device.type, enabled=False): + wav = F.pad(wav, (pad_size_l, pad_size_r)).float() + spec = torch.stft( + wav, + n_fft, + hop_length=hop_length, + win_length=win_length, + window=window_fn(win_length).to(wav), + center=False, + normalized=False, + onesided=True, + return_complex=True, + ) + return spec + + def _create_downsample_layer( + self: "SpectrogramConvNeXtEncoder", + in_channels: int, + out_channels: int, + stride: int, + causal: bool, + padding_mode: str, + ) -> nn.Module: + if causal: + downsample_layer = SConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + causal=True, + norm="weight_norm", + ) + else: # original non-causal implementation + downsample_layer = WNConv1d( + in_channels=in_channels, + out_channels=out_channels, + kernel_size=2 * stride, + stride=stride, + padding=math.ceil(stride / 2), + padding_mode=padding_mode, + ) + return downsample_layer + + def forward(self, x: torch.Tensor) -> torch.Tensor: + """ + Forward pass: + x: waveform in [batch, in_channels, length] (mono: in_channels=1, stereo: in_channels=2) + Returns: encoder output in [batch, length_out, dim_latent], + where the spectrogram's magnitude and phase are concatenated along the channel dimension. + """ + + # Handle stereo input by merging channel dim into batch dim + batch, channels, length = x.shape + if channels > 1: # Stereo case + x = x.reshape(batch * channels, 1, length) # [batch * channels, 1, length] + + # Compute the spectrogram + with torch.autocast(device_type=x.device.type, enabled=False): + spec = self.spectrogram_fn( + x.float().squeeze(1) + ) # Remove the channel dimension for STFT + mag, ph = torch.view_as_real(spec).chunk(2, dim=-1) # Split real and imaginary parts + spectrogram = torch.cat([mag, ph], dim=1).squeeze( + -1 + ) # Concatenate along channel dim: [batch * channels, freq, frame] + + # Cast spectrogram back to original dtype + spectrogram = spectrogram.to(x.dtype) + + # Restore stereo structure if needed + if channels > 1: # Stereo case + freq = spectrogram.shape[1] # Get the frequency dimension + spectrogram = spectrogram.reshape( + batch, channels * freq, *spectrogram.shape[2:] + ) # [batch, freq * channels, frame] + + # forward pass the encoder + output = self.layers(spectrogram) + + return output.transpose(1, 2) # [B, T, C] + + def remove_weight_norm(self: "SpectrogramConvNeXtEncoder") -> None: + print("INFO: Removing all weight norm from SpectrogramConvNeXtEncoder") + for module in self.modules(): + if hasattr( + module, "parametrizations" + ): # for new WN implementation using parameterizations + try: + remove_parametrizations(module, "weight") + except ValueError: + print( + f"[WARNING] No weight norm found in {module} with parameterizations. " + "You can ignore this if you know that this module does not apply weight norm." + ) + elif hasattr(module, "weight"): + try: + remove_weight_norm(module) + except ValueError: + pass + + +class LatentAutoEncoderV2(nn.Module): + """ + A Latent AutoEncoder class with cleaner implementation to generalize using bottleneck.py + + Attributes: + h: Configuration object containing model hyperparameters. + encoder (nn.Module): The encoder module based on configuration. + bottleneck (VAEBottleneck): VAE Bottleneck module. + decoder (nn.Module): The decoder module based on configuration. + """ + + def __init__(self, model_config: Dict[str, Any]) -> None: + super().__init__() + self.model_config = model_config + + # Set up basic model properties + self.stereo = model_config.get("stereo", False) + + # Determine input type + self.input_type = None + if model_config.get("use_wav_as_input", False): + print("INFO: Encoder's input feature is waveform") + self.input_type = "waveform" + model_config["input_channels"] = 1 + elif model_config.get("use_linear_spec_as_input", False): + print("INFO: Encoder's input feature is linear") + self.input_type = "linear" + model_config["input_channels"] = model_config["num_linears"] + elif model_config.get("use_discrete_code_as_input", False): + print("INFO: Encoder's input feature is discrete_code") + self.input_type = "discrete_code" + model_config["input_channels"] = 1 + else: + print("INFO: Encoder's input feature is mel") + self.input_type = "mel" + model_config["input_channels"] = model_config["num_mels"] + + # hop_size defines the down/up sampling factor of the autoencoder + self.hop_size = model_config["hop_size"] + + # Initialize encoder + self.enc_type = model_config.get("enc_type", "convnext") + print(f"INFO: Using {self.enc_type} as encoder") + + # Define encoder (only spec_convnext supported in cleaned version) + if self.enc_type == "spec_convnext": + self.encoder = SpectrogramConvNeXtEncoder(model_config) + else: + raise NotImplementedError( + f"Encoder type '{self.enc_type}' not supported in cleaned AVAE. Only 'spec_convnext' is supported." + ) + + # Initialize encoder projector (Identity for spec_convnext) + self.encoder_proj = nn.Identity() + + if "bottleneck" in model_config: + self.bottleneck = VAEBottleneck() + print(f"INFO: Created bottleneck of type {model_config['bottleneck']['type']}") + else: + raise ValueError("Bottleneck configuration must be specified") + + # Check for encoder-only mode + self.encoder_only = model_config.get("encoder_only", False) + + if not self.encoder_only: + # Initialize decoder + self.dec_type = model_config.get("dec_type", "oobleck") + print(f"INFO: Using {self.dec_type} as decoder") + if self.dec_type == "oobleck": + self.decoder = OobleckDecoder(model_config) + else: + raise NotImplementedError( + f"Decoder type '{self.dec_type}' not supported in cleaned AVAE. Only 'oobleck' is supported." + ) + else: + # Skip decoder initialization + self.decoder = None + print("INFO: Running in encoder-only mode, decoder is set to None") + + # Whether to freeze encoder + self.freeze_encoder = model_config.get("freeze_encoder", False) + if self.freeze_encoder: + print( + "WARNING: freeze_encoder set to true. The encoder will not be updated during training!" + ) + for param in self.encoder.parameters(): + param.requires_grad = False + + # Optional latent normalisation (from cosmos3-internal AVAEModel) + 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) + + # --- weight loading (mirrors cosmos3-internal AVAEModel._load_avae_model) --- + state_dict: Optional[Dict[str, Any]] = None + + # 1. safetensors (standard TRT-LLM / HF format) + sft_candidates = ["model.safetensors", "diffusion_pytorch_model.safetensors"] + for name in sft_candidates: + 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 + + # 2. PyTorch bin + if state_dict is None: + bin_candidates = ["pytorch_model.bin", "diffusion_pytorch_model.bin"] + for name in bin_candidates: + path = os.path.join(checkpoint_dir, name) + if os.path.exists(path): + state_dict = torch.load(path, map_location="cpu", weights_only=True) + break + + if state_dict is None: + raise FileNotFoundError( + f"No weight file found in '{checkpoint_dir}'. " + "Expected model.safetensors, pytorch_model.bin, or *.ckpt." + ) + + missing, unexpected = model.load_state_dict(state_dict, strict=False) + if missing: + print(f"[WARNING] Missing keys when loading sound tokenizer: {missing}") + if unexpected: + print(f"[WARNING] 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 calculate_latent_lengths( + self: "LatentAutoEncoderV2", audio_lengths: torch.Tensor + ) -> torch.Tensor: + """ + Calculates the latent lengths given the original audio lengths. + + Args: + audio_lengths (torch.Tensor): A tensor of shape [B] containing the lengths of the original audio samples. + + Returns: + torch.Tensor: A tensor of shape [B] containing the corresponding latent lengths. + """ + if self.input_type == "waveform": + # The latent length is the audio length divided by the hop_size + latent_lengths = torch.ceil(audio_lengths.float() / self.hop_size).long() + else: + # The latent length is same as audio_lengths + latent_lengths = audio_lengths + + return latent_lengths + + def forward(self: "LatentAutoEncoderV2", x: torch.Tensor) -> dict[str, torch.Tensor]: + """ + Forward pass through the model. + + Args: + x (torch.Tensor): Input tensor to the model with shape [B, C, T]. + + Returns: + dict[str, torch.Tensor]: Dictionary of output tensors including: + - encoder_out: Raw encoder output + - latent: Bottleneck latent representation + - decoder_out: Decoded output (if decoder exists) + - Additional outputs specific to the bottleneck type + """ + return_dict = {} + + # Encoder + encoder_out = self.encoder(x) # Shape: [B, T_frame, encoder_out_dim] + encoder_out_proj = self.encoder_proj(encoder_out) # Shape: [B, T_frame, encoder_proj_dim] + + # Apply bottleneck after reshaping to [B, C, T] again + latent, bottleneck_enc_info = self.bottleneck.encode( + encoder_out_proj.transpose(1, 2), return_info=True + ) + + # Update return dictionary + return_dict.update({"encoder_out": encoder_out.transpose(1, 2), "latent": latent}) + # Add bottleneck-specific info to return dict + for k, v in bottleneck_enc_info.items(): + return_dict[k] = v + + # Decode (if decoder exists) + if self.decoder is not None: + # Apply bottleneck decode + decoded_latent, bottleneck_dec_info = self.bottleneck.decode(latent, return_info=True) + # Apply decoder + decoder_out = self.decoder(decoded_latent) + + # Update return dictionary + return_dict["decoder_out"] = decoder_out + # Add bottleneck-specific info to return dict + for k, v in bottleneck_dec_info.items(): + return_dict[k] = v + + return return_dict + + def encode(self: "LatentAutoEncoderV2", x: torch.Tensor) -> torch.Tensor: + """ + Encode waveform to latent tokens. + + Args: + x: Input tensor with shape [B, C, T]. + + Returns: + Latent tensor [B, latent_ch, T_latent]. Normalised by latent_mean/std + when configured (mirrors cosmos3-internal AVAEModel.encode). + """ + encoder_out = self.encoder(x) + encoder_out_proj = self.encoder_proj(encoder_out) + latent = self.bottleneck.encode(encoder_out_proj.transpose(1, 2)) + + if self.latent_mean is not None and self.latent_std is not None: + latent = (latent - self.latent_mean) / self.latent_std + + return latent + + 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 + + decoded_latent = self.bottleneck.decode(latent) + return self.decoder(decoded_latent) + + @property + def audio_channels(self) -> int: + """Number of output audio channels (1 = mono, 2 = stereo).""" + return self.model_config.get("dec_out_channels", 1) + + def remove_weight_norm(self: "LatentAutoEncoderV2") -> None: + """Remove weight normalization from all components.""" + self.encoder.remove_weight_norm() + 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..6b3734535089 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -665,6 +665,8 @@ class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) pretrained_config = model_config.pretrained_config + self.sound_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 +686,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.sound_gen: + self.sound_dim = pretrained_config.sound_dim + self.sound_latent_fps = pretrained_config.sound_latent_fps + self.temporal_compression_factor_sound = ( + 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" From bc975348f92b1628e40493e1a7dabc13c0bb91f8 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 15 May 2026 12:57:07 -0700 Subject: [PATCH 02/20] full sound pipeline Signed-off-by: Shreyas Misra --- .../models/cosmos3/pipeline_cosmos3.py | 84 +++++++++-- .../models/cosmos3/transformer_cosmos3.py | 131 +++++++++++++++++- 2 files changed, 199 insertions(+), 16 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index f6ed5cf7ea2a..0a2a992aa300 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -132,6 +132,10 @@ def load_standard_components( checkpoint_dir, subfolder=PipelineComponent.SCHEDULER, ) + if self.sound_gen: + # Separate instance so video and sound scheduler states don't collide + # (UniPC mutates internal correction buffers on every .step() call). + self.sound_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" @@ -478,12 +482,12 @@ def decode_sound(self, latent: torch.Tensor) -> torch.Tensor: """Decode sound latent tokens back to waveform. Args: - latent: Sound latent tensor of shape (C, T). A batch dim is added/removed - internally since AVAE expects (B, C, T). + latent: Sound latent tensor of shape (B, C, T). + + Returns: + Waveform tensor of shape (B, audio_channels, N_samples). """ - # AVAE expects (B, C, T) - waveform = self.sound_tokenizer.decode(latent.unsqueeze(0)) # [1,audio_channels,N_samples] - return waveform.squeeze(0) # [audio_channels,N_samples] + return self.sound_tokenizer.decode(latent) # [B, audio_channels, N_samples] # ========================================================================= # Forward (main generation entry point) @@ -621,6 +625,27 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) + # 3b. Sound noise init + # T_sound = ceil(duration_s * sound_latent_fps / temporal_compression_factor_sound) + # Duration derived from num_frames / frame_rate; matches cosmos3-internal. + do_sound = self.sound_gen and hasattr(self, "sound_tokenizer") + sound_latents = None + if do_sound: + duration_s = num_frames / frame_rate + T_sound = math.ceil( + duration_s + * self.transformer.sound_latent_fps + / self.transformer.temporal_compression_factor_sound + ) + sound_latents = randn_tensor( + (1, self.transformer.sound_dim, T_sound), + generator=generator, + device=self.device, + dtype=latents.dtype, + ) + # Sound uses the same scheduler type/config as video. + self.sound_scheduler.set_timesteps(num_inference_steps, device=self.device) + # 4. Build forward_fn for the denoise loop def forward_fn( latent_input, @@ -635,7 +660,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_sound = extra_stream_latents.get("sound") if extra_stream_latents else None + + result = self.transformer( hidden_states=latent_input, timestep=timestep, attention_timestep=timestep / self.scheduler.config.num_train_timesteps, @@ -644,10 +671,18 @@ def forward_fn( video_shape=video_shape, fps=frame_rate, noisy_frame_mask=velocity_mask, + sound_latents=current_sound, ) + + video_noise_pred = result.video + sound_noise_pred = result.sound + if velocity_mask is not None: - noise_pred = noise_pred * velocity_mask - return noise_pred + video_noise_pred = video_noise_pred * velocity_mask + + if sound_noise_pred is not None: + return video_noise_pred, {"sound": sound_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 @@ -661,7 +696,8 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - latents = self.denoise( + extra_streams = {"sound": (sound_latents, self.sound_scheduler)} if do_sound else None + denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, prompt_embeds=cond_ids, # placeholder — actual conditioning via extra_cfg_tensors @@ -669,10 +705,19 @@ def forward_fn( guidance_scale=guidance_scale, forward_fn=forward_fn, extra_cfg_tensors=extra_cfg_tensors, + extra_streams=extra_streams, ) + + if extra_streams is not None: + latents, extra_latents = denoise_result + sound_latents = extra_latents.get("sound") + else: + latents = denoise_result + sound_latents = None + timer.mark_post_start() - # 7. Decode + # 7. Decode video logger.info("Decoding video...") decode_start = time.time() @@ -682,7 +727,13 @@ def forward_fn( video = self.decode_latents(latents, self._decode_latents) - # Video guardrails + # 7b. Decode sound + waveform = None + if do_sound and sound_latents is not None: + logger.info("Decoding sound...") + waveform = self.decode_sound(sound_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") @@ -691,4 +742,13 @@ def forward_fn( video = check_video_safety(video, self.safety_checker) timer.mark_end() - return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate)) + return timer.fill( + PipelineOutput( + video=video, + frame_rate=frame_rate, + audio=waveform, + audio_sample_rate=self.sound_tokenizer.model_config["sampling_rate"] + if waveform is not None + else None, + ) + ) 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 6b3734535089..3c77d53a7fbe 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).""" + + sound: Optional[torch.Tensor] = None + """[B, sound_dim, T_sound] sound 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, @@ -731,6 +749,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.sound_gen: + # Projections for sound modality (mirrors cosmos3-internal Cosmos3VFMNetwork) + self.sound2llm = nn.Linear(self.sound_dim, self.hidden_size) + self.llm2sound = nn.Linear(self.hidden_size, self.sound_dim) + self.sound_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) @@ -867,6 +891,59 @@ def _compute_rope_freqs( freqs_gen = (cos_gen.unsqueeze(2), sin_gen.unsqueeze(2)) return freqs_und, freqs_gen + # ------------------------------------------------------------------------- + # Sound helpers + # ------------------------------------------------------------------------- + + def _compute_sound_rope_freqs( + self, + T_sound: int, + text_mask: torch.Tensor, + fps_sound: float, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute mRoPE cos/sin for sound tokens. + + Sound tokens use a 1×1 spatial grid (H=W=1) aligned with the vision + temporal axis at the sound latent rate. This mirrors the cosmos3-internal + ``sequence_packing.py`` treatment where sound 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() + + sound_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) + # Sound tokens share the vision temporal space; use modality margin offset. + s_pos, _ = compute_mrope_position_ids_vision( + T_sound, + 1, # grid_h + 1, # grid_w + temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, + fps=fps_sound, + base_fps=self.base_fps, + temporal_compression_factor=1, # sound latent is already at sound_latent_fps + enable_fps_modulation=self.enable_fps_modulation, + ) + sound_pos_list.append(s_pos) + + sound_pos_ids = torch.stack(sound_pos_list, dim=1).to(device) # [3, B, T_sound] + rotary_emb = self.language_model.rotary_emb + _dummy = torch.tensor([], dtype=dtype, device=device) + cos_s, sin_s = rotary_emb(_dummy, position_ids=sound_pos_ids) + return cos_s.unsqueeze(2), sin_s.unsqueeze(2) # [B, T_sound, 1, head_dim] + + def pack_sound_latents(self, sound_latents: torch.Tensor) -> torch.Tensor: + """[B, sound_dim, T_sound] → [B, T_sound, sound_dim].""" + return sound_latents.permute(0, 2, 1) + + def unpack_sound_latents(self, hidden_sound: torch.Tensor) -> torch.Tensor: + """[B, T_sound, sound_dim] → [B, sound_dim, T_sound].""" + return hidden_sound.permute(0, 2, 1) + def reset_cache(self): self.cached_kv = None self.cached_freqs_gen = None @@ -881,8 +958,9 @@ def forward( video_shape: Optional[Tuple[int, int, int]] = None, fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, + sound_latents: Optional[torch.Tensor] = None, **kwargs, - ) -> torch.Tensor: + ) -> "TransformerOutput": """ Forward pass for parallel denoising. @@ -900,9 +978,15 @@ def forward( timestep embedding, predict velocity) and 0=conditioned (clean context, skip timestep embedding). None means all frames noisy (T2V mode). + sound_latents: Optional [B, sound_dim, T_sound] noisy sound latents. + When provided, sound tokens are appended to the generation + sequence and a sound velocity is returned alongside the video + velocity. Requires ``sound_gen=True`` in the pretrained config. Returns: - [B, C, T, H, W] velocity prediction + TransformerOutput with video (and image alias) always set. + sound is set to the predicted sound velocity when sound_latents is + provided; otherwise None. action is always None for now. """ del kwargs # Kept for diffusers API compatibility. T, H, W = video_shape @@ -966,9 +1050,35 @@ def forward( else: self.cached_kv = cached_kv_full + # --- Sound token injection ------------------------------------------------- + T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp + T_sound = 0 + if sound_latents is not None and self.sound_gen: + T_sound = sound_latents.shape[2] + hidden_sound = self.pack_sound_latents(sound_latents).to(hidden_gen.dtype) + hidden_sound = self.sound2llm(hidden_sound) + self.sound_modality_embed + hidden_sound = hidden_sound + time_embed.unsqueeze(1) + cos_s, sin_s = self._compute_sound_rope_freqs( + T_sound, + text_mask, + float(self.sound_latent_fps), + hidden_states.device, + hidden_gen.dtype, + ) + # [B, T_vid+T_sound, hidden_size] + hidden_gen = torch.cat([hidden_gen, hidden_sound], dim=1) + cos_v, sin_v = self.cached_freqs_gen + freqs_gen_combined = ( + torch.cat([cos_v, cos_s], dim=1), + torch.cat([sin_v, sin_s], 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) @@ -989,7 +1099,20 @@ def forward( 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 sound velocity (if requested) --------------------------------- + sound_vel = None + if T_sound > 0 and sound_latents is not None and self.sound_gen: + # hidden_gen[:, T_vid_tokens:] → [B, T_sound, hidden_size] + # → llm2sound → [B, T_sound, sound_dim] → unpack → [B, sound_dim, T_sound] + sound_vel = self.unpack_sound_latents( + self.llm2sound(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_sound]) + ) + + return TransformerOutput(video=video_vel, image=video_vel, sound=sound_vel) def load_weights(self, weights: dict) -> None: """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. From 006eea62738ea603a56dac09d8b65b5cc6e38926 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 15 May 2026 20:44:08 +0000 Subject: [PATCH 03/20] working implementation with sound Signed-off-by: Shreyas Misra --- .../models/cosmos3/pipeline_cosmos3.py | 15 ++++--- .../models/cosmos3/sound_tokenizer.py | 42 +++---------------- .../models/cosmos3/transformer_cosmos3.py | 9 +++- tensorrt_llm/_torch/visual_gen/pipeline.py | 4 +- 4 files changed, 25 insertions(+), 45 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 0a2a992aa300..46979e7966dd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -69,11 +69,11 @@ def __init__(self, pipeline_config): self.sound_gen = False self.action_gen = False - if model_config.pretrained_config.sound_gen: + if getattr(model_config.pretrained_config, "sound_gen", False): logger.info("Initializing Cosmos3OmniMoTPipeline with sound generation.") self.sound_gen = True - if model_config.pretrained_config.action_gen: + if getattr(model_config.pretrained_config, "action_gen", False): logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") self.action_gen = True @@ -94,9 +94,14 @@ def load_standard_components( if self.sound_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: logger.info("Loading sound tokenizer...") - self.sound_tokenizer = LatentAutoEncoderV2.from_pretrained( - checkpoint_dir, - subfolder=PipelineComponent.SOUND_TOKENIZER, + self.sound_tokenizer = ( + LatentAutoEncoderV2.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.SOUND_TOKENIZER, + ) + .to(device) + .to(self.dtype) + .eval() ) if PipelineComponent.TOKENIZER not in skip_components: diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py index b3c97b9c5814..a8e337a00e1f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -25,6 +25,8 @@ 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 ( ConvNeXtBlock, SConv1d, @@ -406,19 +408,11 @@ def forward(self: "OobleckDecoder", x: torch.Tensor) -> torch.Tensor: return x def remove_weight_norm(self: "OobleckDecoder") -> None: - print("INFO: Removing all weight norm from OobleckDecoder") for module in self.modules(): if hasattr( module, "parametrizations" ): # for new WN implementation using parameterizations - try: - remove_parametrizations(module, "weight") - except ValueError: - msg = ( - f"[WARNING] No weight norm found in {module} with parameterizations. " - "You can ignore this if you know that this module does not apply weight norm." - ) - print(msg) + remove_parametrizations(module, "weight") elif hasattr(module, "weight"): try: remove_weight_norm(module) @@ -664,18 +658,11 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return output.transpose(1, 2) # [B, T, C] def remove_weight_norm(self: "SpectrogramConvNeXtEncoder") -> None: - print("INFO: Removing all weight norm from SpectrogramConvNeXtEncoder") for module in self.modules(): if hasattr( module, "parametrizations" ): # for new WN implementation using parameterizations - try: - remove_parametrizations(module, "weight") - except ValueError: - print( - f"[WARNING] No weight norm found in {module} with parameterizations. " - "You can ignore this if you know that this module does not apply weight norm." - ) + remove_parametrizations(module, "weight") elif hasattr(module, "weight"): try: remove_weight_norm(module) @@ -704,19 +691,15 @@ def __init__(self, model_config: Dict[str, Any]) -> None: # Determine input type self.input_type = None if model_config.get("use_wav_as_input", False): - print("INFO: Encoder's input feature is waveform") self.input_type = "waveform" model_config["input_channels"] = 1 elif model_config.get("use_linear_spec_as_input", False): - print("INFO: Encoder's input feature is linear") self.input_type = "linear" model_config["input_channels"] = model_config["num_linears"] elif model_config.get("use_discrete_code_as_input", False): - print("INFO: Encoder's input feature is discrete_code") self.input_type = "discrete_code" model_config["input_channels"] = 1 else: - print("INFO: Encoder's input feature is mel") self.input_type = "mel" model_config["input_channels"] = model_config["num_mels"] @@ -725,7 +708,6 @@ def __init__(self, model_config: Dict[str, Any]) -> None: # Initialize encoder self.enc_type = model_config.get("enc_type", "convnext") - print(f"INFO: Using {self.enc_type} as encoder") # Define encoder (only spec_convnext supported in cleaned version) if self.enc_type == "spec_convnext": @@ -740,7 +722,6 @@ def __init__(self, model_config: Dict[str, Any]) -> None: if "bottleneck" in model_config: self.bottleneck = VAEBottleneck() - print(f"INFO: Created bottleneck of type {model_config['bottleneck']['type']}") else: raise ValueError("Bottleneck configuration must be specified") @@ -750,7 +731,6 @@ def __init__(self, model_config: Dict[str, Any]) -> None: if not self.encoder_only: # Initialize decoder self.dec_type = model_config.get("dec_type", "oobleck") - print(f"INFO: Using {self.dec_type} as decoder") if self.dec_type == "oobleck": self.decoder = OobleckDecoder(model_config) else: @@ -760,16 +740,6 @@ def __init__(self, model_config: Dict[str, Any]) -> None: else: # Skip decoder initialization self.decoder = None - print("INFO: Running in encoder-only mode, decoder is set to None") - - # Whether to freeze encoder - self.freeze_encoder = model_config.get("freeze_encoder", False) - if self.freeze_encoder: - print( - "WARNING: freeze_encoder set to true. The encoder will not be updated during training!" - ) - for param in self.encoder.parameters(): - param.requires_grad = False # Optional latent normalisation (from cosmos3-internal AVAEModel) self.latent_mean = model_config.get("latent_mean", None) @@ -822,9 +792,9 @@ def from_pretrained( missing, unexpected = model.load_state_dict(state_dict, strict=False) if missing: - print(f"[WARNING] Missing keys when loading sound tokenizer: {missing}") + logger.warning(f"Missing keys when loading sound tokenizer: {missing}") if unexpected: - print(f"[WARNING] Unexpected keys when loading sound tokenizer: {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() 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 3c77d53a7fbe..48e436f2bf2d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1136,7 +1136,9 @@ def load_weights(self, weights: dict) -> None: if k.startswith(skip_prefixes): continue - if k.startswith(("vae2llm.", "llm2vae.")): + if k.startswith( + ("vae2llm.", "llm2vae.", "sound2llm.", "llm2sound.", "sound_modality_embed") + ): remapped[k] = value continue @@ -1277,6 +1279,11 @@ def post_load_weights(self) -> None: self.vae2llm.to(target_dtype) self.llm2vae.to(target_dtype) + if self.sound_gen: + self.sound2llm.to(target_dtype) + self.llm2sound.to(target_dtype) + self.sound_modality_embed.data = self.sound_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..f97ab68481d0 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.model_config.torch_dtype @property def device(self): From 05cb320aa4462a01c7e0ce4d495ed663e93bb4b2 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 19 May 2026 15:04:48 +0000 Subject: [PATCH 04/20] simplify - remove encoder Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/modules.py | 280 +-------- .../models/cosmos3/pipeline_cosmos3.py | 18 - .../models/cosmos3/sound_tokenizer.py | 586 +++--------------- 3 files changed, 103 insertions(+), 781 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py index d95c916cb849..c0018b1d945b 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py @@ -1,86 +1,12 @@ import math -import warnings -from typing import Any, Dict, Optional, Tuple +from typing import Any, Dict import torch import torch.nn as nn import torch.nn.functional as F -from torch import Tensor -from torch.cuda import amp from torch.nn import Parameter from torch.nn.utils import spectral_norm, weight_norm -# --------------------------------------------------------------------------- -# VAE Bottleneck -# --------------------------------------------------------------------------- - - -class VAEBottleneck(nn.Module): - """ - Variational Autoencoder (VAE) bottleneck. - - Applies VAE reparameterization trick during encoding. - """ - - def __init__(self) -> None: - super().__init__() - - def sample(self, mean: torch.Tensor, scale: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: - stdev = nn.functional.softplus(scale) + 1e-4 - var = stdev * stdev - logvar = torch.log(var) - latents = torch.randn_like(mean) * stdev + mean - - kl = (mean * mean + var - logvar - 1).sum(1).mean() - - return latents, kl - - def encode( - self, x: torch.Tensor, return_info: bool = False - ) -> Tuple[torch.Tensor, Dict[str, Any]]: - """ - Encode input through VAE bottleneck. - - Args: - x: Input tensor with shape [B, C*2, T] where C*2 contains - concatenated mean and scale parameters - return_info: Whether to return additional info dict - - Returns: - Sampled latents (and optionally info dict with KL divergence) - """ - info = {} - - mean, scale = x.chunk(2, dim=1) - x, kl = self.sample(mean, scale) - - info["kl"] = kl - - if return_info: - return x, info - else: - return x - - def decode( - self, x: torch.Tensor, return_info: bool = False - ) -> Tuple[torch.Tensor, Dict[str, Any]]: - """ - Decode from latents (identity operation for VAE). - - Args: - x: Latent tensor - return_info: Whether to return additional info dict - - Returns: - Latents (and optionally empty info dict) - """ - info = {} - if return_info: - return x, info - else: - return x - - # --------------------------------------------------------------------------- # Activations # --------------------------------------------------------------------------- @@ -116,19 +42,21 @@ def __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(in_features) * alpha) - self.beta = Parameter(torch.zeros(in_features) * alpha) + self.alpha = Parameter(torch.zeros(param_shape) * alpha) + self.beta = Parameter(torch.zeros(param_shape) * alpha) else: - self.alpha = Parameter(torch.ones(in_features) * alpha) - self.beta = Parameter(torch.ones(in_features) * alpha) + 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: - alpha = self.alpha.unsqueeze(0).unsqueeze(-1) # [B, C, T] - beta = self.beta.unsqueeze(0).unsqueeze(-1) + # 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) @@ -136,46 +64,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x + (1.0 / (beta + 1e-9)) * pow(torch.sin(x * alpha), 2) -# --------------------------------------------------------------------------- -# LayerNorm (fp32-safe) -# --------------------------------------------------------------------------- - - -class LayerNorm(nn.Module): - """LayerNorm with optional bias. Forces fp32 to avoid numerical issues.""" - - def __init__(self, size: int, eps: float = 1e-5, use_bias: bool = False) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(size)) - self.bias = nn.Parameter(torch.zeros(size)) if use_bias else None - self.eps = eps - - def forward(self, tensor: Tensor) -> Tensor: - dtype = tensor.dtype - with amp.autocast(enabled=True, dtype=torch.float32): - tensor = F.layer_norm(tensor, self.weight.shape, self.weight, self.bias, self.eps) - return tensor.to(dtype) - - -# --------------------------------------------------------------------------- -# ConvNeXt helpers -# --------------------------------------------------------------------------- - - -def zero_module(module: nn.Module) -> nn.Module: - """Zero out all parameters of a module (identity-friendly init).""" - for p in module.parameters(): - p.detach().zero_() - return module - - -def may_mask(x: Tensor, mask: Optional[Tensor] = None) -> Tensor: - """Apply optional mask tensor to activations.""" - if mask is not None: - x = x * mask - return x - - # --------------------------------------------------------------------------- # WN wrappers # --------------------------------------------------------------------------- @@ -191,71 +79,6 @@ def WNConvTranspose1d(*args: Any, **kwargs: Any) -> nn.ConvTranspose1d: return weight_norm(nn.ConvTranspose1d(*args, **kwargs)) -# --------------------------------------------------------------------------- -# ConvNeXt block -# --------------------------------------------------------------------------- - - -class ConvNeXtBlock(nn.Module): - """ - ConvNeXt 1D Block adapted from https://github.com/charactr-platform/vocos - which is adapted from https://github.com/facebookresearch/ConvNeXt to 1D audio signal. - Supports causal and non-causal mode. - - Args: - dim (int): Number of input channels. - intermediate_dim (int): Dimensionality of the intermediate layer. - identity_init (bool): If True, initializes the 1x1 conv in residual paths to zero (identity-friendly). - use_snake (bool): If True, uses SnakeBeta activation; otherwise, GELU. - causal (bool): If True, applies causal padding; otherwise, applies symmetric padding for non-causal. - """ - - def __init__( - self, - dim: int, - intermediate_dim: int, - identity_init: bool = False, - use_snake: bool = False, - causal: bool = False, - ): - super().__init__() - self.causal = causal - - if causal: - self.dwconv = nn.Sequential( - nn.ConstantPad1d((6, 0), 0), - nn.Conv1d(dim, dim, kernel_size=7, groups=dim), - ) - else: - self.dwconv = nn.Sequential( - nn.ConstantPad1d((3, 3), 0), - nn.Conv1d(dim, dim, kernel_size=7, groups=dim), - ) - - self.norm = LayerNorm(dim) - self.pwconv1 = nn.Conv1d(dim, intermediate_dim, 1) - self.act = SnakeBeta(intermediate_dim) if use_snake else nn.GELU() - - if identity_init: - self.pwconv2 = zero_module(nn.Conv1d(intermediate_dim, dim, 1)) - else: - self.pwconv2 = nn.Conv1d(intermediate_dim, dim, 1) - - def forward(self, x: Tensor, mask: Optional[Tensor] = None) -> Tensor: - residual = x - x = self.dwconv(may_mask(x, mask)) - x = self.norm(x.permute(0, 2, 1)).permute(0, 2, 1) - x = self.pwconv1(x) - x = self.act(x) - x = self.pwconv2(x) - x = residual + x - return may_mask(x, mask) - - def remove_weight_norm(self) -> None: - """No weight norm is applied in ConvNeXtBlock.""" - pass - - # --------------------------------------------------------------------------- # EnCodec-style conv helpers (SConv1d / SConvTranspose1d) # --------------------------------------------------------------------------- @@ -289,15 +112,6 @@ def get_norm_module( return nn.Identity() -def get_extra_padding_for_conv1d( - x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0 -) -> int: - length = x.shape[-1] - n_frames = (length - kernel_size + padding_total) / stride + 1 - ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total) - return ideal_length - length - - def pad1d(x: torch.Tensor, paddings: tuple, mode: str = "zero", value: float = 0.0) -> torch.Tensor: """Tiny wrapper around F.pad that handles reflect padding on short inputs.""" length = x.shape[-1] @@ -324,28 +138,6 @@ def unpad1d(x: torch.Tensor, paddings: tuple) -> torch.Tensor: return x[..., padding_left:end] -class NormConv1d(nn.Module): - """Conv1d 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.conv = apply_parametrization_norm(nn.Conv1d(*args, **kwargs), norm) - self.norm = get_norm_module(self.conv, causal, norm, **norm_kwargs) - self.norm_type = norm - - def forward(self, x: torch.Tensor) -> torch.Tensor: - x = self.conv(x) - x = self.norm(x) - return x - - class NormConvTranspose1d(nn.Module): """ConvTranspose1d with optional weight_norm / spectral_norm.""" @@ -368,60 +160,6 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: return x -class SConv1d(nn.Module): - """Conv1d with builtin asymmetric/causal padding and normalization.""" - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: int, - stride: int = 1, - dilation: int = 1, - groups: int = 1, - bias: bool = True, - causal: bool = False, - norm: str = "none", - norm_kwargs: Dict[str, Any] = {}, - pad_mode: str = "reflect", - ): - super().__init__() - if stride > 1 and dilation > 1: - warnings.warn( - "SConv1d has been initialized with stride > 1 and dilation > 1" - f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})." - ) - self.conv = NormConv1d( - in_channels, - out_channels, - kernel_size, - stride, - dilation=dilation, - groups=groups, - bias=bias, - causal=causal, - norm=norm, - norm_kwargs=norm_kwargs, - ) - self.causal = causal - self.pad_mode = pad_mode - - def forward(self, x: torch.Tensor) -> torch.Tensor: - kernel_size = self.conv.conv.kernel_size[0] - stride = self.conv.conv.stride[0] - dilation = self.conv.conv.dilation[0] - kernel_size = (kernel_size - 1) * dilation + 1 - padding_total = kernel_size - stride - extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total) - if self.causal: - x = pad1d(x, (padding_total, extra_padding), mode=self.pad_mode) - else: - padding_right = padding_total // 2 - padding_left = padding_total - padding_right - x = pad1d(x, (padding_left, padding_right + extra_padding), mode=self.pad_mode) - return self.conv(x) - - class SConvTranspose1d(nn.Module): """ConvTranspose1d with builtin asymmetric/causal padding and normalization.""" 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 46979e7966dd..345a4a00f1ee 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -465,24 +465,6 @@ def _decode_latents(self, latents): # Sound generation # ========================================================================= - def encode_sound(self, waveform: torch.Tensor) -> torch.Tensor: - """Encode audio waveform into latent tokens. - - Args: - waveform: Audio tensor of shape (C, N). A batch dim is added/removed - internally since AVAE expects (B, C, N). - Mono audio is duplicated to stereo if the tokenizer expects 2 channels. - """ - # Ensure correct number of channels (AVAE typically expects stereo) - expected_channels = self.sound_tokenizer.audio_channels - if waveform.shape[0] == 1 and expected_channels == 2: - waveform = waveform.repeat(2, 1) # mono → stereo - elif waveform.shape[0] > expected_channels: - waveform = waveform[:expected_channels] - # AVAE expects (B, C, N) - latent = self.sound_tokenizer.encode(waveform.unsqueeze(0)) # [1,sound_channels,T_sound] - return latent.squeeze(0) # [sound_channels,T_sound] - def decode_sound(self, latent: torch.Tensor) -> torch.Tensor: """Decode sound latent tokens back to waveform. diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py index a8e337a00e1f..26ef0bfb9a50 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -16,26 +16,16 @@ import json import math import os -from functools import partial -from typing import Any, Callable, Dict, Literal, Optional +from typing import Any, Dict, Literal, Optional import torch -import torch.nn.functional as F 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 ( - ConvNeXtBlock, - SConv1d, - SConvTranspose1d, - SnakeBeta, - VAEBottleneck, - WNConv1d, - WNConvTranspose1d, -) +from .modules import SConvTranspose1d, SnakeBeta, WNConv1d, WNConvTranspose1d def get_activation( @@ -111,26 +101,26 @@ def __init__( self.padding_mode = padding_mode - self.layers = nn.Sequential( - get_activation( - "snake" if use_snake else "elu", - antialias=antialias_activation, - channels=out_channels, - ), - WNConv1d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=kernel_size, - dilation=dilation, - padding=self.padding, - padding_mode=self.padding_mode, - ), - get_activation( - "snake" if use_snake else "elu", - antialias=antialias_activation, - channels=out_channels, - ), - WNConv1d(in_channels=out_channels, out_channels=out_channels, kernel_size=1, padding=0), + self.snake1 = get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ) + 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( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=out_channels, + ) + self.conv2 = WNConv1d( + in_channels=out_channels, out_channels=out_channels, kernel_size=1, padding=0 ) def forward(self, x: Tensor) -> Tensor: @@ -145,8 +135,8 @@ def forward(self, x: Tensor) -> Tensor: """ res = x - # apply conv layers - x = self.layers(x) + x = self.conv1(self.snake1(x)) + x = self.conv2(self.snake2(x)) if self.causal: # Trim right padding to get the causal output @@ -185,39 +175,37 @@ def __init__( self.causal = causal - self.layers = nn.Sequential( - get_activation( - "snake" if use_snake else "elu", - antialias=antialias_activation, - channels=in_channels, - ), - self._create_upsample_layer( - in_channels, out_channels, stride, use_nearest_upsample, causal, padding_mode - ), - ResidualUnit( - in_channels=out_channels, - out_channels=out_channels, - dilation=1, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, - ), - ResidualUnit( - in_channels=out_channels, - out_channels=out_channels, - dilation=3, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, - ), - ResidualUnit( - in_channels=out_channels, - out_channels=out_channels, - dilation=9, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, - ), + self.snake1 = get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=in_channels, + ) + self.conv_t1 = self._create_upsample_layer( + in_channels, out_channels, stride, use_nearest_upsample, causal, padding_mode + ) + self.res_unit1 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=1, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, + ) + self.res_unit2 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=3, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, + ) + self.res_unit3 = ResidualUnit( + in_channels=out_channels, + out_channels=out_channels, + dilation=9, + use_snake=use_snake, + causal=causal, + padding_mode=padding_mode, ) def _create_upsample_layer( @@ -286,16 +274,18 @@ def forward(self, x: Tensor) -> Tensor: Returns: Output tensor of shape (B, C, T_upsampled) """ - return self.layers(x) + x = self.conv_t1(self.snake1(x)) + x = self.res_unit1(x) + x = self.res_unit2(x) + x = self.res_unit3(x) + return x def remove_weight_norm(self) -> None: """Remove weight normalization from all layers.""" - - for layer in self.layers: + for layer in [self.conv_t1, self.res_unit1, self.res_unit2, self.res_unit3]: try: remove_weight_norm(layer) except (ValueError, AttributeError): - # Layer doesn't have weight norm or is not a module with weight norm pass @@ -348,23 +338,19 @@ def __init__( self.depth = len(c_mults) - # Padding for the first convolution layer self.first_padding = 6 if causal else 3 - first_conv = WNConv1d( + 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() - if causal: - first_conv = nn.Sequential(first_conv, TrimPadding(self.first_padding)) - - layers = [first_conv] - + blocks = [] for i in range(self.depth - 1, 0, -1): - layers += [ + blocks += [ OobleckDecoderBlock( in_channels=c_mults[i] * channels, out_channels=c_mults[i - 1] * channels, @@ -376,10 +362,15 @@ def __init__( padding_mode=padding_mode, ) ] + self.block = nn.ModuleList(blocks) - # Padding for the final convolution layer self.final_padding = 6 if causal else 3 - final_conv = WNConv1d( + self.snake1 = get_activation( + "snake" if use_snake else "elu", + antialias=antialias_activation, + channels=c_mults[0] * channels, + ) + self.conv2 = WNConv1d( in_channels=c_mults[0] * channels, out_channels=out_channels, kernel_size=7, @@ -387,24 +378,18 @@ def __init__( padding_mode=padding_mode, bias=False, ) - - if causal: - final_conv = nn.Sequential(final_conv, TrimPadding(self.final_padding)) - - layers += [ - get_activation( - "snake" if use_snake else "elu", - antialias=antialias_activation, - channels=c_mults[0] * channels, - ), - final_conv, - nn.Tanh() if final_tanh else nn.Identity(), - ] - - self.layers = nn.Sequential(*layers) + 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: - x = self.layers(x) + 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) + x = self.final_activation(x) return x def remove_weight_norm(self: "OobleckDecoder") -> None: @@ -420,264 +405,12 @@ def remove_weight_norm(self: "OobleckDecoder") -> None: pass -class SpectrogramConvNeXtEncoder(nn.Module): - """ - Spectrogram Encoder with ConvNeXtBlocks - - This encoder processes input waveforms by converting them into spectrograms - (magnitude and phase concatenated along the channel dimension) and encodes them - using a sequence of ConvNeXtBlocks and downsampling layers. - - Args (mapped from h): - in_channels (int): Number of input audio channels (1 for mono, 2 for stereo). - channels (int): Base number of channels for the encoder. - latent_dim (int): Dimensionality of the final latent representation. - c_mults (List[int]): Channel multipliers at each depth of the encoder. - strides (List[int]): Downsampling strides for each depth. - num_blocks (int): Number of ConvNeXtBlocks to stack per depth. - identity_init (bool): Whether to initialize the 1x1 convs in residual paths as zeros. - n_fft (int): Number of FFT points for spectrogram computation. - hop_length (int): Hop length for the STFT. - use_snake (bool): Whether to use Snake activation in ConvNeXtBlocks. - causal (bool): If True, uses causal convolutions. - padding_mode (str): Padding mode for convolutions (default: 'zeros'). - - Inputs: - x (torch.Tensor): Input waveform tensor of shape `[batch, in_channels, time]`. - - Outputs: - torch.Tensor: Encoded representation of shape `[batch, time_out, latent_dim]`. - - Forward Pass: - - Converts waveform input into spectrograms (concatenates magnitude and phase). - - Processes the spectrogram through stacked ConvNeXtBlocks and downsampling layers. - - Outputs the final latent representation of specified dimensionality. - - Example: - encoder = SpectrogramConvNeXtEncoder( - in_channels=2, channels=256, latent_dim=128, c_mults=[1, 2, 4], strides=[4, 4, 8] - ) - waveform = torch.randn(8, 2, 65536) # [batch, channels, time] - encoded = encoder(waveform) # Output: [8, time_out, 128] - - NOTE: output is in [B, T, C] to be consistent with other encoders - """ - - def __init__(self, model_config: Dict[str, Any]) -> None: - super().__init__() - self.model_config = model_config - - self.in_channels = model_config["input_channels"] - if model_config.get("stereo", False): - self.in_channels *= 2 - - # if "enc_latent_dim" is found in v2 config, set it as latent_dim - if "enc_latent_dim" in model_config: - self.latent_dim = model_config["enc_latent_dim"] - else: - # if not found, fallback to v1 logic - self.latent_dim = model_config["vocoder_input_dim"] - if model_config["model_type"] == "vae": - self.latent_dim *= 2 - - self.channels = model_config["enc_dim"] - - self.c_mults = model_config["enc_c_mults"] - self.strides = model_config["enc_strides"] - self.num_blocks = model_config["enc_num_blocks"] - self.identity_init = model_config["enc_identity_init"] - self.causal = model_config["causal"] - self.padding_mode = model_config["padding_mode"] - - self.use_snake = model_config["enc_use_snake"] - - # Basic checks - assert len(self.c_mults) == len(self.strides), ( - f"The length of c_mults and strides must match. Got {len(self.c_mults)} vs {len(self.strides)}." - ) - - # Spectrogram function - self.n_fft = model_config["enc_n_fft"] - self.hop_length = model_config["enc_hop_length"] - self.spectrogram_fn = partial( - self.spectrogram, - n_fft=self.n_fft, - hop_length=self.hop_length, - win_length=self.n_fft, - window_fn=torch.hann_window, - ) - - # --------------------------------------------------------------------- - # 1) Initial projection (similar to the first_conv in OobleckEncoder), - # but here we typically use a 1x1 conv for a "spectrogram style" input. - # --------------------------------------------------------------------- - layers = [] - layers.append( - WNConv1d( - (self.n_fft + 2) * self.in_channels, - self.c_mults[0] * self.channels, - kernel_size=1, - bias=False, - ) - ) - - # --------------------------------------------------------------------- - # 2) Stages: For each i in range(len(c_mults)): - # - Stack num_blocks of ConvNeXtBlock - # - Downsample via stride convolution - # --------------------------------------------------------------------- - for i in range(len(self.c_mults)): - dim_in = self.c_mults[i] * self.channels - # Determine output dimension for the block - if i < len(self.c_mults) - 1: # If not the last block - dim_out = self.c_mults[i + 1] * self.channels - else: # For the last block, dim_out is c_mults[-1] * channels - dim_out = self.c_mults[-1] * self.channels - ds_rate = self.strides[i] - - # (a) Repeated ConvNeXtBlocks - for _ in range(self.num_blocks): - layers.append( - ConvNeXtBlock( - dim=dim_in, - intermediate_dim=dim_in * 4, - identity_init=self.identity_init, - use_snake=self.use_snake, - causal=self.causal, - ) - ) - - # (b) Downsampling convolution - layers.append( - self._create_downsample_layer( - dim_in, dim_out, ds_rate, self.causal, self.padding_mode - ) - ) - - # --------------------------------------------------------------------- - # 3) Final projection from the last channel dimension to latent_dim. - # --------------------------------------------------------------------- - layers.append( - WNConv1d(self.c_mults[-1] * self.channels, self.latent_dim, kernel_size=1, bias=False) - ) - - self.layers = nn.Sequential(*layers) - - def spectrogram( - self: "SpectrogramConvNeXtEncoder", - wav: Tensor, - n_fft: int, - hop_length: int, - win_length: int, - window_fn: Callable[[int], torch.Tensor] = torch.hann_window, - ) -> Tensor: - """ - wav: [batch_size?, time_steps], where batch_size? is an optional batch dimension - """ - pad_size_l = (n_fft - hop_length) // 2 - pad_size_r = (n_fft - hop_length) - pad_size_l - with torch.autocast(device_type=wav.device.type, enabled=False): - wav = F.pad(wav, (pad_size_l, pad_size_r)).float() - spec = torch.stft( - wav, - n_fft, - hop_length=hop_length, - win_length=win_length, - window=window_fn(win_length).to(wav), - center=False, - normalized=False, - onesided=True, - return_complex=True, - ) - return spec - - def _create_downsample_layer( - self: "SpectrogramConvNeXtEncoder", - in_channels: int, - out_channels: int, - stride: int, - causal: bool, - padding_mode: str, - ) -> nn.Module: - if causal: - downsample_layer = SConv1d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=2 * stride, - stride=stride, - causal=True, - norm="weight_norm", - ) - else: # original non-causal implementation - downsample_layer = WNConv1d( - in_channels=in_channels, - out_channels=out_channels, - kernel_size=2 * stride, - stride=stride, - padding=math.ceil(stride / 2), - padding_mode=padding_mode, - ) - return downsample_layer - - def forward(self, x: torch.Tensor) -> torch.Tensor: - """ - Forward pass: - x: waveform in [batch, in_channels, length] (mono: in_channels=1, stereo: in_channels=2) - Returns: encoder output in [batch, length_out, dim_latent], - where the spectrogram's magnitude and phase are concatenated along the channel dimension. - """ - - # Handle stereo input by merging channel dim into batch dim - batch, channels, length = x.shape - if channels > 1: # Stereo case - x = x.reshape(batch * channels, 1, length) # [batch * channels, 1, length] - - # Compute the spectrogram - with torch.autocast(device_type=x.device.type, enabled=False): - spec = self.spectrogram_fn( - x.float().squeeze(1) - ) # Remove the channel dimension for STFT - mag, ph = torch.view_as_real(spec).chunk(2, dim=-1) # Split real and imaginary parts - spectrogram = torch.cat([mag, ph], dim=1).squeeze( - -1 - ) # Concatenate along channel dim: [batch * channels, freq, frame] - - # Cast spectrogram back to original dtype - spectrogram = spectrogram.to(x.dtype) - - # Restore stereo structure if needed - if channels > 1: # Stereo case - freq = spectrogram.shape[1] # Get the frequency dimension - spectrogram = spectrogram.reshape( - batch, channels * freq, *spectrogram.shape[2:] - ) # [batch, freq * channels, frame] - - # forward pass the encoder - output = self.layers(spectrogram) - - return output.transpose(1, 2) # [B, T, C] - - def remove_weight_norm(self: "SpectrogramConvNeXtEncoder") -> 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): """ A Latent AutoEncoder class with cleaner implementation to generalize using bottleneck.py Attributes: - h: Configuration object containing model hyperparameters. - encoder (nn.Module): The encoder module based on configuration. - bottleneck (VAEBottleneck): VAE Bottleneck module. + model_config: Configuration object containing model hyperparameters. decoder (nn.Module): The decoder module based on configuration. """ @@ -703,44 +436,20 @@ def __init__(self, model_config: Dict[str, Any]) -> None: self.input_type = "mel" model_config["input_channels"] = model_config["num_mels"] - # hop_size defines the down/up sampling factor of the autoencoder - self.hop_size = model_config["hop_size"] + # Check for encoder-only mode + self.encoder_only = model_config.get("encoder_only", False) - # Initialize encoder - self.enc_type = model_config.get("enc_type", "convnext") + if self.encoder_only: + raise NotImplementedError("Encoder-only mode not supported") - # Define encoder (only spec_convnext supported in cleaned version) - if self.enc_type == "spec_convnext": - self.encoder = SpectrogramConvNeXtEncoder(model_config) + self.dec_type = model_config.get("dec_type", "oobleck") + if self.dec_type == "oobleck": + self.decoder = OobleckDecoder(model_config) else: raise NotImplementedError( - f"Encoder type '{self.enc_type}' not supported in cleaned AVAE. Only 'spec_convnext' is supported." + f"Decoder type '{self.dec_type}' not supported in cleaned AVAE. Only 'oobleck' is supported." ) - # Initialize encoder projector (Identity for spec_convnext) - self.encoder_proj = nn.Identity() - - if "bottleneck" in model_config: - self.bottleneck = VAEBottleneck() - else: - raise ValueError("Bottleneck configuration must be specified") - - # Check for encoder-only mode - self.encoder_only = model_config.get("encoder_only", False) - - if not self.encoder_only: - # Initialize decoder - self.dec_type = model_config.get("dec_type", "oobleck") - if self.dec_type == "oobleck": - self.decoder = OobleckDecoder(model_config) - else: - raise NotImplementedError( - f"Decoder type '{self.dec_type}' not supported in cleaned AVAE. Only 'oobleck' is supported." - ) - else: - # Skip decoder initialization - self.decoder = None - # Optional latent normalisation (from cosmos3-internal AVAEModel) self.latent_mean = model_config.get("latent_mean", None) self.latent_std = model_config.get("latent_std", None) @@ -761,13 +470,9 @@ def from_pretrained( config = json.load(f) model = cls(config) - - # --- weight loading (mirrors cosmos3-internal AVAEModel._load_avae_model) --- state_dict: Optional[Dict[str, Any]] = None - # 1. safetensors (standard TRT-LLM / HF format) - sft_candidates = ["model.safetensors", "diffusion_pytorch_model.safetensors"] - for name in sft_candidates: + 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 @@ -775,19 +480,10 @@ def from_pretrained( state_dict = load_file(path, device="cpu") break - # 2. PyTorch bin - if state_dict is None: - bin_candidates = ["pytorch_model.bin", "diffusion_pytorch_model.bin"] - for name in bin_candidates: - path = os.path.join(checkpoint_dir, name) - if os.path.exists(path): - state_dict = torch.load(path, map_location="cpu", weights_only=True) - break - if state_dict is None: raise FileNotFoundError( f"No weight file found in '{checkpoint_dir}'. " - "Expected model.safetensors, pytorch_model.bin, or *.ckpt." + "Expected diffusion_pytorch_model.safetensors." ) missing, unexpected = model.load_state_dict(state_dict, strict=False) @@ -810,93 +506,6 @@ def from_pretrained( return model - def calculate_latent_lengths( - self: "LatentAutoEncoderV2", audio_lengths: torch.Tensor - ) -> torch.Tensor: - """ - Calculates the latent lengths given the original audio lengths. - - Args: - audio_lengths (torch.Tensor): A tensor of shape [B] containing the lengths of the original audio samples. - - Returns: - torch.Tensor: A tensor of shape [B] containing the corresponding latent lengths. - """ - if self.input_type == "waveform": - # The latent length is the audio length divided by the hop_size - latent_lengths = torch.ceil(audio_lengths.float() / self.hop_size).long() - else: - # The latent length is same as audio_lengths - latent_lengths = audio_lengths - - return latent_lengths - - def forward(self: "LatentAutoEncoderV2", x: torch.Tensor) -> dict[str, torch.Tensor]: - """ - Forward pass through the model. - - Args: - x (torch.Tensor): Input tensor to the model with shape [B, C, T]. - - Returns: - dict[str, torch.Tensor]: Dictionary of output tensors including: - - encoder_out: Raw encoder output - - latent: Bottleneck latent representation - - decoder_out: Decoded output (if decoder exists) - - Additional outputs specific to the bottleneck type - """ - return_dict = {} - - # Encoder - encoder_out = self.encoder(x) # Shape: [B, T_frame, encoder_out_dim] - encoder_out_proj = self.encoder_proj(encoder_out) # Shape: [B, T_frame, encoder_proj_dim] - - # Apply bottleneck after reshaping to [B, C, T] again - latent, bottleneck_enc_info = self.bottleneck.encode( - encoder_out_proj.transpose(1, 2), return_info=True - ) - - # Update return dictionary - return_dict.update({"encoder_out": encoder_out.transpose(1, 2), "latent": latent}) - # Add bottleneck-specific info to return dict - for k, v in bottleneck_enc_info.items(): - return_dict[k] = v - - # Decode (if decoder exists) - if self.decoder is not None: - # Apply bottleneck decode - decoded_latent, bottleneck_dec_info = self.bottleneck.decode(latent, return_info=True) - # Apply decoder - decoder_out = self.decoder(decoded_latent) - - # Update return dictionary - return_dict["decoder_out"] = decoder_out - # Add bottleneck-specific info to return dict - for k, v in bottleneck_dec_info.items(): - return_dict[k] = v - - return return_dict - - def encode(self: "LatentAutoEncoderV2", x: torch.Tensor) -> torch.Tensor: - """ - Encode waveform to latent tokens. - - Args: - x: Input tensor with shape [B, C, T]. - - Returns: - Latent tensor [B, latent_ch, T_latent]. Normalised by latent_mean/std - when configured (mirrors cosmos3-internal AVAEModel.encode). - """ - encoder_out = self.encoder(x) - encoder_out_proj = self.encoder_proj(encoder_out) - latent = self.bottleneck.encode(encoder_out_proj.transpose(1, 2)) - - if self.latent_mean is not None and self.latent_std is not None: - latent = (latent - self.latent_mean) / self.latent_std - - return latent - def decode(self: "LatentAutoEncoderV2", latent: torch.Tensor) -> torch.Tensor: """ Decode latent tokens back to waveform. @@ -910,16 +519,9 @@ def decode(self: "LatentAutoEncoderV2", latent: torch.Tensor) -> torch.Tensor: if self.latent_mean is not None and self.latent_std is not None: latent = latent * self.latent_std + self.latent_mean - decoded_latent = self.bottleneck.decode(latent) - return self.decoder(decoded_latent) - - @property - def audio_channels(self) -> int: - """Number of output audio channels (1 = mono, 2 = stereo).""" - return self.model_config.get("dec_out_channels", 1) + return self.decoder(latent) def remove_weight_norm(self: "LatentAutoEncoderV2") -> None: """Remove weight normalization from all components.""" - self.encoder.remove_weight_norm() if self.decoder is not None: self.decoder.remove_weight_norm() From 4d970f68d0cff4ee5ca49fb1f1f9f91eb5d92895 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 19 May 2026 16:14:18 +0000 Subject: [PATCH 05/20] add enable_sound req param Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py | 5 +++++ .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 5 ++++- 2 files changed, 9 insertions(+), 1 deletion(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index e54e818356c4..458a42a67340 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -56,4 +56,9 @@ default=True, description="Whether to use the guardrails.", ), + "enable_sound": ExtraParamSchema( + type="bool", + default=False, + description="Whether to enable sound generation.", + ), } 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 345a4a00f1ee..0d438ef434bc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -201,6 +201,7 @@ 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_sound=False, ) def infer(self, req): @@ -220,6 +221,7 @@ def infer(self, req): 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), + enable_sound=req.params.extra_params.get("enable_sound", False), ) def _format_prompt_with_template( @@ -499,6 +501,7 @@ 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_sound: bool = COSMOS3_EXTRA_SPECS["enable_sound"].default, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -615,7 +618,7 @@ def forward( # 3b. Sound noise init # T_sound = ceil(duration_s * sound_latent_fps / temporal_compression_factor_sound) # Duration derived from num_frames / frame_rate; matches cosmos3-internal. - do_sound = self.sound_gen and hasattr(self, "sound_tokenizer") + do_sound = enable_sound and self.sound_gen and hasattr(self, "sound_tokenizer") sound_latents = None if do_sound: duration_s = num_frames / frame_rate From f25540c33b42928181699aaf76f6b393206446cc Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 29 May 2026 17:03:46 +0000 Subject: [PATCH 06/20] working with new checkpoint Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/defaults.py | 4 +- .../models/cosmos3/pipeline_cosmos3.py | 94 +++++----- .../models/cosmos3/sound_tokenizer.py | 173 ++++++++++++------ .../models/cosmos3/transformer_cosmos3.py | 161 ++++++++-------- .../_torch/visual_gen/pipeline_registry.py | 1 + 5 files changed, 250 insertions(+), 183 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 458a42a67340..60598d9aeb9d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -56,9 +56,9 @@ default=True, description="Whether to use the guardrails.", ), - "enable_sound": ExtraParamSchema( + "enable_audio": ExtraParamSchema( type="bool", default=False, - description="Whether to enable sound generation.", + description="Whether to enable audio generation.", ), } 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 0d438ef434bc..24abb85daaf3 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -34,8 +34,8 @@ from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS from .guardrails import check_video_safety, download_guardrail_checkpoint -from .transformer_cosmos3 import Cosmos3VFMTransformer 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, " @@ -67,11 +67,15 @@ class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): super().__init__(pipeline_config) - self.sound_gen = False + self.audio_gen = False self.action_gen = False - if getattr(model_config.pretrained_config, "sound_gen", False): - logger.info("Initializing Cosmos3OmniMoTPipeline with sound generation.") - self.sound_gen = True + if getattr( + model_config.pretrained_config, + "audio_gen", + getattr(model_config.pretrained_config, "sound_gen", False), + ): + logger.info("Initializing Cosmos3OmniMoTPipeline with audio generation.") + self.audio_gen = True if getattr(model_config.pretrained_config, "action_gen", False): logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") @@ -91,10 +95,10 @@ def load_standard_components( self, checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = [] ) -> None: skip_components = skip_components or [] - - if self.sound_gen and PipelineComponent.SOUND_TOKENIZER not in skip_components: - logger.info("Loading sound tokenizer...") - self.sound_tokenizer = ( + + 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, @@ -137,10 +141,10 @@ def load_standard_components( checkpoint_dir, subfolder=PipelineComponent.SCHEDULER, ) - if self.sound_gen: - # Separate instance so video and sound scheduler states don't collide + 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.sound_scheduler = UniPCMultistepScheduler.from_config(self.scheduler.config) + 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" @@ -201,7 +205,7 @@ 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_sound=False, + enable_audio=False, ) def infer(self, req): @@ -221,7 +225,7 @@ def infer(self, req): 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), - enable_sound=req.params.extra_params.get("enable_sound", False), + enable_audio=req.params.extra_params.get("enable_audio", False), ) def _format_prompt_with_template( @@ -464,19 +468,19 @@ def _decode_latents(self, latents): return video # ========================================================================= - # Sound generation + # Audio generation # ========================================================================= - def decode_sound(self, latent: torch.Tensor) -> torch.Tensor: - """Decode sound latent tokens back to waveform. + def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: + """Decode audio latent tokens back to waveform. Args: - latent: Sound latent tensor of shape (B, C, T). + latent: Audio latent tensor of shape (B, C, T). Returns: Waveform tensor of shape (B, audio_channels, N_samples). """ - return self.sound_tokenizer.decode(latent) # [B, audio_channels, N_samples] + return self.audio_tokenizer.decode(latent) # [B, audio_channels, N_samples] # ========================================================================= # Forward (main generation entry point) @@ -501,7 +505,7 @@ 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_sound: bool = COSMOS3_EXTRA_SPECS["enable_sound"].default, + enable_audio: bool = COSMOS3_EXTRA_SPECS["enable_audio"].default, ): pipeline_start = time.time() timer = CudaPhaseTimer() @@ -615,26 +619,26 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) - # 3b. Sound noise init - # T_sound = ceil(duration_s * sound_latent_fps / temporal_compression_factor_sound) + # 3b. Audio noise init + # T_audio = ceil(duration_s * audio_latent_fps / temporal_compression_factor_audio) # Duration derived from num_frames / frame_rate; matches cosmos3-internal. - do_sound = enable_sound and self.sound_gen and hasattr(self, "sound_tokenizer") - sound_latents = None - if do_sound: + do_audio = enable_audio and self.audio_gen and hasattr(self, "audio_tokenizer") + audio_latents = None + if do_audio: duration_s = num_frames / frame_rate - T_sound = math.ceil( + T_audio = math.ceil( duration_s - * self.transformer.sound_latent_fps - / self.transformer.temporal_compression_factor_sound + * self.transformer.audio_latent_fps + / self.transformer.temporal_compression_factor_audio ) - sound_latents = randn_tensor( - (1, self.transformer.sound_dim, T_sound), + audio_latents = randn_tensor( + (1, self.transformer.audio_dim, T_audio), generator=generator, device=self.device, dtype=latents.dtype, ) - # Sound uses the same scheduler type/config as video. - self.sound_scheduler.set_timesteps(num_inference_steps, device=self.device) + # 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( @@ -650,7 +654,7 @@ def forward_fn( Since Cosmos3 embeds text internally, we pass token IDs via extra_tensors rather than through encoder_hidden_states. """ - current_sound = extra_stream_latents.get("sound") if extra_stream_latents else None + current_audio = extra_stream_latents.get("audio") if extra_stream_latents else None result = self.transformer( hidden_states=latent_input, @@ -661,17 +665,17 @@ def forward_fn( video_shape=video_shape, fps=frame_rate, noisy_frame_mask=velocity_mask, - sound_latents=current_sound, + audio_latents=current_audio, ) video_noise_pred = result.video - sound_noise_pred = result.sound + audio_noise_pred = result.audio if velocity_mask is not None: video_noise_pred = video_noise_pred * velocity_mask - if sound_noise_pred is not None: - return video_noise_pred, {"sound": sound_noise_pred} + 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 @@ -686,7 +690,7 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - extra_streams = {"sound": (sound_latents, self.sound_scheduler)} if do_sound else None + extra_streams = {"audio": (audio_latents, self.audio_scheduler)} if do_audio else None denoise_result = self.denoise( latents=latents, scheduler=self.scheduler, @@ -700,10 +704,10 @@ def forward_fn( if extra_streams is not None: latents, extra_latents = denoise_result - sound_latents = extra_latents.get("sound") + audio_latents = extra_latents.get("audio") else: latents = denoise_result - sound_latents = None + audio_latents = None timer.mark_post_start() @@ -717,11 +721,11 @@ def forward_fn( video = self.decode_latents(latents, self._decode_latents) - # 7b. Decode sound + # 7b. Decode audio waveform = None - if do_sound and sound_latents is not None: - logger.info("Decoding sound...") - waveform = self.decode_sound(sound_latents) # [B, audio_channels, N_samples] + 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: @@ -737,7 +741,7 @@ def forward_fn( video=video, frame_rate=frame_rate, audio=waveform, - audio_sample_rate=self.sound_tokenizer.model_config["sampling_rate"] + 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 index 26ef0bfb9a50..593a8489f807 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -28,33 +28,74 @@ 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", "none"], + 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', 'snake', or 'none') + 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 == "snake": - act = SnakeBeta(channels) + 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") @@ -87,6 +128,9 @@ def __init__( 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__() @@ -100,11 +144,14 @@ def __init__( self.padding = (dilation * (kernel_size - 1)) // 2 self.padding_mode = padding_mode + activation_name = activation if use_snake else "elu" self.snake1 = get_activation( - "snake" if use_snake else "elu", + 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, @@ -115,9 +162,11 @@ def __init__( padding_mode=self.padding_mode, ) self.snake2 = get_activation( - "snake" if use_snake else "elu", + 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 @@ -170,42 +219,51 @@ def __init__( 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( - "snake" if use_snake else "elu", + 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, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, + **res_unit_kwargs, ) self.res_unit2 = ResidualUnit( in_channels=out_channels, out_channels=out_channels, dilation=3, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, + **res_unit_kwargs, ) self.res_unit3 = ResidualUnit( in_channels=out_channels, out_channels=out_channels, dilation=9, - use_snake=use_snake, - causal=causal, - padding_mode=padding_mode, + **res_unit_kwargs, ) def _create_upsample_layer( @@ -319,10 +377,7 @@ def __init__( self.model_config = model_config latent_dim = model_config["vocoder_input_dim"] - - out_channels = model_config["input_channels"] - if model_config.get("stereo", False): - out_channels *= 2 + out_channels = _resolve_decoder_out_channels(model_config) channels = model_config["dec_dim"] c_mults = model_config["dec_c_mults"] @@ -333,6 +388,19 @@ def __init__( 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] @@ -355,20 +423,18 @@ def __init__( in_channels=c_mults[i] * channels, out_channels=c_mults[i - 1] * channels, stride=strides[i - 1], - use_snake=use_snake, - antialias_activation=antialias_activation, - use_nearest_upsample=use_nearest_upsample, - causal=causal, - padding_mode=padding_mode, + **block_kwargs, ) ] self.block = nn.ModuleList(blocks) self.final_padding = 6 if causal else 3 self.snake1 = get_activation( - "snake" if use_snake else "elu", + 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, @@ -407,50 +473,27 @@ def remove_weight_norm(self: "OobleckDecoder") -> None: class LatentAutoEncoderV2(nn.Module): """ - A Latent AutoEncoder class with cleaner implementation to generalize using bottleneck.py + Decoder-only autoencoder_v2 wrapper for Cosmos3 sound generation. - Attributes: - model_config: Configuration object containing model hyperparameters. - decoder (nn.Module): The decoder module based on configuration. + 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 - - # Set up basic model properties self.stereo = model_config.get("stereo", False) - # Determine input type - self.input_type = None - if model_config.get("use_wav_as_input", False): - self.input_type = "waveform" - model_config["input_channels"] = 1 - elif model_config.get("use_linear_spec_as_input", False): - self.input_type = "linear" - model_config["input_channels"] = model_config["num_linears"] - elif model_config.get("use_discrete_code_as_input", False): - self.input_type = "discrete_code" - model_config["input_channels"] = 1 - else: - self.input_type = "mel" - model_config["input_channels"] = model_config["num_mels"] - - # Check for encoder-only mode - self.encoder_only = model_config.get("encoder_only", False) - - if self.encoder_only: + if model_config.get("encoder_only", False): raise NotImplementedError("Encoder-only mode not supported") - self.dec_type = model_config.get("dec_type", "oobleck") - if self.dec_type == "oobleck": - self.decoder = OobleckDecoder(model_config) - else: + dec_type = model_config.get("dec_type", "oobleck") + if dec_type != "oobleck": raise NotImplementedError( - f"Decoder type '{self.dec_type}' not supported in cleaned AVAE. Only 'oobleck' is supported." + f"Decoder type '{dec_type}' not supported. Only 'oobleck' is supported." ) - # Optional latent normalisation (from cosmos3-internal AVAEModel) + self.decoder = OobleckDecoder(model_config) self.latent_mean = model_config.get("latent_mean", None) self.latent_std = model_config.get("latent_std", None) @@ -486,9 +529,19 @@ def from_pretrained( "Expected diffusion_pytorch_model.safetensors." ) - missing, unexpected = model.load_state_dict(state_dict, strict=False) - if missing: - logger.warning(f"Missing keys when loading sound tokenizer: {missing}") + 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}") 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 48e436f2bf2d..3f05252fc5e5 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -69,8 +69,8 @@ class TransformerOutput: image: torch.Tensor """[B, C, 1, H, W] alias of video for image generation (same tensor).""" - sound: Optional[torch.Tensor] = None - """[B, sound_dim, T_sound] sound velocity prediction, or None.""" + 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.""" @@ -683,7 +683,7 @@ class Cosmos3VFMTransformer(BaseDiffusionModel): def __init__(self, model_config: DiffusionModelConfig): super().__init__(model_config) pretrained_config = model_config.pretrained_config - self.sound_gen = getattr(pretrained_config, "sound_gen", False) + 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 @@ -704,11 +704,15 @@ 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.sound_gen: - self.sound_dim = pretrained_config.sound_dim - self.sound_latent_fps = pretrained_config.sound_latent_fps - self.temporal_compression_factor_sound = ( - pretrained_config.temporal_compression_factor_sound + if self.audio_gen: + self.audio_dim = getattr(pretrained_config, "audio_dim", pretrained_config.sound_dim) + self.audio_latent_fps = getattr( + pretrained_config, "audio_latent_fps", pretrained_config.sound_latent_fps + ) + self.temporal_compression_factor_audio = getattr( + pretrained_config, + "temporal_compression_factor_audio", + pretrained_config.temporal_compression_factor_sound, ) if pretrained_config.position_embedding_type != "unified_3d_mrope": @@ -749,11 +753,11 @@ 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.sound_gen: - # Projections for sound modality (mirrors cosmos3-internal Cosmos3VFMNetwork) - self.sound2llm = nn.Linear(self.sound_dim, self.hidden_size) - self.llm2sound = nn.Linear(self.hidden_size, self.sound_dim) - self.sound_modality_embed = nn.Parameter(torch.zeros(self.hidden_size)) + 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) @@ -892,57 +896,57 @@ def _compute_rope_freqs( return freqs_und, freqs_gen # ------------------------------------------------------------------------- - # Sound helpers + # Audio helpers # ------------------------------------------------------------------------- - def _compute_sound_rope_freqs( + def _compute_audio_rope_freqs( self, - T_sound: int, + T_audio: int, text_mask: torch.Tensor, - fps_sound: float, + fps_audio: float, device: torch.device, dtype: torch.dtype, ) -> Tuple[torch.Tensor, torch.Tensor]: - """Compute mRoPE cos/sin for sound tokens. + """Compute mRoPE cos/sin for audio tokens. - Sound tokens use a 1×1 spatial grid (H=W=1) aligned with the vision - temporal axis at the sound latent rate. This mirrors the cosmos3-internal - ``sequence_packing.py`` treatment where sound mRoPE uses + 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() - sound_pos_list = [] + 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) - # Sound tokens share the vision temporal space; use modality margin offset. + # Audio tokens share the vision temporal space; use modality margin offset. s_pos, _ = compute_mrope_position_ids_vision( - T_sound, + T_audio, 1, # grid_h 1, # grid_w temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, - fps=fps_sound, + fps=fps_audio, base_fps=self.base_fps, - temporal_compression_factor=1, # sound latent is already at sound_latent_fps + temporal_compression_factor=1, # audio latent is already at audio_latent_fps enable_fps_modulation=self.enable_fps_modulation, ) - sound_pos_list.append(s_pos) + audio_pos_list.append(s_pos) - sound_pos_ids = torch.stack(sound_pos_list, dim=1).to(device) # [3, B, T_sound] + 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_s, sin_s = rotary_emb(_dummy, position_ids=sound_pos_ids) - return cos_s.unsqueeze(2), sin_s.unsqueeze(2) # [B, T_sound, 1, head_dim] + 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_sound_latents(self, sound_latents: torch.Tensor) -> torch.Tensor: - """[B, sound_dim, T_sound] → [B, T_sound, sound_dim].""" - return sound_latents.permute(0, 2, 1) + 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_sound_latents(self, hidden_sound: torch.Tensor) -> torch.Tensor: - """[B, T_sound, sound_dim] → [B, sound_dim, T_sound].""" - return hidden_sound.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 @@ -958,7 +962,7 @@ def forward( video_shape: Optional[Tuple[int, int, int]] = None, fps: float | None = None, noisy_frame_mask: torch.Tensor | None = None, - sound_latents: Optional[torch.Tensor] = None, + audio_latents: Optional[torch.Tensor] = None, **kwargs, ) -> "TransformerOutput": """ @@ -978,14 +982,14 @@ def forward( timestep embedding, predict velocity) and 0=conditioned (clean context, skip timestep embedding). None means all frames noisy (T2V mode). - sound_latents: Optional [B, sound_dim, T_sound] noisy sound latents. - When provided, sound tokens are appended to the generation - sequence and a sound velocity is returned alongside the video - velocity. Requires ``sound_gen=True`` in the pretrained config. + 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: TransformerOutput with video (and image alias) always set. - sound is set to the predicted sound velocity when sound_latents is + 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. @@ -1050,27 +1054,27 @@ def forward( else: self.cached_kv = cached_kv_full - # --- Sound token injection ------------------------------------------------- + # --- Audio token injection ------------------------------------------------- T_vid_tokens = hidden_gen.shape[1] # T * Hp * Wp - T_sound = 0 - if sound_latents is not None and self.sound_gen: - T_sound = sound_latents.shape[2] - hidden_sound = self.pack_sound_latents(sound_latents).to(hidden_gen.dtype) - hidden_sound = self.sound2llm(hidden_sound) + self.sound_modality_embed - hidden_sound = hidden_sound + time_embed.unsqueeze(1) - cos_s, sin_s = self._compute_sound_rope_freqs( - T_sound, + 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.sound_latent_fps), + float(self.audio_latent_fps), hidden_states.device, hidden_gen.dtype, ) - # [B, T_vid+T_sound, hidden_size] - hidden_gen = torch.cat([hidden_gen, hidden_sound], dim=1) + # [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_s], dim=1), - torch.cat([sin_v, sin_s], dim=1), + torch.cat([cos_v, cos_a], dim=1), + torch.cat([sin_v, sin_a], dim=1), ) else: freqs_gen_combined = self.cached_freqs_gen @@ -1103,16 +1107,16 @@ def forward( # --- Decode video velocity ------------------------------------------------ video_vel = self.unpatchify(self.llm2vae(hidden_gen[:, :T_vid_tokens]), T, H, W) - # --- Decode sound velocity (if requested) --------------------------------- - sound_vel = None - if T_sound > 0 and sound_latents is not None and self.sound_gen: - # hidden_gen[:, T_vid_tokens:] → [B, T_sound, hidden_size] - # → llm2sound → [B, T_sound, sound_dim] → unpack → [B, sound_dim, T_sound] - sound_vel = self.unpack_sound_latents( - self.llm2sound(hidden_gen[:, T_vid_tokens : T_vid_tokens + T_sound]) + # --- 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, sound=sound_vel) + 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. @@ -1126,8 +1130,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(): @@ -1136,19 +1138,26 @@ def load_weights(self, weights: dict) -> None: if k.startswith(skip_prefixes): continue - if k.startswith( - ("vae2llm.", "llm2vae.", "sound2llm.", "llm2sound.", "sound_modality_embed") - ): - remapped[k] = value - continue - 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.") @@ -1279,10 +1288,10 @@ def post_load_weights(self) -> None: self.vae2llm.to(target_dtype) self.llm2vae.to(target_dtype) - if self.sound_gen: - self.sound2llm.to(target_dtype) - self.llm2sound.to(target_dtype) - self.sound_modality_embed.data = self.sound_modality_embed.data.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): 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 From c0ee3a4d9741d03d61f2cb400ede200498acacc5 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 1 Jun 2026 09:21:41 -0700 Subject: [PATCH 07/20] Fix latent dims for audio Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 24abb85daaf3..8b0834b93217 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -619,18 +619,14 @@ def forward( # 3. Set up scheduler self.scheduler.set_timesteps(num_inference_steps, device=self.device) - # 3b. Audio noise init - # T_audio = ceil(duration_s * audio_latent_fps / temporal_compression_factor_audio) - # Duration derived from num_frames / frame_rate; matches cosmos3-internal. + # 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: - duration_s = num_frames / frame_rate - T_audio = math.ceil( - duration_s - * self.transformer.audio_latent_fps - / self.transformer.temporal_compression_factor_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, From eb69ef8162714e06c8321d583dd51fda63a0ef23 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 5 Jun 2026 11:26:23 -0700 Subject: [PATCH 08/20] tests with audio enabled Signed-off-by: Shreyas Misra --- .../test_cosmos3_transformer_parallel.py | 86 +++++++++- .../visual_gen/test_cosmos3_transformer.py | 155 +++++++++++++++++- 2 files changed, 234 insertions(+), 7 deletions(-) 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_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() From 913d9ed1157c4c7edeb537ab7ac19f9adc371080 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 5 Jun 2026 11:51:03 -0700 Subject: [PATCH 09/20] address coderabbit comments Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/modules.py | 24 ++++++++++++++++--- .../models/cosmos3/transformer_cosmos3.py | 21 ++++++++-------- 2 files changed, 31 insertions(+), 14 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py index c0018b1d945b..d5ea93ae288a 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py @@ -97,13 +97,29 @@ def apply_parametrization_norm(module: nn.Module, norm: str = "none") -> nn.Modu 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 == "layer_norm": + if norm in ("layer_norm", "time_layer_norm"): assert isinstance(module, nn.modules.conv._ConvNd) - return nn.LayerNorm(module.out_channels, **norm_kwargs) + return ConvLayerNorm(module.out_channels, **norm_kwargs) elif norm == "time_group_norm": if causal: raise ValueError("GroupNorm doesn't support causal evaluation.") @@ -112,7 +128,9 @@ def get_norm_module( return nn.Identity() -def pad1d(x: torch.Tensor, paddings: tuple, mode: str = "zero", value: float = 0.0) -> torch.Tensor: +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 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 3f05252fc5e5..937cf2c4f49c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -705,14 +705,10 @@ def __init__(self, model_config: DiffusionModelConfig): self.enable_fps_modulation = pretrained_config.enable_fps_modulation if self.audio_gen: - self.audio_dim = getattr(pretrained_config, "audio_dim", pretrained_config.sound_dim) - self.audio_latent_fps = getattr( - pretrained_config, "audio_latent_fps", pretrained_config.sound_latent_fps - ) - self.temporal_compression_factor_audio = getattr( - pretrained_config, - "temporal_compression_factor_audio", - pretrained_config.temporal_compression_factor_sound, + 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": @@ -1138,6 +1134,12 @@ def load_weights(self, weights: dict) -> None: if k.startswith(skip_prefixes): continue + # Normalize a leading "model." prefix up front so every remap below + # matches whether or not the checkpoint namespaces top-level tensors + # (e.g. "model.audio_proj_in.weight") under "model.". + if k.startswith("model."): + k = k[len("model.") :] + if k.startswith("proj_in."): remapped[k.replace("proj_in.", "vae2llm.", 1)] = value continue @@ -1164,9 +1166,6 @@ def load_weights(self, weights: dict) -> None: 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 From 924dfab87331b0c2da19683bc263589e32ec7aa0 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 8 Jun 2026 08:28:56 -0700 Subject: [PATCH 10/20] update defaults Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/defaults.py | 4 +- .../models/cosmos3/pipeline_cosmos3.py | 81 ++++++++++++------- 2 files changed, 52 insertions(+), 33 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 60598d9aeb9d..7daae28f6c3f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -38,12 +38,12 @@ COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", - default=True, + default=False, description="Whether to use the duration template.", ), "use_resolution_template": ExtraParamSchema( type="bool", - default=True, + default=False, description="Whether to use the resolution template.", ), "use_system_prompt": ExtraParamSchema( 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 8b0834b93217..28d3982e344f 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -37,19 +37,20 @@ 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_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." + +# Inverse templates are appended to the negative prompt so the unconditional +# branch is steered away from the requested duration/resolution. +COSMOS3_INVERSE_DURATION_TEMPLATE = ( + "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." +) +COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." + TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -221,14 +222,14 @@ def infer(self, req): 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_duration_template=req.params.extra_params.get("use_duration_template", False), + use_resolution_template=req.params.extra_params.get("use_resolution_template", False), use_system_prompt=req.params.extra_params.get("use_system_prompt", False), use_guardrails=req.params.extra_params.get("use_guardrails", True), enable_audio=req.params.extra_params.get("enable_audio", False), ) - def _format_prompt_with_template( + def _apply_metadata_templates( self, prompt: str, *, @@ -236,22 +237,29 @@ 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 prompt. - if use_duration_template and num_frames > 1: + ``duration_template`` / ``resolution_template`` of ``None`` disables that + template. The positive prompt uses the forward templates; the negative + prompt uses the inverse templates with ``force_duration_template=True`` + so the duration clause is appended even for single-frame requests. + """ + 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 _resize_and_center_crop_image( self, image: PIL.Image.Image, height: int, width: int @@ -530,7 +538,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) @@ -555,25 +563,36 @@ 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. + dur_tmpl = COSMOS3_DURATION_TEMPLATE if use_duration_template else None + res_tmpl = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE if use_resolution_template else None + + # Negative prompt: inverse templates, gated on the same flags as the + # positive templates, with the duration clause forced on so it is present + # even for single-frame requests. + inv_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if use_duration_template else None + inv_res_tmpl = COSMOS3_INVERSE_RESOLUTION_TEMPLATE if use_resolution_template else None + + negative_prompt = self._apply_metadata_templates( 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=inv_dur_tmpl, + resolution_template=inv_res_tmpl, + force_duration_template=True, ) prompt = [ - self._format_prompt_with_template( + self._apply_metadata_templates( 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 ] From e3fc639ff94c60cab14886ba6ee0866a85610645 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 8 Jun 2026 09:17:40 -0700 Subject: [PATCH 11/20] t2i updates Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/defaults.py | 20 +- .../models/cosmos3/pipeline_cosmos3.py | 265 +++++++++++++++--- 2 files changed, 249 insertions(+), 36 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 7daae28f6c3f..3ce9e08c564a 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", @@ -61,4 +74,9 @@ default=False, description="Whether to enable audio generation.", ), + "output_type": ExtraParamSchema( + type="str", + default="video", + description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", + ), } 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 28d3982e344f..3002f30f6899 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -16,7 +16,7 @@ import math import os import time -from typing import List, Optional, Union +from typing import List, Optional, Tuple, Union import PIL.Image import torch @@ -32,7 +32,7 @@ from tensorrt_llm._utils import nvtx_range 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 @@ -41,8 +41,12 @@ 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." # Inverse templates are appended to the negative prompt so the unconditional # branch is steered away from the requested duration/resolution. @@ -50,6 +54,7 @@ "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." ) COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." +COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution." TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -142,6 +147,14 @@ 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). @@ -176,6 +189,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)] @@ -209,24 +239,62 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N 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", False), - use_resolution_template=req.params.extra_params.get("use_resolution_template", False), - use_system_prompt=req.params.extra_params.get("use_system_prompt", False), - use_guardrails=req.params.extra_params.get("use_guardrails", True), - enable_audio=req.params.extra_params.get("enable_audio", False), + use_duration_template=extra_params.get("use_duration_template", False), + use_resolution_template=extra_params.get("use_resolution_template", False), + 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 _apply_metadata_templates( @@ -278,14 +346,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 [] ) @@ -490,6 +562,67 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent) # [B, audio_channels, N_samples] + @nvtx_range("Cosmos3OmniMoTPipeline._denoise_t2i", color="blue") + def _denoise_t2i( + self, + latents: torch.Tensor, + cond_ids: torch.Tensor, + cond_mask: torch.Tensor, + uncond_ids: torch.Tensor, + uncond_mask: torch.Tensor, + guidance_scale: float, + video_shape: Tuple[int, int, int], + frame_rate: float, + guidance_interval: Optional[Tuple[float, float]], + ) -> torch.Tensor: + """Denoise loop for text-to-image. + + The only T2I-specific behavior is the guidance interval — CFG is applied + only when the scheduler timestep falls inside ``guidance_interval``; + outside it the conditional prediction is used directly. + """ + do_cfg = guidance_scale > 1.0 + if guidance_interval is not None: + interval_lo, interval_hi = guidance_interval + else: + interval_lo, interval_hi = float("-inf"), float("inf") + + if do_cfg: + vgm = self.model_config.visual_gen_mapping + if vgm is not None and getattr(vgm, "cfg_size", 1) >= 2 and self.rank == 0: + raise RuntimeError("Cosmos3 T2I does not use CFG-parallel. Use cfg_size=1.") + text_ids = torch.cat([uncond_ids, cond_ids], dim=0) + text_mask = torch.cat([uncond_mask, cond_mask], dim=0) + else: + text_ids = cond_ids + text_mask = cond_mask + + for t in self.scheduler.timesteps: + latent_input = torch.cat([latents] * 2) if do_cfg else latents + timestep = t.expand(latent_input.shape[0]) + + noise_pred = self.transformer( + hidden_states=latent_input, + timestep=timestep, + text_ids=text_ids, + text_mask=text_mask, + video_shape=video_shape, + fps=frame_rate, + noisy_frame_mask=None, + ).video + + if do_cfg: + noise_uncond, noise_cond = noise_pred.chunk(2) + t_val = t.item() if t.dim() == 0 else t[0].item() + if interval_lo <= t_val <= interval_hi: + noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond) + else: + noise_pred = noise_cond + + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + return latents + # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -514,6 +647,7 @@ def forward( 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() @@ -521,6 +655,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) @@ -563,15 +716,30 @@ def forward( if negative_prompt is None: negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT - # Positive prompt: forward duration/resolution templates. + # 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 - res_tmpl = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE if use_resolution_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: inverse templates, gated on the same flags as the # positive templates, with the duration clause forced on so it is present # even for single-frame requests. inv_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if use_duration_template else None - inv_res_tmpl = COSMOS3_INVERSE_RESOLUTION_TEMPLATE if use_resolution_template else None + if use_resolution_template: + inv_res_tmpl = ( + COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE + if is_t2i + else COSMOS3_INVERSE_RESOLUTION_TEMPLATE + ) + else: + inv_res_tmpl = None negative_prompt = self._apply_metadata_templates( negative_prompt, @@ -602,9 +770,12 @@ 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 @@ -705,24 +876,42 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - 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 - neg_prompt_embeds=uncond_ids, - guidance_scale=guidance_scale, - forward_fn=forward_fn, - extra_cfg_tensors=extra_cfg_tensors, - extra_streams=extra_streams, - ) - - if extra_streams is not None: - latents, extra_latents = denoise_result - audio_latents = extra_latents.get("audio") - else: - latents = denoise_result + if is_t2i: + # T2I uses a dedicated loop with sequential CFG (separate cond/uncond + # K/V caches) and a CFG guidance interval. The shared BasePipeline + # denoise() batches [uncond, cond] and applies CFG at every step, + # which cannot express the guidance interval. + latents = self._denoise_t2i( + latents=latents, + cond_ids=cond_ids, + cond_mask=cond_mask, + uncond_ids=uncond_ids, + uncond_mask=uncond_mask, + guidance_scale=guidance_scale, + video_shape=video_shape, + frame_rate=frame_rate, + guidance_interval=guidance_interval, + ) audio_latents = None + else: + 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 + neg_prompt_embeds=uncond_ids, + guidance_scale=guidance_scale, + forward_fn=forward_fn, + extra_cfg_tensors=extra_cfg_tensors, + extra_streams=extra_streams, + ) + + 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() @@ -751,6 +940,12 @@ def forward_fn( video = check_video_safety(video, self.safety_checker) timer.mark_end() + + 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, From f2a42017d94277fc3ed33878274031c926684524 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 8 Jun 2026 12:37:39 -0700 Subject: [PATCH 12/20] use per batch attention with sliced KV Signed-off-by: Shreyas Misra --- .../visual_gen/configs/cosmos3-nano-1gpu.yaml | 10 +-- .../configs/cosmos3-super-4gpu.yaml | 6 +- examples/visual_gen/models/cosmos3_ti2v.py | 63 +++++++++++++++++- .../models/cosmos3/transformer_cosmos3.py | 64 ++++++++++++++----- 4 files changed, 111 insertions(+), 32 deletions(-) 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_ti2v.py b/examples/visual_gen/models/cosmos3_ti2v.py index b9efb9cdec81..f4ee5927b715 100644 --- a/examples/visual_gen/models/cosmos3_ti2v.py +++ b/examples/visual_gen/models/cosmos3_ti2v.py @@ -45,7 +45,7 @@ - ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE Usage: - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \\ + 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. " \\ @@ -66,11 +66,38 @@ "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." \\ + "in a retail setting." \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + --prompt "A low-angle tracking shot follows a man riding a vintage black motorcycle " \\ + "across a lush green grassy yard. Sunlight filters through overhead trees, casting " \\ + "dappled shadows across the vibrating chrome exhaust and the rider's leather jacket. " \\ + "He kicks up small blades of grass as he maneuvers the bike. He gradually decelerates, " \\ + "the front fork compressing slightly as he brakes to a smooth halt beside another " \\ + "individual standing in the shade. The camera settles into a medium two-shot, capturing " \\ + "the rider lifting his visor to speak, his face framed by a matte helmet. The video is " \\ + "8 seconds long and is of 24 FPS. This video is of 1280x720 resolution. Audio description: " \\ + "The rhythmic, mechanical chugging of a four-stroke motorcycle engine dominates the " \\ + "foreground, characterized by a throaty, guttural timbre. Periodic high-pitched revs " \\ + "punctuate the steady idle as the throttle is twisted. The sound of tires crunching " \\ + "softly over dry grass and twigs provides a textured background layer. As the vehicle " \\ + "slows, the engine note drops to a low-frequency rumble before clicking into neutral. " \\ + "A muffled, mid-range male voice begins speaking, accompanied by the metallic clink of " \\ + "a helmet visor snapping upward and the faint chirping of distant birds in an open-air " \\ + "environment." \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --enable_audio + + python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + --prompt "A cute puppy playing with a ball in a park" \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ + --output_type image \ + --output_path output.png """ import argparse +import json from tensorrt_llm import VisualGen, VisualGenArgs @@ -108,6 +135,26 @@ def main(): default="cosmos3_ti2v_output.mp4", help="Path to save the output video", ) + parser.add_argument( + "--enable_duration_template", action="store_true", help="Enable duration template in prompt" + ) + parser.add_argument( + "--enable_resolution_template", + action="store_true", + help="Enable resolution template in prompt", + ) + 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() # Engine config from shared YAML (optional); model-specific defaults apply otherwise. @@ -120,6 +167,17 @@ def main(): if args.image_path is not None: params.image = args.image_path + negative_prompt = json.load(open("neg_prompt.json")) + + params.extra_params["use_duration_template"] = args.enable_duration_template + params.extra_params["use_resolution_template"] = args.enable_resolution_template + params.extra_params["use_system_prompt"] = args.use_system_prompt + params.extra_params["enable_audio"] = args.enable_audio + params.extra_params["use_guardrails"] = not args.disable_guardrails + params.extra_params["output_type"] = args.output_type + + params.negative_prompt = json.dumps(negative_prompt) + output = visual_gen.generate( inputs=args.prompt, params=params, @@ -127,6 +185,7 @@ def main(): output.save(args.output_path) print(f"Saved: {args.output_path}") + print(output.metrics) if __name__ == "__main__": 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 937cf2c4f49c..91f86ed978c6 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -351,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: @@ -374,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) @@ -503,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) @@ -515,6 +534,7 @@ def forward( freqs_cos=cos, freqs_sin=sin, timestep=timestep, + real_text_lens=real_text_lens, ) hidden_states = residual + hidden_states @@ -992,6 +1012,7 @@ def forward( 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)) @@ -1088,13 +1109,22 @@ 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) From 3d3cfe917fe36f0a314bdb4dde0d4cecbe533c01 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 9 Jun 2026 08:42:40 -0700 Subject: [PATCH 13/20] guidance interval to base pipeline Signed-off-by: Shreyas Misra --- .../models/cosmos3/pipeline_cosmos3.py | 134 ++++-------------- tensorrt_llm/_torch/visual_gen/pipeline.py | 44 ++++-- 2 files changed, 67 insertions(+), 111 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 3002f30f6899..35acca1432ed 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -16,7 +16,7 @@ import math import os import time -from typing import List, Optional, Tuple, Union +from typing import List, Optional, Union import PIL.Image import torch @@ -71,25 +71,31 @@ ) class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, pipeline_config): - super().__init__(pipeline_config) - + primary_pretrained_config = pipeline_config.primary_pretrained_config self.audio_gen = False self.action_gen = False if getattr( - model_config.pretrained_config, + primary_pretrained_config, "audio_gen", - getattr(model_config.pretrained_config, "sound_gen", False), + getattr(primary_pretrained_config, "sound_gen", False), ): logger.info("Initializing Cosmos3OmniMoTPipeline with audio generation.") self.audio_gen = True - if getattr(model_config.pretrained_config, "action_gen", False): + if getattr(primary_pretrained_config, "action_gen", False): logger.info("Initializing Cosmos3OmniMoTPipeline with action generation.") self.action_gen = True + super().__init__(pipeline_config) + + @property + def dtype(self): + return self.pipeline_config.torch_dtype + 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"): @@ -562,67 +568,6 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: """ return self.audio_tokenizer.decode(latent) # [B, audio_channels, N_samples] - @nvtx_range("Cosmos3OmniMoTPipeline._denoise_t2i", color="blue") - def _denoise_t2i( - self, - latents: torch.Tensor, - cond_ids: torch.Tensor, - cond_mask: torch.Tensor, - uncond_ids: torch.Tensor, - uncond_mask: torch.Tensor, - guidance_scale: float, - video_shape: Tuple[int, int, int], - frame_rate: float, - guidance_interval: Optional[Tuple[float, float]], - ) -> torch.Tensor: - """Denoise loop for text-to-image. - - The only T2I-specific behavior is the guidance interval — CFG is applied - only when the scheduler timestep falls inside ``guidance_interval``; - outside it the conditional prediction is used directly. - """ - do_cfg = guidance_scale > 1.0 - if guidance_interval is not None: - interval_lo, interval_hi = guidance_interval - else: - interval_lo, interval_hi = float("-inf"), float("inf") - - if do_cfg: - vgm = self.model_config.visual_gen_mapping - if vgm is not None and getattr(vgm, "cfg_size", 1) >= 2 and self.rank == 0: - raise RuntimeError("Cosmos3 T2I does not use CFG-parallel. Use cfg_size=1.") - text_ids = torch.cat([uncond_ids, cond_ids], dim=0) - text_mask = torch.cat([uncond_mask, cond_mask], dim=0) - else: - text_ids = cond_ids - text_mask = cond_mask - - for t in self.scheduler.timesteps: - latent_input = torch.cat([latents] * 2) if do_cfg else latents - timestep = t.expand(latent_input.shape[0]) - - noise_pred = self.transformer( - hidden_states=latent_input, - timestep=timestep, - text_ids=text_ids, - text_mask=text_mask, - video_shape=video_shape, - fps=frame_rate, - noisy_frame_mask=None, - ).video - - if do_cfg: - noise_uncond, noise_cond = noise_pred.chunk(2) - t_val = t.item() if t.dim() == 0 else t[0].item() - if interval_lo <= t_val <= interval_hi: - noise_pred = noise_uncond + guidance_scale * (noise_cond - noise_uncond) - else: - noise_pred = noise_cond - - latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] - - return latents - # ========================================================================= # Forward (main generation entry point) # ========================================================================= @@ -876,42 +821,25 @@ def forward_fn( # 6. Denoise timer.mark_denoise_start() - if is_t2i: - # T2I uses a dedicated loop with sequential CFG (separate cond/uncond - # K/V caches) and a CFG guidance interval. The shared BasePipeline - # denoise() batches [uncond, cond] and applies CFG at every step, - # which cannot express the guidance interval. - latents = self._denoise_t2i( - latents=latents, - cond_ids=cond_ids, - cond_mask=cond_mask, - uncond_ids=uncond_ids, - uncond_mask=uncond_mask, - guidance_scale=guidance_scale, - video_shape=video_shape, - frame_rate=frame_rate, - guidance_interval=guidance_interval, - ) - audio_latents = None - else: - 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 - neg_prompt_embeds=uncond_ids, - guidance_scale=guidance_scale, - forward_fn=forward_fn, - extra_cfg_tensors=extra_cfg_tensors, - extra_streams=extra_streams, - ) + 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 + neg_prompt_embeds=uncond_ids, + 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 + 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() diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index f97ab68481d0..cb0f8c9693af 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -729,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 ): @@ -887,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 = { @@ -922,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) @@ -987,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. @@ -1017,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. @@ -1055,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"] @@ -1088,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}"): @@ -1132,6 +1159,7 @@ def denoise( current_guidance_scale, guidance_rescale, local_extras, + do_cfg=do_cfg, ) # Scheduler step for all streams From dfc092ad60e6db4cc97cd507fce2ee4f5f653b76 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 9 Jun 2026 09:04:48 -0700 Subject: [PATCH 14/20] add neg prompt example Signed-off-by: Shreyas Misra --- .../models/cosmos3_negative_prompt.json | 108 ++++++++++++++++++ examples/visual_gen/models/cosmos3_ti2v.py | 15 ++- 2 files changed, 122 insertions(+), 1 deletion(-) create mode 100644 examples/visual_gen/models/cosmos3_negative_prompt.json diff --git a/examples/visual_gen/models/cosmos3_negative_prompt.json b/examples/visual_gen/models/cosmos3_negative_prompt.json new file mode 100644 index 000000000000..44dff2693426 --- /dev/null +++ b/examples/visual_gen/models/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_ti2v.py b/examples/visual_gen/models/cosmos3_ti2v.py index f4ee5927b715..d81b4c3df70d 100644 --- a/examples/visual_gen/models/cosmos3_ti2v.py +++ b/examples/visual_gen/models/cosmos3_ti2v.py @@ -98,6 +98,7 @@ import argparse import json +import os from tensorrt_llm import VisualGen, VisualGenArgs @@ -123,6 +124,12 @@ def main(): required=True, help="Text prompt for generation", ) + 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, @@ -167,7 +174,13 @@ def main(): if args.image_path is not None: params.image = args.image_path - negative_prompt = json.load(open("neg_prompt.json")) + if args.negative_prompt is not None: + if os.path.isfile(args.negative_prompt) and args.negative_prompt.endswith(".json"): + negative_prompt = json.load(open(args.negative_prompt)) + else: + negative_prompt = args.negative_prompt + else: + negative_prompt = None params.extra_params["use_duration_template"] = args.enable_duration_template params.extra_params["use_resolution_template"] = args.enable_resolution_template From 1f29600e02074c29c438dd75bbc97e2d29f0ab46 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 9 Jun 2026 09:30:33 -0700 Subject: [PATCH 15/20] tests Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3_ti2v.py | 7 +- tensorrt_llm/_torch/visual_gen/pipeline.py | 2 +- .../visual_gen/test_cosmos3_pipeline.py | 105 ++++++++++++++---- 3 files changed, 89 insertions(+), 25 deletions(-) diff --git a/examples/visual_gen/models/cosmos3_ti2v.py b/examples/visual_gen/models/cosmos3_ti2v.py index d81b4c3df70d..b45a8a4ccf5b 100644 --- a/examples/visual_gen/models/cosmos3_ti2v.py +++ b/examples/visual_gen/models/cosmos3_ti2v.py @@ -189,7 +189,12 @@ def main(): params.extra_params["use_guardrails"] = not args.disable_guardrails params.extra_params["output_type"] = args.output_type - params.negative_prompt = json.dumps(negative_prompt) + 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=args.prompt, diff --git a/tensorrt_llm/_torch/visual_gen/pipeline.py b/tensorrt_llm/_torch/visual_gen/pipeline.py index cb0f8c9693af..13cbe77bdf12 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline.py @@ -216,7 +216,7 @@ def world_size(self): @property def dtype(self): - return self.model_config.torch_dtype + return self.pipeline_config.torch_dtype @property def device(self): diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 61bf5dfd1f90..6e175fa67f98 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -4,13 +4,17 @@ """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 Override checkpoint: DIFFUSION_MODEL_PATH_COSMOS3=/path/to/Cosmos3-Nano \\ @@ -28,7 +32,7 @@ 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.pipeline_loader import PipelineLoader from tensorrt_llm.visual_gen.args import TorchCompileConfig, VisualGenArgs @@ -79,6 +83,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, @@ -141,6 +151,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): @@ -158,20 +213,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 +239,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 +296,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) From bc083267e629661369bcb0ce65379bfdf8edd4f0 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 12 Jun 2026 11:19:00 -0700 Subject: [PATCH 16/20] address comments Signed-off-by: Shreyas Misra --- .../cosmos3_negative_prompt.json | 0 .../models/{ => cosmos3}/cosmos3_ti2v.py | 174 ++++++++++++------ .../models/cosmos3/prompts/i2v.json | 5 + .../models/cosmos3/prompts/t2av.json | 5 + .../models/cosmos3/prompts/t2i.json | 4 + .../models/cosmos3/prompts/t2v.json | 4 + .../visual_gen/models/cosmos3/defaults.py | 4 +- .../visual_gen/models/cosmos3/modules.py | 15 ++ .../models/cosmos3/pipeline_cosmos3.py | 102 ++++++---- .../models/cosmos3/sound_tokenizer.py | 41 +++-- .../visual_gen/test_cosmos3_pipeline.py | 161 ++++++++++++++++ 11 files changed, 404 insertions(+), 111 deletions(-) rename examples/visual_gen/models/{ => cosmos3}/cosmos3_negative_prompt.json (100%) rename examples/visual_gen/models/{ => cosmos3}/cosmos3_ti2v.py (53%) create mode 100644 examples/visual_gen/models/cosmos3/prompts/i2v.json create mode 100644 examples/visual_gen/models/cosmos3/prompts/t2av.json create mode 100644 examples/visual_gen/models/cosmos3/prompts/t2i.json create mode 100644 examples/visual_gen/models/cosmos3/prompts/t2v.json diff --git a/examples/visual_gen/models/cosmos3_negative_prompt.json b/examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json similarity index 100% rename from examples/visual_gen/models/cosmos3_negative_prompt.json rename to examples/visual_gen/models/cosmos3/cosmos3_negative_prompt.json diff --git a/examples/visual_gen/models/cosmos3_ti2v.py b/examples/visual_gen/models/cosmos3/cosmos3_ti2v.py similarity index 53% rename from examples/visual_gen/models/cosmos3_ti2v.py rename to examples/visual_gen/models/cosmos3/cosmos3_ti2v.py index b45a8a4ccf5b..8b8e92f7046f 100644 --- a/examples/visual_gen/models/cosmos3_ti2v.py +++ b/examples/visual_gen/models/cosmos3/cosmos3_ti2v.py @@ -17,7 +17,7 @@ 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. +reference frame, or use ``prompts/i2v.json`` which includes a ``vision_path``. Checkpoints (pass the Hub ID or local path via ``--model``): @@ -44,64 +44,102 @@ - ``cosmos3-nano-1gpu.yaml`` — 1 GPU, FP8 dynamic quant - ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE -Usage: +Example prompts live under ``prompts/`` (mirroring ``cosmos3-internal/inputs/omni``). + +Usage:: + + # Text-to-video 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." \ + --prompt_file prompts/t2v.json \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # Image-to-video (vision_path is read from the prompt file) python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ - --prompt "A low-angle tracking shot follows a man riding a vintage black motorcycle " \\ - "across a lush green grassy yard. Sunlight filters through overhead trees, casting " \\ - "dappled shadows across the vibrating chrome exhaust and the rider's leather jacket. " \\ - "He kicks up small blades of grass as he maneuvers the bike. He gradually decelerates, " \\ - "the front fork compressing slightly as he brakes to a smooth halt beside another " \\ - "individual standing in the shade. The camera settles into a medium two-shot, capturing " \\ - "the rider lifting his visor to speak, his face framed by a matte helmet. The video is " \\ - "8 seconds long and is of 24 FPS. This video is of 1280x720 resolution. Audio description: " \\ - "The rhythmic, mechanical chugging of a four-stroke motorcycle engine dominates the " \\ - "foreground, characterized by a throaty, guttural timbre. Periodic high-pitched revs " \\ - "punctuate the steady idle as the throttle is twisted. The sound of tires crunching " \\ - "softly over dry grass and twigs provides a textured background layer. As the vehicle " \\ - "slows, the engine note drops to a low-frequency rumble before clicking into neutral. " \\ - "A muffled, mid-range male voice begins speaking, accompanied by the metallic clink of " \\ - "a helmet visor snapping upward and the faint chirping of distant birds in an open-air " \\ - "environment." \ - --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --enable_audio + --prompt_file prompts/i2v.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + # Text-to-video with audio python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ - --prompt "A cute puppy playing with a ball in a park" \ + --prompt_file prompts/t2av.json \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml + + # Text-to-image + python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + --prompt_file prompts/t2i.json \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml \ - --output_type image \ --output_path output.png + + # Inline prompt (``--prompt`` or a JSON file path) + python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + --prompt "A cute puppy playing with a ball in a park" \ + --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml """ 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 example") @@ -121,8 +159,14 @@ def main(): parser.add_argument( "--prompt", type=str, - required=True, - help="Text prompt for generation", + 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", @@ -134,7 +178,7 @@ def main(): "--image_path", type=str, default=None, - help="Optional conditioning image path for I2V/TI2V", + help="Optional conditioning image path or URL for I2V/TI2V", ) parser.add_argument( "--output_path", @@ -143,12 +187,14 @@ def main(): help="Path to save the output video", ) parser.add_argument( - "--enable_duration_template", action="store_true", help="Enable duration template in prompt" + "--disable_duration_template", + action="store_true", + help="Disable duration metadata template (enabled by default, matching cosmos-framework CLI)", ) parser.add_argument( - "--enable_resolution_template", + "--disable_resolution_template", action="store_true", - help="Enable resolution template in prompt", + 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" @@ -164,6 +210,14 @@ def main(): ) 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) @@ -171,23 +225,27 @@ def main(): # --- 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 + 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(args.negative_prompt) and args.negative_prompt.endswith(".json"): - negative_prompt = json.load(open(args.negative_prompt)) + 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 - params.extra_params["use_duration_template"] = args.enable_duration_template - params.extra_params["use_resolution_template"] = args.enable_resolution_template + 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"] = args.enable_audio + params.extra_params["enable_audio"] = enable_audio params.extra_params["use_guardrails"] = not args.disable_guardrails - params.extra_params["output_type"] = args.output_type + params.extra_params["output_type"] = output_type if negative_prompt is None: params.negative_prompt = None @@ -197,7 +255,7 @@ def main(): params.negative_prompt = json.dumps(negative_prompt) output = visual_gen.generate( - inputs=args.prompt, + inputs=prompt, params=params, ) 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/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 3ce9e08c564a..10cdcb268f59 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -51,12 +51,12 @@ COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { "use_duration_template": ExtraParamSchema( type="bool", - default=False, + default=True, description="Whether to use the duration template.", ), "use_resolution_template": ExtraParamSchema( type="bool", - default=False, + default=True, description="Whether to use the resolution template.", ), "use_system_prompt": ExtraParamSchema( diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py index d5ea93ae288a..1ec9f64b57ee 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/modules.py @@ -1,3 +1,18 @@ +# 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 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 35acca1432ed..ad680ece47e8 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 @@ -48,14 +49,6 @@ COSMOS3_DEFAULT_RESOLUTION_TEMPLATE = "This video is of {height}x{width} resolution." COSMOS3_IMAGE_RESOLUTION_TEMPLATE = "This image is of {height}x{width} resolution." -# Inverse templates are appended to the negative prompt so the unconditional -# branch is steered away from the requested duration/resolution. -COSMOS3_INVERSE_DURATION_TEMPLATE = ( - "The video is not {duration:.1f} seconds long and is not of {fps:.0f} FPS." -) -COSMOS3_INVERSE_RESOLUTION_TEMPLATE = "This video is not of {height}x{width} resolution." -COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE = "This image is not of {height}x{width} resolution." - TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" @@ -88,10 +81,6 @@ def __init__(self, pipeline_config): super().__init__(pipeline_config) - @property - def dtype(self): - return self.pipeline_config.torch_dtype - def _init_transformer(self) -> None: logger.info("Initializing Cosmos3VFMTransformer") model_config = self.pipeline_config.model_configs["transformer"] @@ -295,8 +284,14 @@ def infer(self, req): seed=req.params.seed, max_sequence_length=req.params.max_sequence_length, frame_rate=req.params.frame_rate, - use_duration_template=extra_params.get("use_duration_template", False), - use_resolution_template=extra_params.get("use_resolution_template", False), + 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), @@ -315,12 +310,10 @@ def _apply_metadata_templates( resolution_template: Optional[str] = COSMOS3_DEFAULT_RESOLUTION_TEMPLATE, force_duration_template: bool = False, ) -> str: - """Append duration and resolution metadata to a prompt. + """Append duration and resolution metadata to a plain-text prompt. ``duration_template`` / ``resolution_template`` of ``None`` disables that - template. The positive prompt uses the forward templates; the negative - prompt uses the inverse templates with ``force_duration_template=True`` - so the duration clause is appended even for single-frame requests. + template. JSON prompts are handled by ``_format_prompt_with_metadata``. """ parts: List[str] = [] head = prompt.rstrip(".").strip() @@ -335,6 +328,52 @@ def _apply_metadata_templates( 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 ) -> PIL.Image.Image: @@ -566,7 +605,7 @@ def decode_audio(self, latent: torch.Tensor) -> torch.Tensor: Returns: Waveform tensor of shape (B, audio_channels, N_samples). """ - return self.audio_tokenizer.decode(latent) # [B, audio_channels, N_samples] + return self.audio_tokenizer.decode(latent).float() # [B, audio_channels, N_samples] # ========================================================================= # Forward (main generation entry point) @@ -673,32 +712,21 @@ def forward( else: res_tmpl = None - # Negative prompt: inverse templates, gated on the same flags as the - # positive templates, with the duration clause forced on so it is present - # even for single-frame requests. - inv_dur_tmpl = COSMOS3_INVERSE_DURATION_TEMPLATE if use_duration_template else None - if use_resolution_template: - inv_res_tmpl = ( - COSMOS3_INVERSE_IMAGE_RESOLUTION_TEMPLATE - if is_t2i - else COSMOS3_INVERSE_RESOLUTION_TEMPLATE - ) - else: - inv_res_tmpl = None - - negative_prompt = self._apply_metadata_templates( + # 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, - duration_template=inv_dur_tmpl, - resolution_template=inv_res_tmpl, - force_duration_template=True, + duration_template=dur_tmpl, + resolution_template=res_tmpl, + force_duration_template=False, ) prompt = [ - self._apply_metadata_templates( + self._format_prompt_with_metadata( p, height=height, width=width, diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py index 593a8489f807..6dc3ac898150 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/sound_tokenizer.py @@ -182,16 +182,19 @@ def forward(self, x: Tensor) -> Tensor: Returns: Output tensor of shape (B, C, T) """ - res = x - - x = self.conv1(self.snake1(x)) - x = self.conv2(self.snake2(x)) + output_tensor = self.conv1(self.snake1(x)) + output_tensor = self.conv2(self.snake2(output_tensor)) if self.causal: - # Trim right padding to get the causal output - x = x[:, :, : -self.padding] + output_tensor = output_tensor[:, :, : -self.padding] + res = x[:, :, : -self.padding] + return res + output_tensor - return x + res + 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): @@ -332,11 +335,11 @@ def forward(self, x: Tensor) -> Tensor: Returns: Output tensor of shape (B, C, T_upsampled) """ - x = self.conv_t1(self.snake1(x)) + x = self.snake1(x) + x = self.conv_t1(x) x = self.res_unit1(x) x = self.res_unit2(x) - x = self.res_unit3(x) - return x + return self.res_unit3(x) def remove_weight_norm(self) -> None: """Remove weight normalization from all layers.""" @@ -448,14 +451,24 @@ def __init__( 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) - 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) - x = self.final_activation(x) + if not isinstance(self.final_activation, nn.Identity): + x = self.final_activation(x) return x def remove_weight_norm(self: "OobleckDecoder") -> None: @@ -572,7 +585,7 @@ def decode(self: "LatentAutoEncoderV2", latent: torch.Tensor) -> torch.Tensor: 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) + return self.decoder(latent).clamp(-1.0, 1.0) def remove_weight_norm(self: "LatentAutoEncoderV2") -> None: """Remove weight normalization from all components.""" diff --git a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index 6e175fa67f98..a72e425df0d1 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -16,12 +16,16 @@ 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 \\ pytest tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py -v -s """ import gc +import json import os from pathlib import Path @@ -33,6 +37,12 @@ import torch 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 @@ -203,6 +213,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() From ea6a2545ee087b820dab09ff7ff0414817466a62 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 15 Jun 2026 07:51:30 -0700 Subject: [PATCH 17/20] image url support Signed-off-by: Shreyas Misra --- .../_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) 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 ad680ece47e8..1ab9acd30f62 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -31,6 +31,7 @@ from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.inputs.utils import load_image from tensorrt_llm.logger import logger from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS, COSMOS3_T2I_PARAMS @@ -754,7 +755,7 @@ def forward( # 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") From 83e7ca2aa60abd988c53636f210aaa1856e0249f Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 15 Jun 2026 07:58:03 -0700 Subject: [PATCH 18/20] generalize example script and update docstring Signed-off-by: Shreyas Misra --- .../cosmos3/{cosmos3_ti2v.py => cosmos3.py} | 47 ++++++++++++------- 1 file changed, 31 insertions(+), 16 deletions(-) rename examples/visual_gen/models/cosmos3/{cosmos3_ti2v.py => cosmos3.py} (84%) diff --git a/examples/visual_gen/models/cosmos3/cosmos3_ti2v.py b/examples/visual_gen/models/cosmos3/cosmos3.py similarity index 84% rename from examples/visual_gen/models/cosmos3/cosmos3_ti2v.py rename to examples/visual_gen/models/cosmos3/cosmos3.py index 8b8e92f7046f..09cdcec98217 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3_ti2v.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -13,11 +13,19 @@ # 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. +r"""Cosmos3 Text(+Image)-to-Video(+Audio) 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, or use ``prompts/i2v.json`` which includes a ``vision_path``. +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``): @@ -41,36 +49,43 @@ Deployment configs (``examples/visual_gen/configs/``): -- ``cosmos3-nano-1gpu.yaml`` — 1 GPU, FP8 dynamic quant +- ``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:: - # Text-to-video - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + # T2V: text-to-video + python cosmos3.py --model nvidia/Cosmos3-Nano \ --prompt_file prompts/t2v.json \ --visual_gen_args ../configs/cosmos3-nano-1gpu.yaml - # Image-to-video (vision_path is read from the prompt file) - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + # 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 - # Text-to-video with audio - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + # 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 - # Text-to-image - python cosmos3_ti2v.py --model nvidia/Cosmos3-Nano \ + # 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_ti2v.py --model nvidia/Cosmos3-Nano \ + 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 """ @@ -142,7 +157,7 @@ def resolve_prompt_and_options( def main(): - parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video example") + parser = argparse.ArgumentParser(description="Cosmos3 Text(+Image)-to-Video(+Audio) example") parser.add_argument( "--model", type=str, @@ -183,7 +198,7 @@ def main(): parser.add_argument( "--output_path", type=str, - default="cosmos3_ti2v_output.mp4", + default="cosmos3_output.mp4", help="Path to save the output video", ) parser.add_argument( From 0601b25eb45985e5e2dc4850efd453a384fcb802 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 30 Jun 2026 08:28:19 -0700 Subject: [PATCH 19/20] address comments Signed-off-by: Shreyas Misra --- examples/visual_gen/models/cosmos3/README.md | 77 +++++++++++++++++++ examples/visual_gen/models/cosmos3/cosmos3.py | 76 ------------------ .../visual_gen/models/cosmos3/defaults.py | 2 +- 3 files changed, 78 insertions(+), 77 deletions(-) create mode 100644 examples/visual_gen/models/cosmos3/README.md 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 index 09cdcec98217..de9e9e5010ad 100644 --- a/examples/visual_gen/models/cosmos3/cosmos3.py +++ b/examples/visual_gen/models/cosmos3/cosmos3.py @@ -13,82 +13,6 @@ # 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(+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 `_ -- `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 -- ``cosmos3-super-4gpu.yaml`` — 4 GPU, CFG + Ulysses + parallel VAE - -Example prompts live under ``prompts/`` (mirroring ``cosmos3-internal/inputs/omni``). - -Usage:: - - # 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 -""" import argparse import json diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py index 10cdcb268f59..f5747544946d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -75,7 +75,7 @@ description="Whether to enable audio generation.", ), "output_type": ExtraParamSchema( - type="str", + type="Literal['video', 'image']", default="video", description="Output modality: 'video' (T2V/I2V) or 'image' (text-to-image).", ), From 5b9dd9e089b58f24ea04e94f6c91da45ff4b464b Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 1 Jul 2026 12:34:55 -0700 Subject: [PATCH 20/20] [TRTLLM-13120][test] fix cosmos3 example path and t2i_smoke kwarg collision - test_cosmos3_example: point script_path at the reorganized examples/visual_gen/models/cosmos3/cosmos3.py (was flat cosmos3_ti2v.py) - _run_forward: accept height/width/guidance_scale overrides so test_t2i_smoke no longer passes duplicate keyword arguments to forward() Signed-off-by: Shreyas Misra --- .../defs/examples/visual_gen/test_visual_gen.py | 6 ++++-- .../_torch/visual_gen/test_cosmos3_pipeline.py | 17 +++++++++++++---- 2 files changed, 17 insertions(+), 6 deletions(-) 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/test_cosmos3_pipeline.py b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py index a72e425df0d1..c0329b71a060 100644 --- a/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py +++ b/tests/unittest/_torch/visual_gen/test_cosmos3_pipeline.py @@ -123,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,