safety check

edict_functions.py

@@ -17,6 +17,8 @@ import os
 from torchvision import datasets
 import pickle
 
+
+
 # StableDiffusion P2P implementation originally from https://github.com/bloc97/CrossAttentionControl
 use_half_prec = True
 if use_half_prec:
@@ -66,7 +68,30 @@ else:
     clip.double().to(device)
 print("Loaded all models")
 
-
+from diffusers.pipelines.stable_diffusion.safety_checker import StableDiffusionSafetyChecker
+from transformers import AutoFeatureExtractor
+# load safety model
+safety_model_id = "CompVis/stable-diffusion-safety-checker"
+safety_feature_extractor = AutoFeatureExtractor.from_pretrained(safety_model_id)
+safety_checker = StableDiffusionSafetyChecker.from_pretrained(safety_model_id)
+def load_replacement(x):
+    try:
+        hwc = x.shape
+        y = Image.open("assets/rick.jpeg").convert("RGB").resize((hwc[1], hwc[0]))
+        y = (np.array(y)/255.0).astype(x.dtype)
+        assert y.shape == x.shape
+        return y
+    except Exception:
+        return x
+def check_safety(x_image):
+    safety_checker_input = safety_feature_extractor(numpy_to_pil(x_image), return_tensors="pt")
+    x_checked_image, has_nsfw_concept = safety_checker(images=x_image, clip_input=safety_checker_input.pixel_values)
+    assert x_checked_image.shape[0] == len(has_nsfw_concept)
+    for i in range(len(has_nsfw_concept)):
+        if has_nsfw_concept[i]:
+            # x_checked_image[i] = load_replacement(x_checked_image[i])
+            x_checked_image[i] *= 0 # load_replacement(x_checked_image[i])
+    return x_checked_image, has_nsfw_concept
     
     
 def EDICT_editing(im_path,
@@ -597,6 +622,9 @@ def baseline_stablediffusion(prompt="",
 
     image = (image / 2 + 0.5).clamp(0, 1)
     image = image.cpu().permute(0, 2, 3, 1).numpy()
+
+    image, _ = check_safety(image)
+
     image = (image[0] * 255).round().astype("uint8")
     return Image.fromarray(image)
 ####################################