Update app.py

app.py

@@ -1,5 +1,4 @@
-from fastapi import FastAPI
-from flask import Flask, request, jsonify
+from fastapi import FastAPI, File, UploadFile
 from pydantic import BaseModel
 from transformers import pipeline
 from PIL import Image
@@ -7,15 +6,18 @@ import joblib
 import re
 import string
 import io
+import os
 import uvicorn
-from threading import Thread
 
-# Initialize Flask & FastAPI
-app = Flask(__name__)
-api = FastAPI()
+# ✅ Set Hugging Face Cache Directory (Fixes Permission Error)
+CACHE_DIR = "/tmp/hf_cache"
+os.makedirs(CACHE_DIR, exist_ok=True)
 
-# ✅ Load NSFW Image Classification Model
-pipe = pipeline("image-classification", model="LukeJacob2023/nsfw-image-detector")
+# ✅ Initialize FastAPI
+app = FastAPI()
+
+# ✅ Load NSFW Image Classification Model (with custom cache directory)
+pipe = pipeline("image-classification", model="LukeJacob2023/nsfw-image-detector", cache_dir=CACHE_DIR)
 
 # ✅ Load Toxic Text Classification Model
 try:
@@ -37,24 +39,23 @@ def preprocess_text(text):
     text = text.translate(str.maketrans('', '', string.punctuation))  # Remove punctuation
     return text.strip()
 
-# 📌 NSFW Image Classification API (Flask)
-@app.route('/classify_image', methods=['POST'])
-def classify_image():
-    if 'file' not in request.files:
-        return jsonify({"error": "No file uploaded"}), 400
-    
-    file = request.files['file']
-    image = Image.open(io.BytesIO(file.read()))
-    results = pipe(image)
-    
-    classification_label = max(results, key=lambda x: x['score'])['label']
-    nsfw_labels = {"sexy", "porn", "hentai"}
-    nsfw_status = "NSFW" if classification_label in nsfw_labels else "SFW"
+# 📌 NSFW Image Classification API
+@app.post("/classify_image/")
+async def classify_image(file: UploadFile = File(...)):
+    try:
+        image = Image.open(io.BytesIO(await file.read()))
+        results = pipe(image)
+
+        classification_label = max(results, key=lambda x: x['score'])['label']
+        nsfw_labels = {"sexy", "porn", "hentai"}
+        nsfw_status = "NSFW" if classification_label in nsfw_labels else "SFW"
 
-    return jsonify({"status": nsfw_status, "results": results})
+        return {"status": nsfw_status, "results": results}
+    except Exception as e:
+        return {"error": str(e)}
 
-# 📌 Toxic Text Classification API (FastAPI)
-@api.post("/classify_text/")
+# 📌 Toxic Text Classification API
+@app.post("/classify_text/")
 async def classify_text(data: TextInput):
     try:
         processed_text = preprocess_text(data.text)
@@ -65,14 +66,6 @@ async def classify_text(data: TextInput):
     except Exception as e:
         return {"error": str(e)}
 
-# 🔥 Run both servers using Gunicorn
-def run_flask():
-    app.run(host="0.0.0.0", port=5000)
-
-def run_fastapi():
-    uvicorn.run(api, host="0.0.0.0", port=8000)
-
+# ✅ Run FastAPI using Uvicorn (Hugging Face requires port 7860)
 if __name__ == "__main__":
-    Thread(target=run_flask).start()
-    Thread(target=run_fastapi).start()
-
+    uvicorn.run(app, host="0.0.0.0", port=7860)