Update app.py

app.py

@@ -15,6 +15,7 @@ from flask_cors import CORS
 # =====================
 # CONFIG
 # =====================
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf_cache"  # tránh vượt storage limit
 TEXT_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Text_Image"
 IMAGE_MODEL_REPO = "phuongsuga/PBL6_AI_Model_Image"
 DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
@@ -35,33 +36,35 @@ text_model = AutoModelForSequenceClassification.from_pretrained(
     TEXT_MODEL_REPO,
     subfolder="text_model/checkpoint-3390"
 )
-
 text_model.to(DEVICE).eval()
 
-# Load label2id.json từ repo
+# load label2id.json
 label_url = f"https://huggingface.co/{TEXT_MODEL_REPO}/resolve/main/text_model/label2id.json"
 label2id = requests.get(label_url).json()
 id2label_text = {i: l for l, i in label2id.items()}
 
-
 # =====================
 # LOAD IMAGE MODEL
 # =====================
 print("🔹 Downloading image model...")
+
 class_names = ["an_toan", "bao_luc", "khieu_dam_doi_truy", "nhay_cam_chinh_tri"]
 
 def build_model(num_classes=4):
     return timm.create_model("efficientnet_b3", pretrained=False, num_classes=num_classes)
 
 image_model = build_model()
-image_model_path = f"https://huggingface.co/{IMAGE_MODEL_REPO}/resolve/main/image_model/efficientnet_b3.pth"
 
-torch.hub.download_url_to_file(image_model_path, "efficientnet_b3.pth")
-image_model.load_state_dict(torch.load("efficientnet_b3.pth", map_location=DEVICE))
-image_model.to(DEVICE).eval()
+# tải model tạm trong /tmp để không chiếm storage
+image_model_path = "/tmp/efficientnet_b3.pth"
+if not os.path.exists(image_model_path):
+    url = f"https://huggingface.co/{IMAGE_MODEL_REPO}/resolve/main/image_model/efficientnet_b3.pth"
+    torch.hub.download_url_to_file(url, image_model_path)
 
+image_model.load_state_dict(torch.load(image_model_path, map_location=DEVICE))
+image_model.to(DEVICE).eval()
 
-# Chuẩn hoá ảnh đầu vào
+# chuẩn hóa ảnh
 val_transforms = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -76,7 +79,7 @@ def seg_pyvi(text: str) -> str:
     try:
         seg = ViTokenizer.tokenize(text)
         seg = seg.replace(" ", "_")
-    except:
+    except Exception:
         seg = text
     return seg
 
@@ -129,6 +132,10 @@ def predict_image(pil_image: Image.Image):
 app = Flask(__name__)
 CORS(app)
 
+@app.route("/")
+def home():
+    return jsonify({"message": "✅ AI moderation API is running!"})
+
 @app.route("/analyze", methods=["POST"])
 def analyze():
     result = {"text_result": [], "image_result": []}
@@ -152,6 +159,5 @@ def analyze():
 # RUN APP
 # =====================
 if __name__ == "__main__":
-    import os
-    port = int(os.environ.get("PORT", 7860))  # Hugging Face sẽ truyền PORT vào đây
+    port = int(os.environ.get("PORT", 7860))  # Hugging Face Space truyền PORT vào
     app.run(host="0.0.0.0", port=port)