[Diffusion] Use current_platform instead of hardcoded "cuda" in cosmos3 guardrails (#34612)

Co-authored-by: ronnie_zheng <zl19940307@163.com>
This commit is contained in:
Arseniy Mironov
2026-08-19 15:26:28 +03:00
committed by GitHub
co-authored by ronnie_zheng
parent 574274660f
commit c57ada81e1
2 changed files with 8 additions and 1 deletions
@@ -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