Adding NSFW_simple function of code (not loaded)

model/model_manager.py

@@ -10,6 +10,7 @@ from .fetch_museum_results import draw_from_imagen_museum, draw2_from_imagen_mus
 from .pre_download import pre_download_all_models, pre_download_video_models
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
+import re
 
 def debug_packages():
     import pkg_resources
@@ -17,7 +18,59 @@ def debug_packages():
     installed_packages = pkg_resources.working_set
     for package in installed_packages:
         print(f"{package.key}=={package.version}")
-        
+
+def fetch_unsafe_words(file_path):
+    """
+    Loads unsafe words from a file and returns them as a list.
+    """
+    try:
+        with open(file_path, 'r') as file:
+            # Read lines from file and strip any extra whitespace
+            unsafe_words = [line.strip() for line in file.readlines()]
+            # Remove any empty strings that may result from empty lines
+            unsafe_words = [word for word in unsafe_words if word]
+            return unsafe_words
+    except Exception as e:
+        print(f"Error loading file: {e}. Using default unsafe words.")
+        # Default unsafe words list
+        return [
+            "anal", "anus", "arse", "ass", "ballsack", "bastard", "bdsm", "bitch", "bimbo", 
+            "blow job", "blowjob", "blue waffle", "boob", "booobs", "breasts", "booty call", 
+            "boner", "bondage", "bullshit", "busty", "butthole", "cawk", "chink", "clit", 
+            "cnut", "cock", "cokmuncher", "cowgirl", "crap", "crotch", "cum", "cunt", "damn", 
+            "dick", "dildo", "dink", "deepthroat", "deep throat", "dog style", "doggie style", 
+            "doggy style", "doosh", "douche", "duche", "ejaculate", "ejaculating", 
+            "ejaculation", "ejakulate", "erotic", "erotism", "fag", "fatass", "femdom", 
+            "fingering", "footjob", "foot job", "fuck", "fcuk", "fingerfuck", "fistfuck", 
+            "fook", "fooker", "fuk", "gangbang", "gang bang", "gaysex", "handjob", 
+            "hand job", "hentai", "hooker", "hoer", "homo", "horny", "incest", "jackoff", 
+            "jack off", "jerkoff", "jerk off", "jizz", "masturbate", "mofo", "mothafuck", 
+            "motherfuck", "milf", "muff", "nigga", "nigger", "nipple", "nob", "numbnuts", 
+            "nutsack", "nude", "orgy", "orgasm", "panty", "panties", "penis", "playboy", 
+            "porn", "pussy", "pussies", "rape", "raping", "rapist", "rectum", "retard", 
+            "rimming", "sadist", "sadism", "scrotum", "sex", "semen", "shemale", "she male", 
+            "shit", "slut", "spunk", "strip club", "stripclub", "tit", "threesome", 
+            "three some", "throating", "twat", "viagra", "vagina", "wank", "whore", "whoar", 
+            "xxx"
+        ]
+
+def check_prompt_safety(prompt, unsafe_words_file='./profanity_words.txt'):
+    """
+    Checking prompt safety. Returns boolean (Not Safe = False, Safe = True)
+    """
+    # Load unsafe words from the provided file or use default if loading fails
+    unsafe_words = fetch_unsafe_words(unsafe_words_file)
+    
+    # Convert input string to lowercase to ensure case-insensitive matching
+    prompt = prompt.lower()
+
+    # Check if any unsafe word is in the input string
+    for word in unsafe_words:
+        # Use regex to match whole words only
+        if re.search(r'\b' + re.escape(word) + r'\b', prompt):
+            return False
+    return True
+
 class ModelManager:
     def __init__(self, enable_nsfw=False, do_pre_download=False, do_debug_packages=False):
         self.model_ig_list = IMAGE_GENERATION_MODELS
@@ -52,6 +105,13 @@ class ModelManager:
         else:
             self.guard_tokenizer = None
             self.guard = None
+    
+    def NSFW_filter_simple(self, prompt):
+        is_safe = check_prompt_safety(prompt)
+        if is_safe:
+            return "safe"
+        else:
+            return "unsafe"
         
     def NSFW_filter(self, prompt):
         chat = [{"role": "user", "content": prompt}]

model/profanity_words.txt

@@ -0,0 +1,130 @@
+anal
+anus
+arse
+ass
+ballsack
+bastard
+bdsm
+bitch
+bimbo
+blow job
+blowjob
+blue waffle
+boob
+booobs
+breasts
+booty call
+boner
+bondage
+bullshit
+busty
+butthole
+cawk
+chink
+clit
+cnut
+cock
+cokmuncher
+cowgirl
+crap
+crotch
+cum
+cunt
+damn
+dick
+dildo
+dink
+deepthroat
+deep throat
+dog style
+doggie style
+doggy style
+doosh
+douche
+duche
+ejaculate
+ejaculating
+ejaculation
+ejakulate
+erotic
+erotism
+fag
+fatass
+femdom
+fingering
+footjob
+foot job
+fuck
+fcuk
+fingerfuck
+fistfuck
+fook
+fooker
+fuk
+gangbang
+gang bang
+gaysex
+handjob
+hand job
+hentai
+hooker
+hoer
+homo
+horny
+incest
+jackoff
+jack off
+jerkoff
+jerk off
+jizz
+masturbate
+mofo
+mothafuck
+motherfuck
+milf
+muff
+nigga
+nigger
+nipple
+nob
+numbnuts
+nutsack
+nude
+orgy
+orgasm
+panty
+panties
+penis
+playboy
+porn
+pussy
+pussies
+rape
+raping
+rapist
+rectum
+retard
+rimming
+sadist
+sadism
+scrotum
+sex
+semen
+shemale
+she male
+shit
+slut
+spunk
+strip club
+stripclub
+tit
+threesome
+three some
+throating
+twat
+viagra
+vagina
+wank
+whore
+whoar
+xxx