StableDiffusion系のnsfwチェカーをましにするには？

はじめに、StableDiffusionにはNSFWチェックというものがある。
これは、不健全な画像をユーザーに見せたいためのものだが、すこし過剰に反応すぎる。
オフにするのは簡単だが、一般サービスとして使う場合そうはいかない。

画像判定モデル

Falconsai/nsfw_image_detection

商用可能なApache2.0ライセンス
https://huggingface.co/Falconsai/nsfw_image_detection
画像がNSFWか、判定してくれます。

沢山あるデモの一つ（停止している場合、https://huggingface.co/spaces からnsfwで検索)
https://huggingface.co/spaces/5m4ck3r/nsfw_image_detection

これを組み込んだコード

stablediffusion側では、nsfwを向うにして、ユーザーに返す直前に以下を呼び出せば

return classificate_safe_check(image)

すれば、完全にアウトな時だけ、黒塗りを返す。実用性の高いサービスになります。

以下コード

from PIL import Image
from transformers import pipeline
import torch

device = 0 if torch.cuda.is_available() else -1
#device = -1
def is_nsfw(img):
    classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection",device=device)
    result = classifier(img)
    #print(result[0])
    return result[0]["label"] == "nsfw"
def classificate_safe_check(img):
    nsfw = is_nsfw(img)
    if nsfw:
        black_image = Image.new("RGB", img.size, "black")
        return black_image
    else:
        return img

img = Image.open("test.png")
out = classificate_safe_check(img)
out.save("out.jpg")

その他

MichalMlodawski/nsfw-image-detection-large
より精度はいいが、ライセンスが、CC-BY-NC-SA-4.0 で商用不可
https://huggingface.co/MichalMlodawski/nsfw-image-detection-large