From e4a4f4bf91cf9cfe533b84a621103ab9de7ec72b Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 5 May 2026 19:00:45 -0700 Subject: [PATCH 01/17] initial commit Signed-off-by: Shreyas Misra --- requirements.txt | 2 + .../_torch/visual_gen/models/__init__.py | 2 + .../visual_gen/models/cosmos3/__init__.py | 3 + .../visual_gen/models/cosmos3/guardrails.py | 360 ++++++ .../models/cosmos3/pipeline_cosmos3.py | 493 ++++++++ .../models/cosmos3/transformer_cosmos3.py | 1092 +++++++++++++++++ .../_torch/visual_gen/pipeline_registry.py | 4 + 7 files changed, 1956 insertions(+) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py diff --git a/requirements.txt b/requirements.txt index 3a1de400c452..6ad927652f99 100644 --- a/requirements.txt +++ b/requirements.txt @@ -89,3 +89,5 @@ etcd-sdk-python==0.0.7 python-multipart smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 +nltk==3.9.4 +better_profanity==0.7.0 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..5120f1fc8a34 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py @@ -0,0 +1,3 @@ +from .pipeline_cosmos3 import Cosmos3OmniMoTPipeline + +__all__ = ["Cosmos3OmniMoTPipeline"] 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..05f281fc4286 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -0,0 +1,360 @@ +from __future__ import annotations + +import os +import warnings +from typing import Callable + +import cv2 +import numpy as np +import torch +import torch.nn as nn + +from tensorrt_llm.logger import logger + +TextGuardrailFn = Callable[[str], tuple[bool, str]] +VideoGuardrailFn = Callable[[np.ndarray], np.ndarray] + +GUARDRAIL_HF_REPO = "nvidia/Cosmos-Guardrail1" +GUARDRAIL_HF_REVISION = "d6d4bfa899a71454a700907664f3e88f503950cf" +CUTOFF_UNSAFE_FRAMES_PERCENT = 10 + + +# --------------------------------------------------------------------------- +# Video safety classifier (matches reference: SigLIP so400m + 3-layer head) +# --------------------------------------------------------------------------- +class SafetyClassifier(nn.Module): + """3-layer classifier with BatchNorm (1152 → 512 → 256 → 7).""" + + def __init__(self, input_size: int = 1152, num_classes: int = 7): + super().__init__() + self.layers = nn.Sequential( + nn.Linear(input_size, 512), + nn.BatchNorm1d(512), + nn.ReLU(), + nn.Linear(512, 256), + nn.BatchNorm1d(256), + nn.ReLU(), + nn.Linear(256, num_classes), + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.layers(x) + + +CLASS_IDX_TO_NAME = { + 0: "Safe", + 1: "Sexual_Content", + 3: "Drugs", + 4: "Child_Abuse", + 5: "Hate_and_Harassment", + 6: "Self-Harm", +} + + +# --------------------------------------------------------------------------- +# Face pixelation utility +# --------------------------------------------------------------------------- +def _pixelate_face(face_img: np.ndarray, blocks: int = 5) -> np.ndarray: + h, w = face_img.shape[:2] + if h == 0 or w == 0: + return face_img + temp = cv2.resize(face_img, (blocks, blocks), interpolation=cv2.INTER_LINEAR) + return cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST) + + +# --------------------------------------------------------------------------- +# Default guardrail builders +# --------------------------------------------------------------------------- +def download_guardrail_checkpoint() -> str: + from huggingface_hub import snapshot_download + + try: + return snapshot_download( + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_HF_REVISION, + local_files_only=True, + ) + except FileNotFoundError: + logger.warning( + f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO} {GUARDRAIL_HF_REVISION}" + ) + return snapshot_download( + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_HF_REVISION, + ) + + +def build_text_guardrail(guardrail_ckpt_dir: str) -> TextGuardrailFn: + checkers: list[Callable[[str], tuple[bool, str]]] = [] + + # 1. Blocklist + try: + import nltk + from better_profanity import profanity as profanity_filter + + blocklist_dir = os.path.join(guardrail_ckpt_dir, "blocklist") + nltk.data.path.append(os.path.join(blocklist_dir, "nltk_data")) + + def _read_keywords(dirpath: str) -> list[str]: + words: list[str] = [] + if not os.path.isdir(dirpath): + return words + for fname in sorted(os.listdir(dirpath)): + fpath = os.path.join(dirpath, fname) + if os.path.isfile(fpath): + with open(fpath) as f: + words.extend(line.strip() for line in f if line.strip()) + return words + + blocklist_words = _read_keywords(os.path.join(blocklist_dir, "custom")) + whitelist_words = _read_keywords(os.path.join(blocklist_dir, "whitelist")) + profanity_filter.load_censor_words( + custom_words=blocklist_words, whitelist_words=whitelist_words + ) + + def _blocklist_check(prompt: str) -> tuple[bool, str]: + if profanity_filter.contains_profanity(prompt): + return False, "Blocked by keyword filter" + return True, "" + + checkers.append(_blocklist_check) + logger.info("Blocklist guardrail loaded (%d keywords)", len(blocklist_words)) + except ImportError: + logger.warning("better-profanity or nltk not installed; skipping blocklist guardrail") + + # 2. Qwen3Guard + try: + from transformers import AutoModelForCausalLM, AutoTokenizer + + model_id = "Qwen/Qwen3Guard-Gen-0.6B" + qwen_tokenizer = AutoTokenizer.from_pretrained(model_id) + qwen_model = ( + AutoModelForCausalLM.from_pretrained( + model_id, + torch_dtype=torch.bfloat16, + ) + .to("cuda") + .eval() + ) + + def _qwen_check(prompt: str) -> tuple[bool, str]: + conversations = [{"role": "user", "content": prompt}] + input_ids = qwen_tokenizer.apply_chat_template( + conversations, + tokenize=True, + return_tensors="pt", + add_generation_prompt=True, + ).to("cuda") + with torch.no_grad(): + output_ids = qwen_model.generate(input_ids, max_new_tokens=128) + response = qwen_tokenizer.decode( + output_ids[0][input_ids.shape[1] :], + skip_special_tokens=True, + ) + if "unsafe" in response.lower(): + return False, f"Qwen3Guard: {response.strip()}" + return True, "" + + checkers.append(_qwen_check) + logger.info("Qwen3Guard guardrail loaded") + except ImportError: + logger.warning("transformers not installed; skipping Qwen3Guard") + + def text_guardrail(prompt: str) -> None: + for checker in checkers: + is_safe, msg = checker(prompt) + if not is_safe: + return is_safe, msg + return True, "" + + return text_guardrail + + +def build_video_guardrail(guardrail_ckpt_dir: str) -> VideoGuardrailFn: + safety_checker: Callable[[np.ndarray], tuple[bool, str]] | None = None + face_blurrer: Callable[[np.ndarray], np.ndarray] | None = None + + # 1. Video content safety filter: SigLIP so400m + SafetyClassifier + try: + from PIL import Image + from transformers import SiglipModel, SiglipProcessor + + siglip_id = "google/siglip-so400m-patch14-384" + siglip_model = SiglipModel.from_pretrained(siglip_id).to("cuda", dtype=torch.float32).eval() + siglip_processor = SiglipProcessor.from_pretrained(siglip_id) + + classifier = SafetyClassifier(input_size=1152, num_classes=7) + ckpt_path = os.path.join( + guardrail_ckpt_dir, "video_content_safety_filter", "safety_filter.pt" + ) + checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True) + state = {k.removeprefix("network."): v for k, v in checkpoint["model"].items()} + classifier.load_state_dict(state) + classifier = classifier.to("cuda", dtype=torch.float32).eval() + + def _safety_check(frames: np.ndarray) -> tuple[bool, str]: + nonlocal siglip_model, classifier + + unsafe_count = 0 + total = len(frames) + for frame in frames: + img = Image.fromarray(frame) + inputs = siglip_processor(images=img, return_tensors="pt").to( + "cuda", dtype=torch.float32 + ) + with torch.no_grad(): + features = siglip_model.get_image_features(**inputs) + features = features / features.norm(dim=-1, keepdim=True) + logits = classifier(features) + pred = logits.argmax(dim=-1).item() + class_name = CLASS_IDX_TO_NAME.get(pred, "Unknown") + if class_name != "Safe": + unsafe_count += 1 + + if unsafe_count / total > CUTOFF_UNSAFE_FRAMES_PERCENT / 100: + return False, f"Video content safety: {unsafe_count}/{total} frames unsafe" + return True, "" + + safety_checker = _safety_check + logger.info("Video content safety filter loaded (SigLIP so400m + classifier)") + except (ImportError, FileNotFoundError) as e: + logger.warning("Could not load video safety filter: %s", e) + + # 2. Face blur: RetinaFace + pixelation + try: + from retinaface.data import cfg_re50 + from retinaface.layers.functions.prior_box import PriorBox + from retinaface.models.retinaface import RetinaFace + from retinaface.utils.nms.py_cpu_nms import py_cpu_nms + + face_ckpt = os.path.join(guardrail_ckpt_dir, "face_blur_filter", "Resnet50_Final.pth") + if not os.path.exists(face_ckpt): + raise FileNotFoundError(face_ckpt) + + cfg = dict(cfg_re50) + cfg["pretrain"] = False + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + retinaface_net = RetinaFace(cfg=cfg, phase="test") + + # Load weights (strip 'module.' prefix if present) + pretrained_dict = torch.load(face_ckpt, map_location="cpu", weights_only=True) + if "state_dict" in pretrained_dict: + pretrained_dict = pretrained_dict["state_dict"] + pretrained_dict = { + k.replace("module.", "", 1) if k.startswith("module.") else k: v + for k, v in pretrained_dict.items() + } + retinaface_net.load_state_dict(pretrained_dict, strict=False) + retinaface_device = "cuda" + retinaface_net = retinaface_net.to(retinaface_device, dtype=torch.float32).eval() + + CONF_THRESH = 0.7 + NMS_THRESH = 0.4 + TOP_K = 5000 + KEEP_TOP_K = 750 + + def _decode_batch(loc, priors, variances): + batch_size = loc.size(0) + p = priors.unsqueeze(0).expand(batch_size, -1, -1) + boxes = torch.cat( + ( + p[:, :, :2] + loc[:, :, :2] * variances[0] * p[:, :, 2:], + p[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1]), + ), + dim=2, + ) + boxes[:, :, :2] -= boxes[:, :, 2:] / 2 + boxes[:, :, 2:] += boxes[:, :, :2] + return boxes + + def _face_blur(frames: np.ndarray) -> np.ndarray: + nonlocal retinaface_net + + prior_data = None + scale = None + result_frames = [] + + for frame in frames: + frame_t = torch.from_numpy(frame).to("cuda", dtype=torch.float32) + frame_t = frame_t.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W] + frame_t = frame_t[:, [2, 1, 0], :, :] # RGB → BGR + means = torch.tensor( + [104.0, 117.0, 123.0], device="cuda", dtype=torch.float32 + ).view(1, 3, 1, 1) + frame_t = frame_t - means + + h, w = frame_t.shape[2], frame_t.shape[3] + if prior_data is None: + priorbox = PriorBox(cfg, image_size=(h, w)) + prior_data = priorbox.forward().to("cuda", dtype=torch.float32) + if scale is None: + scale = torch.tensor([w, h, w, h], device="cuda", dtype=torch.float32) + + with torch.no_grad(): + loc, conf, _ = retinaface_net(frame_t) + + boxes = _decode_batch(loc, prior_data, cfg["variance"]) + boxes = (boxes * scale).squeeze(0).cpu().numpy() + scores = conf.squeeze(0)[:, 1].cpu().numpy() + + # Filter by confidence + inds = np.where(scores > CONF_THRESH)[0] + boxes_f = boxes[inds] + scores_f = scores[inds] + order = scores_f.argsort()[::-1][:TOP_K] + boxes_f = boxes_f[order] + scores_f = scores_f[order] + + # NMS + dets = np.hstack((boxes_f, scores_f[:, np.newaxis])).astype(np.float32) + keep = py_cpu_nms(dets, NMS_THRESH) + dets = dets[keep][:KEEP_TOP_K] + + out_frame = frame.copy() + for det in dets: + x1, y1, x2, y2 = map(int, det[:4]) + if x2 - x1 < 20 or y2 - y1 < 20: + continue + max_h, max_w = out_frame.shape[:2] + y1c, y2c = max(y1, 0), min(y2, max_h) + x1c, x2c = max(x1, 0), min(x2, max_w) + out_frame[y1c:y2c, x1c:x2c] = _pixelate_face(out_frame[y1c:y2c, x1c:x2c]) + + result_frames.append(out_frame) + + return np.array(result_frames) + + face_blurrer = _face_blur + logger.info("Face blur filter loaded (RetinaFace Resnet50)") + except (ImportError, FileNotFoundError) as e: + logger.warning("Could not load face blur filter: %s", e) + + def video_guardrail(frames: np.ndarray) -> np.ndarray | None: + if safety_checker is not None: + is_safe, msg = safety_checker(frames) + if not is_safe: + logger.warning(f"Video content safety: {msg}") + return None + if face_blurrer is not None: + frames = face_blurrer(frames) + return frames + + return video_guardrail + + +def check_video_safety( + video_tensor: torch.Tensor, video_guardrail: VideoGuardrailFn +) -> torch.Tensor | None: + v = video_tensor.detach().cpu() + if v.dim() == 5: + v = v[0] + frames_np = v.numpy() + frames_np = video_guardrail(frames_np) + if frames_np is None: + return None + + result = torch.from_numpy(frames_np) + if video_tensor.dim() == 4: + 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..652eb9ed6202 --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -0,0 +1,493 @@ +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.config import PipelineComponent +from tensorrt_llm._torch.visual_gen.output import MediaOutput +from tensorrt_llm._torch.visual_gen.pipeline import BasePipeline +from tensorrt_llm._torch.visual_gen.pipeline_registry import 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 .guardrails import ( + build_text_guardrail, + build_video_guardrail, + check_video_safety, + download_guardrail_checkpoint, +) +from .transformer_cosmos3 import Cosmos3VFMTransformer + +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} FPS." +TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" + + +@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 + ) -> None: + if PipelineComponent.TOKENIZER not in skip_components: + logger.info("Loading tokenizer...") + self.tokenizer = Qwen2Tokenizer.from_pretrained( + checkpoint_dir, + subfolder="text_tokenizer", + ) + + 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", 4) + self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16) + 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, + ) + self.scheduler = UniPCMultistepScheduler.from_config( + self.scheduler.config, flow_shift=5.0 + ) # for 720p trained checkpoint + + # load guardrails by default + if ( + not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + and PipelineComponent.TEXT_GUARDRAIL not in skip_components + and PipelineComponent.VIDEO_GUARDRAIL not in skip_components + ): + guardrail_ckpt_dir = download_guardrail_checkpoint() + self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) + self.video_guardrail = build_video_guardrail(guardrail_ckpt_dir) + + 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 [61, 81] + + 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=4.0, + seed=42, + max_sequence_length=256, + 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_system_prompt=req.params.extra_params.get("use_system_prompt", False), + use_guardrails=req.params.extra_params.get("use_guardrails", True), + ) + + @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, + ) + token_ids = token_ids[:max_sequence_length] + 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 = 720, + width: int = 1280, + num_frames: int = 81, + num_inference_steps: int = 35, + guidance_scale: float = 4.0, + seed: int = 42, + max_sequence_length: int = 256, + frame_rate: float = 24.0, + use_duration_template: bool = True, + use_system_prompt: bool = False, + use_guardrails: bool = True, + ): + pipeline_start = time.time() + 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." + ) + + if self.rank == 0 and use_guardrails: + for p in prompt: + is_safe, msg = self.text_guardrail(p) + if not is_safe: + logger.warning(f"Text guardrail blocked prompt: {msg}") + return MediaOutput() + + generator = torch.Generator(device=self.device).manual_seed(seed) + + if negative_prompt is None: + negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT + + if use_duration_template and num_frames > 1: + duration = num_frames / frame_rate + suffix = COSMOS3_DURATION_TEMPLATE.format(duration=duration, fps=frame_rate) + prompt = [f"{p} {suffix}" for p in prompt] + logger.info(f"Prompt with duration: '{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 = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + 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 + 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, + ) + + # 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) + + 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: + video = check_video_safety(video, self.video_guardrail) + if video is None: + logger.warning("Video guardrail blocked video generation") + return MediaOutput() + + return MediaOutput(video=video) 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..3b7fea26f99e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -0,0 +1,1092 @@ +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.parallelism import setup_sequence_parallelism +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: + 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, + ): + 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, + ) + 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.hidden_size + // model_config.pretrained_config.num_attention_heads, + 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.hidden_size + // model_config.pretrained_config.num_attention_heads, + 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.hidden_size + // model_config.pretrained_config.num_attention_heads + ) + 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" + ) + + self.use_ulysses, self.ulysses_size, self.ulysses_pg, self.ulysses_rank = ( + setup_sequence_parallelism( + model_config=model_config, + num_attention_heads=self.num_attention_heads, + num_kv_heads=self.num_kv_heads, + ) + ) + + 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.ulysses_size > 1: + rank = self.ulysses_rank + # Round max_real_len up to next multiple of ulysses_size. + # At most ulysses_size-1 extra positions, negligible softmax dilution. + val = (self.ulysses_size - max_real_len % self.ulysses_size) % self.ulysses_size + S_text_shard_total = int(max_real_len) + val + S_text_shard = S_text_shard_total // self.ulysses_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.ulysses_size > 1: + S_gen = hidden_gen.shape[1] + pad = (self.ulysses_size - S_gen % self.ulysses_size) % self.ulysses_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 + self.cached_freqs_gen = ( + F.pad(cos, (0, 0, 0, pad)), + F.pad(sin, (0, 0, 0, pad)), + ) + S_shard = S_gen // self.ulysses_size + hidden_gen = hidden_gen[ + :, self.ulysses_rank * S_shard : (self.ulysses_rank + 1) * S_shard + ] + # Shard freqs_gen to match + cos, sin = self.cached_freqs_gen + freqs_gen = ( + cos[:, self.ulysses_rank * S_shard : (self.ulysses_rank + 1) * S_shard], + sin[:, self.ulysses_rank * S_shard : (self.ulysses_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.ulysses_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.ulysses_size > 1: + hidden_gen = hidden_gen.contiguous() + parts = [torch.empty_like(hidden_gen) for _ in range(self.ulysses_size)] + dist.all_gather(parts, hidden_gen, group=self.ulysses_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 Diffusers/HF Cosmos3 checkpoints. + + Expects tensor names under ``model.*`` (e.g. ``model.layers.{i}.self_attn.q_proj.weight``). + Maps UND vs GEN blocks into this module's layout (causal self-attn vs cross-attn + MLPs). + """ + remapped = {} + + for key, value in weights.items(): + k = key + + if k.startswith(("vae2llm.", "llm2vae.")): + remapped[k] = value + continue + + if k.startswith("time_embedder.mlp."): + k = k.replace("time_embedder.mlp.0.", "time_embedder.mlp.linear_1.") + k = k.replace("time_embedder.mlp.2.", "time_embedder.mlp.linear_2.") + remapped[k] = value + continue + + if k.startswith("lm_head."): + continue + + if not k.startswith("model."): + logger.warning(f"Skipping unknown checkpoint key: {key}") + continue + 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 LM 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.q_proj.": f"{und_lp}.self_attn.to_q.", + "self_attn.k_proj.": f"{und_lp}.self_attn.to_k.", + "self_attn.v_proj.": f"{und_lp}.self_attn.to_v.", + "self_attn.o_proj.": f"{und_lp}.self_attn.to_out.0.", + "self_attn.q_norm.": f"{und_lp}.self_attn.norm_q.", + "self_attn.k_norm.": f"{und_lp}.self_attn.norm_k.", + } + + # --- GEN attention → gen_layers.{i}.cross_attention.* --- + attn_gen_map = { + "self_attn.q_proj_moe_gen.": f"{gen_lp}.cross_attention.to_q.", + "self_attn.k_proj_moe_gen.": f"{gen_lp}.cross_attention.to_k.", + "self_attn.v_proj_moe_gen.": f"{gen_lp}.cross_attention.to_v.", + "self_attn.o_proj_moe_gen.": f"{gen_lp}.cross_attention.to_out.0.", + "self_attn.q_norm_moe_gen.": f"{gen_lp}.cross_attention.norm_q.", + "self_attn.k_norm_moe_gen.": 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/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: From ecf4b46b2810ac51866f9ddf1b81b9d900d7bf80 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 7 May 2026 23:21:08 +0000 Subject: [PATCH 02/17] add defaults, templates, fix image resizing Signed-off-by: Shreyas Misra --- .../visual_gen/attention_backend/vanilla.py | 9 +- .../visual_gen/models/cosmos3/defaults.py | 60 +++++++++ .../models/cosmos3/pipeline_cosmos3.py | 121 ++++++++++++++---- 3 files changed, 166 insertions(+), 24 deletions(-) create mode 100644 tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py 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/cosmos3/defaults.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py new file mode 100644 index 000000000000..b4bc20502b2e --- /dev/null +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -0,0 +1,60 @@ +# 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 Wan pipelines. + +Deduction cascade: model version (2.1/2.2) → model size → model name. +Shared by WanPipeline (T2V) and WanImageToVideoPipeline (I2V). +""" + +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/pipeline_cosmos3.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py index 652eb9ed6202..cbf5e7d44df8 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1,3 +1,4 @@ +import math import os import time from typing import List, Optional, Union @@ -17,6 +18,7 @@ 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 ( build_text_guardrail, build_video_guardrail, @@ -25,11 +27,19 @@ ) from .transformer_cosmos3 import Cosmos3VFMTransformer -COSMOS3_DEFAULT_NEGATIVE_PROMPT = "" +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} FPS." +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" @@ -76,9 +86,6 @@ def load_standard_components( checkpoint_dir, subfolder=PipelineComponent.SCHEDULER, ) - self.scheduler = UniPCMultistepScheduler.from_config( - self.scheduler.config, flow_shift=5.0 - ) # for 720p trained checkpoint # load guardrails by default if ( @@ -100,6 +107,14 @@ def default_warmup_resolutions(self): def default_warmup_num_frames(self): return [61, 81] + @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( @@ -109,9 +124,9 @@ def _run_warmup(self, height: int, width: int, num_frames: int, steps: int) -> N width=width, num_frames=num_frames, num_inference_steps=steps, - guidance_scale=4.0, + guidance_scale=COSMOS3_720P_PARAMS["guidance_scale"], seed=42, - max_sequence_length=256, + max_sequence_length=COSMOS3_720P_PARAMS["max_sequence_length"], use_guardrails=False, image=None, ) @@ -130,10 +145,51 @@ def infer(self, req): 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 @@ -338,17 +394,18 @@ def forward( prompt: Union[str, List[str]], negative_prompt: Optional[str] = None, image: Optional[Union[PIL.Image.Image, torch.Tensor, str]] = None, - height: int = 720, - width: int = 1280, - num_frames: int = 81, - num_inference_steps: int = 35, - guidance_scale: float = 4.0, + 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 = 256, - frame_rate: float = 24.0, - use_duration_template: bool = True, - use_system_prompt: bool = False, - use_guardrails: bool = True, + 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() use_guardrails = use_guardrails and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS @@ -381,11 +438,29 @@ def forward( if negative_prompt is None: negative_prompt = COSMOS3_DEFAULT_NEGATIVE_PROMPT - if use_duration_template and num_frames > 1: - duration = num_frames / frame_rate - suffix = COSMOS3_DURATION_TEMPLATE.format(duration=duration, fps=frame_rate) - prompt = [f"{p} {suffix}" for p in prompt] - logger.info(f"Prompt with duration: '{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] @@ -403,7 +478,7 @@ def forward( if isinstance(image, PIL.Image.Image): image = image.convert("RGB") - image = image.resize((width, height), PIL.Image.Resampling.LANCZOS) + image = self._resize_and_center_crop_image(image, height=height, width=width) image = self.video_processor.preprocess( image, height=height, From ac49291ee83bee7e2c42c4bcfbafd0dc64a4cce7 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Thu, 7 May 2026 23:54:32 +0000 Subject: [PATCH 03/17] enable ulysses for cross attention Signed-off-by: Shreyas Misra --- .../models/cosmos3/transformer_cosmos3.py | 77 +++++++++++++------ .../models/ltx2/transformer_ltx2.py | 1 + .../visual_gen/models/wan/transformer_wan.py | 1 + .../_torch/visual_gen/modules/attention.py | 3 +- 4 files changed, 59 insertions(+), 23 deletions(-) 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 3b7fea26f99e..0a809e4bb509 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -13,7 +13,6 @@ 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.parallelism import setup_sequence_parallelism from tensorrt_llm._torch.visual_gen.quantization.loader import DynamicLinearWeightLoader from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig @@ -652,13 +651,45 @@ def __init__(self, model_config: DiffusionModelConfig): f"Position embedding type {pretrained_config.position_embedding_type} not supported" ) - self.use_ulysses, self.ulysses_size, self.ulysses_pg, self.ulysses_rank = ( - setup_sequence_parallelism( - model_config=model_config, - num_attention_heads=self.num_attention_heads, - num_kv_heads=self.num_kv_heads, + 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 + and 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: + self.use_seq_parallel = True + self.seq_parallel_size = attn2d_mesh_size + self.seq_parallel_pg = vgm.attn2d_mesh_group + self.seq_parallel_rank = vgm.attn2d_mesh_rank + 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) @@ -872,13 +903,15 @@ def forward( cached_kv_full = self.language_model(text_ids, text_mask, freqs_und) self.cached_freqs_gen = freqs_gen - if self.ulysses_size > 1: - rank = self.ulysses_rank + if self.use_seq_parallel: + rank = self.seq_parallel_rank # Round max_real_len up to next multiple of ulysses_size. - # At most ulysses_size-1 extra positions, negligible softmax dilution. - val = (self.ulysses_size - max_real_len % self.ulysses_size) % self.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.ulysses_size + S_text_shard = S_text_shard_total // self.seq_parallel_size self.cached_kv = [] for k, v in cached_kv_full: @@ -897,9 +930,9 @@ def forward( else: self.cached_kv = cached_kv_full - if self.ulysses_size > 1: + if self.use_seq_parallel: S_gen = hidden_gen.shape[1] - pad = (self.ulysses_size - S_gen % self.ulysses_size) % self.ulysses_size + 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)) @@ -908,30 +941,30 @@ def forward( F.pad(cos, (0, 0, 0, pad)), F.pad(sin, (0, 0, 0, pad)), ) - S_shard = S_gen // self.ulysses_size + S_shard = S_gen // self.seq_parallel_size hidden_gen = hidden_gen[ - :, self.ulysses_rank * S_shard : (self.ulysses_rank + 1) * S_shard + :, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard ] # Shard freqs_gen to match cos, sin = self.cached_freqs_gen freqs_gen = ( - cos[:, self.ulysses_rank * S_shard : (self.ulysses_rank + 1) * S_shard], - sin[:, self.ulysses_rank * S_shard : (self.ulysses_rank + 1) * S_shard], + cos[:, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard], + sin[:, 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.ulysses_size <= 1: + 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.ulysses_size > 1: + if self.use_seq_parallel: hidden_gen = hidden_gen.contiguous() - parts = [torch.empty_like(hidden_gen) for _ in range(self.ulysses_size)] - dist.all_gather(parts, hidden_gen, group=self.ulysses_pg) + 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) 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 From c30b9ef92b0ac64a43e53d4714978734d653f53f Mon Sep 17 00:00:00 2001 From: Shreyas Mista Date: Mon, 11 May 2026 18:37:36 +0000 Subject: [PATCH 04/17] make it work with latest code changes Signed-off-by: Shreyas Mista --- .../models/cosmos3/pipeline_cosmos3.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 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 cbf5e7d44df8..a2c463c2ff4e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -11,7 +11,7 @@ from transformers import Qwen2Tokenizer from tensorrt_llm._torch.visual_gen.config import PipelineComponent -from tensorrt_llm._torch.visual_gen.output import MediaOutput +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 register_pipeline from tensorrt_llm._torch.visual_gen.utils import postprocess_video_tensor @@ -105,7 +105,7 @@ def default_warmup_resolutions(self): @property def default_warmup_num_frames(self): - return [61, 81] + return [189] @property def default_generation_params(self): @@ -210,6 +210,7 @@ def _tokenize_prompt( conversations, tokenize=True, add_generation_prompt=True, + return_dict=False, ) token_ids = token_ids[:max_sequence_length] token_ids.append(self.tokenizer.eos_token_id) # 151645 @@ -408,6 +409,9 @@ def forward( 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): @@ -431,7 +435,8 @@ def forward( is_safe, msg = self.text_guardrail(p) if not is_safe: logger.warning(f"Text guardrail blocked prompt: {msg}") - return MediaOutput() + timer.mark_end() + return timer.fill(PipelineOutput()) generator = torch.Generator(device=self.device).manual_seed(seed) @@ -535,6 +540,7 @@ def forward_fn( self.transformer.reset_cache() # 6. Denoise + timer.mark_denoise_start() latents = self.denoise( latents=latents, scheduler=self.scheduler, @@ -544,6 +550,7 @@ def forward_fn( forward_fn=forward_fn, extra_cfg_tensors=extra_cfg_tensors, ) + timer.mark_post_start() # 7. Decode logger.info("Decoding video...") @@ -563,6 +570,8 @@ def forward_fn( video = check_video_safety(video, self.video_guardrail) if video is None: logger.warning("Video guardrail blocked video generation") - return MediaOutput() + timer.mark_end() + return timer.fill(PipelineOutput()) - return MediaOutput(video=video) + timer.mark_end() + return timer.fill(PipelineOutput(video=video, frame_rate=frame_rate)) From efaa2334a6a2c39341ef1e3404b18c3178d9037c Mon Sep 17 00:00:00 2001 From: Shreyas Mista Date: Mon, 11 May 2026 19:26:08 +0000 Subject: [PATCH 05/17] guardrail fixes Signed-off-by: Shreyas Mista --- requirements.txt | 1 + tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py | 4 +++- 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 6ad927652f99..2caa35608c39 100644 --- a/requirements.txt +++ b/requirements.txt @@ -91,3 +91,4 @@ smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 nltk==3.9.4 better_profanity==0.7.0 +retinaface @ git+https://github.com/NVShreyas/retinaface.git@main diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 05f281fc4286..37aeab1b1699 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -144,6 +144,7 @@ def _qwen_check(prompt: str) -> tuple[bool, str]: tokenize=True, return_tensors="pt", add_generation_prompt=True, + return_dict=False, ).to("cuda") with torch.no_grad(): output_ids = qwen_model.generate(input_ids, max_new_tokens=128) @@ -203,7 +204,8 @@ def _safety_check(frames: np.ndarray) -> tuple[bool, str]: "cuda", dtype=torch.float32 ) with torch.no_grad(): - features = siglip_model.get_image_features(**inputs) + siglip_out = siglip_model.get_image_features(**inputs) + features = siglip_out.pooler_output features = features / features.norm(dim=-1, keepdim=True) logits = classifier(features) pred = logits.argmax(dim=-1).item() From 136e20fdb8674d256cd18b61045dd6e92bd48908 Mon Sep 17 00:00:00 2001 From: Shreyas Mista Date: Mon, 11 May 2026 22:18:08 +0000 Subject: [PATCH 06/17] add guardrail-checkpoint-dir argument Signed-off-by: Shreyas Mista --- tensorrt_llm/_torch/visual_gen/config.py | 3 +++ .../visual_gen/models/cosmos3/guardrails.py | 15 ++++++++++++--- .../visual_gen/models/cosmos3/pipeline_cosmos3.py | 8 +++++++- .../models/cosmos3/transformer_cosmos3.py | 11 +++-------- 4 files changed, 25 insertions(+), 12 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 3f7009bef569..2158daad9df3 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -402,6 +402,9 @@ def from_pretrained( if value: extra_attrs[key] = value + if args and args.guardrail_checkpoint_dir: + extra_attrs["guardrail_checkpoint_dir"] = args.guardrail_checkpoint_dir + # Discover pipeline components (diffusers layout) components = discover_pipeline_components(checkpoint_path) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 37aeab1b1699..548e47d85898 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -180,9 +180,18 @@ def build_video_guardrail(guardrail_ckpt_dir: str) -> VideoGuardrailFn: from PIL import Image from transformers import SiglipModel, SiglipProcessor - siglip_id = "google/siglip-so400m-patch14-384" - siglip_model = SiglipModel.from_pretrained(siglip_id).to("cuda", dtype=torch.float32).eval() - siglip_processor = SiglipProcessor.from_pretrained(siglip_id) + siglip_dir = os.path.join( + guardrail_ckpt_dir, + "video_content_safety_filter", + "models--google--siglip-so400m-patch14-384/snapshots/9fdffc58afc957d1a03a25b10dba0329ab15c2a3", + ) + if not os.path.exists(siglip_dir): + raise FileNotFoundError(siglip_dir) + + siglip_model = ( + SiglipModel.from_pretrained(siglip_dir).to("cuda", dtype=torch.float32).eval() + ) + siglip_processor = SiglipProcessor.from_pretrained(siglip_dir) classifier = SafetyClassifier(input_size=1152, num_classes=7) ckpt_path = os.path.join( 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 a2c463c2ff4e..2091cd83afdd 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -93,7 +93,13 @@ def load_standard_components( and PipelineComponent.TEXT_GUARDRAIL not in skip_components and PipelineComponent.VIDEO_GUARDRAIL not in skip_components ): - guardrail_ckpt_dir = download_guardrail_checkpoint() + if self.model_config.extra_attrs.get("guardrail_checkpoint_dir", None) is not None: + logger.info( + f"Loading guardrails from {self.model_config.extra_attrs['guardrail_checkpoint_dir']}" + ) + guardrail_ckpt_dir = self.model_config.extra_attrs["guardrail_checkpoint_dir"] + else: + guardrail_ckpt_dir = download_guardrail_checkpoint() self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) self.video_guardrail = build_video_guardrail(guardrail_ckpt_dir) 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 0a809e4bb509..19b59e19413c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -352,8 +352,7 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): 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.hidden_size - // model_config.pretrained_config.num_attention_heads, + head_dim=model_config.pretrained_config.head_dim, model_config=model_config, layer_idx=layer_idx, ) @@ -415,8 +414,7 @@ def __init__(self, model_config: DiffusionModelConfig, layer_idx: int): 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.hidden_size - // model_config.pretrained_config.num_attention_heads, + head_dim=model_config.pretrained_config.head_dim, model_config=model_config, layer_idx=layer_idx, ) @@ -499,10 +497,7 @@ def _compute_default_rope_parameters( """ base = model_config.pretrained_config.rope_theta partial_rotary_factor = 1 - head_dim = ( - model_config.pretrained_config.hidden_size - // model_config.pretrained_config.num_attention_heads - ) + 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 From 83c55c7c7613e80e153dd62c12498223434cbe90 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 12 May 2026 12:30:12 -0700 Subject: [PATCH 07/17] address coderabbit comments Signed-off-by: Shreyas Misra --- requirements.txt | 2 +- .../visual_gen/models/cosmos3/__init__.py | 15 ++++ .../visual_gen/models/cosmos3/defaults.py | 5 +- .../visual_gen/models/cosmos3/guardrails.py | 36 ++++++---- .../models/cosmos3/pipeline_cosmos3.py | 69 +++++++++++++++---- .../models/cosmos3/transformer_cosmos3.py | 23 +++++-- 6 files changed, 114 insertions(+), 36 deletions(-) diff --git a/requirements.txt b/requirements.txt index 2caa35608c39..d38c7c25fad2 100644 --- a/requirements.txt +++ b/requirements.txt @@ -91,4 +91,4 @@ smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 nltk==3.9.4 better_profanity==0.7.0 -retinaface @ git+https://github.com/NVShreyas/retinaface.git@main +retinaface @ git+https://github.com/NVShreyas/retinaface.git@1bb2e4589b5f1cf4e1d01a56b7b0a7259ac563de diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py index 5120f1fc8a34..98f83cee9b08 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/__init__.py @@ -1,3 +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 index b4bc20502b2e..e54e818356c4 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/defaults.py @@ -12,10 +12,9 @@ # 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 Wan pipelines. +"""Per-model default generation parameters for Cosmos3 pipelines. -Deduction cascade: model version (2.1/2.2) → model size → model name. -Shared by WanPipeline (T2V) and WanImageToVideoPipeline (I2V). +Shared by the Cosmos3 OmniMoT text-to-video and image-to-video generation paths. """ from typing import Dict diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 548e47d85898..96c2c9256bbc 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -1,3 +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 __future__ import annotations import os @@ -119,8 +134,8 @@ def _blocklist_check(prompt: str) -> tuple[bool, str]: checkers.append(_blocklist_check) logger.info("Blocklist guardrail loaded (%d keywords)", len(blocklist_words)) - except ImportError: - logger.warning("better-profanity or nltk not installed; skipping blocklist guardrail") + except (ImportError, OSError, RuntimeError, ValueError) as e: + logger.warning("Could not load blocklist guardrail: %s", e) # 2. Qwen3Guard try: @@ -158,8 +173,8 @@ def _qwen_check(prompt: str) -> tuple[bool, str]: checkers.append(_qwen_check) logger.info("Qwen3Guard guardrail loaded") - except ImportError: - logger.warning("transformers not installed; skipping Qwen3Guard") + except (ImportError, OSError, RuntimeError, ValueError) as e: + logger.warning("Could not load Qwen3Guard guardrail: %s", e) def text_guardrail(prompt: str) -> None: for checker in checkers: @@ -203,8 +218,6 @@ def build_video_guardrail(guardrail_ckpt_dir: str) -> VideoGuardrailFn: classifier = classifier.to("cuda", dtype=torch.float32).eval() def _safety_check(frames: np.ndarray) -> tuple[bool, str]: - nonlocal siglip_model, classifier - unsafe_count = 0 total = len(frames) for frame in frames: @@ -228,7 +241,7 @@ def _safety_check(frames: np.ndarray) -> tuple[bool, str]: safety_checker = _safety_check logger.info("Video content safety filter loaded (SigLIP so400m + classifier)") - except (ImportError, FileNotFoundError) as e: + except (ImportError, FileNotFoundError, OSError, RuntimeError, ValueError) as e: logger.warning("Could not load video safety filter: %s", e) # 2. Face blur: RetinaFace + pixelation @@ -280,8 +293,6 @@ def _decode_batch(loc, priors, variances): return boxes def _face_blur(frames: np.ndarray) -> np.ndarray: - nonlocal retinaface_net - prior_data = None scale = None result_frames = [] @@ -338,7 +349,7 @@ def _face_blur(frames: np.ndarray) -> np.ndarray: face_blurrer = _face_blur logger.info("Face blur filter loaded (RetinaFace Resnet50)") - except (ImportError, FileNotFoundError) as e: + except (ImportError, FileNotFoundError, OSError, RuntimeError, ValueError) as e: logger.warning("Could not load face blur filter: %s", e) def video_guardrail(frames: np.ndarray) -> np.ndarray | None: @@ -358,7 +369,8 @@ def check_video_safety( video_tensor: torch.Tensor, video_guardrail: VideoGuardrailFn ) -> torch.Tensor | None: v = video_tensor.detach().cpu() - if v.dim() == 5: + was_batched = v.dim() == 5 + if was_batched: v = v[0] frames_np = v.numpy() frames_np = video_guardrail(frames_np) @@ -366,6 +378,6 @@ def check_video_safety( return None result = torch.from_numpy(frames_np) - if video_tensor.dim() == 4: + 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 index 2091cd83afdd..a4731af3dd55 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -1,3 +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. + import math import os import time @@ -59,7 +74,7 @@ def load_weights(self, weights: dict) -> None: self.transformer.eval() def load_standard_components( - self, checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = None + self, checkpoint_dir: str, device: torch.device, skip_components: Optional[list] = [] ) -> None: if PipelineComponent.TOKENIZER not in skip_components: logger.info("Loading tokenizer...") @@ -87,11 +102,11 @@ def load_standard_components( subfolder=PipelineComponent.SCHEDULER, ) - # load guardrails by default - if ( - not TRTLLM_DISABLE_COSMOS3_GUARDRAILS - and PipelineComponent.TEXT_GUARDRAIL not in skip_components - and PipelineComponent.VIDEO_GUARDRAIL not in skip_components + self.text_guardrail = None + self.video_guardrail = None + if not TRTLLM_DISABLE_COSMOS3_GUARDRAILS and ( + PipelineComponent.TEXT_GUARDRAIL not in skip_components + or PipelineComponent.VIDEO_GUARDRAIL not in skip_components ): if self.model_config.extra_attrs.get("guardrail_checkpoint_dir", None) is not None: logger.info( @@ -100,8 +115,10 @@ def load_standard_components( guardrail_ckpt_dir = self.model_config.extra_attrs["guardrail_checkpoint_dir"] else: guardrail_ckpt_dir = download_guardrail_checkpoint() - self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) - self.video_guardrail = build_video_guardrail(guardrail_ckpt_dir) + if PipelineComponent.TEXT_GUARDRAIL not in skip_components: + self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) + if PipelineComponent.VIDEO_GUARDRAIL not in skip_components: + self.video_guardrail = build_video_guardrail(guardrail_ckpt_dir) self.video_processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) @@ -218,7 +235,12 @@ def _tokenize_prompt( add_generation_prompt=True, return_dict=False, ) - token_ids = token_ids[:max_sequence_length] + 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) @@ -436,13 +458,22 @@ def forward( f"use a single image with multiple prompts instead." ) - if self.rank == 0 and use_guardrails: + # Text guardrail + text_blocked = torch.zeros((), device=self.device, dtype=torch.int32) + if self.rank == 0 and use_guardrails and self.text_guardrail is not None: for p in prompt: is_safe, msg = self.text_guardrail(p) if not is_safe: logger.warning(f"Text guardrail blocked prompt: {msg}") - timer.mark_end() - return timer.fill(PipelineOutput()) + 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) @@ -568,16 +599,24 @@ def forward_fn( video = self.decode_latents(latents, self._decode_latents) + # Video guardrail + video_blocked = torch.zeros((), device=self.device, dtype=torch.int32) 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: + if use_guardrails and self.video_guardrail is not None: video = check_video_safety(video, self.video_guardrail) if video is None: logger.warning("Video guardrail blocked video generation") - timer.mark_end() - return timer.fill(PipelineOutput()) + video_blocked.fill_(1) + + if torch.distributed.is_available() and torch.distributed.is_initialized(): + torch.distributed.broadcast(video_blocked, src=0) + + if video_blocked.item(): + timer.mark_end() + return timer.fill(PipelineOutput()) 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 index 19b59e19413c..de6278cb749c 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1,3 +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. + import math from typing import Tuple @@ -82,7 +97,7 @@ def compute_mrope_position_ids_vision( Returns: (position_ids [3, grid_t * grid_h * grid_w], next_temporal_offset) """ - if enable_fps_modulation: + 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) @@ -659,10 +674,8 @@ def __init__(self, model_config: DiffusionModelConfig): f"Cosmos3 does not support tensor parallelism. Got tp_size={vgm.tp_size}" ) - if ( - use_ulysses - and self.num_attention_heads % ulysses_size != 0 - and self.num_kv_heads % ulysses_size != 0 + 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 " From 0ddd2e90bfd1e8d4be341c4cbcacce1e6f7b4618 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 12 May 2026 12:51:19 -0700 Subject: [PATCH 08/17] fix seq padding for ulysses Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/transformer_cosmos3.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) 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 de6278cb749c..348ff0f7b165 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -946,10 +946,11 @@ def forward( hidden_gen = F.pad(hidden_gen, (0, 0, 0, pad)) cos, sin = self.cached_freqs_gen self.cached_freqs_gen = ( - F.pad(cos, (0, 0, 0, pad)), - F.pad(sin, (0, 0, 0, pad)), + F.pad(cos, (0, 0, 0, 0, 0, pad)), + F.pad(sin, (0, 0, 0, 0, 0, pad)), ) - S_shard = S_gen // self.seq_parallel_size + 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 ] From fe1533dc00de61ad46b40a3dc1cf3b80356d3675 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Fri, 15 May 2026 19:07:49 +0000 Subject: [PATCH 09/17] raise proper error for guardrails Signed-off-by: Shreyas Misra --- .../_torch/visual_gen/models/cosmos3/guardrails.py | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 96c2c9256bbc..657c4c93f578 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -82,6 +82,7 @@ def _pixelate_face(face_img: np.ndarray, blocks: int = 5) -> np.ndarray: # --------------------------------------------------------------------------- def download_guardrail_checkpoint() -> str: from huggingface_hub import snapshot_download + from huggingface_hub.errors import GatedRepoError try: return snapshot_download( @@ -93,10 +94,15 @@ def download_guardrail_checkpoint() -> str: logger.warning( f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO} {GUARDRAIL_HF_REVISION}" ) - return snapshot_download( - GUARDRAIL_HF_REPO, - revision=GUARDRAIL_HF_REVISION, - ) + try: + return snapshot_download( + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_HF_REVISION, + ) + except GatedRepoError: + raise ValueError( + f"{GUARDRAIL_HF_REPO} requires an approved HF_TOKEN. Please set HF_TOKEN and try again." + ) def build_text_guardrail(guardrail_ckpt_dir: str) -> TextGuardrailFn: From d3aa6c83107804bc64045553d65023ac93d665f7 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 18 May 2026 11:37:05 -0700 Subject: [PATCH 10/17] address comments Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/guardrails.py | 45 ++++++++++++------- .../models/cosmos3/pipeline_cosmos3.py | 45 +++++++++++++------ .../models/cosmos3/transformer_cosmos3.py | 36 ++++++++++----- 3 files changed, 84 insertions(+), 42 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 657c4c93f578..533bd0c2253e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -80,24 +80,22 @@ def _pixelate_face(face_img: np.ndarray, blocks: int = 5) -> np.ndarray: # --------------------------------------------------------------------------- # Default guardrail builders # --------------------------------------------------------------------------- -def download_guardrail_checkpoint() -> str: +def download_guardrail_checkpoint(repo_url: str, revision: str | None = None) -> str: from huggingface_hub import snapshot_download from huggingface_hub.errors import GatedRepoError try: return snapshot_download( - GUARDRAIL_HF_REPO, - revision=GUARDRAIL_HF_REVISION, + repo_url, + revision=revision, local_files_only=True, ) except FileNotFoundError: - logger.warning( - f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO} {GUARDRAIL_HF_REVISION}" - ) + logger.warning(f"Guardrail checkpoint not found, downloading from {repo_url} {revision}") try: return snapshot_download( - GUARDRAIL_HF_REPO, - revision=GUARDRAIL_HF_REVISION, + repo_url, + revision=revision, ) except GatedRepoError: raise ValueError( @@ -147,11 +145,14 @@ def _blocklist_check(prompt: str) -> tuple[bool, str]: try: from transformers import AutoModelForCausalLM, AutoTokenizer - model_id = "Qwen/Qwen3Guard-Gen-0.6B" - qwen_tokenizer = AutoTokenizer.from_pretrained(model_id) + model_dir = download_guardrail_checkpoint( + "Qwen/Qwen3Guard-Gen-0.6B", + revision="main", + ) + qwen_tokenizer = AutoTokenizer.from_pretrained(model_dir) qwen_model = ( AutoModelForCausalLM.from_pretrained( - model_id, + model_dir, torch_dtype=torch.bfloat16, ) .to("cuda") @@ -182,7 +183,13 @@ def _qwen_check(prompt: str) -> tuple[bool, str]: except (ImportError, OSError, RuntimeError, ValueError) as e: logger.warning("Could not load Qwen3Guard guardrail: %s", e) - def text_guardrail(prompt: str) -> None: + if not checkers: + raise RuntimeError( + "All text guardrail components failed to load. " + "Set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 to explicitly disable guardrails." + ) + + def text_guardrail(prompt: str) -> tuple[bool, str]: for checker in checkers: is_safe, msg = checker(prompt) if not is_safe: @@ -359,11 +366,15 @@ def _face_blur(frames: np.ndarray) -> np.ndarray: logger.warning("Could not load face blur filter: %s", e) def video_guardrail(frames: np.ndarray) -> np.ndarray | None: - if safety_checker is not None: - is_safe, msg = safety_checker(frames) - if not is_safe: - logger.warning(f"Video content safety: {msg}") - return None + if safety_checker is None: + raise RuntimeError( + "Video content safety classifier failed to load. " + "Set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 to explicitly disable guardrails." + ) + is_safe, msg = safety_checker(frames) + if not is_safe: + logger.warning(f"Video content safety: {msg}") + return None if face_blurrer is not None: frames = face_blurrer(frames) return frames 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 a4731af3dd55..5f586095f217 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,8 @@ from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS from .guardrails import ( + GUARDRAIL_HF_REPO, + GUARDRAIL_HF_REVISION, build_text_guardrail, build_video_guardrail, check_video_safety, @@ -83,6 +85,10 @@ def load_standard_components( 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( @@ -91,8 +97,12 @@ def load_standard_components( 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", 4) - self.vae_scale_factor_spatial = getattr(self.vae.config, "scale_factor_spatial", 16) + 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: @@ -104,17 +114,20 @@ def load_standard_components( self.text_guardrail = None self.video_guardrail = None - if not TRTLLM_DISABLE_COSMOS3_GUARDRAILS and ( - PipelineComponent.TEXT_GUARDRAIL not in skip_components - or PipelineComponent.VIDEO_GUARDRAIL not in skip_components + # 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 + and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS + and ( + PipelineComponent.TEXT_GUARDRAIL not in skip_components + or PipelineComponent.VIDEO_GUARDRAIL not in skip_components + ) ): - if self.model_config.extra_attrs.get("guardrail_checkpoint_dir", None) is not None: - logger.info( - f"Loading guardrails from {self.model_config.extra_attrs['guardrail_checkpoint_dir']}" - ) - guardrail_ckpt_dir = self.model_config.extra_attrs["guardrail_checkpoint_dir"] - else: - guardrail_ckpt_dir = download_guardrail_checkpoint() + guardrail_ckpt_dir = self.model_config.extra_attrs.get( + "guardrail_checkpoint_dir" + ) or download_guardrail_checkpoint(GUARDRAIL_HF_REPO, GUARDRAIL_HF_REVISION) + logger.info(f"Loading guardrails from {guardrail_ckpt_dir}") if PipelineComponent.TEXT_GUARDRAIL not in skip_components: self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) if PipelineComponent.VIDEO_GUARDRAIL not in skip_components: @@ -458,10 +471,14 @@ def forward( f"use a single image with multiple prompts instead." ) - # Text guardrail + # 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.text_guardrail is not None: - for p in prompt: + 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, msg = self.text_guardrail(p) if not is_safe: logger.warning(f"Text guardrail blocked prompt: {msg}") 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 348ff0f7b165..8d4b510341fe 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -289,6 +289,13 @@ def __init__( model_config: DiffusionModelConfig, layer_idx: int = 0, ): + original_backend = model_config.attention.backend + if model_config.attention.backend == "TRTLLM": + logger.warning( + "TRTLLM backend is not supported for Cosmos3CrossAttention. Falling back to VANILLA." + ) + model_config.attention.backend = "VANILLA" + super().__init__( hidden_size=hidden_size, num_attention_heads=num_attention_heads, @@ -302,6 +309,8 @@ def __init__( 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) @@ -684,10 +693,12 @@ def __init__(self, model_config: DiffusionModelConfig): ) if use_attn2d: - self.use_seq_parallel = True - self.seq_parallel_size = attn2d_mesh_size - self.seq_parallel_pg = vgm.attn2d_mesh_group - self.seq_parallel_rank = vgm.attn2d_mesh_rank + # 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 @@ -945,20 +956,23 @@ def forward( # 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 - self.cached_freqs_gen = ( - F.pad(cos, (0, 0, 0, 0, 0, pad)), - F.pad(sin, (0, 0, 0, 0, 0, pad)), - ) + 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 - cos, sin = self.cached_freqs_gen freqs_gen = ( - cos[:, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard], - sin[:, self.seq_parallel_rank * S_shard : (self.seq_parallel_rank + 1) * S_shard], + 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 From 3c4f9fe96af3d03507db4ef33201eeb0fbc522fd Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Mon, 18 May 2026 19:02:16 +0000 Subject: [PATCH 11/17] remove log Signed-off-by: Shreyas Misra --- .../_torch/visual_gen/models/cosmos3/transformer_cosmos3.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) 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 8d4b510341fe..665fc864d70d 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -291,9 +291,7 @@ def __init__( ): original_backend = model_config.attention.backend if model_config.attention.backend == "TRTLLM": - logger.warning( - "TRTLLM backend is not supported for Cosmos3CrossAttention. Falling back to VANILLA." - ) + # TRTLLM backend is not supported for Cosmos3CrossAttention model_config.attention.backend = "VANILLA" super().__init__( From f643e5c582a4a41413281ce495bb72258eb7dfb5 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 19 May 2026 08:05:43 -0700 Subject: [PATCH 12/17] remove siglip guardrail as requested by team Signed-off-by: Shreyas Misra --- .../visual_gen/models/cosmos3/guardrails.py | 105 +----------------- .../models/cosmos3/pipeline_cosmos3.py | 13 +-- .../models/cosmos3/transformer_cosmos3.py | 6 +- 3 files changed, 9 insertions(+), 115 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index 533bd0c2253e..e5ee178979ae 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -22,7 +22,6 @@ import cv2 import numpy as np import torch -import torch.nn as nn from tensorrt_llm.logger import logger @@ -31,39 +30,6 @@ GUARDRAIL_HF_REPO = "nvidia/Cosmos-Guardrail1" GUARDRAIL_HF_REVISION = "d6d4bfa899a71454a700907664f3e88f503950cf" -CUTOFF_UNSAFE_FRAMES_PERCENT = 10 - - -# --------------------------------------------------------------------------- -# Video safety classifier (matches reference: SigLIP so400m + 3-layer head) -# --------------------------------------------------------------------------- -class SafetyClassifier(nn.Module): - """3-layer classifier with BatchNorm (1152 → 512 → 256 → 7).""" - - def __init__(self, input_size: int = 1152, num_classes: int = 7): - super().__init__() - self.layers = nn.Sequential( - nn.Linear(input_size, 512), - nn.BatchNorm1d(512), - nn.ReLU(), - nn.Linear(512, 256), - nn.BatchNorm1d(256), - nn.ReLU(), - nn.Linear(256, num_classes), - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.layers(x) - - -CLASS_IDX_TO_NAME = { - 0: "Safe", - 1: "Sexual_Content", - 3: "Drugs", - 4: "Child_Abuse", - 5: "Hate_and_Harassment", - 6: "Self-Harm", -} # --------------------------------------------------------------------------- @@ -200,64 +166,9 @@ def text_guardrail(prompt: str) -> tuple[bool, str]: def build_video_guardrail(guardrail_ckpt_dir: str) -> VideoGuardrailFn: - safety_checker: Callable[[np.ndarray], tuple[bool, str]] | None = None face_blurrer: Callable[[np.ndarray], np.ndarray] | None = None - # 1. Video content safety filter: SigLIP so400m + SafetyClassifier - try: - from PIL import Image - from transformers import SiglipModel, SiglipProcessor - - siglip_dir = os.path.join( - guardrail_ckpt_dir, - "video_content_safety_filter", - "models--google--siglip-so400m-patch14-384/snapshots/9fdffc58afc957d1a03a25b10dba0329ab15c2a3", - ) - if not os.path.exists(siglip_dir): - raise FileNotFoundError(siglip_dir) - - siglip_model = ( - SiglipModel.from_pretrained(siglip_dir).to("cuda", dtype=torch.float32).eval() - ) - siglip_processor = SiglipProcessor.from_pretrained(siglip_dir) - - classifier = SafetyClassifier(input_size=1152, num_classes=7) - ckpt_path = os.path.join( - guardrail_ckpt_dir, "video_content_safety_filter", "safety_filter.pt" - ) - checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True) - state = {k.removeprefix("network."): v for k, v in checkpoint["model"].items()} - classifier.load_state_dict(state) - classifier = classifier.to("cuda", dtype=torch.float32).eval() - - def _safety_check(frames: np.ndarray) -> tuple[bool, str]: - unsafe_count = 0 - total = len(frames) - for frame in frames: - img = Image.fromarray(frame) - inputs = siglip_processor(images=img, return_tensors="pt").to( - "cuda", dtype=torch.float32 - ) - with torch.no_grad(): - siglip_out = siglip_model.get_image_features(**inputs) - features = siglip_out.pooler_output - features = features / features.norm(dim=-1, keepdim=True) - logits = classifier(features) - pred = logits.argmax(dim=-1).item() - class_name = CLASS_IDX_TO_NAME.get(pred, "Unknown") - if class_name != "Safe": - unsafe_count += 1 - - if unsafe_count / total > CUTOFF_UNSAFE_FRAMES_PERCENT / 100: - return False, f"Video content safety: {unsafe_count}/{total} frames unsafe" - return True, "" - - safety_checker = _safety_check - logger.info("Video content safety filter loaded (SigLIP so400m + classifier)") - except (ImportError, FileNotFoundError, OSError, RuntimeError, ValueError) as e: - logger.warning("Could not load video safety filter: %s", e) - - # 2. Face blur: RetinaFace + pixelation + # Face blur: RetinaFace + pixelation try: from retinaface.data import cfg_re50 from retinaface.layers.functions.prior_box import PriorBox @@ -365,19 +276,13 @@ def _face_blur(frames: np.ndarray) -> np.ndarray: except (ImportError, FileNotFoundError, OSError, RuntimeError, ValueError) as e: logger.warning("Could not load face blur filter: %s", e) - def video_guardrail(frames: np.ndarray) -> np.ndarray | None: - if safety_checker is None: + def video_guardrail(frames: np.ndarray) -> np.ndarray: + if face_blurrer is None: raise RuntimeError( - "Video content safety classifier failed to load. " + "Face blur filter not loaded. " "Set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 to explicitly disable guardrails." ) - is_safe, msg = safety_checker(frames) - if not is_safe: - logger.warning(f"Video content safety: {msg}") - return None - if face_blurrer is not None: - frames = face_blurrer(frames) - return frames + return face_blurrer(frames) return video_guardrail 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 5f586095f217..acce60194775 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -616,24 +616,13 @@ def forward_fn( video = self.decode_latents(latents, self._decode_latents) - # Video guardrail - video_blocked = torch.zeros((), device=self.device, dtype=torch.int32) + # 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.video_guardrail is not None: video = check_video_safety(video, self.video_guardrail) - if video is None: - logger.warning("Video guardrail blocked video generation") - video_blocked.fill_(1) - - if torch.distributed.is_available() and torch.distributed.is_initialized(): - torch.distributed.broadcast(video_blocked, src=0) - - if video_blocked.item(): - timer.mark_end() - return timer.fill(PipelineOutput()) 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 index 665fc864d70d..9796292d3170 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -1006,9 +1006,9 @@ def load_weights(self, weights: dict) -> None: remapped[k] = value continue - if k.startswith("time_embedder.mlp."): - k = k.replace("time_embedder.mlp.0.", "time_embedder.mlp.linear_1.") - k = k.replace("time_embedder.mlp.2.", "time_embedder.mlp.linear_2.") + 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 From 66344d52e41c9e99500f7b463a1729ea54c55827 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 20 May 2026 14:25:26 -0700 Subject: [PATCH 13/17] update repo and add note about upstream PR Signed-off-by: Shreyas Misra --- requirements.txt | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index d38c7c25fad2..21e40fa1bf73 100644 --- a/requirements.txt +++ b/requirements.txt @@ -91,4 +91,5 @@ smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 nltk==3.9.4 better_profanity==0.7.0 -retinaface @ git+https://github.com/NVShreyas/retinaface.git@1bb2e4589b5f1cf4e1d01a56b7b0a7259ac563de +# NOTE: Should be updated to official release once PR in upstream is merged - https://github.com/andresprados/Pytorch_Retinaface/pull/2 +retinaface-py @ git+https://github.com/NVShreyas/Pytorch_Retinaface.git@a7c9c6bec4943cf5d2120128ad78f5542e7e317d From 28bcbc1cd4822add19e1d06a43ebdde3f81480a4 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 26 May 2026 17:13:05 +0000 Subject: [PATCH 14/17] new checkpoint format Signed-off-by: Shreyas Misra --- tensorrt_llm/_torch/visual_gen/config.py | 5 +- .../models/cosmos3/pipeline_cosmos3.py | 8 ++- .../models/cosmos3/transformer_cosmos3.py | 57 ++++++++++++------- .../_torch/visual_gen/pipeline_registry.py | 2 + 4 files changed, 43 insertions(+), 29 deletions(-) diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 2158daad9df3..9b65baf888db 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -397,14 +397,11 @@ def from_pretrained( resolved_pipeline_config = kwargs.pop("pipeline_config", None) if resolved_pipeline_config is None: resolved_pipeline_config = dict(args.pipeline_config) if args else {} - for key in ("spatial_upsampler_path", "distilled_lora_path"): + for key in ("spatial_upsampler_path", "distilled_lora_path", "guardrail_checkpoint_dir"): value = resolved_pipeline_config.get(key) if value: extra_attrs[key] = value - if args and args.guardrail_checkpoint_dir: - extra_attrs["guardrail_checkpoint_dir"] = args.guardrail_checkpoint_dir - # Discover pipeline components (diffusers layout) components = discover_pipeline_components(checkpoint_path) 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 acce60194775..1d3dedbb4d1e 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -25,10 +25,9 @@ from diffusers.video_processor import VideoProcessor from transformers import Qwen2Tokenizer -from tensorrt_llm._torch.visual_gen.config import PipelineComponent 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 register_pipeline +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 @@ -60,7 +59,8 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" -@register_pipeline("Cosmos3OmniMoTPipeline") +# TODO: add hf_ids +@register_pipeline("Cosmos3OmniMoTPipeline", defaults={"guardrail_checkpoint_dir": None}) class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, model_config): super().__init__(model_config) @@ -78,6 +78,8 @@ def load_weights(self, weights: dict) -> None: 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( 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 9796292d3170..d260b8e10ea1 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/transformer_cosmos3.py @@ -992,33 +992,46 @@ def forward( return self.unpatchify(self.llm2vae(hidden_gen), T, H, W) def load_weights(self, weights: dict) -> None: - """Load weights with key remapping from Diffusers/HF Cosmos3 checkpoints. + """Load weights with key remapping from Cosmos3-Nano / Diffusers checkpoints. - Expects tensor names under ``model.*`` (e.g. ``model.layers.{i}.self_attn.q_proj.weight``). + 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("lm_head."): - continue - - if not k.startswith("model."): - logger.warning(f"Skipping unknown checkpoint key: {key}") - continue - k = k[len("model.") :] + if k.startswith("model."): + k = k[len("model.") :] # embed_tokens and norm → language_model.* if k.startswith("embed_tokens.") or k.startswith("norm."): @@ -1031,7 +1044,7 @@ def load_weights(self, weights: dict) -> None: continue if not k.startswith("layers."): - logger.warning(f"Skipping unknown LM key: {key}") + logger.warning(f"Skipping unknown checkpoint key: {key}") continue parts = k.split(".", 2) # ['layers', '{i}', '{rest}'] @@ -1045,22 +1058,22 @@ def load_weights(self, weights: dict) -> None: # --- UND attention → language_model.layers.{i}.self_attn.* --- attn_und_map = { - "self_attn.q_proj.": f"{und_lp}.self_attn.to_q.", - "self_attn.k_proj.": f"{und_lp}.self_attn.to_k.", - "self_attn.v_proj.": f"{und_lp}.self_attn.to_v.", - "self_attn.o_proj.": f"{und_lp}.self_attn.to_out.0.", - "self_attn.q_norm.": f"{und_lp}.self_attn.norm_q.", - "self_attn.k_norm.": f"{und_lp}.self_attn.norm_k.", + "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.q_proj_moe_gen.": f"{gen_lp}.cross_attention.to_q.", - "self_attn.k_proj_moe_gen.": f"{gen_lp}.cross_attention.to_k.", - "self_attn.v_proj_moe_gen.": f"{gen_lp}.cross_attention.to_v.", - "self_attn.o_proj_moe_gen.": f"{gen_lp}.cross_attention.to_out.0.", - "self_attn.q_norm_moe_gen.": f"{gen_lp}.cross_attention.norm_q.", - "self_attn.k_norm_moe_gen.": f"{gen_lp}.cross_attention.norm_k.", + "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 --- diff --git a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index 8d804b8dc1a9..c96340f5ac55 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -56,6 +56,8 @@ class PipelineComponent(str, Enum): SCHEDULER = "scheduler" IMAGE_ENCODER = "image_encoder" IMAGE_PROCESSOR = "image_processor" + TEXT_GUARDRAIL = "text_guardrail" + VIDEO_GUARDRAIL = "video_guardrail" @dataclass From e956985ba6e204f534d9abd78a3e7ddac60ed676 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 26 May 2026 18:25:20 +0000 Subject: [PATCH 15/17] revamp guardrails Signed-off-by: Shreyas Misra --- requirements.txt | 4 - tensorrt_llm/_torch/visual_gen/config.py | 2 +- .../visual_gen/models/cosmos3/guardrails.py | 268 +----------------- .../models/cosmos3/pipeline_cosmos3.py | 59 ++-- .../_torch/visual_gen/pipeline_registry.py | 2 - 5 files changed, 43 insertions(+), 292 deletions(-) diff --git a/requirements.txt b/requirements.txt index 21e40fa1bf73..3a1de400c452 100644 --- a/requirements.txt +++ b/requirements.txt @@ -89,7 +89,3 @@ etcd-sdk-python==0.0.7 python-multipart smg-grpc-proto>=0.4.2 cache-dit>=1.3.5 -nltk==3.9.4 -better_profanity==0.7.0 -# NOTE: Should be updated to official release once PR in upstream is merged - https://github.com/andresprados/Pytorch_Retinaface/pull/2 -retinaface-py @ git+https://github.com/NVShreyas/Pytorch_Retinaface.git@a7c9c6bec4943cf5d2120128ad78f5542e7e317d diff --git a/tensorrt_llm/_torch/visual_gen/config.py b/tensorrt_llm/_torch/visual_gen/config.py index 9b65baf888db..3f7009bef569 100644 --- a/tensorrt_llm/_torch/visual_gen/config.py +++ b/tensorrt_llm/_torch/visual_gen/config.py @@ -397,7 +397,7 @@ def from_pretrained( resolved_pipeline_config = kwargs.pop("pipeline_config", None) if resolved_pipeline_config is None: resolved_pipeline_config = dict(args.pipeline_config) if args else {} - for key in ("spatial_upsampler_path", "distilled_lora_path", "guardrail_checkpoint_dir"): + for key in ("spatial_upsampler_path", "distilled_lora_path"): value = resolved_pipeline_config.get(key) if value: extra_attrs[key] = value diff --git a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py index e5ee178979ae..881d830450fa 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/guardrails.py @@ -15,287 +15,49 @@ from __future__ import annotations -import os -import warnings -from typing import Callable +from typing import Any -import cv2 -import numpy as np import torch from tensorrt_llm.logger import logger -TextGuardrailFn = Callable[[str], tuple[bool, str]] -VideoGuardrailFn = Callable[[np.ndarray], np.ndarray] +GUARDRAIL_HF_REPO = "nvidia/Cosmos-1.0-Guardrail" +GUARDRAIL_REVISION = "cf03c0395fac8c4de386c0bdab12cc4fc8d66362" -GUARDRAIL_HF_REPO = "nvidia/Cosmos-Guardrail1" -GUARDRAIL_HF_REVISION = "d6d4bfa899a71454a700907664f3e88f503950cf" - -# --------------------------------------------------------------------------- -# Face pixelation utility -# --------------------------------------------------------------------------- -def _pixelate_face(face_img: np.ndarray, blocks: int = 5) -> np.ndarray: - h, w = face_img.shape[:2] - if h == 0 or w == 0: - return face_img - temp = cv2.resize(face_img, (blocks, blocks), interpolation=cv2.INTER_LINEAR) - return cv2.resize(temp, (w, h), interpolation=cv2.INTER_NEAREST) - - -# --------------------------------------------------------------------------- -# Default guardrail builders -# --------------------------------------------------------------------------- -def download_guardrail_checkpoint(repo_url: str, revision: str | None = None) -> str: +def download_guardrail_checkpoint() -> str: from huggingface_hub import snapshot_download from huggingface_hub.errors import GatedRepoError try: return snapshot_download( - repo_url, - revision=revision, + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_REVISION, local_files_only=True, ) except FileNotFoundError: - logger.warning(f"Guardrail checkpoint not found, downloading from {repo_url} {revision}") + logger.warning(f"Guardrail checkpoint not found, downloading from {GUARDRAIL_HF_REPO}") try: return snapshot_download( - repo_url, - revision=revision, + GUARDRAIL_HF_REPO, + revision=GUARDRAIL_REVISION, ) except GatedRepoError: raise ValueError( - f"{GUARDRAIL_HF_REPO} requires an approved HF_TOKEN. Please set HF_TOKEN and try again." - ) - - -def build_text_guardrail(guardrail_ckpt_dir: str) -> TextGuardrailFn: - checkers: list[Callable[[str], tuple[bool, str]]] = [] - - # 1. Blocklist - try: - import nltk - from better_profanity import profanity as profanity_filter - - blocklist_dir = os.path.join(guardrail_ckpt_dir, "blocklist") - nltk.data.path.append(os.path.join(blocklist_dir, "nltk_data")) - - def _read_keywords(dirpath: str) -> list[str]: - words: list[str] = [] - if not os.path.isdir(dirpath): - return words - for fname in sorted(os.listdir(dirpath)): - fpath = os.path.join(dirpath, fname) - if os.path.isfile(fpath): - with open(fpath) as f: - words.extend(line.strip() for line in f if line.strip()) - return words - - blocklist_words = _read_keywords(os.path.join(blocklist_dir, "custom")) - whitelist_words = _read_keywords(os.path.join(blocklist_dir, "whitelist")) - profanity_filter.load_censor_words( - custom_words=blocklist_words, whitelist_words=whitelist_words - ) - - def _blocklist_check(prompt: str) -> tuple[bool, str]: - if profanity_filter.contains_profanity(prompt): - return False, "Blocked by keyword filter" - return True, "" - - checkers.append(_blocklist_check) - logger.info("Blocklist guardrail loaded (%d keywords)", len(blocklist_words)) - except (ImportError, OSError, RuntimeError, ValueError) as e: - logger.warning("Could not load blocklist guardrail: %s", e) - - # 2. Qwen3Guard - try: - from transformers import AutoModelForCausalLM, AutoTokenizer - - model_dir = download_guardrail_checkpoint( - "Qwen/Qwen3Guard-Gen-0.6B", - revision="main", - ) - qwen_tokenizer = AutoTokenizer.from_pretrained(model_dir) - qwen_model = ( - AutoModelForCausalLM.from_pretrained( - model_dir, - torch_dtype=torch.bfloat16, - ) - .to("cuda") - .eval() - ) - - def _qwen_check(prompt: str) -> tuple[bool, str]: - conversations = [{"role": "user", "content": prompt}] - input_ids = qwen_tokenizer.apply_chat_template( - conversations, - tokenize=True, - return_tensors="pt", - add_generation_prompt=True, - return_dict=False, - ).to("cuda") - with torch.no_grad(): - output_ids = qwen_model.generate(input_ids, max_new_tokens=128) - response = qwen_tokenizer.decode( - output_ids[0][input_ids.shape[1] :], - skip_special_tokens=True, - ) - if "unsafe" in response.lower(): - return False, f"Qwen3Guard: {response.strip()}" - return True, "" - - checkers.append(_qwen_check) - logger.info("Qwen3Guard guardrail loaded") - except (ImportError, OSError, RuntimeError, ValueError) as e: - logger.warning("Could not load Qwen3Guard guardrail: %s", e) - - if not checkers: - raise RuntimeError( - "All text guardrail components failed to load. " - "Set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 to explicitly disable guardrails." - ) - - def text_guardrail(prompt: str) -> tuple[bool, str]: - for checker in checkers: - is_safe, msg = checker(prompt) - if not is_safe: - return is_safe, msg - return True, "" - - return text_guardrail - - -def build_video_guardrail(guardrail_ckpt_dir: str) -> VideoGuardrailFn: - face_blurrer: Callable[[np.ndarray], np.ndarray] | None = None - - # Face blur: RetinaFace + pixelation - try: - from retinaface.data import cfg_re50 - from retinaface.layers.functions.prior_box import PriorBox - from retinaface.models.retinaface import RetinaFace - from retinaface.utils.nms.py_cpu_nms import py_cpu_nms - - face_ckpt = os.path.join(guardrail_ckpt_dir, "face_blur_filter", "Resnet50_Final.pth") - if not os.path.exists(face_ckpt): - raise FileNotFoundError(face_ckpt) - - cfg = dict(cfg_re50) - cfg["pretrain"] = False - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - retinaface_net = RetinaFace(cfg=cfg, phase="test") - - # Load weights (strip 'module.' prefix if present) - pretrained_dict = torch.load(face_ckpt, map_location="cpu", weights_only=True) - if "state_dict" in pretrained_dict: - pretrained_dict = pretrained_dict["state_dict"] - pretrained_dict = { - k.replace("module.", "", 1) if k.startswith("module.") else k: v - for k, v in pretrained_dict.items() - } - retinaface_net.load_state_dict(pretrained_dict, strict=False) - retinaface_device = "cuda" - retinaface_net = retinaface_net.to(retinaface_device, dtype=torch.float32).eval() - - CONF_THRESH = 0.7 - NMS_THRESH = 0.4 - TOP_K = 5000 - KEEP_TOP_K = 750 - - def _decode_batch(loc, priors, variances): - batch_size = loc.size(0) - p = priors.unsqueeze(0).expand(batch_size, -1, -1) - boxes = torch.cat( - ( - p[:, :, :2] + loc[:, :, :2] * variances[0] * p[:, :, 2:], - p[:, :, 2:] * torch.exp(loc[:, :, 2:] * variances[1]), - ), - dim=2, + "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" ) - boxes[:, :, :2] -= boxes[:, :, 2:] / 2 - boxes[:, :, 2:] += boxes[:, :, :2] - return boxes - - def _face_blur(frames: np.ndarray) -> np.ndarray: - prior_data = None - scale = None - result_frames = [] - - for frame in frames: - frame_t = torch.from_numpy(frame).to("cuda", dtype=torch.float32) - frame_t = frame_t.permute(2, 0, 1).unsqueeze(0) # [1, C, H, W] - frame_t = frame_t[:, [2, 1, 0], :, :] # RGB → BGR - means = torch.tensor( - [104.0, 117.0, 123.0], device="cuda", dtype=torch.float32 - ).view(1, 3, 1, 1) - frame_t = frame_t - means - - h, w = frame_t.shape[2], frame_t.shape[3] - if prior_data is None: - priorbox = PriorBox(cfg, image_size=(h, w)) - prior_data = priorbox.forward().to("cuda", dtype=torch.float32) - if scale is None: - scale = torch.tensor([w, h, w, h], device="cuda", dtype=torch.float32) - - with torch.no_grad(): - loc, conf, _ = retinaface_net(frame_t) - - boxes = _decode_batch(loc, prior_data, cfg["variance"]) - boxes = (boxes * scale).squeeze(0).cpu().numpy() - scores = conf.squeeze(0)[:, 1].cpu().numpy() - - # Filter by confidence - inds = np.where(scores > CONF_THRESH)[0] - boxes_f = boxes[inds] - scores_f = scores[inds] - order = scores_f.argsort()[::-1][:TOP_K] - boxes_f = boxes_f[order] - scores_f = scores_f[order] - - # NMS - dets = np.hstack((boxes_f, scores_f[:, np.newaxis])).astype(np.float32) - keep = py_cpu_nms(dets, NMS_THRESH) - dets = dets[keep][:KEEP_TOP_K] - - out_frame = frame.copy() - for det in dets: - x1, y1, x2, y2 = map(int, det[:4]) - if x2 - x1 < 20 or y2 - y1 < 20: - continue - max_h, max_w = out_frame.shape[:2] - y1c, y2c = max(y1, 0), min(y2, max_h) - x1c, x2c = max(x1, 0), min(x2, max_w) - out_frame[y1c:y2c, x1c:x2c] = _pixelate_face(out_frame[y1c:y2c, x1c:x2c]) - - result_frames.append(out_frame) - - return np.array(result_frames) - - face_blurrer = _face_blur - logger.info("Face blur filter loaded (RetinaFace Resnet50)") - except (ImportError, FileNotFoundError, OSError, RuntimeError, ValueError) as e: - logger.warning("Could not load face blur filter: %s", e) - - def video_guardrail(frames: np.ndarray) -> np.ndarray: - if face_blurrer is None: - raise RuntimeError( - "Face blur filter not loaded. " - "Set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 to explicitly disable guardrails." - ) - return face_blurrer(frames) - - return video_guardrail -def check_video_safety( - video_tensor: torch.Tensor, video_guardrail: VideoGuardrailFn -) -> torch.Tensor | None: +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 = video_guardrail(frames_np) + frames_np = safety_checker.check_video_safety(frames_np) if frames_np is None: return None 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 1d3dedbb4d1e..958ec9c23b06 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -33,14 +33,7 @@ from tensorrt_llm.logger import logger from .defaults import COSMOS3_720P_PARAMS, COSMOS3_EXTRA_SPECS -from .guardrails import ( - GUARDRAIL_HF_REPO, - GUARDRAIL_HF_REVISION, - build_text_guardrail, - build_video_guardrail, - check_video_safety, - download_guardrail_checkpoint, -) +from .guardrails import check_video_safety, download_guardrail_checkpoint from .transformer_cosmos3 import Cosmos3VFMTransformer COSMOS3_DEFAULT_NEGATIVE_PROMPT = ( @@ -59,8 +52,22 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" +if not TRTLLM_DISABLE_COSMOS3_GUARDRAILS: + 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 " + "disable guardrails by setting TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1." + "`pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python`" + ) + + # TODO: add hf_ids -@register_pipeline("Cosmos3OmniMoTPipeline", defaults={"guardrail_checkpoint_dir": None}) +@register_pipeline("Cosmos3OmniMoTPipeline") class Cosmos3OmniMoTPipeline(BasePipeline): def __init__(self, model_config): super().__init__(model_config) @@ -114,26 +121,14 @@ def load_standard_components( subfolder=PipelineComponent.SCHEDULER, ) - self.text_guardrail = None - self.video_guardrail = None # 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 - and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS - and ( - PipelineComponent.TEXT_GUARDRAIL not in skip_components - or PipelineComponent.VIDEO_GUARDRAIL not in skip_components - ) - ): - guardrail_ckpt_dir = self.model_config.extra_attrs.get( - "guardrail_checkpoint_dir" - ) or download_guardrail_checkpoint(GUARDRAIL_HF_REPO, GUARDRAIL_HF_REVISION) - logger.info(f"Loading guardrails from {guardrail_ckpt_dir}") - if PipelineComponent.TEXT_GUARDRAIL not in skip_components: - self.text_guardrail = build_text_guardrail(guardrail_ckpt_dir) - if PipelineComponent.VIDEO_GUARDRAIL not in skip_components: - self.video_guardrail = build_video_guardrail(guardrail_ckpt_dir) + if self.rank == 0 and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS: + # 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) @@ -476,14 +471,14 @@ 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. text_blocked = torch.zeros((), device=self.device, dtype=torch.int32) - if self.rank == 0 and use_guardrails and self.text_guardrail is not None: + 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, msg = self.text_guardrail(p) + is_safe = self.safety_checker.check_text_safety(p) if not is_safe: - logger.warning(f"Text guardrail blocked prompt: {msg}") + logger.warning("Text guardrail blocked prompt") text_blocked.fill_(1) break @@ -623,8 +618,8 @@ def forward_fn( 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.video_guardrail is not None: - video = check_video_safety(video, self.video_guardrail) + 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/pipeline_registry.py b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py index c96340f5ac55..8d804b8dc1a9 100644 --- a/tensorrt_llm/_torch/visual_gen/pipeline_registry.py +++ b/tensorrt_llm/_torch/visual_gen/pipeline_registry.py @@ -56,8 +56,6 @@ class PipelineComponent(str, Enum): SCHEDULER = "scheduler" IMAGE_ENCODER = "image_encoder" IMAGE_PROCESSOR = "image_processor" - TEXT_GUARDRAIL = "text_guardrail" - VIDEO_GUARDRAIL = "video_guardrail" @dataclass From acd100b9a94a2c873a71c3946af14d4826bb9a61 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Tue, 26 May 2026 15:25:56 -0700 Subject: [PATCH 16/17] set TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1 in CI Signed-off-by: Shreyas Misra --- jenkins/L0_Test.groovy | 2 ++ tests/integration/defs/conftest.py | 5 +++++ tests/unittest/conftest.py | 5 +++++ 3 files changed, 12 insertions(+) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index ecc967d59aca..19b6a1cb3a6f 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -854,6 +854,8 @@ def getPytestBaseCommandLine( extraInternalEnv += " NCCL_DEBUG=INFO" // Pass stage name to perf sanity tests for OpenSearch tracking extraInternalEnv += " stageName=${stageName}" + // CI images do not ship cosmos_guardrail; must be set before conftest imports tensorrt_llm. + extraInternalEnv += " TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1" // Container port allocation environment variables for avoiding port conflicts def portEnvVars = "" diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 007998ee9b41..12058984f9fa 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -18,6 +18,11 @@ import gc import logging import os + +# CI images do not ship cosmos_guardrail. Set before tensorrt_llm import (visual_gen loads cosmos3 at import time). +if "JENKINS_HOME" in os.environ: + os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + import platform import re import shutil diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index 13c51d0336ee..b2c58cd68fb5 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -14,6 +14,11 @@ # limitations under the License. # # Force resource release after test import os + +# CI images do not ship cosmos_guardrail. Set before any tensorrt_llm import. +if "JENKINS_HOME" in os.environ: + os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") + import signal import sys import traceback From 06b027ee8bf013bb4117c9ba404f5d6011cdc8a1 Mon Sep 17 00:00:00 2001 From: Shreyas Misra Date: Wed, 27 May 2026 07:24:03 -0700 Subject: [PATCH 17/17] lazy import cosmos_guardrails Signed-off-by: Shreyas Misra --- jenkins/L0_Test.groovy | 2 - .../models/cosmos3/pipeline_cosmos3.py | 44 +++++++++---------- tests/integration/defs/conftest.py | 5 --- tests/unittest/conftest.py | 5 --- 4 files changed, 22 insertions(+), 34 deletions(-) diff --git a/jenkins/L0_Test.groovy b/jenkins/L0_Test.groovy index 19b6a1cb3a6f..ecc967d59aca 100644 --- a/jenkins/L0_Test.groovy +++ b/jenkins/L0_Test.groovy @@ -854,8 +854,6 @@ def getPytestBaseCommandLine( extraInternalEnv += " NCCL_DEBUG=INFO" // Pass stage name to perf sanity tests for OpenSearch tracking extraInternalEnv += " stageName=${stageName}" - // CI images do not ship cosmos_guardrail; must be set before conftest imports tensorrt_llm. - extraInternalEnv += " TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1" // Container port allocation environment variables for avoiding port conflicts def portEnvVars = "" 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 958ec9c23b06..800006216bf2 100644 --- a/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py +++ b/tensorrt_llm/_torch/visual_gen/models/cosmos3/pipeline_cosmos3.py @@ -52,20 +52,6 @@ TRTLLM_DISABLE_COSMOS3_GUARDRAILS = os.environ.get("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "0") == "1" -if not TRTLLM_DISABLE_COSMOS3_GUARDRAILS: - 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 " - "disable guardrails by setting TRTLLM_DISABLE_COSMOS3_GUARDRAILS=1." - "`pip install cosmos_guardrail==0.3.0 && pip uninstall opencv-python`" - ) - - # TODO: add hf_ids @register_pipeline("Cosmos3OmniMoTPipeline") class Cosmos3OmniMoTPipeline(BasePipeline): @@ -121,14 +107,28 @@ def load_standard_components( subfolder=PipelineComponent.SCHEDULER, ) - # 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 and not TRTLLM_DISABLE_COSMOS3_GUARDRAILS: - # 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) + 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) diff --git a/tests/integration/defs/conftest.py b/tests/integration/defs/conftest.py index 12058984f9fa..007998ee9b41 100644 --- a/tests/integration/defs/conftest.py +++ b/tests/integration/defs/conftest.py @@ -18,11 +18,6 @@ import gc import logging import os - -# CI images do not ship cosmos_guardrail. Set before tensorrt_llm import (visual_gen loads cosmos3 at import time). -if "JENKINS_HOME" in os.environ: - os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") - import platform import re import shutil diff --git a/tests/unittest/conftest.py b/tests/unittest/conftest.py index b2c58cd68fb5..13c51d0336ee 100644 --- a/tests/unittest/conftest.py +++ b/tests/unittest/conftest.py @@ -14,11 +14,6 @@ # limitations under the License. # # Force resource release after test import os - -# CI images do not ship cosmos_guardrail. Set before any tensorrt_llm import. -if "JENKINS_HOME" in os.environ: - os.environ.setdefault("TRTLLM_DISABLE_COSMOS3_GUARDRAILS", "1") - import signal import sys import traceback