From 3ea9607d1c6ef9b5d92475b4f323b42941895349 Mon Sep 17 00:00:00 2001 From: Dawid Majchrowski Date: Fri, 29 May 2026 06:08:45 +0200 Subject: [PATCH] [diffusion] model: update to new model format (#26492) --- .../configs/models/dits/cosmos3video.py | 87 ++-- .../runtime/models/dits/cosmos3video.py | 16 +- .../cosmos3_guardrails.py | 459 ++---------------- .../multimodal_gen/test/unit/test_cosmos3.py | 76 +-- 4 files changed, 143 insertions(+), 495 deletions(-) diff --git a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py index 6efb44987..da7476135 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/configs/models/dits/cosmos3video.py @@ -15,101 +15,100 @@ def _build_cosmos3_param_names_mapping() -> dict: """Map diffusers-format Cosmos3 weights to the sglang model namespace. Source keys (diffusers transformer ckpt) → target keys (sglang model): - model.embed_tokens.weight -> language_model.embed_tokens.weight - model.layers.X.input_layernorm.weight -> language_model.layers.X.input_layernorm.weight - model.layers.X.input_layernorm_moe_gen.weight -> gen_layers.X.input_layernorm.weight - model.layers.X.self_attn.{q,k,v}_proj.weight -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0) - model.layers.X.self_attn.{q,k,v}_proj_moe_gen.weight -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0) - model.layers.X.mlp.{gate,up}_proj.weight -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0) - model.layers.X.mlp_moe_gen.{gate,up}_proj.weight -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0) - model.norm_moe_gen.weight -> norm_moe_gen.weight - time_embedder.mlp.{0,2}.weight -> time_embedder.linear_{1,2}.weight - vae2llm.weight, llm2vae.weight -> (pass-through) + embed_tokens.weight -> language_model.embed_tokens.weight + layers.X.input_layernorm.weight -> language_model.layers.X.input_layernorm.weight + layers.X.input_layernorm_moe_gen.weight -> gen_layers.X.input_layernorm.weight + layers.X.self_attn.{to_q,to_k,to_v}.weight -> language_model.layers.X.self_attn.to_qkv.weight (concat dim 0) + layers.X.self_attn.{add_q,add_k,add_v}_proj.weight -> gen_layers.X.cross_attention.to_qkv.weight (concat dim 0) + layers.X.mlp.{gate,up}_proj.weight -> language_model.layers.X.mlp.gate_up_proj.weight (concat dim 0) + layers.X.mlp_moe_gen.{gate,up}_proj.weight -> gen_layers.X.mlp.gate_up_proj.weight (concat dim 0) + norm_moe_gen.weight -> norm_moe_gen.weight + time_embedder.linear_{1,2}.weight -> (pass-through) + proj_in.weight, proj_out.weight -> (pass-through) - GEN patterns (`*_moe_gen`) must precede the UND catch-all so the - catch-all can't claim GEN keys. `model.norm.weight` and `lm_head.weight` - are inherited from Qwen3-VL pretraining and not used at inference, so - they are skipped via empty-string replacement. + GEN patterns (`*_moe_gen`, `add_*`, `to_add_out`, `norm_added_*`) must + precede the UND catch-all so the catch-all can't claim GEN keys. + `norm.weight` and `lm_head.weight` are inherited from Qwen3-VL + pretraining and not used at inference; audio/action keys are reserved + for a future modality extension — all skipped via empty-string replacement. """ return { # Inherited from Qwen3-VL pretraining; unused at diffusion inference. r"^lm_head\.weight$": "", - r"^model\.norm\.weight$": "", - # Top-level norms / heads. - r"^model\.norm_moe_gen\.(.*)$": r"norm_moe_gen.\1", - r"^model\.embed_tokens\.(.*)$": r"language_model.embed_tokens.\1", - # Time embedder: mlp.0 -> linear_1, mlp.2 -> linear_2 (SiLU at index 1). - r"^time_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1", - r"^time_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\1", + r"^norm\.weight$": "", + # Audio / action modalities — not yet wired; skip to avoid load warnings. + r"^audio_.*$": "", + r"^action_.*$": "", + # Top-level norms / embeddings. + r"^norm_moe_gen\.(.*)$": r"norm_moe_gen.\1", + r"^embed_tokens\.(.*)$": r"language_model.embed_tokens.\1", # GEN pathway: per-layer (must run before the UND catch-all below). # Q/K/V merge into MergedColumnParallelLinear to_qkv (concat order: Q, K, V). - r"^model\.layers\.(\d+)\.self_attn\.q_proj_moe_gen\.(.*)$": ( + r"^layers\.(\d+)\.self_attn\.add_q_proj\.(.*)$": ( r"gen_layers.\1.cross_attention.to_qkv.\2", 0, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.k_proj_moe_gen\.(.*)$": ( + r"^layers\.(\d+)\.self_attn\.add_k_proj\.(.*)$": ( r"gen_layers.\1.cross_attention.to_qkv.\2", 1, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.v_proj_moe_gen\.(.*)$": ( + r"^layers\.(\d+)\.self_attn\.add_v_proj\.(.*)$": ( r"gen_layers.\1.cross_attention.to_qkv.\2", 2, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.o_proj_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.to_out.\2", - r"^model\.layers\.(\d+)\.self_attn\.q_norm_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.norm_q.\2", - r"^model\.layers\.(\d+)\.self_attn\.k_norm_moe_gen\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2", - r"^model\.layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2", - r"^model\.layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2", + r"^layers\.(\d+)\.self_attn\.to_add_out\.(.*)$": r"gen_layers.\1.cross_attention.to_out.\2", + r"^layers\.(\d+)\.self_attn\.norm_added_q\.(.*)$": r"gen_layers.\1.cross_attention.norm_q.\2", + r"^layers\.(\d+)\.self_attn\.norm_added_k\.(.*)$": r"gen_layers.\1.cross_attention.norm_k.\2", + r"^layers\.(\d+)\.input_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.input_layernorm.\2", + r"^layers\.(\d+)\.post_attention_layernorm_moe_gen\.(.*)$": r"gen_layers.\1.post_attention_layernorm.\2", # GEN MLP gate/up merge into MergedColumnParallelLinear gate_up_proj. # Must precede the mlp_moe_gen catch-all below. - r"^model\.layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": ( + r"^layers\.(\d+)\.mlp_moe_gen\.gate_proj\.(.*)$": ( r"gen_layers.\1.mlp.gate_up_proj.\2", 0, 2, ), - r"^model\.layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": ( + r"^layers\.(\d+)\.mlp_moe_gen\.up_proj\.(.*)$": ( r"gen_layers.\1.mlp.gate_up_proj.\2", 1, 2, ), - r"^model\.layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2", - # UND pathway: per-layer attention rename (q/k/v_proj -> to_qkv merged, - # q_norm/k_norm -> norm_q/k, o_proj -> to_out). - r"^model\.layers\.(\d+)\.self_attn\.q_proj\.(.*)$": ( + r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2", + # UND pathway: Q/K/V merge into to_qkv; remaining attention keys + # (to_out, norm_q, norm_k) and layernorms pass through the catch-all. + r"^layers\.(\d+)\.self_attn\.to_q\.(.*)$": ( r"language_model.layers.\1.self_attn.to_qkv.\2", 0, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.k_proj\.(.*)$": ( + r"^layers\.(\d+)\.self_attn\.to_k\.(.*)$": ( r"language_model.layers.\1.self_attn.to_qkv.\2", 1, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.v_proj\.(.*)$": ( + r"^layers\.(\d+)\.self_attn\.to_v\.(.*)$": ( r"language_model.layers.\1.self_attn.to_qkv.\2", 2, 3, ), - r"^model\.layers\.(\d+)\.self_attn\.o_proj\.(.*)$": r"language_model.layers.\1.self_attn.to_out.\2", - r"^model\.layers\.(\d+)\.self_attn\.q_norm\.(.*)$": r"language_model.layers.\1.self_attn.norm_q.\2", - r"^model\.layers\.(\d+)\.self_attn\.k_norm\.(.*)$": r"language_model.layers.\1.self_attn.norm_k.\2", # UND MLP gate/up merge into MergedColumnParallelLinear gate_up_proj. # Must precede the layers catch-all below. - r"^model\.layers\.(\d+)\.mlp\.gate_proj\.(.*)$": ( + r"^layers\.(\d+)\.mlp\.gate_proj\.(.*)$": ( r"language_model.layers.\1.mlp.gate_up_proj.\2", 0, 2, ), - r"^model\.layers\.(\d+)\.mlp\.up_proj\.(.*)$": ( + r"^layers\.(\d+)\.mlp\.up_proj\.(.*)$": ( r"language_model.layers.\1.mlp.gate_up_proj.\2", 1, 2, ), - # UND pathway: layernorms + remaining mlp keys pass through unchanged. - r"^model\.layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2", + # UND pathway: layernorms + remaining attention/mlp keys pass through + # under language_model.layers namespace. + r"^layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2", } diff --git a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py index 8d6192b3f..d0a65bb97 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/cosmos3video.py @@ -863,19 +863,19 @@ class Cosmos3OmniTransformer(CachableDiT): ) # Latent projection layers - ReplicatedLinear for quantization support - self.vae2llm = ReplicatedLinear( + self.proj_in = ReplicatedLinear( self.patch_latent_dim, self.hidden_size, bias=True, quant_config=quant_config, - prefix="vae2llm", + prefix="proj_in", ) - self.llm2vae = ReplicatedLinear( + self.proj_out = ReplicatedLinear( self.hidden_size, self.patch_latent_dim, bias=True, quant_config=quant_config, - prefix="llm2vae", + prefix="proj_out", ) # Timestep embedder @@ -1084,7 +1084,7 @@ class Cosmos3OmniTransformer(CachableDiT): sequence_shard_enabled = self.sp_size > 1 # Patchify and project to hidden dim - hidden_gen, _ = self.vae2llm(self.patchify(hidden_states, T, H, W)) + hidden_gen, _ = self.proj_in(self.patchify(hidden_states, T, H, W)) seq_len_orig = hidden_gen.shape[1] seq_shard_pad = 0 @@ -1195,7 +1195,7 @@ class Cosmos3OmniTransformer(CachableDiT): # this cuts the post-loop SP collective bandwidth ~21x. hidden_gen = hidden_gen + residual hidden_gen = self.norm_moe_gen(hidden_gen) - output, _ = self.llm2vae(hidden_gen) + output, _ = self.proj_out(hidden_gen) if sequence_shard_enabled: output = sequence_model_parallel_all_gather(output, dim=1) @@ -1343,10 +1343,10 @@ class Cosmos3OmniTransformer(CachableDiT): # Ensure embeddings and projections are in target dtype self.language_model.embed_tokens.to(target_dtype) - for module in self.vae2llm.modules(): + for module in self.proj_in.modules(): if not _is_quantized(module): _cast_direct(module, target_dtype) - for module in self.llm2vae.modules(): + for module in self.proj_out.modules(): if not _is_quantized(module): _cast_direct(module, target_dtype) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py index a9815aeaf..808aa1a1c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/cosmos3_guardrails.py @@ -1,28 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 """Cosmos3 guardrail stages. -Text: - 1. Blocklist — ``better_profanity`` + nvidia/Cosmos-Guardrail1 word lists. - 2. Qwen3Guard — 0.6B LLM classifier (Qwen/Qwen3Guard-Gen-0.6B). - -Video: - 1. SigLIP content-safety filter — 7-class frame classifier; blocks if - more than 10% of frames are unsafe. - 2. RetinaFace face blur — detects faces and pixelates them. +Text and video safety checks via the ``cosmos_guardrail`` package. +Install with: pip install cosmos-guardrail==0.3.1 Enabled by default; opt out with ``SGLANG_DISABLE_COSMOS3_GUARDRAILS=1``. """ 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 sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( @@ -34,413 +22,65 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) -# --------------------------------------------------------------------------- -# Constants -# --------------------------------------------------------------------------- -GUARDRAIL_HF_REPO = "nvidia/Cosmos-Guardrail1" -GUARDRAIL_HF_REVISION = "d6d4bfa899a71454a700907664f3e88f503950cf" -CUTOFF_UNSAFE_FRAMES_PERCENT = 10 - -TextGuardrailFn = Callable[[str], None] -VideoGuardrailFn = Callable[[np.ndarray], np.ndarray] - - -# --------------------------------------------------------------------------- -# Video safety classifier (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) - - -# --------------------------------------------------------------------------- -# Checkpoint download helper -# --------------------------------------------------------------------------- -def _download_checkpoint() -> str: - from huggingface_hub import snapshot_download - - return snapshot_download(GUARDRAIL_HF_REPO, revision=GUARDRAIL_HF_REVISION) - - -# --------------------------------------------------------------------------- -# Text guardrail builder -# --------------------------------------------------------------------------- -def _build_text_guardrail(offload_to_cpu: bool) -> TextGuardrailFn: - checkers: list[Callable[[str], tuple[bool, str]]] = [] - - # 1. Blocklist - try: - import nltk - from better_profanity import profanity as profanity_filter - - ckpt_dir = _download_checkpoint() - blocklist_dir = os.path.join(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) - device = "cpu" if offload_to_cpu else "cuda" - qwen_model = ( - AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16) - .to(device) - .eval() - ) - - def _qwen_check(prompt: str) -> tuple[bool, str]: - conversations = [{"role": "user", "content": prompt}] - inputs = qwen_tokenizer.apply_chat_template( - conversations, - tokenize=True, - return_tensors="pt", - add_generation_prompt=True, - return_dict=True, - ).to(device) - input_len = inputs["input_ids"].shape[1] - with torch.no_grad(): - output_ids = qwen_model.generate(**inputs, max_new_tokens=128) - response = qwen_tokenizer.decode( - output_ids[0][input_len:], 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: - raise ValueError(f"Guardrail blocked prompt: {msg}") - - return text_guardrail - - -# --------------------------------------------------------------------------- -# Video guardrail builder -# --------------------------------------------------------------------------- -def _build_video_guardrail(offload_to_cpu: bool) -> VideoGuardrailFn: - ckpt_dir = _download_checkpoint() - 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 - - device = "cpu" if offload_to_cpu else "cuda" - siglip_id = "google/siglip-so400m-patch14-384" - siglip_model = ( - SiglipModel.from_pretrained(siglip_id) - .to(device, dtype=torch.float32) - .eval() - ) - siglip_processor = SiglipProcessor.from_pretrained(siglip_id) - - classifier = SafetyClassifier(input_size=1152, num_classes=7) - ckpt_path = os.path.join( - ckpt_dir, "video_content_safety_filter", "safety_filter.pt" - ) - checkpoint = torch.load(ckpt_path, map_location="cpu", weights_only=True) - # Checkpoint keys have "network." prefix from the VideoSafetyModel wrapper. - state = {k.removeprefix("network."): v for k, v in checkpoint["model"].items()} - classifier.load_state_dict(state) - classifier = classifier.to(device, dtype=torch.float32).eval() - - def _safety_check(frames: np.ndarray) -> tuple[bool, str]: - nonlocal siglip_model, classifier - if offload_to_cpu: - siglip_model = siglip_model.to("cuda") - classifier = classifier.to("cuda") - - unsafe_count = 0 - total = len(frames) - for frame in frames: - if frame.dtype != np.uint8: - frame = (np.clip(frame, 0.0, 1.0) * 255.0).astype(np.uint8) - 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) - if hasattr(features, "pooler_output"): - features = features.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 offload_to_cpu: - siglip_model = siglip_model.to("cpu") - classifier = classifier.to("cpu") - - 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(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") - - 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 = "cpu" if offload_to_cpu else "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 - if offload_to_cpu: - retinaface_net = retinaface_net.to("cuda") - - prior_data = None - scale = None - result_frames = [] - - for frame in frames: - frame_t = torch.from_numpy(frame).to("cuda", dtype=torch.float32) - if frame.dtype != np.uint8: - frame_t = frame_t * 255.0 - 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) - - if offload_to_cpu: - retinaface_net = retinaface_net.to("cpu") - - 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: - if safety_checker is not None: - is_safe, msg = safety_checker(frames) - if not is_safe: - raise ValueError(f"Guardrail blocked video: {msg}") - if face_blurrer is not None: - frames = face_blurrer(frames) - return frames - - return video_guardrail - - -# --------------------------------------------------------------------------- -# Singleton initialization -# --------------------------------------------------------------------------- -_text_guardrail: TextGuardrailFn | None = None -_video_guardrail: VideoGuardrailFn | None = None -_initialized = False +_checker = None def _init_guardrails(offload_to_cpu: bool = False) -> None: - global _text_guardrail, _video_guardrail, _initialized - if _initialized: + global _checker + if _checker is not None: return + try: + from cosmos_guardrail import CosmosSafetyChecker + except ImportError: + raise ImportError( + "cosmos_guardrail is required for Cosmos3 safety checks. " + "Install it with: pip install cosmos-guardrail==0.3.1" + ) logger.info( - "Initializing Cosmos3 guardrails (offload_to_cpu=%s)...", offload_to_cpu + "Initializing Cosmos3 guardrails (offload_to_cpu=%s) ...", offload_to_cpu ) - _text_guardrail = _build_text_guardrail(offload_to_cpu) - _video_guardrail = _build_video_guardrail(offload_to_cpu) - _initialized = True + _checker = CosmosSafetyChecker() + idle_device = "cpu" if offload_to_cpu else "cuda" + for runner in (_checker.text_guardrail, _checker.video_guardrail): + if runner is None or not hasattr(runner, "models"): + continue + for m in runner.models: + if isinstance(m, torch.nn.Module): + m.to(idle_device) logger.info("Cosmos3 guardrails initialized.") -# --------------------------------------------------------------------------- -# Public API — video guardrail function for use inside Cosmos3DecodingStage -# --------------------------------------------------------------------------- +def check_text_safety(prompt: str) -> None: + if _checker is None: + return + if not _checker.check_text_safety(prompt): + raise ValueError("Guardrail blocked prompt.") + + def check_video_safety(video: np.ndarray) -> np.ndarray: - """Run video guardrails on decoded frames (numpy [B, T, H, W, C] or [T, H, W, C]). + """Apply video guardrails to decoded frames. - Raises ``ValueError`` if content is blocked. - Returns (potentially face-blurred) frames. + Args: + video: numpy [B, T, H, W, C] or [T, H, W, C], uint8. + + Returns: + Processed frames in the same shape, or raises ValueError if blocked. """ - if _video_guardrail is None: + if _checker is None: return video - frames = video[0] if video.ndim == 5 else video - frames = _video_guardrail(frames) if video.ndim == 5: - frames = frames[np.newaxis] - return frames + processed = [] + for frames in video: + result = _checker.check_video_safety(frames) + processed.append(result if result is not None else frames) + return np.stack(processed) + result = _checker.check_video_safety(video) + return result if result is not None else video -# --------------------------------------------------------------------------- -# Pipeline stage — text guardrail (runs before generation) -# --------------------------------------------------------------------------- class Cosmos3TextGuardrailStage(PipelineStage): """Check prompt text against safety policies before generation. - Runs blocklist keyword matching and Qwen3Guard LLM classifier. Raises ``ValueError`` if the prompt is blocked. """ @@ -451,11 +91,12 @@ class Cosmos3TextGuardrailStage(PipelineStage): _init_guardrails(offload_to_cpu) def forward(self, batch: Req, server_args: ServerArgs) -> Req: - if _text_guardrail is not None and batch.prompt is not None: - prompt = batch.prompt - if isinstance(prompt, list): - for p in prompt: - _text_guardrail(p) - else: - _text_guardrail(prompt) + prompt = batch.prompt + if prompt is None: + return batch + if isinstance(prompt, list): + for p in prompt: + check_text_safety(p) + else: + check_text_safety(prompt) return batch diff --git a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py index db021484f..31d3b4296 100644 --- a/python/sglang/multimodal_gen/test/unit/test_cosmos3.py +++ b/python/sglang/multimodal_gen/test/unit/test_cosmos3.py @@ -32,123 +32,131 @@ class TestCosmos3ParamNamesMapping(unittest.TestCase): key, idx, total = _apply(self.fn, "lm_head.weight") self.assertEqual(key, "") - def test_model_norm_dropped(self): - key, idx, total = _apply(self.fn, "model.norm.weight") + def test_norm_dropped(self): + key, idx, total = _apply(self.fn, "norm.weight") + self.assertEqual(key, "") + + def test_audio_proj_in_dropped(self): + key, *_ = _apply(self.fn, "audio_proj_in.weight") + self.assertEqual(key, "") + + def test_action_proj_in_dropped(self): + key, *_ = _apply(self.fn, "action_proj_in.weight") self.assertEqual(key, "") # --- top-level pass-through --- def test_embed_tokens(self): - key, *_ = _apply(self.fn, "model.embed_tokens.weight") + key, *_ = _apply(self.fn, "embed_tokens.weight") self.assertEqual(key, "language_model.embed_tokens.weight") def test_norm_moe_gen(self): - key, *_ = _apply(self.fn, "model.norm_moe_gen.weight") + key, *_ = _apply(self.fn, "norm_moe_gen.weight") self.assertEqual(key, "norm_moe_gen.weight") - # --- time embedder --- + # --- time embedder (pass-through: checkpoint already uses linear_1/2) --- - def test_time_embedder_mlp_0(self): - key, *_ = _apply(self.fn, "time_embedder.mlp.0.weight") + def test_time_embedder_linear_1_passthrough(self): + key, *_ = _apply(self.fn, "time_embedder.linear_1.weight") self.assertEqual(key, "time_embedder.linear_1.weight") - def test_time_embedder_mlp_2(self): - key, *_ = _apply(self.fn, "time_embedder.mlp.2.bias") + def test_time_embedder_linear_2_passthrough(self): + key, *_ = _apply(self.fn, "time_embedder.linear_2.bias") self.assertEqual(key, "time_embedder.linear_2.bias") # --- GEN pathway: Q/K/V merge (must not be claimed by UND catch-all) --- def test_gen_q_proj_key_and_merge_index(self): - key, idx, total = _apply( - self.fn, "model.layers.3.self_attn.q_proj_moe_gen.weight" - ) + key, idx, total = _apply(self.fn, "layers.3.self_attn.add_q_proj.weight") self.assertEqual(key, "gen_layers.3.cross_attention.to_qkv.weight") self.assertEqual(idx, 0) self.assertEqual(total, 3) def test_gen_k_proj_merge_index(self): - _, idx, total = _apply( - self.fn, "model.layers.0.self_attn.k_proj_moe_gen.weight" - ) + _, idx, total = _apply(self.fn, "layers.0.self_attn.add_k_proj.weight") self.assertEqual(idx, 1) self.assertEqual(total, 3) def test_gen_v_proj_merge_index(self): - _, idx, total = _apply( - self.fn, "model.layers.0.self_attn.v_proj_moe_gen.weight" - ) + _, idx, total = _apply(self.fn, "layers.0.self_attn.add_v_proj.weight") self.assertEqual(idx, 2) self.assertEqual(total, 3) def test_gen_o_proj(self): - key, idx, total = _apply( - self.fn, "model.layers.5.self_attn.o_proj_moe_gen.weight" - ) + key, idx, total = _apply(self.fn, "layers.5.self_attn.to_add_out.weight") self.assertEqual(key, "gen_layers.5.cross_attention.to_out.weight") self.assertIsNone(idx) + def test_gen_norm_added_q(self): + key, idx, _ = _apply(self.fn, "layers.2.self_attn.norm_added_q.weight") + self.assertEqual(key, "gen_layers.2.cross_attention.norm_q.weight") + self.assertIsNone(idx) + + def test_gen_norm_added_k(self): + key, idx, _ = _apply(self.fn, "layers.2.self_attn.norm_added_k.weight") + self.assertEqual(key, "gen_layers.2.cross_attention.norm_k.weight") + self.assertIsNone(idx) + def test_gen_mlp_gate_proj(self): - key, idx, total = _apply(self.fn, "model.layers.2.mlp_moe_gen.gate_proj.weight") + key, idx, total = _apply(self.fn, "layers.2.mlp_moe_gen.gate_proj.weight") self.assertEqual(key, "gen_layers.2.mlp.gate_up_proj.weight") self.assertEqual(idx, 0) self.assertEqual(total, 2) def test_gen_mlp_up_proj(self): - key, idx, total = _apply(self.fn, "model.layers.2.mlp_moe_gen.up_proj.weight") + key, idx, total = _apply(self.fn, "layers.2.mlp_moe_gen.up_proj.weight") self.assertEqual(key, "gen_layers.2.mlp.gate_up_proj.weight") self.assertEqual(idx, 1) self.assertEqual(total, 2) def test_gen_mlp_down_proj_passthrough(self): - key, idx, _ = _apply(self.fn, "model.layers.2.mlp_moe_gen.down_proj.weight") + key, idx, _ = _apply(self.fn, "layers.2.mlp_moe_gen.down_proj.weight") self.assertEqual(key, "gen_layers.2.mlp.down_proj.weight") self.assertIsNone(idx) # --- UND pathway: Q/K/V merge --- def test_und_q_proj_key_and_merge_index(self): - key, idx, total = _apply(self.fn, "model.layers.7.self_attn.q_proj.weight") + key, idx, total = _apply(self.fn, "layers.7.self_attn.to_q.weight") self.assertEqual(key, "language_model.layers.7.self_attn.to_qkv.weight") self.assertEqual(idx, 0) self.assertEqual(total, 3) def test_und_k_proj_merge_index(self): - _, idx, total = _apply(self.fn, "model.layers.0.self_attn.k_proj.weight") + _, idx, total = _apply(self.fn, "layers.0.self_attn.to_k.weight") self.assertEqual(idx, 1) self.assertEqual(total, 3) def test_und_v_proj_merge_index(self): - _, idx, total = _apply(self.fn, "model.layers.0.self_attn.v_proj.weight") + _, idx, total = _apply(self.fn, "layers.0.self_attn.to_v.weight") self.assertEqual(idx, 2) self.assertEqual(total, 3) def test_und_mlp_gate_proj(self): - key, idx, total = _apply(self.fn, "model.layers.1.mlp.gate_proj.weight") + key, idx, total = _apply(self.fn, "layers.1.mlp.gate_proj.weight") self.assertEqual(key, "language_model.layers.1.mlp.gate_up_proj.weight") self.assertEqual(idx, 0) self.assertEqual(total, 2) def test_und_mlp_up_proj(self): - _, idx, total = _apply(self.fn, "model.layers.1.mlp.up_proj.weight") + _, idx, total = _apply(self.fn, "layers.1.mlp.up_proj.weight") self.assertEqual(idx, 1) self.assertEqual(total, 2) def test_und_layernorm_catch_all(self): - key, idx, _ = _apply(self.fn, "model.layers.0.input_layernorm.weight") + key, idx, _ = _apply(self.fn, "layers.0.input_layernorm.weight") self.assertEqual(key, "language_model.layers.0.input_layernorm.weight") self.assertIsNone(idx) # --- ordering: GEN patterns must not be swallowed by UND catch-all --- def test_gen_layernorm_not_mapped_to_und(self): - key, *_ = _apply(self.fn, "model.layers.0.input_layernorm_moe_gen.weight") + key, *_ = _apply(self.fn, "layers.0.input_layernorm_moe_gen.weight") self.assertIn("gen_layers", key) self.assertNotIn("language_model", key) def test_gen_post_attention_layernorm_not_mapped_to_und(self): - key, *_ = _apply( - self.fn, "model.layers.4.post_attention_layernorm_moe_gen.weight" - ) + key, *_ = _apply(self.fn, "layers.4.post_attention_layernorm_moe_gen.weight") self.assertIn("gen_layers", key) self.assertNotIn("language_model", key)