Vision Safety Checker

from safety_checker import StableDiffusionSafetyChecker
from transformers import CLIPFeatureExtractor

safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker").to(device)
feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")

def check_nsfw_images(images: list[Image.Image]) -> list[bool]:
    safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)
    has_nsfw_concepts = safety_checker(images=[images], clip_input=safety_checker_input.pixel_values.to(device))
    return has_nsfw_concepts