Update safty checker

backend/safety_check.py

@@ -1,30 +1,16 @@
-from transformers import CLIPProcessor, CLIPModel
-from PIL import Image
-
-
-# model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-# processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+from transformers import pipeline
 
 
 def is_safe_image(
-    model,
-    processor,
+    classifier,
     image,
 ):
-    # Load image
-    # image = Image.open(
-    #     r"F:\om\2025\fastsdcpumcp\fastsdcpu\results\829a2123-92c8-4957-ad2f-06365a19665a-1.png"
-    # )
-    categories = ["safe", "nsfw"]
-    inputs = processor(
-        text=categories,
-        images=image,
-        return_tensors="pt",
-        padding=True,
-    )
-    outputs = model(**inputs)
-    logits_per_image = outputs.logits_per_image
-    probs = logits_per_image.softmax(dim=1)
-    safe_prob = dict(zip(categories, probs[0].tolist()))
-    print(safe_prob)
-    return safe_prob["safe"] > safe_prob["nsfw"]
+    pred = classifier(image)
+    nsfw_score = 0
+    normal_score = 0
+    for label in pred:
+        if label["label"] == "nsfw":
+            nsfw_score = label["score"]
+        elif label["label"] == "normal":
+            normal_score = label["score"]
+    return normal_score > nsfw_score

frontend/webui/hf_demo.py

@@ -13,15 +13,17 @@ from PIL import Image
 from backend.models.lcmdiffusion_setting import DiffusionTask
 from backend.safety_check import is_safe_image
 from pprint import pprint
-from transformers import CLIPProcessor, CLIPModel
+from transformers import pipeline
 
 lcm_text_to_image = LCMTextToImage()
 lcm_lora = LCMLora(
     base_model_id="Lykon/dreamshaper-7",
     lcm_lora_id="latent-consistency/lcm-lora-sdv1-5",
 )
-model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
-processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
+classifier = pipeline(
+    "image-classification",
+    model="Falconsai/nsfw_image_detection",
+)
 
 
 # https://github.com/gradio-app/gradio/issues/2635#issuecomment-1423531319
@@ -69,7 +71,10 @@ def predict(
     latency = perf_counter() - start
     print(f"Latency: {latency:.2f} seconds")
     result = images[0]
-    if is_safe_image(model, processor, result):
+    if is_safe_image(
+        classifier,
+        result,
+    ):
         return result  # .resize([512, 512], PIL.Image.ANTIALIAS)
     else:
         print("Unsafe image detected")