[diffusion] model: update to new model format (#26492)
This commit is contained in:
Dawid Majchrowski
@@ -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",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
+50
-409
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user