How to enable safety checker in stable-diffusion-2-1 pipeline?

Safety checker is not specified in stabilityai/stable-diffusion-2-1 pipeline:

How can I enable safety checker? What Python packages need to be installed? Thanks.

hi @samulik ,

Since the safety checker is not on the model config, you can do it manually

from transformers import CLIPFeatureExtractor
from diffusers import DiffusionPipeline, DPMSolverMultistepScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker

repo_id = "stabilityai/stable-diffusion-2"
pipe = DiffusionPipeline.from_pretrained(repo_id,
                                         safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker"),
                                         feature_extractor=CLIPFeatureExtractor.from_pretrained("openai/clip-vit-base-patch32"))
pipe.scheduler = DPMSolverMultistepScheduler.from_config(pipe.scheduler.config)
pipe = pipe.to("cuda")

prompt = "naked austronaut in the beach"
results = pipe(prompt, guidance_scale=9, num_inference_steps=25, num_images_per_prompt=1)

if not results.nsfw_content_detected[0]:
  print("Not NSFW")
  results.images[0].save("astronaut.png")
else:
  print("NSFW Detected")

Hello, I tried this and everything seems to run fine, only nsfw_content_detected=None is being returned for very NSFW images. I pretty much replicated your code. Not sure why it isn’t returning correctly.

hi @caldwecg , if you’re running the code above, you should expected an array of Booleans for the nsfw_content_detected=[True, False, False ...] depending the number of generated images per prompt. So if you’re getting nsfw_content_detected=None something might be missing

@radames thanks for the response. I’m using a dreambooth trained model, would that affect the results? I copied your code verbatim and when printing the result object it gave the following:

StableDiffusionPipelineOutput(images=[<PIL.Image.Image image mode=RGB size=512x512 at 0x7F2CDC646440>], nsfw_content_detected=None)

just tested the above code with dreamboothed model `repo_id = “nitrosocke/mo-di-diffusion” and the safety filter works. Could you please share a bit more?

@radames I’m using Shivam Srirao’s implementation of DreamBooth (Google Colab) and it looks like it uses stable-diffusion-v1-5. Would the safety checker be different for this version? Thanks again for the help.

@caldwecg I don’t believe the safety checker is trained on a specific stable diffusion version

2 Likes

hi @caldwecg in your notebook, just replace the inference cell with this

import torch
from torch import autocast
from diffusers import StableDiffusionPipeline, DDIMScheduler
from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
from IPython.display import display

model_path = WEIGHTS_DIR             # If you want to use previously trained model saved in gdrive, replace this with the full path of model in gdrive

pipe = StableDiffusionPipeline.from_pretrained(model_path,
                                               safety_checker=StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-safety-checker", torch_dtype=torch.float16),
                                               torch_dtype=torch.float16).to("cuda")
pipe.scheduler = DDIMScheduler.from_config(pipe.scheduler.config)
pipe.enable_xformers_memory_efficient_attention()
g_cuda = None

hi @radames I ran it your recent solution and then again on the nitrosocke/mo-di-diffusion with your original solution and still ran into nsfw_content_detected=None. Not sure what else could be different between ours… maybe something now working with the feature_extractor?

I also looked at StableDiffusionPipeline and there’s a “requires_safety_checker: bool = False” field, do I need to specify this?

could you please share a reproducible code? Spaces or on Colab thanks

Hi @radames it ended up being an issue with the version of diffusers I had installed (0.9.0). Upgrading to 0.13.1 worked. Thanks for the help.

1 Like

hi @radames does this safety feature take additional time, i have a startup, we generate a lots of images, im thinking of trying the safety feature, will it comsume more time. i cannot afford to experiement. its expensive. can you atleast make a guess what extra time it may take? tnks