Update app.py

app.py

@@ -94,6 +94,7 @@ detoxify_model = Detoxify('multilingual')
 print("Loading NSFW image classification model...")
 nsfw_model = AutoModelForImageClassification.from_pretrained("Falconsai/nsfw_image_detection")
 nsfw_processor = ViTImageProcessor.from_pretrained('Falconsai/nsfw_image_detection')
+nsfw_model.eval()  # Set to evaluation mode
 print("NSFW image classification model loaded.")
 
 MODERATION_SYSTEM_PROMPT = (
@@ -305,11 +306,15 @@ def classify_text_with_detoxify(text):
 
 def classify_image(image_data):
     try:
+        # Open and convert the image
         img = Image.open(io.BytesIO(image_data)).convert("RGB")
         
-        # Use the model and processor directly as shown in the example
+        # Process the image with the NSFW model
         with torch.no_grad():
             inputs = nsfw_processor(images=img, return_tensors="pt")
+            # Move to the same device as the model
+            inputs = {k: v.to(nsfw_model.device) for k, v in inputs.items()}
+            
             outputs = nsfw_model(**inputs)
             logits = outputs.logits
             
@@ -336,6 +341,7 @@ def classify_image(image_data):
             "nsfw_score": nsfw_score
         }
     except Exception as e:
+        print(f"Error in classify_image: {str(e)}")
         return {
             "classification": "s",
             "label": "ERROR",
@@ -345,6 +351,7 @@ def classify_image(image_data):
         }
 
 def process_content_item(item, text_model="gemma"):
+    # Handle string input (simple text)
     if isinstance(item, str):
         if text_model == "gemma":
             gemma_result = classify_text_with_gemma(item)
@@ -459,8 +466,11 @@ def process_content_item(item, text_model="gemma"):
                 "text": item
             }
     
+    # Handle dictionary input (structured content)
     elif isinstance(item, dict):
-        if item.get("type") == "text":
+        content_type = item.get("type")
+        
+        if content_type == "text":
             text = item.get("text", "")
             
             if text_model == "gemma":
@@ -576,89 +586,35 @@ def process_content_item(item, text_model="gemma"):
                     "text": text
                 }
         
-        elif item.get("type") == "image":
+        elif content_type == "image":
             image_data = None
+            image_url = item.get("url")
+            image_base64 = item.get("base64")
             
-            if item.get("url"):
+            # Get image data from URL
+            if image_url:
                 try:
-                    response = requests.get(item.get("url"))
-                    image_data = response.content
-                except Exception:
-                    return {
-                        "flagged": False,
-                        "categories": {
-                            "hate": False,
-                            "hate/threatening": False,
-                            "harassment": False,
-                            "harassment/threatening": False,
-                            "self-harm": False,
-                            "self-harm/intent": False,
-                            "self-harm/instructions": False,
-                            "sexual": False,
-                            "sexual/minors": False,
-                            "violence": False,
-                            "violence/graphic": False,
-                            "nsfw": False
-                        },
-                        "category_scores": {
-                            "hate": 0.1,
-                            "hate/threatening": 0.1,
-                            "harassment": 0.1,
-                            "harassment/threatening": 0.1,
-                            "self-harm": 0.1,
-                            "self-harm/intent": 0.1,
-                            "self-harm/instructions": 0.1,
-                            "sexual": 0.1,
-                            "sexual/minors": 0.1,
-                            "violence": 0.1,
-                            "violence/graphic": 0.1,
-                            "nsfw": 0.1
-                        },
-                        "image_url": item.get("url")
-                    }
+                    response = requests.get(image_url)
+                    if response.status_code == 200:
+                        image_data = response.content
+                    else:
+                        print(f"Failed to fetch image from URL: {image_url}, status code: {response.status_code}")
+                except Exception as e:
+                    print(f"Error fetching image from URL: {str(e)}")
             
-            elif item.get("base64"):
+            # Get image data from base64
+            elif image_base64:
                 try:
-                    if item.get("base64").startswith("data:image"):
-                        base64_data = item.get("base64").split(",")[1]
+                    if image_base64.startswith("data:image"):
+                        base64_data = image_base64.split(",")[1]
                     else:
-                        base64_data = item.get("base64")
+                        base64_data = image_base64
                     
                     image_data = base64.b64decode(base64_data)
-                except Exception:
-                    return {
-                        "flagged": False,
-                        "categories": {
-                            "hate": False,
-                            "hate/threatening": False,
-                            "harassment": False,
-                            "harassment/threatening": False,
-                            "self-harm": False,
-                            "self-harm/intent": False,
-                            "self-harm/instructions": False,
-                            "sexual": False,
-                            "sexual/minors": False,
-                            "violence": False,
-                            "violence/graphic": False,
-                            "nsfw": False
-                        },
-                        "category_scores": {
-                            "hate": 0.1,
-                            "hate/threatening": 0.1,
-                            "harassment": 0.1,
-                            "harassment/threatening": 0.1,
-                            "self-harm": 0.1,
-                            "self-harm/intent": 0.1,
-                            "self-harm/instructions": 0.1,
-                            "sexual": 0.1,
-                            "sexual/minors": 0.1,
-                            "violence": 0.1,
-                            "violence/graphic": 0.1,
-                            "nsfw": 0.1
-                        },
-                        "image_base64": item.get("base64")[:50] + "..." if len(item.get("base64", "")) > 50 else item.get("base64", "")
-                    }
+                except Exception as e:
+                    print(f"Error decoding base64 image: {str(e)}")
             
+            # Process the image if we have data
             if image_data:
                 image_result = classify_image(image_data)
                 flagged = image_result["classification"] == "u"
@@ -693,10 +649,46 @@ def process_content_item(item, text_model="gemma"):
                         "violence/graphic": 0.1,
                         "nsfw": image_result["nsfw_score"]
                     },
-                    "image_url": item.get("url"),
-                    "image_base64": item.get("base64")[:50] + "..." if item.get("base64") and len(item.get("base64", "")) > 50 else item.get("base64", "")
+                    "image_url": image_url,
+                    "image_base64": image_base64[:50] + "..." if image_base64 and len(image_base64) > 50 else image_base64
+                }
+            else:
+                # Return error if no image data
+                return {
+                    "flagged": False,
+                    "categories": {
+                        "hate": False,
+                        "hate/threatening": False,
+                        "harassment": False,
+                        "harassment/threatening": False,
+                        "self-harm": False,
+                        "self-harm/intent": False,
+                        "self-harm/instructions": False,
+                        "sexual": False,
+                        "sexual/minors": False,
+                        "violence": False,
+                        "violence/graphic": False,
+                        "nsfw": False
+                    },
+                    "category_scores": {
+                        "hate": 0.1,
+                        "hate/threatening": 0.1,
+                        "harassment": 0.1,
+                        "harassment/threatening": 0.1,
+                        "self-harm": 0.1,
+                        "self-harm/intent": 0.1,
+                        "self-harm/instructions": 0.1,
+                        "sexual": 0.1,
+                        "sexual/minors": 0.1,
+                        "violence": 0.1,
+                        "violence/graphic": 0.1,
+                        "nsfw": 0.1
+                    },
+                    "image_url": image_url,
+                    "image_base64": image_base64[:50] + "..." if image_base64 and len(image_base64) > 50 else image_base64
                 }
     
+    # Default return for invalid items
     return {
         "flagged": False,
         "categories": {
@@ -763,10 +755,15 @@ async def moderate_content(
         input_data = request.input
         text_model = request.model or "gemma"
         
+        # Normalize input to a list of items
+        items = []
+        
         if isinstance(input_data, str):
+            # Single string input
             items = [input_data]
             total_tokens += count_tokens(input_data)
         elif isinstance(input_data, list):
+            # List of items
             items = input_data
             for item in items:
                 if isinstance(item, str):
@@ -779,6 +776,7 @@ async def moderate_content(
         if len(items) > 10:
             raise HTTPException(status_code=400, detail="Too many input items. Maximum 10 allowed.")
         
+        # Process each item individually
         results = []
         for item in items:
             result = process_content_item(item, text_model)