Update app.py

app.py

@@ -4,6 +4,7 @@ import os
 from huggingface_hub import InferenceClient
 from io import BytesIO
 from PIL import Image
+from transformers import pipeline
 
 # Initialize the Flask app
 app = Flask(__name__)
@@ -19,6 +20,9 @@ mutated fingers, poorly drawn fingers, disfigured fingers,
 too many fingers, deformed hands, extra hands, malformed hands, 
 blurry hands, disproportionate fingers"""
 
+# Initialize the NSFW detection pipeline
+nsfw_classifier = pipeline("image-classification", model="Falconsai/nsfw_image_detection")
+
 @app.route('/')
 def home():
     return "Welcome to the Image Background Remover!"
@@ -31,6 +35,14 @@ def is_prompt_explicit(prompt):
             return True
     return False
 
+# NSFW detection function
+def is_nsfw_image(image):
+    results = nsfw_classifier(image)
+    for result in results:
+        if result['label'] == 'nsfw' and result['score'] > 0.5:
+            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:
@@ -71,7 +83,6 @@ def generate_api():
     try:
         # Check for explicit content
         if is_prompt_explicit(prompt):
-            # Return the pre-defined "thinkgood.png" image
             return send_file(
                 "thinkgood.jpeg",
                 mimetype='image/png',
@@ -79,10 +90,19 @@ def generate_api():
                 download_name='thinkgood.png'
             )
 
-        # Call the generate_image function with the provided parameters
+        # Generate the image
         image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
 
         if image:
+            # Check for NSFW content
+            if is_nsfw_image(image):
+                return send_file(
+                    "nsfw.jpg", 
+                    mimetype='image/jpeg', 
+                    as_attachment=False, 
+                    download_name='nsfw.jpg'
+                )
+
             # Save the image to a BytesIO object
             img_byte_arr = BytesIO()
             image.save(img_byte_arr, format='PNG')  # Convert the image to PNG