nsfw filter

StableDiffuser.py

@@ -5,10 +5,12 @@ from baukit import TraceDict
 from diffusers import AutoencoderKL, UNet2DConditionModel
 from PIL import Image
 from tqdm.auto import tqdm
-from transformers import CLIPTextModel, CLIPTokenizer
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPFeatureExtractor
+from diffusers.schedulers import EulerAncestralDiscreteScheduler
 from diffusers.schedulers.scheduling_ddim import DDIMScheduler
 from diffusers.schedulers.scheduling_ddpm import DDPMScheduler
 from diffusers.schedulers.scheduling_lms_discrete import LMSDiscreteScheduler
+from diffusers.pipelines.stable_diffusion import StableDiffusionSafetyChecker
 import util
 
 
@@ -53,6 +55,9 @@ class StableDiffuser(torch.nn.Module):
         self.unet = UNet2DConditionModel.from_pretrained(
             "CompVis/stable-diffusion-v1-4", subfolder="unet")
         
+        self.feature_extractor = CLIPFeatureExtractor.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="feature_extractor")
+        self.safety_checker = StableDiffusionSafetyChecker.from_pretrained("CompVis/stable-diffusion-v1-4", subfolder="safety_checker")
+
         if scheduler == 'LMS':
             self.scheduler = LMSDiscreteScheduler(beta_start=0.00085, beta_end=0.012, beta_schedule="scaled_linear", num_train_timesteps=1000)
         elif scheduler == 'DDIM':
@@ -237,6 +242,15 @@ class StableDiffuser(torch.nn.Module):
         latents_steps = [self.decode(latents.to(self.unet.device)) for latents in latents_steps]
         images_steps = [self.to_image(latents) for latents in latents_steps]
 
+        for i in range(len(images_steps)):
+            self.safety_checker = self.safety_checker.float()
+            safety_checker_input = self.feature_extractor(images_steps[i], return_tensors="pt").to(latents_steps[0].device)
+            image, has_nsfw_concept = self.safety_checker(
+                images=latents_steps[i].float().cpu().numpy(), clip_input=safety_checker_input.pixel_values.float()
+            )
+
+            images_steps[i][0] = self.to_image(torch.from_numpy(image))[0]
+
         images_steps = list(zip(*images_steps))
 
         if trace_steps: