Update app.py

app.py

@@ -4,7 +4,6 @@ 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__)
@@ -20,9 +19,6 @@ 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!"
@@ -35,13 +31,23 @@ def is_prompt_explicit(prompt):
             return True
     return False
 
-# NSFW detection function
+# NSFW detection function using InferenceClient
 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
+    # 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())
+        for item in result:
+            if item['label'] == 'nsfw' and item['score'] > 0.5:
+                return True
+        return False
+    except Exception as e:
+        print(f"Error in NSFW detection: {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):