Upload app.py with huggingface_hub

app.py

@@ -1,12 +1,8 @@
 """TRIBE V2 — Brain Response Prediction (Meta)
-Predicts fMRI brain responses using Meta's TRIBE V2 model on ZeroGPU.
-"""
-import subprocess, sys
 
-# Install tribev2 with --no-deps to bypass Python 3.11 requirement.
-# The actual code works on 3.10 — the pyproject.toml constraint is conservative.
-# We install sub-deps separately to avoid version conflicts with ZeroGPU base image.
-# No pip installs at runtime — use only what's in requirements.txt + base image
+Full multimodal brain prediction using Meta's TRIBE V2 model.
+Supports video, audio, and text scoring on ZeroGPU (Python 3.12).
+"""
 
 import gradio as gr
 import spaces
@@ -15,102 +11,111 @@ import numpy as np
 import os
 import json
 import tempfile
+import io
 
-# ---- Model ----
+# ---- HF Auth for gated models (LLaMA 3.2-3B) ----
+try:
+    from huggingface_hub import login
+    hf_token = os.environ.get("HF_TOKEN", "")
+    if hf_token:
+        login(token=hf_token, add_to_git_credential=False)
+        print("HF authenticated for gated model access.")
+except Exception as e:
+    print(f"HF auth warning: {e}")
+
+# ---- Model (loads on first GPU call) ----
 model = None
 
 def ensure_model():
-    """Load LLaMA 3.2-3B for text encoding. TRIBE v2's full pipeline requires
-    Python 3.11, so we use the text encoder directly for brain-region scoring."""
     global model
     if model is not None:
         return model
-    # Use a non-gated model for text encoding (LLaMA 3.2 is gated, fails on ZeroGPU workers)
-    # microsoft/phi-2 is a strong 2.7B model that's fully open
-    print("Loading text encoder (Phi-2 2.7B — open, no gating)...")
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-    model_id = "microsoft/phi-2"
-    model = {
-        "tokenizer": AutoTokenizer.from_pretrained(model_id, trust_remote_code=True),
-        "model": AutoModelForCausalLM.from_pretrained(
-            model_id, torch_dtype=torch.float16,
-            output_hidden_states=True, trust_remote_code=True,
-        ),
-    }
-    print("Phi-2 loaded.")
+    print("Loading TRIBE V2 model...")
+    from tribev2 import TribeModel
+    model = TribeModel.from_pretrained("facebook/tribev2")
+    print(f"Model loaded: {type(model)}")
     return model
 
-# Authenticate with HF to access gated models (LLaMA 3.2-3B)
-try:
-    from huggingface_hub import login
-    hf_token = os.environ.get("HF_TOKEN", "")
-    if hf_token:
-        login(token=hf_token, add_to_git_credential=False)
-        print("HF authenticated.")
-    else:
-        print("WARNING: No HF_TOKEN set. Gated models (LLaMA) may fail to download.")
-except Exception as e:
-    print(f"HF login warning: {e}")
-
-print("TRIBE V2 ready. Model loads on first GPU call.")
 
 # ---- ROI Mapping ----
 REGIONS = {
-    "attention": ["S_intrapariet", "G_front_middle", "S_front_sup", "G_pariet_inf-Supramar"],
-    "emotion": ["G_insular", "S_circular_insula", "G_cingul", "G_front_inf-Orbital"],
-    "language": ["G_front_inf-Opercular", "G_front_inf-Triangul", "G_temp_sup-Lateral"],
-    "visual": ["G_occipital", "S_occipital", "G_cuneus", "S_calcarine", "Pole_occipital"],
-    "default_mode": ["G_front_sup", "G_precuneus", "G_cingul-Post"],
+    "attention": ["S_intrapariet", "G_front_middle", "S_front_sup",
+                   "G_pariet_inf-Supramar", "G_temp_sup-G_T_transv"],
+    "emotion": ["G_insular", "S_circular_insula", "G_cingul",
+                 "G_front_inf-Orbital", "G_rectus", "G_subcallosal"],
+    "language": ["G_front_inf-Opercular", "G_front_inf-Triangul",
+                  "G_temp_sup-Lateral", "G_temp_sup-Plan_tempo",
+                  "S_temporal_sup", "G_and_S_subcentral"],
+    "visual": ["G_occipital", "S_occipital", "G_cuneus", "S_calcarine",
+                "Pole_occipital", "G_oc-temp_lat-fusifor",
+                "S_oc_sup_and_transversal", "G_oc-temp_med-Lingual"],
+    "default_mode": ["G_front_sup", "G_precuneus", "G_cingul-Post",
+                      "G_temp_sup-Plan_polar", "G_parietal_sup"],
 }
-_roi = {"labels": None, "names": None}
+
+_roi = {"labels": None, "names": None, "loaded": False}
 
 def _load_roi():
-    if _roi["labels"] is not None:
+    if _roi["loaded"]:
         return _roi["labels"], _roi["names"]
     try:
         from nilearn import datasets
         d = datasets.fetch_atlas_surf_destrieux()
         _roi["labels"] = np.concatenate([d["labels_lh"], d["labels_rh"]])
         _roi["names"] = [n.decode() if isinstance(n, bytes) else str(n) for n in d["label_names"]]
-    except Exception:
-        pass
+    except Exception as e:
+        print(f"ROI atlas warning: {e}")
+    _roi["loaded"] = True
     return _roi["labels"], _roi["names"]
 
-def _sig(val, c=0.008, s=300.0):
-    """Sigmoid normalization calibrated for projected hidden state magnitudes.
-    Projected activations are very small (~0.001-0.02), so center and scale
-    are tuned accordingly to map into 20-80 range."""
-    return float(100.0 / (1.0 + np.exp(-s * (val - c))))
 
-def interpret(preds):
+def _get_mask(labels, names, region_key):
+    if labels is None:
+        return None
+    subs = REGIONS.get(region_key, [])
+    ids = [i for i, n in enumerate(names) if any(s in n for s in subs)]
+    mask = np.isin(labels, ids)
+    return mask if mask.any() else None
+
+
+def _sig(val, center=0.15, scale=20.0):
+    return float(100.0 / (1.0 + np.exp(-scale * (val - center))))
+
+
+def interpret(preds, modalities=None):
+    """Convert (n_timesteps, n_vertices) cortical predictions to scores."""
+    if modalities is None:
+        modalities = ["text"]
     if isinstance(preds, torch.Tensor):
         preds = preds.cpu().numpy()
     if preds.ndim == 1:
         preds = preds.reshape(1, -1)
+
     n_t, n_v = preds.shape
     labels, names = _load_roi()
 
-    scores = {}
-    for key, subs in REGIONS.items():
-        if labels is not None and labels.shape[0] == n_v:
-            ids = [i for i, n in enumerate(names) if any(s in n for s in subs)]
-            mask = np.isin(labels, ids)
-            scores[key] = float(np.abs(preds[:, mask]).mean()) if mask.any() else float(np.abs(preds).mean())
+    region_scores = {}
+    for key in REGIONS:
+        mask = _get_mask(labels, names, key)
+        if mask is not None and mask.shape[0] == n_v:
+            region_scores[key] = float(np.abs(preds[:, mask]).mean())
         else:
-            scores[key] = float(np.abs(preds).mean())
+            region_scores[key] = float(np.abs(preds).mean())
 
-    att = _sig(scores["attention"])
-    emo = _sig(scores["emotion"])
-    lang = _sig(scores["language"])
-    vis = _sig(scores["visual"])
-    dm = _sig(scores["default_mode"])
+    att = _sig(region_scores["attention"])
+    emo = _sig(region_scores["emotion"])
+    lang = _sig(region_scores["language"])
+    vis = _sig(region_scores["visual"])
+    dm = _sig(region_scores["default_mode"])
     overall = (att + emo + lang + vis + dm) / 5.0
     viral = att * 0.4 + emo * 0.4 + vis * 0.2
 
     temporal = [_sig(float(np.abs(preds[t]).mean())) for t in range(n_t)]
-    hook = np.mean(temporal[:2]) if len(temporal) >= 2 else overall
-    body = np.mean(temporal[2:]) if len(temporal) > 2 else overall
+    hook = float(np.mean(temporal[:2])) if len(temporal) >= 2 else overall
+    body = float(np.mean(temporal[2:])) if len(temporal) > 2 else overall
     retention = min(body / max(hook, 1) * 100, 100)
+    peak_tr = int(np.argmax(temporal)) if temporal else 0
+    peak_time = peak_tr * 2.0 + 5.0
 
     return {
         "scores": {
@@ -121,189 +126,217 @@ def interpret(preds):
             "language_processing": round(lang, 1),
             "visual_imagery": round(vis, 1),
             "hook_effectiveness": round(hook, 1),
-            "retention_prediction": round(retention, 1),
+            "retention_prediction": round(min(retention, 100), 1),
         },
         "raw": {
-            "n_timesteps": n_t, "n_vertices": n_v,
+            "n_timesteps": n_t,
+            "n_vertices": n_v,
+            "peak_engagement_time_s": round(peak_time, 1),
             "temporal_profile": [round(v, 1) for v in temporal],
-            "region_raw": {k: round(v, 4) for k, v in scores.items()},
+            "modalities_used": modalities,
+            "region_activations_raw": {k: round(v, 4) for k, v in region_scores.items()},
         },
     }
 
-# ---- GPU Functions ----
-@spaces.GPU(duration=120)
-def predict_text_gpu(text):
-    """Use LLM perplexity + semantic features to predict brain engagement.
-
-    Instead of random projections (which produce uniform scores), we measure:
-    - Perplexity: surprising/novel content → higher attention
-    - Token entropy: diverse vocabulary → higher language processing
-    - Sentiment strength: emotional words → higher emotional valence
-    - Specificity: numbers, names, concrete nouns → higher visual imagery
-    """
+
+# ---- Visualization ----
+def make_radar(scores, title="Brain Engagement"):
+    import matplotlib
+    matplotlib.use("Agg")
+    import matplotlib.pyplot as plt
+
+    cats = ["Attention", "Emotion", "Language", "Visual", "Viral"]
+    vals = [scores.get("attention_capture", 0), scores.get("emotional_valence", 0),
+            scores.get("language_processing", 0), scores.get("visual_imagery", 0),
+            scores.get("viral_potential", 0)]
+    vals += vals[:1]
+    angles = [n / 5.0 * 2 * np.pi for n in range(5)] + [0]
+
+    fig, ax = plt.subplots(figsize=(5, 5), subplot_kw=dict(polar=True))
+    fig.patch.set_facecolor("#0D1B2A")
+    ax.set_facecolor("#0D1B2A")
+    ax.plot(angles, vals, "o-", linewidth=2, color="#FFD166")
+    ax.fill(angles, vals, alpha=0.25, color="#FFD166")
+    ax.set_ylim(0, 100)
+    ax.set_xticks(angles[:-1])
+    ax.set_xticklabels(cats, size=11, color="white")
+    ax.set_yticks([25, 50, 75])
+    ax.set_yticklabels(["25", "50", "75"], size=8, color="grey")
+    ax.tick_params(colors="grey")
+    ax.spines["polar"].set_color("grey")
+    ax.grid(color="grey", alpha=0.3)
+    ax.set_title(title, size=14, color="white", pad=20)
+
+    buf = io.BytesIO()
+    fig.savefig(buf, format="png", bbox_inches="tight", facecolor=fig.get_facecolor(), dpi=100)
+    plt.close(fig)
+    buf.seek(0)
+    return buf
+
+
+def make_summary(scores):
+    o = scores.get("overall_brain_engagement", 50)
+    parts = []
+    if o >= 70:
+        parts.append(f"Strong engagement ({o}/100).")
+    elif o >= 50:
+        parts.append(f"Decent engagement ({o}/100).")
+    else:
+        parts.append(f"Weak engagement ({o}/100).")
+    if scores.get("attention_capture", 50) >= 70:
+        parts.append("Great attention hook.")
+    elif scores.get("attention_capture", 50) < 40:
+        parts.append("Needs stronger opening hook.")
+    if scores.get("emotional_valence", 50) >= 70:
+        parts.append("Strong emotional trigger.")
+    elif scores.get("emotional_valence", 50) < 40:
+        parts.append("Add personal stakes or urgency.")
+    if scores.get("hook_effectiveness", 50) >= 70 and scores.get("retention_prediction", 50) < 50:
+        parts.append("Good hook but drops off mid-section.")
+    return " ".join(parts)
+
+
+def _format(scores):
+    return "\n".join([
+        f"Overall Engagement: {scores['overall_brain_engagement']}/100",
+        f"Viral Potential:    {scores['viral_potential']}/100",
+        f"Attention Capture:  {scores['attention_capture']}/100",
+        f"Emotional Valence:  {scores['emotional_valence']}/100",
+        f"Language Processing: {scores['language_processing']}/100",
+        f"Visual Imagery:     {scores['visual_imagery']}/100",
+        f"Hook Effectiveness: {scores['hook_effectiveness']}/100",
+        f"Retention:          {scores['retention_prediction']}/100",
+    ])
+
+
+# ---- GPU Prediction Functions ----
+
+@spaces.GPU(duration=60)
+def _predict_text(text):
     m = ensure_model()
-    tokenizer = m["tokenizer"]
-    llm = m["model"].cuda().half()
-
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to("cuda")
-
-    with torch.inference_mode():
-        outputs = llm(**inputs)
-
-    logits = outputs.logits  # (1, seq_len, vocab_size)
-    hidden = outputs.hidden_states[-1]  # last layer (1, seq_len, hidden)
-
-    # --- Feature extraction ---
-    # 1. Perplexity (how surprising is the text?)
-    shift_logits = logits[:, :-1, :].contiguous()
-    shift_labels = inputs["input_ids"][:, 1:].contiguous()
-    loss_fn = torch.nn.CrossEntropyLoss(reduction="none")
-    token_losses = loss_fn(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
-    perplexity = float(torch.exp(token_losses.mean()).cpu())
-    # Normalize: perplexity 1-10 = boring, 10-50 = interesting, 50+ = very novel
-    attention_raw = min(perplexity / 30.0, 1.0)  # 0-1 scale
-
-    # 2. Token entropy (vocabulary diversity)
-    token_ids = inputs["input_ids"][0].cpu().numpy()
-    unique_ratio = len(set(token_ids.tolist())) / max(len(token_ids), 1)
-    language_raw = unique_ratio  # 0-1
-
-    # 3. Emotional intensity (variance in hidden states = more expressive)
-    hidden_np = hidden.squeeze().cpu().float().numpy()  # (seq_len, hidden)
-    token_norms = np.linalg.norm(hidden_np, axis=1)
-    emotion_raw = float(np.std(token_norms) / (np.mean(token_norms) + 1e-8))
-
-    # 4. Visual/specificity (presence of numbers, caps, punctuation = concrete)
-    text_lower = text.lower()
-    has_numbers = sum(1 for c in text if c.isdigit()) / max(len(text), 1)
-    has_caps = sum(1 for c in text if c.isupper()) / max(len(text), 1)
-    urgency_words = sum(1 for w in ["now", "shock", "destroy", "change", "secret",
-        "never", "always", "must", "urgent", "breaking", "exclusive", "free"] if w in text_lower)
-    visual_raw = min((has_numbers * 10 + has_caps * 5 + urgency_words * 0.15), 1.0)
-
-    # 5. Default mode (self-referential: I, me, my, you, your)
-    personal_words = sum(1 for w in text_lower.split() if w in
-        ["i", "me", "my", "you", "your", "we", "our", "myself"])
-    dm_raw = min(personal_words / max(len(text_lower.split()), 1) * 5, 1.0)
-
-    # --- Map to 0-100 scores ---
-    def to_score(val, center=0.3, steepness=8.0):
-        clamped = max(0.0, min(1.0, val))
-        return float(100.0 / (1.0 + np.exp(-steepness * (clamped - center))))
-
-    scores = {
-        "attention": to_score(attention_raw, 0.25, 6.0),
-        "emotion": to_score(emotion_raw, 0.15, 10.0),
-        "language": to_score(language_raw, 0.5, 8.0),
-        "visual": to_score(visual_raw, 0.2, 8.0),
-        "default_mode": to_score(dm_raw, 0.2, 6.0),
-    }
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8") as f:
+        f.write(text)
+        path = f.name
+    try:
+        df = m.get_events_dataframe(text_path=path)
+        preds, segs = m.predict(events=df)
+    finally:
+        os.unlink(path)
+    torch.cuda.empty_cache()
+    return interpret(preds, modalities=["text"])
 
-    overall = np.mean(list(scores.values()))
-    viral = scores["attention"] * 0.4 + scores["emotion"] * 0.4 + scores["visual"] * 0.2
-    hook_score = scores["attention"]  # attention IS the hook
-    retention = min(scores["language"] / max(scores["attention"], 1) * 100, 100)
 
+@spaces.GPU(duration=120)
+def _predict_video(video_path):
+    m = ensure_model()
+    df = m.get_events_dataframe(video_path=video_path)
+    preds, segs = m.predict(events=df)
     torch.cuda.empty_cache()
+    return interpret(preds, modalities=["video", "audio", "text"])
 
-    return {
-        "scores": {
-            "overall_brain_engagement": round(overall, 1),
-            "viral_potential": round(viral, 1),
-            "attention_capture": round(scores["attention"], 1),
-            "emotional_valence": round(scores["emotion"], 1),
-            "language_processing": round(scores["language"], 1),
-            "visual_imagery": round(scores["visual"], 1),
-            "hook_effectiveness": round(hook_score, 1),
-            "retention_prediction": round(retention, 1),
-        },
-        "raw": {
-            "perplexity": round(perplexity, 2),
-            "token_unique_ratio": round(unique_ratio, 3),
-            "hidden_state_variance": round(emotion_raw, 4),
-            "specificity": round(visual_raw, 3),
-            "personal_reference": round(dm_raw, 3),
-        },
-    }
 
 # ---- Handlers ----
-def score_text(text):
+
+def handle_text(text):
     if not text or not text.strip():
-        return "Please enter text.", ""
+        return "Enter text to score.", None, ""
     try:
-        r = predict_text_gpu(text.strip())
+        r = _predict_text(text.strip())
         s = r["scores"]
-        lines = [
-            f"Overall Engagement: {s['overall_brain_engagement']}/100",
-            f"Viral Potential:    {s['viral_potential']}/100",
-            f"Attention Capture:  {s['attention_capture']}/100",
-            f"Emotional Valence:  {s['emotional_valence']}/100",
-            f"Language Processing: {s['language_processing']}/100",
-            f"Visual Imagery:     {s['visual_imagery']}/100",
-            f"Hook Effectiveness: {s['hook_effectiveness']}/100",
-            f"Retention:          {s['retention_prediction']}/100",
-        ]
-        # Summary
-        o = s["overall_brain_engagement"]
-        summary = f"{'Strong' if o >= 70 else 'Decent' if o >= 50 else 'Weak'} engagement ({o}/100). "
-        if s["attention_capture"] < 40:
-            summary += "Needs stronger opening hook. "
-        if s["emotional_valence"] >= 70:
-            summary += "Great emotional trigger. "
-        elif s["emotional_valence"] < 40:
-            summary += "Add personal stakes or urgency. "
-        if s["hook_effectiveness"] >= 70 and s["retention_prediction"] < 50:
-            summary += "Good hook but drops off mid-section. "
-        return "\n".join(lines), summary
+        chart = make_radar(s)
+        return _format(s), chart, make_summary(s)
     except Exception as e:
         import traceback
-        return f"Error: {type(e).__name__}: {e}\n{traceback.format_exc()}", ""
+        return f"Error: {e}\n{traceback.format_exc()}", None, ""
 
-def score_json(text):
-    if not text or not text.strip():
-        return '{"error": "No text provided"}'
+
+def handle_video(video):
+    if video is None:
+        return "Upload a video.", None, ""
     try:
-        r = predict_text_gpu(text.strip())
-        return json.dumps(r, indent=2)
+        r = _predict_video(video)
+        s = r["scores"]
+        chart = make_radar(s, title="Video Brain Engagement")
+        peak = r["raw"].get("peak_engagement_time_s", "N/A")
+        text = _format(s) + f"\nPeak Engagement: {peak}s"
+        return text, chart, make_summary(s)
     except Exception as e:
-        return json.dumps({"error": str(e)})
+        import traceback
+        return f"Error: {e}\n{traceback.format_exc()}", None, ""
 
-def ab_test(a, b):
+
+def handle_ab(a, b):
     if not a or not b:
         return "Enter both versions."
     try:
-        ra = predict_text_gpu(a.strip())
-        rb = predict_text_gpu(b.strip())
+        ra = _predict_text(a.strip())
+        rb = _predict_text(b.strip())
         sa, sb = ra["scores"], rb["scores"]
         va, vb = sa["viral_potential"], sb["viral_potential"]
-        w = f"Version A wins ({va} vs {vb})" if va > vb else f"Version B wins ({vb} vs {va})" if vb > va else "Tie"
-        return f"{w}\n\nA: engagement={sa['overall_brain_engagement']} viral={va} hook={sa['hook_effectiveness']}\nB: engagement={sb['overall_brain_engagement']} viral={vb} hook={sb['hook_effectiveness']}"
+        w = f"Version A wins ({va} vs {vb})" if va > vb else (
+            f"Version B wins ({vb} vs {va})" if vb > va else "Tie")
+        return f"{w}\n\n--- A ---\n{_format(sa)}\n{make_summary(sa)}\n\n--- B ---\n{_format(sb)}\n{make_summary(sb)}"
     except Exception as e:
         return f"Error: {e}"
 
-# ---- Gradio UI (Textbox only — avoids ZeroGPU schema bug) ----
-with gr.Blocks(title="TRIBE V2 Brain Prediction") as demo:
-    gr.Markdown("# TRIBE V2 — Brain Response Prediction\nMeta's fMRI model predicts brain engagement with your content.\n")
-
-    with gr.Tab("Score Text"):
-        inp = gr.Textbox(label="Content", lines=5, placeholder="Paste script or hook...")
-        btn = gr.Button("Analyze", variant="primary")
-        out = gr.Textbox(label="Scores", lines=10)
-        summary = gr.Textbox(label="Insight")
-        btn.click(score_text, [inp], [out, summary], api_name="predict")
-
-    with gr.Tab("A/B Test"):
-        a = gr.Textbox(label="Version A", lines=3)
-        b = gr.Textbox(label="Version B", lines=3)
-        btn2 = gr.Button("Compare", variant="primary")
-        res = gr.Textbox(label="Result", lines=6)
-        btn2.click(ab_test, [a, b], [res], api_name="ab_test")
-
-    with gr.Tab("API (JSON)"):
-        gr.Markdown("Returns raw JSON for programmatic use.")
+
+def handle_api(text):
+    if not text or not text.strip():
+        return '{"error": "No text"}'
+    try:
+        r = _predict_text(text.strip())
+        return json.dumps(r, indent=2)
+    except Exception as e:
+        return json.dumps({"error": str(e)})
+
+
+# ---- Gradio UI ----
+
+with gr.Blocks(title="TRIBE V2 Brain Prediction", theme=gr.themes.Base(
+    primary_hue="amber", secondary_hue="cyan", neutral_hue="slate",
+    font=gr.themes.GoogleFont("Inter"),
+)) as demo:
+
+    gr.Markdown("# 🧠 TRIBE V2 — Brain Response Prediction\n"
+                "Meta's fMRI model predicts how your content activates the brain.\n")
+
+    with gr.Tab("📝 Text Scorer"):
+        gr.Markdown("Score a script, hook, or post. ~30s on GPU.")
+        t_in = gr.Textbox(label="Content", lines=5, placeholder="Paste script or hook...")
+        t_btn = gr.Button("🧠 Analyze", variant="primary")
+        with gr.Row():
+            t_scores = gr.Textbox(label="Scores", lines=10)
+            t_chart = gr.Image(label="Brain Radar", type="filepath")
+        t_summary = gr.Textbox(label="Insight")
+        t_btn.click(handle_text, [t_in], [t_scores, t_chart, t_summary], api_name="predict")
+
+    with gr.Tab("🎬 Video Scorer"):
+        gr.Markdown("Upload a video for full multimodal brain analysis. ~2-5 min on GPU.")
+        v_in = gr.Video(label="Upload Video")
+        v_btn = gr.Button("🧠 Analyze Video", variant="primary")
+        with gr.Row():
+            v_scores = gr.Textbox(label="Scores", lines=10)
+            v_chart = gr.Image(label="Brain Radar", type="filepath")
+        v_summary = gr.Textbox(label="Insight")
+        v_btn.click(handle_video, [v_in], [v_scores, v_chart, v_summary], api_name="predict_video")
+
+    with gr.Tab("⚔️ A/B Test"):
+        gr.Markdown("Compare two hooks head-to-head.")
+        with gr.Row():
+            ab_a = gr.Textbox(label="Version A", lines=3)
+            ab_b = gr.Textbox(label="Version B", lines=3)
+        ab_btn = gr.Button("⚔️ Compare", variant="primary")
+        ab_out = gr.Textbox(label="Result", lines=10)
+        ab_btn.click(handle_ab, [ab_a, ab_b], [ab_out], api_name="ab_test")
+
+    with gr.Tab("🔌 API"):
+        gr.Markdown("Returns raw JSON for `score_script.py` compatibility.")
         api_in = gr.Textbox(label="Text", lines=3)
         api_btn = gr.Button("Get JSON")
-        api_out = gr.Textbox(label="Response", lines=15)
-        api_btn.click(score_json, [api_in], [api_out], api_name="api_predict")
+        api_out = gr.Textbox(label="JSON", lines=15)
+        api_btn.click(handle_api, [api_in], [api_out], api_name="api_predict")
+
+    gr.Markdown("---\n*[Meta TRIBE V2](https://github.com/facebookresearch/tribev2) | "
+                "ZeroGPU (A10G) | Python 3.12 | Built by somebeast*")
 
 demo.queue().launch()