File size: 2,826 Bytes
226c7c9 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import os
import numpy as np
from cosmos_transfer1.auxiliary.guardrail.blocklist.blocklist import Blocklist
from cosmos_transfer1.auxiliary.guardrail.common.core import GuardrailRunner
from cosmos_transfer1.auxiliary.guardrail.face_blur_filter.face_blur_filter import RetinaFaceFilter
from cosmos_transfer1.auxiliary.guardrail.llamaGuard3.llamaGuard3 import LlamaGuard3
from cosmos_transfer1.auxiliary.guardrail.video_content_safety_filter.video_content_safety_filter import (
VideoContentSafetyFilter,
)
from cosmos_transfer1.utils import log
def create_text_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
"""Create the text guardrail runner."""
return GuardrailRunner(safety_models=[Blocklist(checkpoint_dir), LlamaGuard3(checkpoint_dir)])
def create_video_guardrail_runner(checkpoint_dir: str) -> GuardrailRunner:
"""Create the video guardrail runner."""
return GuardrailRunner(
safety_models=[VideoContentSafetyFilter(checkpoint_dir)],
postprocessors=[RetinaFaceFilter(checkpoint_dir)],
)
def run_text_guardrail(prompt: str, guardrail_runner: GuardrailRunner) -> bool:
"""Run the text guardrail on the prompt, checking for content safety.
Args:
prompt: The text prompt.
guardrail_runner: The text guardrail runner.
Returns:
bool: Whether the prompt is safe.
"""
is_safe, message = guardrail_runner.run_safety_check(prompt)
if not is_safe:
log.critical(f"GUARDRAIL BLOCKED: {message}")
return is_safe
def run_video_guardrail(frames: np.ndarray, guardrail_runner: GuardrailRunner) -> np.ndarray | None:
"""Run the video guardrail on the frames, checking for content safety and applying face blur.
Args:
frames: The frames of the generated video.
guardrail_runner: The video guardrail runner.
Returns:
The processed frames if safe, otherwise None.
"""
is_safe, message = guardrail_runner.run_safety_check(frames)
if not is_safe:
log.critical(f"GUARDRAIL BLOCKED: {message}")
return None
frames = guardrail_runner.postprocess(frames)
return frames
|