Merge pull request #65 from anotherjesse/nsfw_checker

enable safety_checker for SD 2.1

Author: Jesse Andrews

cog.yaml

@@ -3,11 +3,11 @@ build:
   cuda: "11.6"
   python_version: "3.10"
   python_packages:
-    - "diffusers==0.10.0"
+    - "diffusers==0.11.1"
     - "torch==1.13.0"
     - "ftfy==6.1.1"
     - "scipy==1.9.3"
     - "transformers==4.25.1"
-    - "accelerate==0.14.0"
+    - "accelerate==0.15.0"
 
 predict: "predict.py:Predictor"

predict.py

@@ -12,17 +12,28 @@
     EulerAncestralDiscreteScheduler,
     DPMSolverMultistepScheduler,
 )
+from diffusers.pipelines.stable_diffusion.safety_checker import (
+    StableDiffusionSafetyChecker,
+)
+
 
 MODEL_ID = "stabilityai/stable-diffusion-2-1"
 MODEL_CACHE = "diffusers-cache"
+SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
 
 
 class Predictor(BasePredictor):
     def setup(self):
         """Load the model into memory to make running multiple predictions efficient"""
         print("Loading pipeline...")
+        safety_checker = StableDiffusionSafetyChecker.from_pretrained(
+            SAFETY_MODEL_ID,
+            cache_dir=MODEL_CACHE,
+            local_files_only=True,
+        )
         self.pipe = StableDiffusionPipeline.from_pretrained(
             MODEL_ID,
+            safety_checker=safety_checker,
             cache_dir=MODEL_CACHE,
             local_files_only=True,
         ).to("cuda")
@@ -107,9 +118,15 @@ def predict(
 
         output_paths = []
         for i, sample in enumerate(output.images):
-            output_path = f"/tmp/out-{i}.png"
-            sample.save(output_path)
-            output_paths.append(Path(output_path))
+            if output.nsfw_content_detected and not output.nsfw_content_detected[i]:
+                output_path = f"/tmp/out-{i}.png"
+                sample.save(output_path)
+                output_paths.append(Path(output_path))
+
+        if len(output_paths) == 0:
+            raise Exception(
+                f"NSFW content detected. Try running it again, or try a different prompt."
+            )
 
         return output_paths
 

script/download-weights

@@ -3,15 +3,25 @@
 import os
 import shutil
 from diffusers import StableDiffusionPipeline
+from diffusers.pipelines.stable_diffusion.safety_checker import (
+    StableDiffusionSafetyChecker,
+)
 
 
-model_id = "stabilityai/stable-diffusion-2-1"
+MODEL_ID = "stabilityai/stable-diffusion-2-1"
 MODEL_CACHE = "diffusers-cache"
+SAFETY_MODEL_ID = "CompVis/stable-diffusion-safety-checker"
+
 if os.path.exists(MODEL_CACHE):
     shutil.rmtree(MODEL_CACHE)
 os.makedirs(MODEL_CACHE, exist_ok=True)
 
+saftey_checker = StableDiffusionSafetyChecker.from_pretrained(
+    SAFETY_MODEL_ID,
+    cache_dir=MODEL_CACHE,
+)
+
 pipe = StableDiffusionPipeline.from_pretrained(
-    model_id,
+    MODEL_ID,
     cache_dir=MODEL_CACHE,
 )