diff --git a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py b/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py index bcd16893638a..ea6c3947fc78 100644 --- a/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py +++ b/tensorrt_llm/_torch/visual_gen/attention_backend/vanilla.py @@ -99,7 +99,14 @@ def forward( f"Invalid v shape: expected [B={q.shape[0]}, H_kv, S_kv, D={self.head_dim}], got {v.shape}" ) - return F.scaled_dot_product_attention(q, k, v, is_causal=is_causal, scale=self.scale) + return F.scaled_dot_product_attention( + q, + k, + v, + is_causal=is_causal, + scale=self.scale, + enable_gqa=self.num_heads != self.num_kv_heads, + ) @property def preferred_layout(self) -> AttentionTensorLayout: diff --git a/tensorrt_llm/_torch/visual_gen/models/__init__.py b/tensorrt_llm/_torch/visual_gen/models/__init__.py index 235b5a65f3d5..4b425d2e7d89 100644 --- a/tensorrt_llm/_torch/visual_gen/models/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/__init__.py @@ -19,6 +19,7 @@ from ..pipeline import BasePipeline from ..pipeline_registry import AutoPipeline, register_pipeline +from .cosmos3 import Cosmos3OmniMoTPipeline from .flux import Flux2Pipeline, FluxPipeline from .ltx2 import LTX2Pipeline # noqa: F401 from .wan import WanImageToVideoPipeline, WanPipeline @@ -30,5 +31,6 @@ "Flux2Pipeline", "WanPipeline", "WanImageToVideoPipeline", + "Cosmos3OmniMoTPipeline", "register_pipeline", ] diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py new file mode 100644 index 000000000000..98f83cee9b08 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py @@ -0,0 +1,18 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from .pipeline_cosmos3 import Cosmos3OmniMoTPipeline + +__all__ = ["Cosmos3OmniMoTPipeline"] diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py new file mode 100644 index 000000000000..e54e818356c4 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -0,0 +1,59 @@ +# 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. +"""Per-model default generation parameters for Cosmos3 pipelines. + +Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. +""" + +from typing import Dict + +from tensorrt_llm._torch.visual_gen.pipeline import ExtraParamSchema + +# --------------------------------------------------------------------------- +# Constant tables +# --------------------------------------------------------------------------- + +COSMOS3_720P_PARAMS = { + "height": 720, + "width": 1280, + "num_inference_steps": 35, + "guidance_scale": 6.0, + "max_sequence_length": 1024, + "num_frames": 189, + "frame_rate": 24.0, +} + +COSMOS3_EXTRA_SPECS: Dict[str, ExtraParamSchema] = { + "use_duration_template": ExtraParamSchema( + type="bool", + default=True, + description="Whether to use the duration template.", + ), + "use_resolution_template": ExtraParamSchema( + type="bool", + default=True, + description="Whether to use the resolution template.", + ), + "use_system_prompt": ExtraParamSchema( + type="bool", + default=False, + description="Whether to use the system prompt.", + ), + "use_guardrails": ExtraParamSchema( + type="bool", + default=True, + description="Whether to use the guardrails.", + ), +} diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py new file mode 100644 index 000000000000..881d830450fa --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -0,0 +1,67 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from __future__ import annotations + +from typing import Any + +import torch + +from tensorrt_llm.logger import logger + +GUARDRAIL_HF_REPO = "nvidia/Cosmos-1.0-Guardrail" +GUARDRAIL_REVISION = "cf03c0395fac8c4de386c0bdab12cc4fc8d66362" + + +def download_guardrail_checkpoint() -> str: + from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError + + try: + return snapshot_download( + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_REVISION, + local_files_only=True, + ) + except FileNotFoundError: + logger.warning(f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO}") + try: + return snapshot_download( + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_REVISION, + ) + except GatedRepoError: + raise ValueError( + "Cosmos Guardrail checkpoint not found. " + "Please ensure " + "a) you have accepted the terms of use (https://huggingface.co/nvidia/Cosmos-1.0-Guardrail) " + "b) you have set a valid HF_TOKEN environment variable" + ) + + +def check_video_safety(video_tensor: torch.Tensor, safety_checker: Any) -> torch.Tensor | None: + v = video_tensor.detach().cpu() + was_batched = v.dim() == 5 + if was_batched: + v = v[0] + frames_np = v.numpy() + frames_np = safety_checker.check_video_safety(frames_np) + if frames_np is None: + return None + + result = torch.from_numpy(frames_np) + if was_batched: + result = result.unsqueeze(0) + return result.to(video_tensor.device) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py new file mode 100644 index 000000000000..800006216bf2 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -0,0 +1,625 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +import os +import time +from typing import List, Optional, Union + +import PIL.Image +import torch +from diffusers import AutoencoderKLWan, UniPCMultistepScheduler +from diffusers.utils.torch_utils import randn_tensor +from diffusers.video_processor import VideoProcessor +from transformers import Qwen2Tokenizer + +from tensorrt_llm._torch.visual_gen.output import CudaPhaseTimer, PipelineOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import PipelineComponent, register_pipeline +from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor +from tensorrt_llm._utils import nvtx_range +from tensorrt_llm.logger import logger + +from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS +from .guardrails import check_video_safety, download_guardrail_checkpoint +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_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." +TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" + + +# TODO: add hf_ids +@register_pipeline("Cosmos3OmniMoTPipeline") +class Cosmos3OmniMoTPipeline(BasePipeline): + def __init__(self, model_config): + super().__init__(model_config) + + def _init_transformer(self) -> None: + logger.info("Initializing Cosmos3VFMTransformer") + self.transformer = Cosmos3VFMTransformer(self.model_config) + + def load_weights(self, weights: dict) -> None: + if self.transformer is not None and hasattr(self.transformer, "load_weights"): + transformer_weights = weights.get("transformer", weights) + self.transformer.load_weights(transformer_weights) + self.transformer.eval() + + def load_standard_components( + self, checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = [] + ) -> None: + skip_components = skip_components or [] + + if PipelineComponent.TOKENIZER not in skip_components: + logger.info("Loading tokenizer...") + self.tokenizer = Qwen2Tokenizer.from_pretrained( + checkpoint_dir, + subfolder="text_tokenizer", + ) + + # Cosmos3 canonical defaults — overwritten if VAE is loaded + self.vae_scale_factor_temporal = 4 + self.vae_scale_factor_spatial = 16 + + if PipelineComponent.VAE not in skip_components: + logger.info("Loading VAE...") + self.vae = AutoencoderKLWan.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.VAE, + torch_dtype=torch.bfloat16, # load VAE in BF16 for memory saving + ).to(device) + + self.vae_scale_factor_temporal = getattr( + self.vae.config, "scale_factor_temporal", self.vae_scale_factor_temporal + ) + self.vae_scale_factor_spatial = getattr( + self.vae.config, "scale_factor_spatial", self.vae_scale_factor_spatial + ) + self.transformer.temporal_compression_factor = self.vae_scale_factor_temporal + + if PipelineComponent.SCHEDULER not in skip_components: + logger.info("Loading scheduler...") + self.scheduler = UniPCMultistepScheduler.from_pretrained( + checkpoint_dir, + subfolder=PipelineComponent.SCHEDULER, + ) + + if not TRTLLM_DISABLE_COSMOS3_GUARDRAILS: + # lazy import + try: + from cosmos_guardrail import CosmosSafetyChecker + except (ImportError, ModuleNotFoundError): + raise ValueError( + "Cosmos Guardrail is not installed. This is in violation of the " + "[NVIDIA Open Model License Agreement]" + "(https://www.nvidia.com/en-us/agreements/enterprise-software/nvidia-open-model-license). " + "Please run the following installation commands or " + "explicitly disable guardrails by setting TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 " + "(user is responsible for deploying the model without guardrails). " + "- `pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python`" + ) + # Guardrails are only evaluated on rank 0; load them only there to avoid + # dead model weights occupying GPU memory on every other rank. + if self.rank == 0: + # the download guardrail checkpoint will bypass CosmosSafetyChecker's checkpoint download. + # Both will use HF_HOME as the cache directory. + download_guardrail_checkpoint() + self.safety_checker = CosmosSafetyChecker() + self.safety_checker.to(device) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) + + @property + def default_warmup_resolutions(self): + return [(720, 1280)] + + @property + def default_warmup_num_frames(self): + return [189] + + @property + def default_generation_params(self): + return dict(COSMOS3_720P_PARAMS) + + @property + def extra_param_specs(self): + return dict(COSMOS3_EXTRA_SPECS) + + def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> None: + with torch.no_grad(): + self.forward( + prompt="warmup", + negative_prompt="", + height=height, + width=width, + num_frames=num_frames, + num_inference_steps=steps, + guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], + seed=42, + max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], + use_guardrails=False, + image=None, + ) + + def infer(self, req): + return self.forward( + prompt=req.prompt, + negative_prompt=req.params.negative_prompt, + image=req.params.image, + height=req.params.height, + width=req.params.width, + num_frames=req.params.num_frames, + num_inference_steps=req.params.num_inference_steps, + guidance_scale=req.params.guidance_scale, + seed=req.params.seed, + max_sequence_length=req.params.max_sequence_length, + frame_rate=req.params.frame_rate, + use_duration_template=req.params.extra_params.get("use_duration_template", True), + use_resolution_template=req.params.extra_params.get("use_resolution_template", True), + use_system_prompt=req.params.extra_params.get("use_system_prompt", False), + use_guardrails=req.params.extra_params.get("use_guardrails", True), + ) + + def _format_prompt_with_template( + self, + prompt: str, + *, + height: int, + width: int, + num_frames: int, + frame_rate: float, + use_duration_template: bool = True, + use_resolution_template: bool = True, + ) -> str: + prompt = prompt.strip() + + if use_duration_template and num_frames > 1: + 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 + + def _resize_and_center_crop_image( + self, image: PIL.Image.Image, height: int, width: int + ) -> PIL.Image.Image: + """Match Cosmos3 reference preprocessing for conditioning images.""" + orig_w, orig_h = image.size + scaling_ratio = max(width / orig_w, height / orig_h) + resize_w = int(math.ceil(scaling_ratio * orig_w)) + resize_h = int(math.ceil(scaling_ratio * orig_h)) + + image = image.resize((resize_w, resize_h), PIL.Image.Resampling.LANCZOS) + + left = max((resize_w - width) // 2, 0) + top = max((resize_h - height) // 2, 0) + return image.crop((left, top, left + width, top + height)) + + @nvtx_range("_tokenize_prompt", color="blue") + def _tokenize_prompt( + self, text: str, max_sequence_length: int, use_system_prompt: bool = False + ): + """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}] + if use_system_prompt + else [] + ) + conversations.append( + {"role": "user", "content": text}, + ) + token_ids = self.tokenizer.apply_chat_template( + conversations, + tokenize=True, + add_generation_prompt=True, + return_dict=False, + ) + reserved_tokens = 2 + if max_sequence_length < reserved_tokens: + raise ValueError( + f"max_sequence_length must be at least {reserved_tokens}, got {max_sequence_length}" + ) + token_ids = token_ids[: max_sequence_length - reserved_tokens] + token_ids.append(self.tokenizer.eos_token_id) # 151645 + token_ids.append(self.tokenizer.convert_tokens_to_ids("<|vision_start|>")) # 151652 + seq_len = len(token_ids) + + # Pad to max_sequence_length + pad_len = max_sequence_length - seq_len + attention_mask = [1] * seq_len + [0] * pad_len + token_ids = token_ids + [self.tokenizer.pad_token_id or 0] * pad_len + + input_ids = torch.tensor([token_ids], dtype=torch.long, device=self.device) + attention_mask = torch.tensor([attention_mask], dtype=torch.long, device=self.device) + return input_ids, attention_mask + + # ========================================================================= + # Latent preparation + # ========================================================================= + + @nvtx_range("_prepare_latents", color="blue") + def _prepare_latents(self, height, width, num_frames, generator): + num_channels_latents = self.transformer.latent_channel_size + num_latent_frames = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + shape = ( + 1, + num_channels_latents, + num_latent_frames, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ) + return randn_tensor(shape, generator=generator, device=self.device, dtype=self.dtype) + + # -- I2V latent preparation ----------------------------------------------- + + def _encode_conditioning_video( + self, + image_tensor: torch.Tensor, + num_frames: int, + height: int, + width: int, + ) -> torch.Tensor: + """VAE-encode a conditioning image as a full-length video. + + The WAN VAE has temporal compression (factor 4), so encoding a single + frame produces degenerate temporal features. Following imaginaire4's + ``build_conditioned_video_batch``, we fill the entire pixel-space video + with the conditioning image (repeating it across all frames) so the + temporal encoder sees plausible content everywhere. The caller then + keeps only the conditioned latent frame(s) and replaces the rest with + noise. + + Args: + image_tensor: [1, 3, H, W] in [-1, 1] + num_frames: total pixel frames for the video + height: pixel height + width: pixel width + + Returns: + [1, C, T_latent, H_latent, W_latent] normalized latent of the + full conditioning video. + """ + # Build pixel-space video: repeat the conditioning image across all frames + # image_tensor: [1, 3, H, W] -> [1, 3, 1, H, W] -> [1, 3, num_frames, H, W] + video = image_tensor.unsqueeze(2).expand(-1, -1, num_frames, -1, -1).contiguous() + video = video.to(device=self.device, dtype=self.vae.dtype) + + latent = self.vae.encode(video).latent_dist.mode() + + # Normalize (inverse of _decode_latents denormalization) + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(latent.device, latent.dtype) + ) + latent = (latent - latents_mean) / latents_std + else: + scaling_factor = getattr(self.vae.config, "scaling_factor", 1.0) + latent = latent * scaling_factor + + return latent.to(self.dtype) + + def _prepare_latents_i2v( + self, + image_tensor: torch.Tensor, + height: int, + width: int, + num_frames: int, + generator: torch.Generator, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Prepare initial latents with frame 0 conditioned on the input image. + + The conditioning image is repeated across all pixel frames before VAE + encoding so the temporal encoder sees plausible content everywhere + (avoids degenerate single-frame encoding with the WAN VAE's temporal + compression). Only frame 0 of the resulting latent is kept clean; + the rest is replaced with noise. + + Returns: + latents: [1, C, T_lat, H_lat, W_lat] with frame 0 = image, rest = noise + velocity_mask: [1, 1, T_lat, 1, 1] with frame 0 = 0, rest = 1 + image_latent: [1, C, 1, H_lat, W_lat] clean frame 0 for re-injection + """ + C = self.transformer.latent_channel_size + T_lat = (num_frames - 1) // self.vae_scale_factor_temporal + 1 + + # Pure noise + noise = randn_tensor( + ( + 1, + C, + T_lat, + height // self.vae_scale_factor_spatial, + width // self.vae_scale_factor_spatial, + ), + generator=generator, + device=self.device, + dtype=self.dtype, + ) + + # Encode full conditioning video (image repeated across all frames) + cond_latent = self._encode_conditioning_video( + image_tensor, + num_frames, + height, + width, + ) # [1, C, T_lat, H_lat, W_lat] + + # Keep only frame 0 for conditioning; replace rest with noise + image_latent = cond_latent[:, :, 0:1, :, :] # [1, C, 1, H_lat, W_lat] + + condition_mask = torch.zeros(1, 1, T_lat, 1, 1, device=self.device, dtype=self.dtype) + condition_mask[:, :, 0, :, :] = 1.0 + + latents = condition_mask * cond_latent + (1.0 - condition_mask) * noise + + velocity_mask = 1.0 - condition_mask + return latents, velocity_mask, image_latent + + # ========================================================================= + # VAE decode + # ========================================================================= + + @nvtx_range("_decode_latents", color="blue") + def _decode_latents(self, latents): + latents = latents.to(self.vae.dtype) + + if hasattr(self.vae.config, "latents_mean") and hasattr(self.vae.config, "latents_std"): + if not hasattr(self, "_latents_mean"): + self._latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, -1, 1, 1, 1) + .to(self.device, self.vae.dtype) + ) + self._latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, -1, 1, 1, 1) + .to(self.device, self.vae.dtype) + ) + latents = (latents * self._latents_std) + self._latents_mean + else: + scaling_factor = self.vae.config.get("scaling_factor", 1.0) + latents = latents / scaling_factor + + video = self.vae.decode(latents, return_dict=False)[0] + video = postprocess_video_tensor(video) + return video + + # ========================================================================= + # Forward (main generation entry point) + # ========================================================================= + + @nvtx_range("Cosmos3OmniMoTPipeline.forward") + @torch.inference_mode() + def forward( + self, + prompt: Union[str, List[str]], + negative_prompt: Optional[str] = None, + image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, + height: int = COSMOS3_720P_PARAMS["height"], + width: int = COSMOS3_720P_PARAMS["width"], + num_frames: int = COSMOS3_720P_PARAMS["num_frames"], + num_inference_steps: int = COSMOS3_720P_PARAMS["num_inference_steps"], + guidance_scale: float = COSMOS3_720P_PARAMS["guidance_scale"], + seed: int = 42, + max_sequence_length: int = COSMOS3_720P_PARAMS["max_sequence_length"], + frame_rate: float = COSMOS3_720P_PARAMS["frame_rate"], + use_duration_template: bool = COSMOS3_EXTRA_SPECS["use_duration_template"].default, + 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, + ): + pipeline_start = time.time() + timer = CudaPhaseTimer() + timer.mark_pre_start() + + use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + + if isinstance(prompt, str): + prompt = [prompt] + batch_size = len(prompt) + + if batch_size > 1: + # TODO: support batch generation + raise ValueError("Batch generation is not supported for Cosmos3") + + # Validate image input — only single image is supported for batch generation + if image is not None and not isinstance(image, (PIL.Image.Image, torch.Tensor, str)): + raise ValueError( + f"`image` must be a PIL.Image, torch.Tensor, or file path string, " + f"got {type(image)}. Batch of different images is not supported; " + f"use a single image with multiple prompts instead." + ) + + # Text guardrail — check both positive and user-supplied negative prompts. + # None negative_prompt means the hardcoded 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) + if negative_prompt is not None: + prompts_to_check.append(negative_prompt) + for p in prompts_to_check: + is_safe = self.safety_checker.check_text_safety(p) + if not is_safe: + logger.warning("Text guardrail blocked prompt") + text_blocked.fill_(1) + break + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.broadcast(text_blocked, src=0) + + if text_blocked.item(): + timer.mark_end() + return timer.fill(PipelineOutput()) + + generator = torch.Generator(device=self.device).manual_seed(seed) + + if negative_prompt is None: + negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT + + negative_prompt = self._format_prompt_with_template( + 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, + ) + + prompt = [ + self._format_prompt_with_template( + 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, + ) + for p in prompt + ] + logger.info(f"Prompt with metadata: '{prompt}'") + + prompt = prompt[0] + + # 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) + uncond_ids, uncond_mask = self._tokenize_prompt( + negative_prompt, max_sequence_length, use_system_prompt + ) + + # 2. Prepare latents + if image is not None: + if isinstance(image, str): + image = PIL.Image.open(image).convert("RGB") + + if isinstance(image, PIL.Image.Image): + image = image.convert("RGB") + image = self._resize_and_center_crop_image(image, height=height, width=width) + image = self.video_processor.preprocess( + image, + height=height, + width=width, + ) + + latents, velocity_mask, image_latent = self._prepare_latents_i2v( + image, height=height, width=width, num_frames=num_frames, generator=generator + ) + else: + latents = self._prepare_latents(height, width, num_frames, generator) + velocity_mask = None + image_latent = None + + # Compute video shape in latent space + T_latent = latents.shape[2] + H_latent = latents.shape[3] + W_latent = latents.shape[4] + video_shape = (T_latent, H_latent, W_latent) + + # 3. Set up scheduler + self.scheduler.set_timesteps(num_inference_steps, device=self.device) + + # 4. Build forward_fn for the denoise loop + def forward_fn( + latent_input, extra_stream_latents, timestep, encoder_hidden_states, extra_tensors + ): + """Cosmos3 forward function for BasePipeline.denoise(). + + Since Cosmos3 embeds text internally, we pass token IDs via extra_tensors + rather than through encoder_hidden_states. + """ + noise_pred = self.transformer( + hidden_states=latent_input, + timestep=timestep, + text_ids=extra_tensors["text_ids"], + text_mask=extra_tensors["text_mask"], + video_shape=video_shape, + fps=frame_rate, + noisy_frame_mask=velocity_mask, + ) + if velocity_mask is not None: + noise_pred = noise_pred * velocity_mask + return 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 + # We pass text IDs/masks through extra_cfg_tensors so they get split correctly + extra_cfg_tensors = { + "text_ids": (cond_ids, uncond_ids), + "text_mask": (cond_mask, uncond_mask), + } + + self.transformer.reset_cache() + + # 6. Denoise + timer.mark_denoise_start() + latents = 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, + ) + timer.mark_post_start() + + # 7. Decode + logger.info("Decoding video...") + decode_start = time.time() + + if image_latent is not None: + latents = latents.clone() + latents[:, :, 0:1, :, :] = image_latent.to(device=latents.device, dtype=latents.dtype) + + video = self.decode_latents(latents, self._decode_latents) + + # Video guardrails + 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") + + if use_guardrails and self.safety_checker is not None: + video = check_video_safety(video, self.safety_checker) + + timer.mark_end() + return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate)) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py new file mode 100644 index 000000000000..d260b8e10ea1 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -0,0 +1,1159 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import math +from typing import Tuple + +import torch +import torch.distributed as dist +import torch.nn as nn +import torch.nn.functional as F +from diffusers.models.embeddings import TimestepEmbedding + +from tensorrt_llm._torch.attention_backend.interface import PredefinedAttentionMask +from tensorrt_llm._torch.modules.embedding import Embedding +from tensorrt_llm._torch.modules.gated_mlp import GatedMLP +from tensorrt_llm._torch.modules.linear import Linear +from tensorrt_llm._torch.visual_gen.config import DiffusionModelConfig +from tensorrt_llm._torch.visual_gen.modules.attention import Attention, QKVMode +from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader +from tensorrt_llm.logger import logger +from tensorrt_llm.models.modeling_utils import QuantConfig + + +class Qwen3VLTextRMSNorm(nn.Module): + def __init__( + self, hidden_size: int, eps: float = 1e-6, dtype: torch.dtype = torch.bfloat16 + ) -> None: + """ + Qwen3VLTextRMSNorm is equivalent to T5LayerNorm + """ + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + self.dtype = dtype + + def post_load_weights(self): + self.weight.data = self.weight.data.to(self.dtype) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + output = self.weight * hidden_states.to(input_dtype) + return output + + +def compute_mrope_position_ids_text( + num_tokens: int, + temporal_offset: int, +) -> tuple[torch.Tensor, int]: + """Generate 3D mRoPE position IDs for text tokens. + + Text tokens: all three axes (T, H, W) share the same monotonically + increasing position IDs: (0,0,0), (1,1,1), (2,2,2), ... + + Returns: + (position_ids [3, num_tokens], next_temporal_offset) + """ + ids = torch.arange(num_tokens, dtype=torch.long) + temporal_offset + mrope_ids = ids.unsqueeze(0).expand(3, -1).contiguous() + return mrope_ids, temporal_offset + num_tokens + + +def compute_mrope_position_ids_vision( + grid_t: int, + grid_h: int, + grid_w: int, + temporal_offset: int | float, + fps: float | None = None, + base_fps: float = 24.0, + temporal_compression_factor: int = 4, + enable_fps_modulation: bool = False, +) -> tuple[torch.Tensor, int | float]: + """Generate 3D mRoPE position IDs for vision tokens. + + Creates a (T, H, W) position grid. Spatial indices reset to 0 + per vision segment (Qwen3VL-style, reset_spatial_indices=True). + Flattened in T-major order. + + When ``enable_fps_modulation`` is ``True``, temporal positions are scaled + to reflect real time so that videos at different frame rates get comparable + temporal embeddings. + + Returns: + (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset) + """ + if enable_fps_modulation and fps is not None: + tps = fps / temporal_compression_factor + base_tps = base_fps / temporal_compression_factor + frame_indices = torch.arange(grid_t, dtype=torch.float32) + t_index = ( + (frame_indices / tps * base_tps + temporal_offset) + .view(-1, 1) + .expand(-1, grid_h * grid_w) + .flatten() + ) + else: + t_index = torch.arange(grid_t, dtype=torch.long).view(-1, 1).expand( + -1, grid_h * grid_w + ).flatten() + int(temporal_offset) + + h_index = ( + torch.arange(grid_h, dtype=torch.long).view(1, -1, 1).expand(grid_t, -1, grid_w).flatten() + ) + w_index = ( + torch.arange(grid_w, dtype=torch.long).view(1, 1, -1).expand(grid_t, grid_h, -1).flatten() + ) + + if enable_fps_modulation: + mrope_ids = torch.stack( + [t_index, h_index.to(torch.float32), w_index.to(torch.float32)], dim=0 + ) + else: + mrope_ids = torch.stack([t_index, h_index, w_index], dim=0) + + next_offset = math.ceil(mrope_ids.max().item()) + 1 + return mrope_ids, next_offset + + +class TimestepEmbedder(nn.Module): + """ + Embeds scalar timesteps into vector representations. + """ + + def __init__( + self, + hidden_size, + frequency_embedding_size=256, + max_period=10000, + target_dtype=torch.bfloat16, + ): + super().__init__() + self.mlp = TimestepEmbedding( + in_channels=frequency_embedding_size, time_embed_dim=hidden_size, act_fn="silu" + ) + self.frequency_embedding_size = frequency_embedding_size + self.hidden_size = hidden_size + + half = frequency_embedding_size // 2 + freqs = torch.exp( + -math.log(max_period) * torch.arange(start=0, end=half, dtype=target_dtype) / half + ) + self.register_buffer("freqs", freqs, persistent=False) + + def _init_weights(self): + std = 1.0 / math.sqrt(self.frequency_embedding_size) + torch.nn.init.trunc_normal_(self.mlp.linear_1.weight, std=std, a=-3 * std, b=3 * std) + + std = 1.0 / math.sqrt(self.hidden_size) + torch.nn.init.trunc_normal_(self.mlp.linear_2.weight, std=std, a=-3 * std, b=3 * std) + + def forward(self, t): + # use .float() here if acc loss + args = t[:, None] * self.freqs[None] + t_freq = torch.cat([torch.cos(args), torch.sin(args)], dim=-1) + t_emb = self.mlp(t_freq) + return t_emb + + +def qwen3_rotate_half(x: torch.Tensor) -> torch.Tensor: + """Qwen3/Llama-style rotate_half: split first/second half of head_dim.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def qwen3_apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Qwen3-style RoPE: (x * cos) + (rotate_half(x) * sin). + + Args: + q: [B, S, H, D] + k: [B, S, H_kv, D] + cos: [1, S, 1, D] or broadcastable + sin: [1, S, 1, D] or broadcastable + """ + q_embed = (q * cos) + (qwen3_rotate_half(q) * sin) + k_embed = (k * cos) + (qwen3_rotate_half(k) * sin) + return q_embed, k_embed + + +class Cosmos3CausalAttention(Attention): + """Understanding pathway: causal self-attention on text tokens. + + Inherits from Attention for projections (SEPARATE_QKV), per-head QK norms, + and backend. Overrides forward to: + - Reshape to 4D before QK norm (per-head) + - Apply Qwen3-style RoPE (rotate_half, not interleaved) + - Pass causal mask to backend + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + model_config: DiffusionModelConfig, + layer_idx: int = 0, + ): + super().__init__( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + qkv_mode=QKVMode.SEPARATE_QKV, + qk_norm=True, + qk_norm_mode="per_head", + bias=False, + config=model_config, + layer_idx=layer_idx, + enable_ulysses=False, + ) + self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + + def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-head RMSNorm on 4D tensors [B, S, H, D].""" + q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) + k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) + return q, k + + def forward_with_kv( + self, + hidden_states: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> torch.Tensor: + batch_size, seq_len = hidden_states.shape[:2] + + q, k, v = self.get_qkv(hidden_states) + + q = q.view(batch_size, seq_len, self.num_attention_heads, self.head_dim) + k = k.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim) + v = v.view(batch_size, seq_len, self.num_key_value_heads, self.head_dim) + + q, k = self.apply_qk_norm(q, k) + q, k = qwen3_apply_rotary_pos_emb(q, k, freqs_cos, freqs_sin) + + out = self._attn_impl( + q, + k, + v, + attention_mask=PredefinedAttentionMask.CAUSAL, + ) + + return self.to_out[0](out), k, v + + def forward(self): + raise NotImplementedError( + "forward method not implemented for Cosmos3CausalAttention. Use forward_with_kv instead." + ) + + +class Cosmos3CrossAttention(Attention): + """Generation pathway: full attention where visual Q attends to all K/V. + + Inherits from Attention for gen-pathway projections, per-head QK norms, + and backend. Overrides forward to: + - Accept pre-computed und K/V for concatenation + - Reshape to 4D before QK norm (per-head) + - Apply Qwen3-style RoPE + - Full (non-causal) attention with Q_gen attending to [K_und, K_gen] + """ + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + num_key_value_heads: int, + head_dim: int, + model_config: DiffusionModelConfig, + layer_idx: int = 0, + ): + original_backend = model_config.attention.backend + if model_config.attention.backend == "TRTLLM": + # TRTLLM backend is not supported for Cosmos3CrossAttention + model_config.attention.backend = "VANILLA" + + super().__init__( + hidden_size=hidden_size, + num_attention_heads=num_attention_heads, + num_key_value_heads=num_key_value_heads, + head_dim=head_dim, + qkv_mode=QKVMode.FUSE_QKV, + qk_norm=True, + qk_norm_mode="per_head", + bias=False, + config=model_config, + layer_idx=layer_idx, + enable_ulysses=True, + ) + model_config.attention.backend = original_backend + + self.norm_q = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + self.norm_k = Qwen3VLTextRMSNorm(hidden_size=head_dim, dtype=torch.bfloat16) + + def apply_qk_norm(self, q: torch.Tensor, k: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """Per-head RMSNorm on 4D tensors [B, S, H, D].""" + q = F.rms_norm(q, (q.shape[-1],), self.norm_q.weight, self.norm_q.variance_epsilon) + k = F.rms_norm(k, (k.shape[-1],), self.norm_k.weight, self.norm_k.variance_epsilon) + return q, k + + def forward( + self, + hidden_states: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + freqs_cos: torch.Tensor, + freqs_sin: torch.Tensor, + ) -> torch.Tensor: + """ + Args: + hidden_states: [B, S_gen, hidden_size] visual tokens + k_und: [B, S_und, H_kv, D] pre-computed und keys (post-norm, post-RoPE) + v_und: [B, S_und, H_kv, D] pre-computed und values + freqs_cos: [B, S_gen, 1, D] cosine part of RoPE + freqs_sin: [B, S_gen, 1, D] sine part of RoPE + + Returns: + [B, S_gen, hidden_size] cross-attention output + """ + batch_size, seq_len_gen = hidden_states.shape[:2] + + q, k, v = self.get_qkv(hidden_states) + + q = q.view(batch_size, seq_len_gen, self.num_attention_heads, self.head_dim) + k = k.view(batch_size, seq_len_gen, self.num_key_value_heads, self.head_dim) + v = v.view(batch_size, seq_len_gen, self.num_key_value_heads, self.head_dim) + + 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, + ) + + return self.to_out[0](out) + + +class Cosmos3UndDecoderLayer(nn.Module): + """Understanding pathway decoder layer: causal self-attention + MLP.""" + + def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + hidden_size = model_config.pretrained_config.hidden_size + intermediate_size = model_config.pretrained_config.intermediate_size + + self.self_attn = Cosmos3CausalAttention( + hidden_size=hidden_size, + num_attention_heads=model_config.pretrained_config.num_attention_heads, + num_key_value_heads=model_config.pretrained_config.num_key_value_heads, + head_dim=model_config.pretrained_config.head_dim, + model_config=model_config, + layer_idx=layer_idx, + ) + self.input_layernorm = Qwen3VLTextRMSNorm( + hidden_size=hidden_size, + eps=model_config.pretrained_config.rms_norm_eps, + dtype=torch.bfloat16, + ) + self.post_attention_layernorm = Qwen3VLTextRMSNorm( + hidden_size=hidden_size, + eps=model_config.pretrained_config.rms_norm_eps, + dtype=torch.bfloat16, + ) + self.mlp = GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + freqs: Tuple[torch.Tensor, torch.Tensor], + ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """ + Returns: + (hidden_states, K, V) where K/V are post-QKnorm, post-RoPE + for consumption by the GEN cross-attention. + """ + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + cos, sin = freqs + attn_out, k, v = self.self_attn.forward_with_kv(hidden_states, cos, sin) + hidden_states = residual + attn_out + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + B, S, D = hidden_states.shape + hidden_states = self.mlp(hidden_states.view(-1, D)).view(B, S, D) + hidden_states = residual + hidden_states + + return hidden_states, k, v + + +class Cosmos3GenDecoderLayer(nn.Module): + """Generation pathway decoder layer: cross-attention (to UND K/V) + MLP.""" + + def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): + super().__init__() + self.layer_idx = layer_idx + hidden_size = model_config.pretrained_config.hidden_size + intermediate_size = model_config.pretrained_config.intermediate_size + + self.cross_attention = Cosmos3CrossAttention( + hidden_size=hidden_size, + num_attention_heads=model_config.pretrained_config.num_attention_heads, + num_key_value_heads=model_config.pretrained_config.num_key_value_heads, + head_dim=model_config.pretrained_config.head_dim, + model_config=model_config, + layer_idx=layer_idx, + ) + self.input_layernorm = Qwen3VLTextRMSNorm( + hidden_size=hidden_size, + eps=model_config.pretrained_config.rms_norm_eps, + dtype=torch.bfloat16, + ) + self.post_attention_layernorm = Qwen3VLTextRMSNorm( + hidden_size=hidden_size, + eps=model_config.pretrained_config.rms_norm_eps, + dtype=torch.bfloat16, + ) + self.mlp = GatedMLP( + hidden_size=hidden_size, + intermediate_size=intermediate_size, + bias=False, + dtype=torch.bfloat16, + config=model_config, + layer_idx=layer_idx, + ) + + def forward( + self, + hidden_states: torch.Tensor, + k_und: torch.Tensor, + v_und: torch.Tensor, + freqs: Tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + cos, sin = freqs + hidden_states = self.cross_attention( + hidden_states, + k_und=k_und, + v_und=v_und, + freqs_cos=cos, + freqs_sin=sin, + ) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + B, S, D = hidden_states.shape + hidden_states = self.mlp(hidden_states.view(-1, D)).view(B, S, D) + hidden_states = residual + hidden_states + + return hidden_states + + +def _compute_default_rope_parameters( + model_config: DiffusionModelConfig, +) -> tuple["torch.Tensor", float]: + """ + Computes the inverse frequencies according to the original RoPE implementation + Args: + config ([`~transformers.PretrainedConfig`]): + The model configuration. This function assumes that the config will provide at least the following + properties: + + * rope_theta (`float`): The base wavelength from which the inverse frequencies will be derived. + * hidden_size (`int`): The numerator when deriving a head_dim, if not provided directly. + * num_attention_heads (`int`): The denominator when deriving a head_dim, if not provided directly. + + Additionally, this function will make use of the following properties if they are found in the config: + + * head_dim (`int`, *optional*): The size of the key-value heads in the model. If None, this value will be + derived as hidden_size // num_attention_heads. + * partial_rotary_factor (`float`, *optional*): If less than 1.0, inverse frequencies will be returned for + the first fraction of the head_dim. Defaults to 1.0. + device (`torch.device`): + The device to use for initialization of the inverse frequencies. + seq_len (`int`, *optional*): + The current sequence length. Unused for this type of RoPE. + + Returns: + Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the + post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). + """ + base = model_config.pretrained_config.rope_theta + partial_rotary_factor = 1 + head_dim = model_config.pretrained_config.head_dim + dim = int(head_dim * partial_rotary_factor) + + attention_factor = 1.0 # Unused in this type of RoPE + + # Compute the inverse frequencies + inv_freq = 1.0 / ( + base ** (torch.arange(0, dim, 2, dtype=torch.int64).to(dtype=torch.float) / dim) + ) + return inv_freq, attention_factor + + +class Qwen3VLTextRotaryEmbedding(nn.Module): + def __init__(self, model_config: DiffusionModelConfig): + super().__init__() + self.rope_type = model_config.pretrained_config.rope_scaling["rope_type"] + self.max_seq_len_cached = model_config.pretrained_config.max_position_embeddings + self.original_max_seq_len = model_config.pretrained_config.max_position_embeddings + + self.mrope_section = model_config.pretrained_config.rope_scaling["mrope_section"] + + inv_freq, self.attention_scaling = _compute_default_rope_parameters(model_config) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def apply_interleaved_mrope(self, freqs, mrope_section): + """Apply interleaved MRoPE to 3D rotary embeddings. + Reorganizes frequency layout from chunked [TTT...HHH...WWW] to + interleaved [THTHWHTHW...TT], preserving frequency continuity. + args: + x: (3, bs, seq_len, head_dim // 2) + mrope_section: (3,) + returns: + x_t: (bs, seq_len, head_dim // 2) + """ + freqs_t = freqs[0] # just overwrite the first dimension T + for dim, offset in enumerate((1, 2), start=1): # H, W + length = mrope_section[dim] * 3 + idx = slice(offset, length, 3) + freqs_t[..., idx] = freqs[dim, ..., idx] + return freqs_t + + @torch.no_grad() + def forward(self, x, position_ids): + assert self.inv_freq.dtype == torch.float32, ( + f"inv_freq must be float32, but got {self.inv_freq.dtype}" + ) + + # In contrast to other models, Qwen3VL has different position ids for the grids + # So we expand the inv_freq to shape (3, ...) + if position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + inv_freq_expanded = ( + self.inv_freq[None, None, :, None] + .float() + .expand(3, position_ids.shape[1], -1, 1) + .to(x.device) + ) + position_ids_expanded = position_ids[:, :, None, :].float() # shape (3, bs, 1, positions) + + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(2, 3) + freqs = self.apply_interleaved_mrope(freqs, self.mrope_section) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() * self.attention_scaling + sin = emb.sin() * self.attention_scaling + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +class Cosmos3LanguageModel(nn.Module): + """Understanding pathway: a standard causal LM that processes text tokens. + + Returns per-layer K/V tensors for the generation pathway's cross-attention. + The UND pathway is independent of the denoising step, so its K/V can be + computed once and reused across all sampling steps. + """ + + def __init__(self, model_config: DiffusionModelConfig): + super().__init__() + hidden_size = model_config.pretrained_config.hidden_size + num_hidden_layers = model_config.pretrained_config.num_hidden_layers + + self.embed_tokens = Embedding( + model_config.pretrained_config.vocab_size, + hidden_size, + dtype=torch.bfloat16, + gather_output=True, + ) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(model_config) + self.layers = nn.ModuleList( + [Cosmos3UndDecoderLayer(model_config, layer_idx=i) for i in range(num_hidden_layers)] + ) + + def forward( + self, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + freqs: Tuple[torch.Tensor, torch.Tensor], + ) -> list[Tuple[torch.Tensor, torch.Tensor]]: + """ + Args: + text_ids: [B, S] token IDs + text_mask: [B, S] float mask (1=real, 0=pad) + freqs: (cos, sin) each [B, S, 1, D] — precomputed UND RoPE + + Returns: + List of (K, V) per layer. K/V are [B, S, H_kv, D], post-QKnorm + and post-RoPE, ready for GEN cross-attention. + """ + hidden = self.embed_tokens(text_ids) + mask_3d = text_mask.unsqueeze(-1) # [B, S, 1] + + cached_kv: list[Tuple[torch.Tensor, torch.Tensor]] = [] + for layer in self.layers: + hidden = hidden * mask_3d + hidden, k, v = layer(hidden, freqs) + cached_kv.append((k, v)) + + return cached_kv + + +class Cosmos3VFMTransformer(nn.Module): + def __init__(self, model_config: DiffusionModelConfig): + super().__init__() + self.model_config = model_config + pretrained_config = model_config.pretrained_config + + self.hidden_size = pretrained_config.hidden_size + self.num_hidden_layers = pretrained_config.num_hidden_layers + self.latent_patch_size = pretrained_config.latent_patch_size + self.latent_channel_size = pretrained_config.latent_channel + self.patch_latent_dim = (self.latent_patch_size**2) * self.latent_channel_size + self.timestep_scale = pretrained_config.timestep_scale + self.base_fps = pretrained_config.base_fps + + # Comes from VAE. Updated after VAE is loaded. + self.temporal_compression_factor = 4 + + self.unified_3d_mrope_temporal_modality_margin = ( + pretrained_config.unified_3d_mrope_temporal_modality_margin + ) + self.num_attention_heads = pretrained_config.num_attention_heads + self.num_kv_heads = pretrained_config.num_key_value_heads + self.enable_fps_modulation = pretrained_config.enable_fps_modulation + + if pretrained_config.position_embedding_type != "unified_3d_mrope": + raise ValueError( + f"Position embedding type {pretrained_config.position_embedding_type} not supported" + ) + + vgm = model_config.visual_gen_mapping + attn2d_row_size = vgm.attn2d_row_size if vgm else 1 + attn2d_col_size = vgm.attn2d_col_size if vgm else 1 + attn2d_mesh_size = attn2d_row_size * attn2d_col_size + ulysses_size = vgm.ulysses_size if vgm else 1 + use_attn2d = attn2d_mesh_size > 1 + use_ulysses = ulysses_size > 1 + + if vgm is not None and vgm.tp_size > 1: + raise ValueError( + f"Cosmos3 does not support tensor parallelism. Got tp_size={vgm.tp_size}" + ) + + if use_ulysses and ( + self.num_attention_heads % ulysses_size != 0 or self.num_kv_heads % ulysses_size != 0 + ): + raise ValueError( + f"num_attention_heads ({self.num_attention_heads}) and " + f"num_kv_heads ({self.num_kv_heads}) must be divisible by " + f"ulysses_size ({ulysses_size})" + ) + + if use_attn2d: + # Attention2D is not compatible with Cosmos3 cross-attention: its forward() + # TODO: Re-enable once Ring/Attn2D PRs with cross-attention support have landed. + raise NotImplementedError( + "Attention2D (Ring attention) is not supported for Cosmos3. " + "Use Ulysses sequence parallelism instead." + ) + elif use_ulysses: + self.use_seq_parallel = True + self.seq_parallel_size = ulysses_size + self.seq_parallel_pg = vgm.ulysses_group + self.seq_parallel_rank = vgm.ulysses_rank + else: + self.use_seq_parallel = False + self.seq_parallel_size = 1 + self.seq_parallel_pg = None + self.seq_parallel_rank = 0 + + self.language_model = Cosmos3LanguageModel(model_config) + + self.vae2llm = nn.Linear(self.patch_latent_dim, self.hidden_size) + self.llm2vae = nn.Linear(self.hidden_size, self.patch_latent_dim) + + # try timestep embedder in float32 if acc loss + self.time_embedder = TimestepEmbedder(self.hidden_size, target_dtype=torch.bfloat16) + + self.gen_layers = nn.ModuleList( + [ + Cosmos3GenDecoderLayer(model_config, layer_idx=i) + for i in range(self.num_hidden_layers) + ] + ) + + self.norm_moe_gen = Qwen3VLTextRMSNorm( + hidden_size=self.hidden_size, + eps=pretrained_config.rms_norm_eps, + ) + + self.cached_kv = None + self.cached_freqs_gen = None + + self.__post_init__() + + @property + def device(self): + return next(self.parameters()).device + + def __post_init__(self): + # TODO: move this to pipeline loader under meta init so transformers' dont need to know about it here + self.apply_quant_config_exclude_modules() + + for _, module in self.named_modules(): + if callable(getattr(module, "create_weights", None)): + module.create_weights() + + def apply_quant_config_exclude_modules(self): + quant_config = self.model_config.quant_config + if quant_config is None or quant_config.exclude_modules is None: + return + + kv_cache_quant_algo = quant_config.kv_cache_quant_algo if quant_config else None + no_quant_config = QuantConfig(kv_cache_quant_algo=kv_cache_quant_algo) + + for name, module in self.named_modules(): + if isinstance(module, Linear): + is_excluded = quant_config.is_module_excluded_from_quantization(name) + if is_excluded and getattr(module, "quant_config", None) is not None: + module.quant_config = no_quant_config + + def _pad_to_patch_size(self, H: int, W: int) -> Tuple[int, int, int, int]: + """Compute padded spatial dims aligned to patch_size. + + Returns (Hp, Wp, H_padded, W_padded) where Hp/Wp are the patch grid + dimensions and H_padded/W_padded are the padded latent dimensions. + """ + p = self.latent_patch_size + H_padded = ((H + p - 1) // p) * p + W_padded = ((W + p - 1) // p) * p + return H_padded // p, W_padded // p, H_padded, W_padded + + def patchify(self, latents: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor: + """[B, C, T, H, W] -> [B, T*Hp*Wp, p*p*C], padding H/W if needed.""" + B = latents.shape[0] + p = self.latent_patch_size + C = self.latent_channel_size + Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W) + + if H_padded != H or W_padded != W: + latents = F.pad(latents, (0, W_padded - W, 0, H_padded - H)) + + x = latents.reshape(B, C, T, Hp, p, Wp, p) + x = x.permute(0, 2, 3, 5, 4, 6, 1) # [B, T, Hp, Wp, p, p, C] + return x.reshape(B, T * Hp * Wp, p * p * C) + + def unpatchify(self, tokens: torch.Tensor, T: int, H: int, W: int) -> torch.Tensor: + """[B, T*Hp*Wp, p*p*C] -> [B, C, T, H, W], cropping padding if needed.""" + B = tokens.shape[0] + p = self.latent_patch_size + C = self.latent_channel_size + Hp, Wp, H_padded, W_padded = self._pad_to_patch_size(H, W) + + x = tokens.reshape(B, T, Hp, Wp, p, p, C) + x = x.permute(0, 6, 1, 2, 4, 3, 5) # [B, C, T, Hp, p, Wp, p] + x = x.reshape(B, C, T, H_padded, W_padded) + + if H_padded != H or W_padded != W: + x = x[:, :, :, :H, :W] + return x + + def _compute_rope_freqs( + self, + text_mask: torch.Tensor, + T: int, + Hp: int, + Wp: int, + fps: float | None, + device: torch.device, + dtype: torch.dtype, + ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor]]: + """Compute mRoPE cos/sin for UND (text) and GEN (visual) pathways.""" + B = text_mask.shape[0] + S_text = text_mask.shape[1] + text_lengths = text_mask.sum(dim=1).long() + effective_fps = fps if fps is not None and T > 1 else None + + text_pos_list = [] + vis_pos_list = [] + for b in range(B): + real_len = int(text_lengths[b].item()) + t_pos, t_offset = compute_mrope_position_ids_text(real_len, temporal_offset=0) + v_pos, _ = compute_mrope_position_ids_vision( + T, + Hp, + Wp, + temporal_offset=t_offset + self.unified_3d_mrope_temporal_modality_margin, + fps=effective_fps, + base_fps=self.base_fps, + temporal_compression_factor=self.temporal_compression_factor, + enable_fps_modulation=self.enable_fps_modulation, + ) + if real_len < S_text: + t_pos = torch.cat( + [t_pos, torch.zeros(3, S_text - real_len, dtype=t_pos.dtype)], dim=1 + ) + text_pos_list.append(t_pos) + vis_pos_list.append(v_pos) + + text_pos_ids = torch.stack(text_pos_list, dim=1).to(device) # [3, B, S_text] + vis_pos_ids = torch.stack(vis_pos_list, dim=1).to(device) # [3, B, S_vis] + + rotary_emb = self.language_model.rotary_emb + _dummy = torch.tensor([], dtype=dtype, device=device) + cos_und, sin_und = rotary_emb(_dummy, position_ids=text_pos_ids) + cos_gen, sin_gen = rotary_emb(_dummy, position_ids=vis_pos_ids) + + freqs_und = (cos_und.unsqueeze(2), sin_und.unsqueeze(2)) # (B, S, 1, 128) + freqs_gen = (cos_gen.unsqueeze(2), sin_gen.unsqueeze(2)) + return freqs_und, freqs_gen + + def reset_cache(self): + self.cached_kv = None + self.cached_freqs_gen = None + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + text_ids: torch.Tensor, + text_mask: torch.Tensor, + video_shape: Tuple[int, int, int], + fps: float | None = None, + noisy_frame_mask: torch.Tensor | None = None, + **kwargs, + ) -> torch.Tensor: + """ + Forward pass for parallel denoising. + + Args: + hidden_states: [B, C, T, H, W] noisy latents + timestep: [B] diffusion timestep per sample + text_ids: [B, S_text] tokenized text input + text_mask: [B, S_text] attention mask for text (1=real, 0=pad) + video_shape: (T, H, W) in latent space + fps: video frame rate; when provided, temporal mRoPE positions are + scaled to reflect real time (FPS modulation). + noisy_frame_mask: Optional [B, 1, T, 1, 1] mask where 1=noisy (add + timestep embedding, predict velocity) and 0=conditioned (clean + context, skip timestep embedding). None means all frames noisy + (T2V mode). + + Returns: + [B, C, T, H, W] velocity prediction + """ + T, H, W = video_shape + Hp, Wp, _, _ = self._pad_to_patch_size(H, W) + max_real_len = text_mask.sum(dim=1).max().item() + + hidden_gen = self.vae2llm(self.patchify(hidden_states, T, H, W)) + + with torch.autocast("cuda", enabled=True, dtype=torch.float32): + time_embed = self.time_embedder((timestep * self.timestep_scale)) + time_embed = time_embed.to(hidden_states.dtype) + + if noisy_frame_mask is not None: + # Build per-token mask from per-frame mask. + # noisy_frame_mask: [B, 1, T, 1, 1] → token mask: [B, T*Hp*Wp, 1] + noisy_frame_mask = noisy_frame_mask.expand(hidden_gen.shape[0], -1, -1, -1, -1) + token_noisy_mask = ( + noisy_frame_mask[:, 0, :, 0, 0] # [B, T] + .unsqueeze(-1) # [B, T, 1] + .expand(-1, -1, Hp * Wp) # [B, T, Hp*Wp] + .reshape(hidden_gen.shape[0], -1, 1) # [B, T*Hp*Wp, 1] + ) + hidden_gen = hidden_gen + time_embed.unsqueeze(1) * token_noisy_mask + else: + hidden_gen = hidden_gen + time_embed.unsqueeze(1) + + if self.cached_kv is None: + freqs_und, freqs_gen = self._compute_rope_freqs( + text_mask, + T, + Hp, + Wp, + fps, + hidden_states.device, + hidden_states.dtype, + ) + cached_kv_full = self.language_model(text_ids, text_mask, freqs_und) + self.cached_freqs_gen = freqs_gen + + if self.use_seq_parallel: + rank = self.seq_parallel_rank + # Round max_real_len up to next multiple of ulysses_size. + # At most seq_parallel_size-1 extra positions, negligible softmax dilution. + val = ( + self.seq_parallel_size - max_real_len % self.seq_parallel_size + ) % self.seq_parallel_size + S_text_shard_total = int(max_real_len) + val + S_text_shard = S_text_shard_total // self.seq_parallel_size + + self.cached_kv = [] + for k, v in cached_kv_full: + # Slice to S_text_shard_total; zero out the val padding positions + k = k[:, :S_text_shard_total].clone() + v = v[:, :S_text_shard_total].clone() + if val > 0: + k[:, int(max_real_len) :] = 0 + v[:, int(max_real_len) :] = 0 + self.cached_kv.append( + ( + k[:, rank * S_text_shard : (rank + 1) * S_text_shard], + v[:, rank * S_text_shard : (rank + 1) * S_text_shard], + ) + ) + else: + self.cached_kv = cached_kv_full + + if self.use_seq_parallel: + S_gen = hidden_gen.shape[1] + pad = (self.seq_parallel_size - S_gen % self.seq_parallel_size) % self.seq_parallel_size + if pad > 0: + # This will cause minor noise in softmax due to padding. + hidden_gen = F.pad(hidden_gen, (0, 0, 0, pad)) + cos, sin = self.cached_freqs_gen + cos_padded = F.pad(cos, (0, 0, 0, 0, 0, pad)) + sin_padded = F.pad(sin, (0, 0, 0, 0, 0, pad)) + else: + cos_padded, sin_padded = self.cached_freqs_gen + padded_s_gen = S_gen + pad + S_shard = padded_s_gen // self.seq_parallel_size + hidden_gen = hidden_gen[ + :, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard + ] + # Shard freqs_gen to match + freqs_gen = ( + cos_padded[ + :, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard + ], + sin_padded[ + :, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard + ], + ) + else: + freqs_gen = self.cached_freqs_gen + + for i, layer in enumerate(self.gen_layers): + k_und, v_und = self.cached_kv[i] + if self.seq_parallel_size <= 1: + 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) + + if self.use_seq_parallel: + hidden_gen = hidden_gen.contiguous() + parts = [torch.empty_like(hidden_gen) for _ in range(self.seq_parallel_size)] + dist.all_gather(parts, hidden_gen, group=self.seq_parallel_pg) + hidden_gen = torch.cat(parts, dim=1)[:, :S_gen] # [B, S_gen, patch_latent_dim] + + hidden_gen = self.norm_moe_gen(hidden_gen) + return self.unpatchify(self.llm2vae(hidden_gen), T, H, W) + + def load_weights(self, weights: dict) -> None: + """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. + + Expects tensor names as in ``diffusion_pytorch_model.safetensors.index.json`` + (e.g. ``layers.{i}.self_attn.to_q.weight``, ``proj_in.weight``). + Maps UND vs GEN blocks into this module's layout (causal self-attn vs cross-attn + MLPs). + """ + remapped = {} + skip_prefixes = ( + "lm_head.", + "action_modality_embed", + "action_proj_", + "audio_modality_embed", + "audio_proj_", + ) + + for key, value in weights.items(): + k = key + + if k.startswith(skip_prefixes): + continue + + if k.startswith(("vae2llm.", "llm2vae.")): + 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("time_embedder.linear"): + k = k.replace("time_embedder.linear_1.", "time_embedder.mlp.linear_1.") + k = k.replace("time_embedder.linear_2.", "time_embedder.mlp.linear_2.") + remapped[k] = value + continue + + if k.startswith("model."): + k = k[len("model.") :] + + # embed_tokens and norm → language_model.* + if k.startswith("embed_tokens.") or k.startswith("norm."): + remapped[f"language_model.{k}"] = value + continue + + # norm_moe_gen stays at top level + if k.startswith("norm_moe_gen."): + remapped[k] = value + continue + + if not k.startswith("layers."): + logger.warning(f"Skipping unknown checkpoint key: {key}") + continue + + parts = k.split(".", 2) # ['layers', '{i}', '{rest}'] + layer_idx = parts[1] + rest = parts[2] + + # UND (language_model) prefix + und_lp = f"language_model.layers.{layer_idx}" + # GEN prefix + gen_lp = f"gen_layers.{layer_idx}" + + # --- UND attention → language_model.layers.{i}.self_attn.* --- + attn_und_map = { + "self_attn.to_q.": f"{und_lp}.self_attn.to_q.", + "self_attn.to_k.": f"{und_lp}.self_attn.to_k.", + "self_attn.to_v.": f"{und_lp}.self_attn.to_v.", + "self_attn.to_out.": f"{und_lp}.self_attn.to_out.0.", + "self_attn.norm_q.": f"{und_lp}.self_attn.norm_q.", + "self_attn.norm_k.": f"{und_lp}.self_attn.norm_k.", + } + + # --- GEN attention → gen_layers.{i}.cross_attention.* --- + attn_gen_map = { + "self_attn.add_q_proj.": f"{gen_lp}.cross_attention.to_q.", + "self_attn.add_k_proj.": f"{gen_lp}.cross_attention.to_k.", + "self_attn.add_v_proj.": f"{gen_lp}.cross_attention.to_v.", + "self_attn.to_add_out.": f"{gen_lp}.cross_attention.to_out.0.", + "self_attn.norm_added_q.": f"{gen_lp}.cross_attention.norm_q.", + "self_attn.norm_added_k.": f"{gen_lp}.cross_attention.norm_k.", + } + + # --- Norms --- + norm_map = { + "input_layernorm.": f"{und_lp}.input_layernorm.", + "post_attention_layernorm.": f"{und_lp}.post_attention_layernorm.", + "input_layernorm_moe_gen.": f"{gen_lp}.input_layernorm.", + "post_attention_layernorm_moe_gen.": f"{gen_lp}.post_attention_layernorm.", + } + + # --- MLPs --- + mlp_map = { + "mlp.gate_proj.": f"{und_lp}.mlp.gate_proj.", + "mlp.up_proj.": f"{und_lp}.mlp.up_proj.", + "mlp.down_proj.": f"{und_lp}.mlp.down_proj.", + "mlp_moe_gen.gate_proj.": f"{gen_lp}.mlp.gate_proj.", + "mlp_moe_gen.up_proj.": f"{gen_lp}.mlp.up_proj.", + "mlp_moe_gen.down_proj.": f"{gen_lp}.mlp.down_proj.", + } + + matched = False + for mapping in [attn_gen_map, attn_und_map, norm_map, mlp_map]: + for pattern, replacement in mapping.items(): + if rest.startswith(pattern): + suffix = rest[len(pattern) :] + remapped[replacement + suffix] = value + matched = True + break + if matched: + break + + if not matched: + logger.warning(f"Unmatched layer key: {key}") + + # --- Load using DynamicLinearWeightLoader (handles QKV/gate_up fusion + quantization) --- + params_map = { + "qkv_proj": ["to_q", "to_k", "to_v"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + loader = DynamicLinearWeightLoader(self.model_config, params_map=params_map) + + for param_name, param in self._parameters.items(): + if param is not None and param_name in remapped: + param.data.copy_(remapped[param_name].to(param.dtype)) + + loaded_linear = 0 + loaded_other = 0 + skipped_modules = [] + for name, module in self.named_modules(): + if len(module._parameters) == 0: + continue + + if isinstance(module, Linear): + weight_dicts = loader.get_linear_weights(module, name, remapped) + if weight_dicts: + loader.load_linear_weights(module, name, weight_dicts) + loaded_linear += 1 + else: + skipped_modules.append(f"{name}(Linear)") + else: + module_weights = loader.filter_weights(name, remapped) + if module_weights: + loaded_other += 1 + else: + has_params = any(p is not None for p in module._parameters.values()) + if has_params and name: + skipped_modules.append(f"{name}({type(module).__name__})") + for param_name, param in module._parameters.items(): + if param is not None and param_name in module_weights: + param.data.copy_(module_weights[param_name].to(param.dtype)) + + def post_load_weights(self) -> None: + """Post-load processing: dtype conversion and Linear finalization.""" + target_dtype = self.model_config.torch_dtype + + self.time_embedder.to(torch.float32) + self.language_model.embed_tokens.to(target_dtype) + self.vae2llm.to(target_dtype) + self.llm2vae.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/models/ltx2/transformer_ltx2.py b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py index 647e73a554a6..9dbb91059dd1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py +++ b/tensorrt_llm/_torch/visual_gen/models/ltx2/transformer_ltx2.py @@ -120,6 +120,7 @@ def __init__( fuse_qk_norm_rope=True, config=config, layer_idx=layer_idx, + enable_ulysses=use_ulysses and not self._is_cross_attn, ) # For audio self-attention that may need a runtime Ulysses toggle diff --git a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py index 8b5ddaaf84e2..1f6e300c3b51 100644 --- a/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py +++ b/tensorrt_llm/_torch/visual_gen/models/wan/transformer_wan.py @@ -306,6 +306,7 @@ def __init__( eps=eps, config=model_config, layer_idx=_layer_idx, + enable_ulysses=False, ) if cross_attn_norm: diff --git a/tensorrt_llm/_torch/visual_gen/modules/attention.py b/tensorrt_llm/_torch/visual_gen/modules/attention.py index 676f3aec29fa..9cc8f6274bdb 100644 --- a/tensorrt_llm/_torch/visual_gen/modules/attention.py +++ b/tensorrt_llm/_torch/visual_gen/modules/attention.py @@ -51,6 +51,7 @@ def __init__( fuse_qk_norm_rope: Optional[bool] = None, config: Optional[DiffusionModelConfig] = None, layer_idx: Optional[int] = None, + enable_ulysses: bool = True, # make this enable sequence parallelism ): super().__init__() @@ -132,7 +133,7 @@ def __init__( # Currently kept as mutually exclusive. attn2d_size = (vgm.attn2d_row_size * vgm.attn2d_col_size) if vgm else 1 use_attn2d = attn2d_size > 1 and self.qkv_mode != QKVMode.SEPARATE_QKV - use_ulysses = ulysses_size > 1 and self.qkv_mode != QKVMode.SEPARATE_QKV + use_ulysses = ulysses_size > 1 and enable_ulysses # Compute head counts for the backend # Ulysses shards heads across workers; inner backend sees sharded count diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 120bbcd34273..8d804b8dc1a9 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -182,6 +182,10 @@ def _detect_from_checkpoint(checkpoint_dir: str) -> str: if "Flux" in class_name: return "FluxPipeline" + if "Cosmos3" in class_name: + return "Cosmos3OmniMoTPipeline" + + ######################################################### # 2. Single-safetensors with embedded metadata (LTX-2 specific) detected = AutoPipeline._detect_from_single_safetensors(checkpoint_dir) if detected is not None: