Update app.py

app.py

@@ -5,12 +5,12 @@ from huggingface_hub import InferenceClient
 from io import BytesIO
 from PIL import Image
 
-# Initialize the Flask app
+# Initialize Flask app
 app = Flask(__name__)
 CORS(app)  # Enable CORS for all routes
 
-# Initialize the InferenceClient with your Hugging Face token
-HF_TOKEN = os.environ.get("HF_TOKEN")  # Ensure to set your Hugging Face token in the environment
+# Initialize the InferenceClient with Hugging Face token
+HF_TOKEN = os.environ.get("HF_TOKEN")  # Set your Hugging Face token in environment variables
 client = InferenceClient(token=HF_TOKEN)
 
 # Hardcoded negative prompt
@@ -21,54 +21,48 @@ blurry hands, disproportionate fingers"""
 
 @app.route('/')
 def home():
-    return "Welcome to the Image Background Remover!"
+    return "Welcome to the AI Image Generator with NSFW Detection!"
 
-# Simple content moderation function
-def is_prompt_explicit(prompt):
-    explicit_keywords = ["sexual", "nudity", "erotic", "explicit", "porn", "pornographic", "xxx", "hentai", "fetish", "sex", "sensual", "nude", "strip", "stripping", "adult", "lewd", "provocative", "obscene", "vulgar", "intimacy", "intimate", "lust", "arouse", "seductive", "seduction", "kinky", "bdsm", "dominatrix", "bondage", "hardcore", "softcore", "topless", "bottomless", "threesome", "orgy", "incest", "taboo", "masturbation", "genital", "penis", "vagina", "breast", "boob", "nipple", "butt", "anal", "oral", "ejaculation", "climax", "moan", "foreplay", "intercourse", "naked", "exposed", "suicide", "self-harm", "overdose", "poison", "hang", "end life", "kill myself", "noose", "depression", "hopeless", "worthless", "die", "death", "harm myself"]  # Add more keywords as needed
-    for keyword in explicit_keywords:
-        if keyword.lower() in prompt.lower():
-            return True
-    return False
-
-# NSFW detection function using InferenceClient
+# Function for NSFW detection
 def is_nsfw_image(image):
-    # Convert the image to bytes
-    img_byte_arr = BytesIO()
-    image.save(img_byte_arr, format='PNG')
-    img_byte_arr.seek(0)
-
-    # Send the image to the Hugging Face NSFW model
     try:
-        result = client.image_classification("Falconsai/nsfw_image_detection", img_byte_arr.getvalue())
+        # Convert the image to bytes
+        img_byte_arr = BytesIO()
+        image.save(img_byte_arr, format='PNG')
+        img_byte_arr.seek(0)
+
+        # Send the image to Hugging Face for NSFW classification
+        result = client.image_classification(model="Falconsai/nsfw_image_detection", inputs=img_byte_arr.getvalue())
+        
+        # Check if any prediction is NSFW with high confidence
         for item in result:
-            if item['label'] == 'nsfw' and item['score'] > 0.5:
+            if item['label'].lower() == 'nsfw' and item['score'] > 0.5:
                 return True
         return False
     except Exception as e:
-        print(f"Error in NSFW detection: {e}")
+        print(f"NSFW detection error: {e}")
         return False
 
-# Function to generate an image from a text prompt
+# Function to generate an image
 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
+        # Generate the image using Hugging Face's API
         image = client.text_to_image(
-            prompt=prompt, 
-            negative_prompt=NEGATIVE_PROMPT_FINGERS,
-            height=height, 
-            width=width, 
+            prompt=prompt,
+            negative_prompt=negative_prompt or NEGATIVE_PROMPT_FINGERS,
+            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
-            seed=seed  # Control the seed for reproducibility
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            seed=seed
         )
-        return image  # Return the generated image
+        return image
     except Exception as e:
-        print(f"Error generating image: {str(e)}")
+        print(f"Error generating image: {e}")
         return None
 
-# Flask route for the API endpoint to generate an image
+# Flask route for image generation API
 @app.route('/generate_image', methods=['POST'])
 def generate_api():
     data = request.get_json()
@@ -76,26 +70,17 @@ def generate_api():
     # Extract required fields from the request
     prompt = data.get('prompt', '')
     negative_prompt = data.get('negative_prompt', None)
-    height = data.get('height', 1024)  # Default height
-    width = data.get('width', 720)  # Default width
-    num_inference_steps = data.get('num_inference_steps', 50)  # Default number of inference steps
-    guidance_scale = data.get('guidance_scale', 7.5)  # Default guidance scale
-    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')  # Default model
-    seed = data.get('seed', None)  # Seed for reproducibility, default is None
+    height = data.get('height', 512)
+    width = data.get('width', 512)
+    num_inference_steps = data.get('num_inference_steps', 50)
+    guidance_scale = data.get('guidance_scale', 7.5)
+    model_name = data.get('model', 'stabilityai/stable-diffusion-2-1')
+    seed = data.get('seed', None)
 
     if not prompt:
         return jsonify({"error": "Prompt is required"}), 400
 
     try:
-        # Check for explicit content
-        if is_prompt_explicit(prompt):
-            return send_file(
-                "thinkgood.jpeg",
-                mimetype='image/png',
-                as_attachment=False,
-                download_name='thinkgood.png'
-            )
-
         # Generate the image
         image = generate_image(prompt, negative_prompt, height, width, model_name, num_inference_steps, guidance_scale, seed)
 
@@ -103,30 +88,30 @@ def generate_api():
             # Check for NSFW content
             if is_nsfw_image(image):
                 return send_file(
-                    "nsfw.jpg", 
-                    mimetype='image/jpeg', 
-                    as_attachment=False, 
+                    "nsfw.jpg",  # Path to your predefined NSFW placeholder image
+                    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
-            img_byte_arr.seek(0)  # Move to the start of the byte stream
+            image.save(img_byte_arr, format='PNG')
+            img_byte_arr.seek(0)
 
-            # Send the generated image as a response
+            # Send the generated image
             return send_file(
-                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
+                img_byte_arr,
+                mimetype='image/png',
+                as_attachment=False,
+                download_name='generated_image.png'
             )
         else:
             return jsonify({"error": "Failed to generate image"}), 500
     except Exception as e:
-        print(f"Error in generate_api: {str(e)}")  # Log the error
+        print(f"Error in generate_api: {e}")
         return jsonify({"error": str(e)}), 500
 
-# Add this block to make sure your app runs when called
-if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860)  # Run directly if needed for testing
+# Run the Flask app
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=7860)