From c57ada81e1f01f365268f8b5aaffa9f5bf7d28eb Mon Sep 17 00:00:00 2001 From: Arseniy Mironov <98156294+Napkin-AI@users.noreply.github.com> Date: Wed, 19 Aug 2026 15:26:28 +0300 Subject: [PATCH] [Diffusion] Use current_platform instead of hardcoded "cuda" in cosmos3 guardrails (#34612) Co-authored-by: ronnie_zheng --- docs/cookbook/diffusion/Cosmos/Cosmos3.mdx | 6 ++++++ .../stages/model_specific_stages/cosmos3_guardrails.py | 3 ++- 2 files changed, 8 insertions(+), 1 deletion(-) diff --git a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx index 8f513c5d4..3b9c70a13 100644 --- a/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx +++ b/docs/cookbook/diffusion/Cosmos/Cosmos3.mdx @@ -44,6 +44,12 @@ pip install "cosmos-guardrail==0.3.1" `cosmos-guardrail` downloads gated NVIDIA guardrail weights, so pass a Hugging Face token if your environment needs one. If the package is not installed, SGLang skips Cosmos3 guardrails and logs a warning. To disable Cosmos3 guardrails for local experiments, set `SGLANG_DISABLE_COSMOS3_GUARDRAILS=1` before starting the server. +There may be problems loading the Cosmos-1.0-Guardrail weights on Ascend NPU. If the *_pickle.UnpicklingError* error occurs during startup, you should change ```weight_only=True``` to ```weights_only=False``` parameter in *cosmos_guardrail/cosmos_utils.py*: +``` +#!/usr/bin/env bash +COSMOS_GUARDRAIL_DIR="$(dirname "$(python -c 'import cosmos_guardrail; print(cosmos_guardrail.__file__)')")" +sed -i 's/weights_only=True/weights_only=False/g' "$COSMOS_GUARDRAIL_DIR/cosmos_utils.py" +``` ## 3. Serve Cosmos3 Serve `Cosmos3-Nano` directly from the Hugging Face model ID: 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 c3e05e53a..a475aec5f 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 @@ -21,6 +21,7 @@ from sglang.multimodal_gen.runtime.pipelines_core.stages.base import ( PipelineStage, StageParallelismType, ) +from sglang.multimodal_gen.runtime.platforms import current_platform from sglang.multimodal_gen.runtime.server_args import ServerArgs from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger @@ -49,7 +50,7 @@ def _init_guardrails(offload_to_cpu: bool = False) -> None: "Initializing Cosmos3 guardrails (offload_to_cpu=%s) ...", offload_to_cpu ) _checker = CosmosSafetyChecker() - idle_device = "cpu" if offload_to_cpu else "cuda" + idle_device = "cpu" if offload_to_cpu else current_platform.device_type for runner in (_checker.text_guardrail, _checker.video_guardrail): if runner is None or not hasattr(runner, "models"): continue