added safety checker

app.py

@@ -27,6 +27,8 @@ from transformers import CLIPVisionModelWithProjection
 from huggingface_hub import hf_hub_download
 from safetensors.torch import load_file
 
+from safety_checker_improved import maybe_nsfw
+
 prompt_list = [p for p in list(set(
                 pd.read_csv('./twitter_prompts.csv').iloc[:, 1].tolist())) if type(p) == str]
 
@@ -56,6 +58,8 @@ pipe.to(device=DEVICE)
 # TODO put back
 @spaces.GPU
 def compile_em():
+    # TODO Compile
+    return None
     pipe.unet = torch.compile(pipe.unet)
     pipe.vae = torch.compile(pipe.vae, mode='reduce-overhead')
     autoencoder.model.forward = torch.compile(autoencoder.model.forward, backend='inductor', dynamic=True)
@@ -160,6 +164,11 @@ def predict(
         im_emb, _ = pipe.encode_image(
                 image, DEVICE, 1, output_hidden_state
             )
+        
+        nsfw = maybe_nsfw(image)
+        if nsfw:
+            return None, im_emb.to('cpu')
+        
         return image, im_emb.to('cpu')
         
         
@@ -245,10 +254,10 @@ def next_image(embs, img_embs, ys, calibrate_prompts):
             image, img_emb = predict(prompt, im_emb=img_emb)
             img_embs.append(img_emb)
 
-            #if len(embs) > 20:
-            #    embs.pop(0)
-            #    img_embs.pop(0)
-            #    ys.pop(0)
+            if len(embs) > 100:
+                embs.pop(0)
+                img_embs.pop(0)
+                ys.pop(0)
             return image, embs, img_embs, ys, calibrate_prompts
 
 
@@ -274,7 +283,7 @@ def start(_, embs, img_embs, ys, calibrate_prompts):
             ]
 
 
-def choose(choice, embs, img_embs, ys, calibrate_prompts):
+def choose(img, choice, embs, img_embs, ys, calibrate_prompts):
     if choice == 'Like (L)':
         choice = 1
     elif choice == 'Neither (Space)':
@@ -284,6 +293,12 @@ def choose(choice, embs, img_embs, ys, calibrate_prompts):
         return img, embs, img_embs, ys, calibrate_prompts
     else:
         choice = 0
+    
+    print(img, 'img')
+    if img is None:
+        print('NSFW -- choice is disliked')
+        choice = 0
+        
     ys.append(choice)
     img, embs, img_embs, ys, calibrate_prompts = next_image(embs, img_embs, ys, calibrate_prompts)
     return img, embs, img_embs, ys, calibrate_prompts
@@ -363,17 +378,17 @@ with gr.Blocks(css=css, head=js_head) as demo:
         b1 = gr.Button(value='Like (L)', interactive=False, elem_id="like")
         b1.click(
         choose, 
-        [b1, embs, img_embs, ys, calibrate_prompts],
+        [img, b1, embs, img_embs, ys, calibrate_prompts],
         [img, embs, img_embs, ys, calibrate_prompts]
         )
         b2.click(
         choose, 
-        [b2, embs, img_embs, ys, calibrate_prompts],
+        [img, b2, embs, img_embs, ys, calibrate_prompts],
         [img, embs, img_embs, ys, calibrate_prompts]
         )
         b3.click(
         choose, 
-        [b3, embs, img_embs, ys, calibrate_prompts],
+        [img, b3, embs, img_embs, ys, calibrate_prompts],
         [img, embs, img_embs, ys, calibrate_prompts]
         )
     with gr.Row():

nsfweffnetv2-b02-3epochs.h5

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:91422f388d1632c1af21b3d787b4f6c1a8e6114f600162d392b0bf285ff8a433
+size 71027272

safety_checker_improved.py

@@ -0,0 +1,45 @@
+
+# TODO required tensorflow==2.14 for me
+# weights from https://github.com/LAION-AI/safety-pipeline/tree/main
+from PIL import Image
+import tensorflow_hub as hub
+import tensorflow
+import numpy as np
+import sys
+sys.path.append('/home/ryn_mote/Misc/generative_recommender/gradio_video/automl/efficientnetv2/')
+import tensorflow as tf
+from tensorflow.keras import mixed_precision
+physical_devices = tf.config.list_physical_devices('GPU')
+
+tf.config.experimental.set_memory_growth(
+    physical_devices[0], True
+)
+
+model = tf.keras.models.load_model('nsfweffnetv2-b02-3epochs.h5',custom_objects={"KerasLayer":hub.KerasLayer})
+# "The image classifier had been trained on 682550 images from the 5 classes "Drawing" (39026), "Hentai" (28134), "Neutral" (369507), "Porn" (207969) & "Sexy" (37914).
+# ... we created a manually inspected test set that consists of 4900 samples, that contains images & their captions."
+
+# Run prediction
+def maybe_nsfw(pil_image):
+    # Run prediction
+    imm = tensorflow.image.resize(np.array(pil_image)[:, :, :3], (260, 260))
+    imm = (imm / 255)
+    pred = model(tensorflow.expand_dims(imm, 0)).numpy()
+    probs = tensorflow.math.softmax(pred[0]).numpy()
+    print(probs)
+    if all([i < .3 for i in probs[[1, 3, 4]]]):
+        return False
+    return True
+
+# pre-initializing prediction
+maybe_nsfw(Image. new("RGB", (260, 260), 255))
+model.load_weights('nsfweffnetv2-b02-3epochs.h5', by_name=True, )
+
+
+
+
+
+
+
+
+