add prompt safety classifier

app.py

@@ -9,8 +9,11 @@ import os
 from dotenv import load_dotenv
 import json
 from PIL import Image, ImageDraw, ImageFont
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
 import uuid
 import threading
+import functools
 
 # Load environment variables first
 load_dotenv()
@@ -41,6 +44,38 @@ BACKENDS = {
     },
 }
 
+MODEL_URL = "MichalMlodawski/nsfw-text-detection-large"
+TITLE = "🖼️🔍 Image Prompt Safety Classifier 🛡️"
+DESCRIPTION = "✨ Enter an image generation prompt to classify its safety level! ✨"
+
+# Load model and tokenizer
+tokenizer = AutoTokenizer.from_pretrained(MODEL_URL)
+model = AutoModelForSequenceClassification.from_pretrained(MODEL_URL)
+
+# Define class names with emojis and detailed descriptions
+CLASS_NAMES = {
+    0: "✅ SAFE - This prompt is appropriate and harmless.",
+    1: "⚠️ QUESTIONABLE - This prompt may require further review.",
+    2: "🚫 UNSAFE - This prompt is likely to generate inappropriate content."
+}
+
+
+@functools.lru_cache(maxsize=128)
+def classify_text(text):
+    inputs = tokenizer(text,
+                       return_tensors="pt",
+                       truncation=True,
+                       padding=True,
+                       max_length=1024)
+
+    with torch.no_grad():
+        outputs = model(**inputs)
+
+    logits = outputs.logits
+    predicted_class = torch.argmax(logits, dim=1).item()
+
+    return predicted_class, CLASS_NAMES[predicted_class]
+
 
 class BackendStatus:
 
@@ -600,6 +635,24 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
             )
             return
 
+        # Check if the prompt is safe
+        classification, message = classify_text(prompt)
+        if classification != 0:
+            # Handle unsafe prompt case
+            yield (
+                message,
+                message,
+                gr.update(visible=True),
+                gr.update(visible=False),
+                None,
+                None,
+                None,
+                None,
+                session_id,  # Return the session ID
+                None,
+            )
+            return
+
         # Status message
         status_message = f"🔄 PROCESSING: '{prompt}'"
 
@@ -749,10 +802,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
 
 # Launch with increased max_threads
 if __name__ == "__main__":
-    # demo.queue(max_size=50).launch(
-    #     server_name="0.0.0.0",
-    #     max_threads=16,  # Increase thread count for better concurrency
-    # )
+    demo.queue(max_size=50).launch(
+        server_name="0.0.0.0",
+        max_threads=16,  # Increase thread count for better concurrency
+    )
     demo.queue(max_size=4).launch(
         server_name="0.0.0.0",
         max_threads=16,  # Increase thread count for better concurrency

requirements.txt

@@ -3,3 +3,5 @@ aiohttp
 plotly
 python-dotenv
 pydantic==2.8.2
+torch
+transformers==4.37.2