radames/Stable-Diffusion-Safety-Checker-API

1from transformers import CLIPFeatureExtractor2from safety_checker import StableDiffusionSafetyChecker3import torch4from PIL import Image5import gradio as gr6from pathlib import Path7 8device = "cuda" if torch.cuda.is_available() else "cpu"9safety_checker = StableDiffusionSafetyChecker.from_pretrained(10    "CompVis/stable-diffusion-safety-checker"11).to(device)12feature_extractor = CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32")13 14 15import gradio as gr16 17 18def image_classifier(files):19    images = [Image.open(file).convert("RGB").resize((512, 512)) for file in files]20 21    safety_checker_input = feature_extractor(images, return_tensors="pt").to(device)22    has_nsfw_concepts = safety_checker(23        images=[images], clip_input=safety_checker_input.pixel_values.to(torch.float16)24    )25    results = [26        {"has_nsfw": nsfw, "file": Path(file).name}27        for (nsfw, file) in zip(has_nsfw_concepts, files)28    ]29    return {"results": results}30 31 32demo = gr.Interface(33    title="Stable Diffusion Safety Checker API",34    fn=image_classifier,35    inputs=gr.File(file_count="multiple", file_types=["image"]),36    outputs="json",37    api_name="classify",38)39demo.launch()