Update app.py

app.py

@@ -4,9 +4,6 @@ import torch.nn as nn
 import sys
 import os
 
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-
 sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
 
 from models.audio_model import FrozenExtractor, LinearProjection, AGRU, PersonalityRegressor
@@ -36,17 +33,14 @@ REG_HEADS    = 8
 REG_DROPOUT  = 0.1
 PROJ_DROPOUT = 0.1
 
-
 # ── Checkpoint loader helper ───────────────────────────────────────────────────
 def _unwrap_ckpt(ckpt):
-    """Extract state dict from common checkpoint wrapper formats."""
     if not isinstance(ckpt, dict):
         return ckpt
     for key in ("model_state", "model_state_dict", "state_dict", "model"):
         if key in ckpt:
             return ckpt[key]
-    return ckpt  # already a flat state dict
-
+    return ckpt
 
 def _load_ckpt(model, path):
     ckpt = torch.load(path, map_location=DEVICE, weights_only=False)
@@ -54,7 +48,6 @@ def _load_ckpt(model, path):
     model.load_state_dict(state)
     return model
 
-
 # ── Visual model ───────────────────────────────────────────────────────────────
 print("⏳ Loading visual model...")
 visual_model = VisualPersonalityModel(VISUAL_CFG).to(DEVICE)
@@ -62,7 +55,6 @@ visual_model = _load_ckpt(visual_model, "checkpoints/best_model.pt")
 visual_model.eval()
 print("✓ Visual model ready")
 
-
 # ── Fusion model ───────────────────────────────────────────────────────────────
 print("⏳ Loading fusion model...")
 fusion_model = FusionMLP(
@@ -72,10 +64,8 @@ fusion_model = _load_ckpt(fusion_model, "checkpoints/best_fusion_model.pt")
 fusion_model.eval()
 print("✓ Fusion model ready")
 
-
 # ── Audio model ────────────────────────────────────────────────────────────────
-print("⏳ Loading audio model (downloading wav2vec2 if not cached)...")
-
+print("⏳ Loading audio model...")
 frozen_ext = FrozenExtractor("facebook/wav2vec2-base", n_shallow=N_SHALLOW).to(DEVICE)
 frozen_ext.eval()
 
@@ -95,46 +85,35 @@ regressor = PersonalityRegressor(
     dropout=REG_DROPOUT,
 ).to(DEVICE)
 
-# Load audio checkpoint — handle all common formats
-audio_ckpt = torch.load(
-    "checkpoints/best_audio_model.pt", map_location=DEVICE, weights_only=False
-)
+audio_ckpt = torch.load("checkpoints/best_audio_model.pt", map_location=DEVICE, weights_only=False)
 audio_state = _unwrap_ckpt(audio_ckpt)
 
-print(f"Audio checkpoint keys: {list(audio_state.keys())[:8]}")
-
 if "deep_transformer" in audio_state:
-    # Saved as separate component dicts
     deep_transformer.load_state_dict(audio_state["deep_transformer"])
     proj.load_state_dict(audio_state["proj"])
     agru.load_state_dict(audio_state["agru"])
     regressor.load_state_dict(audio_state["regressor"])
 else:
-    # Single flat state dict — split by prefix
     dt_state   = {k[len("deep_transformer."):]: v for k, v in audio_state.items() if k.startswith("deep_transformer.")}
-    proj_state = {k[len("proj."):]: v          for k, v in audio_state.items() if k.startswith("proj.")}
-    agru_state = {k[len("agru."):]: v          for k, v in audio_state.items() if k.startswith("agru.")}
-    reg_state  = {k[len("regressor."):]: v     for k, v in audio_state.items() if k.startswith("regressor.")}
-
+    proj_state = {k[len("proj."):]: v for k, v in audio_state.items() if k.startswith("proj.")}
+    agru_state = {k[len("agru."):]: v for k, v in audio_state.items() if k.startswith("agru.")}
+    reg_state  = {k[len("regressor."):]: v for k, v in audio_state.items() if k.startswith("regressor.")}
+    
     if dt_state:
         deep_transformer.load_state_dict(dt_state)
         proj.load_state_dict(proj_state)
         agru.load_state_dict(agru_state)
         regressor.load_state_dict(reg_state)
-    else:
-        print(f"⚠️ Could not match audio keys. All keys: {list(audio_state.keys())}")
 
 for m in [deep_transformer, proj, agru, regressor]:
     m.eval()
 print("✓ Audio model ready")
 
-
 # ── Face detector ──────────────────────────────────────────────────────────────
 print("⏳ Loading face detector...")
 face_detector = _load_face_detector()
 print("✓ Face detector ready")
 
-
 # ── Inference ──────────────────────────────────────────────────────────────────
 def predict(video_path):
     if video_path is None:
@@ -144,22 +123,15 @@ def predict(video_path):
         with open(video_path, "rb") as f:
             video_bytes = f.read()
 
-        visual_preds, visual_emb = get_visual_embedding(
-            video_bytes, visual_model, DEVICE, face_detector
-        )
-
+        visual_preds, visual_emb = get_visual_embedding(video_bytes, visual_model, DEVICE, face_detector)
         audio_preds, audio_emb = get_audio_embedding(
-            video_bytes, frozen_ext, deep_transformer,
-            proj, agru, regressor, DEVICE
-        )
-
-        fusion_preds = get_fusion_prediction(
-            audio_emb, visual_emb, fusion_model, DEVICE
+            video_bytes, frozen_ext, deep_transformer, proj, agru, regressor, DEVICE
         )
+        fusion_preds = get_fusion_prediction(audio_emb, visual_emb, fusion_model, DEVICE)
 
         results = format_predictions(fusion_preds, visual_preds, audio_preds)
-
         fusion = results["fusion"]
+        
         md = "## 🧠 Fusion Results (Main)\n\n"
         for trait, score in fusion.items():
             bar = "█" * int(score / 5) + "░" * (20 - int(score / 5))
@@ -171,14 +143,10 @@ def predict(video_path):
         import traceback
         return None, f"❌ Error: {str(e)}\n\n```{traceback.format_exc()}```"
 
-
 # ── Gradio UI ──────────────────────────────────────────────────────────────────
-with gr.Blocks(title="Personality Prediction API", theme=gr.themes.Soft()) as demo:
+with gr.Blocks(title="Personality Prediction API") as demo:
 
-    gr.Markdown("""
-    # 🧠 Personality Prediction
-    Upload a video to predict **OCEAN** personality traits using audio + visual fusion.
-    """)
+    gr.Markdown("# 🧠 Personality Prediction\nUpload a video to predict **OCEAN** personality traits.")
 
     with gr.Row():
         with gr.Column(scale=1):
@@ -188,26 +156,14 @@ with gr.Blocks(title="Personality Prediction API", theme=gr.themes.Soft()) as de
         with gr.Column(scale=1):
             summary_out = gr.Markdown(label="Results Summary")
 
-    with gr.Accordion("📊 Full Results (Audio / Visual / Fusion)", open=False):
+    with gr.Accordion("📊 Full Results", open=False):
         json_out = gr.JSON(label="Raw Scores (%)")
 
     predict_btn.click(
         fn=predict,
         inputs=video_input,
         outputs=[json_out, summary_out],
-        api_name="predict"
+        api_name="predict" # This creates the API endpoint
     )
 
-    gr.Markdown("---\n**Traits:** Extraversion · Agreeableness · Conscientiousness · Neuroticism · Openness")
-
-app = FastAPI()
-
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-app = gr.mount_gradio_app(app, demo, path="/")
+demo.launch()