Update app.py

app.py

@@ -14,38 +14,47 @@ CORS(app)  # Enable CORS for all routes
 HF_TOKEN = os.environ.get("HF_TOKEN")  # Ensure to set your Hugging Face token in the environment
 client = InferenceClient(token=HF_TOKEN)
 
-# Initialize NSFW model
-NSFW_MODEL = "MichalMlodawski/nsfw-text-detection-large"
-nsfw_client = InferenceClient(model=NSFW_MODEL, token=HF_TOKEN)
-
 # Hardcoded negative prompt
-NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers, mutated fingers, poorly drawn fingers, disfigured fingers, too many fingers, deformed hands, extra hands, malformed hands, blurry hands, disproportionate fingers"""
+NEGATIVE_PROMPT_FINGERS = """2D,missing fingers, extra fingers, elongated fingers, fused fingers, 
+mutated fingers, poorly drawn fingers, disfigured fingers, 
+too many fingers, deformed hands, extra hands, malformed hands, 
+blurry hands, disproportionate fingers"""
 
-# NSFW detection function
-def is_prompt_explicit(prompt):
-    try:
-        response = nsfw_client(prompt, task="text-classification")
-        if "error" in response:
-            print(f"Error in NSFW detection: {response['error']}")
-            return False
+@app.route('/')
+def home():
+    return "Welcome to the Image Background Remover!"
 
-        # Parse the classification result
-        predicted_class = response[0]["label"]  # E.g., "LABEL_2"
-        class_id = int(predicted_class.split("_")[-1])  # Extract numerical label
-        return class_id == 2  # Class 2 indicates UNSAFE
-    except Exception as e:
-        print(f"Error in NSFW detection: {str(e)}")
-        return False
+# Simple content moderation function
+def is_prompt_explicit(prompt):
+    # Streamlined keyword list to avoid unnecessary restrictions
+ explicit_keywords = [
+    "sexual", "sex", "boobs", "boob", "breasts", "cleavage", "penis", "phallus", "porn", "pornography", "hentai", 
+    "fetish", "nude", "nudity", "provocative", "obscene", "vulgar", "intimate", "kinky", "hardcore", 
+    "threesome", "orgy", "masturbation", "masturbate", "genital", "genitals", "vagina", "vaginal", 
+    "anus", "anal", "butt", "buttocks", "butthole", "ass", "prostate", "erection", "cum", "ejaculation", 
+    "sperm", "semen", "naked", "bare", "lingerie", "thong", "striptease", "stripper", 
+    "seductive", "sensual", "explicit", "lewd", "taboo", "NSFW", "bdsm", "dominatrix", "submission", 
+    "intercourse", "penetration", "orgasm", "fuck", "fucking", "fuckers", "fucker", "slut", "whore", 
+    "prostitute", "hooker", "escort", "camgirl", "camwhore", "sugar daddy", "sugar baby", "adult content", 
+    "sexually explicit", "arousal", "lust", "depraved", "hardcore porn", "softcore", "erotic", "erotica",
+    "roleplay", "incest", "taboo", "voyeur", "exhibitionist", "peeping", "dildo", "sex toy", "vibrator", 
+    "suicide", "self-harm", "depression", "kill myself", "worthless", "abuse", "violence", "rape", 
+    "sexual violence", "molestation", "pedophilia", "child porn", "underage", "illegal content"
+]
+    for keyword in explicit_keywords:
+        if keyword.lower() in prompt.lower():
+            return True
+    return False
 
 # Function to generate an image from a text prompt
 def generate_image(prompt, negative_prompt=None, height=512, width=512, model="stabilityai/stable-diffusion-2-1", num_inference_steps=50, guidance_scale=7.5, seed=None):
     try:
         # Generate the image using Hugging Face's inference API with additional parameters
         image = client.text_to_image(
-            prompt=prompt,
+            prompt=prompt, 
             negative_prompt=NEGATIVE_PROMPT_FINGERS,
-            height=height,
-            width=width,
+            height=height, 
+            width=width, 
             model=model,
             num_inference_steps=num_inference_steps,  # Control the number of inference steps
             guidance_scale=guidance_scale,  # Control the guidance scale
@@ -77,12 +86,12 @@ def generate_api():
     try:
         # Check for explicit content
         if is_prompt_explicit(prompt):
-            # Return the pre-defined "NSFW.jpg" image if the content is explicit
+            # Return the pre-defined "thinkgood.png" image
             return send_file(
-                "nsfw.jpg",  # Make sure this file exists in your project directory
-                mimetype='image/jpeg',
+                "nsfw.jpg",
+                mimetype='image/png',
                 as_attachment=False,
-                download_name='NSFW_detected.jpg'  # The name for download (optional)
+                download_name='thinkgood.png'
             )
 
         # Call the generate_image function with the provided parameters
@@ -96,8 +105,8 @@ def generate_api():
 
             # Send the generated image as a response
             return send_file(
-                img_byte_arr,
-                mimetype='image/png',
+                img_byte_arr, 
+                mimetype='image/png', 
                 as_attachment=False,  # Send the file as an attachment
                 download_name='generated_image.png'  # The file name for download
             )
@@ -110,4 +119,5 @@ def generate_api():
 # Add this block to make sure your app runs when called
 if __name__ == "__main__":
     subprocess.Popen(["python", "wk.py"])  # Start awake.py
+
     app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing