[diffusion] model: update to new model format (#26492)

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
+        embed_tokens.weight                                    -> language_model.embed_tokens.weight
-        model.layers.X.input_layernorm.weight                  -> language_model.layers.X.input_layernorm.weight
+        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
+        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)
+        layers.X.self_attn.{to_q,to_k,to_v}.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)
+        layers.X.self_attn.{add_q,add_k,add_v}_proj.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)
+        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)
+        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
+        norm_moe_gen.weight                                    -> norm_moe_gen.weight
-        time_embedder.mlp.{0,2}.weight                         -> time_embedder.linear_{1,2}.weight
+        time_embedder.linear_{1,2}.weight                      -> (pass-through)
-        vae2llm.weight, llm2vae.weight                         -> (pass-through)
+        proj_in.weight, proj_out.weight                        -> (pass-through)
 
-    GEN patterns (`*_moe_gen`) must precede the UND catch-all so the
+    GEN patterns (`*_moe_gen`, `add_*`, `to_add_out`, `norm_added_*`) must
-    catch-all can't claim GEN keys. `model.norm.weight` and `lm_head.weight`
+    precede the UND catch-all so the catch-all can't claim GEN keys.
-    are inherited from Qwen3-VL pretraining and not used at inference, so
+    `norm.weight` and `lm_head.weight` are inherited from Qwen3-VL
-    they are skipped via empty-string replacement.
+    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$": "",
+        r"^norm\.weight$": "",
-        # Top-level norms / heads.
+        # Audio / action modalities — not yet wired; skip to avoid load warnings.
-        r"^model\.norm_moe_gen\.(.*)$": r"norm_moe_gen.\1",
+        r"^audio_.*$": "",
-        r"^model\.embed_tokens\.(.*)$": r"language_model.embed_tokens.\1",
+        r"^action_.*$": "",
-        # Time embedder: mlp.0 -> linear_1, mlp.2 -> linear_2 (SiLU at index 1).
+        # Top-level norms / embeddings.
-        r"^time_embedder\.mlp\.0\.(.*)$": r"time_embedder.linear_1.\1",
+        r"^norm_moe_gen\.(.*)$": r"norm_moe_gen.\1",
-        r"^time_embedder\.mlp\.2\.(.*)$": r"time_embedder.linear_2.\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"^layers\.(\d+)\.self_attn\.to_add_out\.(.*)$": 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"^layers\.(\d+)\.self_attn\.norm_added_q\.(.*)$": 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"^layers\.(\d+)\.self_attn\.norm_added_k\.(.*)$": 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"^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+)\.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",
+        r"^layers\.(\d+)\.mlp_moe_gen\.(.*)$": r"gen_layers.\1.mlp.\2",
-        # UND pathway: per-layer attention rename (q/k/v_proj -> to_qkv merged,
+        # UND pathway: Q/K/V merge into to_qkv; remaining attention keys
-        # q_norm/k_norm -> norm_q/k, o_proj -> to_out).
+        # (to_out, norm_q, norm_k) and layernorms pass through the catch-all.
-        r"^model\.layers\.(\d+)\.self_attn\.q_proj\.(.*)$": (
+        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.
+        # UND pathway: layernorms + remaining attention/mlp keys pass through
-        r"^model\.layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2",
+        # under language_model.layers namespace.
+        r"^layers\.(\d+)\.(.*)$": r"language_model.layers.\1.\2",
     }
 
 

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)
 

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:
+Text and video safety checks via the ``cosmos_guardrail`` package.
-  1. Blocklist — ``better_profanity`` + nvidia/Cosmos-Guardrail1 word lists.
+Install with: pip install cosmos-guardrail==0.3.1
-  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.
 
 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__)
 
-# ---------------------------------------------------------------------------
+_checker = None
-# 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
 
 
 def _init_guardrails(offload_to_cpu: bool = False) -> None:
-    global _text_guardrail, _video_guardrail, _initialized
+    global _checker
-    if _initialized:
+    if _checker is not None:
         return
-    logger.info(
+    try:
-        "Initializing Cosmos3 guardrails (offload_to_cpu=%s)...", offload_to_cpu
+        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"
         )
-    _text_guardrail = _build_text_guardrail(offload_to_cpu)
+    logger.info(
-    _video_guardrail = _build_video_guardrail(offload_to_cpu)
+        "Initializing Cosmos3 guardrails (offload_to_cpu=%s) ...", 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.")
 
 
-# ---------------------------------------------------------------------------
+def check_text_safety(prompt: str) -> None:
-# Public API — video guardrail function for use inside Cosmos3DecodingStage
+    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.
+    Args:
-    Returns (potentially face-blurred) frames.
+        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]
+        processed = []
-    return frames
+        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 prompt is None:
+            return batch
         if isinstance(prompt, list):
             for p in prompt:
-                    _text_guardrail(p)
+                check_text_safety(p)
         else:
-                _text_guardrail(prompt)
+            check_text_safety(prompt)
         return batch

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):
+    def test_norm_dropped(self):
-        key, idx, total = _apply(self.fn, "model.norm.weight")
+        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):
+    def test_time_embedder_linear_1_passthrough(self):
-        key, *_ = _apply(self.fn, "time_embedder.mlp.0.weight")
+        key, *_ = _apply(self.fn, "time_embedder.linear_1.weight")
         self.assertEqual(key, "time_embedder.linear_1.weight")
 
-    def test_time_embedder_mlp_2(self):
+    def test_time_embedder_linear_2_passthrough(self):
-        key, *_ = _apply(self.fn, "time_embedder.mlp.2.bias")
+        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(
+        key, idx, total = _apply(self.fn, "layers.3.self_attn.add_q_proj.weight")
-            self.fn, "model.layers.3.self_attn.q_proj_moe_gen.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(
+        _, idx, total = _apply(self.fn, "layers.0.self_attn.add_k_proj.weight")
-            self.fn, "model.layers.0.self_attn.k_proj_moe_gen.weight"
-        )
         self.assertEqual(idx, 1)
         self.assertEqual(total, 3)
 
     def test_gen_v_proj_merge_index(self):
-        _, idx, total = _apply(
+        _, idx, total = _apply(self.fn, "layers.0.self_attn.add_v_proj.weight")
-            self.fn, "model.layers.0.self_attn.v_proj_moe_gen.weight"
-        )
         self.assertEqual(idx, 2)
         self.assertEqual(total, 3)
 
     def test_gen_o_proj(self):
-        key, idx, total = _apply(
+        key, idx, total = _apply(self.fn, "layers.5.self_attn.to_add_out.weight")
-            self.fn, "model.layers.5.self_attn.o_proj_moe_gen.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(
+        key, *_ = _apply(self.fn, "layers.4.post_attention_layernorm_moe_gen.weight")
-            self.fn, "model.layers.4.post_attention_layernorm_moe_gen.weight"
-        )
         self.assertIn("gen_layers", key)
         self.assertNotIn("language_model", key)