Update app.py

app.py

@@ -14,38 +14,38 @@ 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)
 
-# 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"""
+# Initialize NSFW model
+NSFW_MODEL = "MichalMlodawski/nsfw-text-detection-large"
+nsfw_client = InferenceClient(model=NSFW_MODEL, token=HF_TOKEN)
 
-@app.route('/')
-def home():
-    return "Welcome to the Image Background Remover!"
+# 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"""
 
-# Simple content moderation function
+# NSFW detection function
 def is_prompt_explicit(prompt):
-    # Streamlined keyword list to avoid unnecessary restrictions
-    explicit_keywords = [
-        "sexual","boobs","boob","penis","porn", "hentai", "fetish", "nude", "provocative", "obscene", "vulgar", "intimate", "kinky", "hardcore",
-        "threesome", "orgy", "masturbation", "genital", "suicide", 
-        "self-harm", "depression", "kill myself", "worthless","vagina","intercourse","fuck","fucking","fuckers","fucker"
-    ]
-    for keyword in explicit_keywords:
-        if keyword.lower() in prompt.lower():
-            return True
-    return False
+    try:
+        response = nsfw_client(prompt, task="text-classification")
+        if "error" in response:
+            print(f"Error in NSFW detection: {response['error']}")
+            return False
+
+        # 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
 
 # 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,7 +77,7 @@ def generate_api():
     try:
         # Check for explicit content
         if is_prompt_explicit(prompt):
-            # Return the pre-defined "thinkgood.png" image
+            # Return the pre-defined "nsfw.jpg" image
             return send_file(
                 "nsfw.jpg",
                 mimetype='image/png',
@@ -96,8 +96,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,5 +110,4 @@ 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