Make Space resilient when LLM endpoint is offline; add cross-atlas demos

app.py

@@ -149,6 +149,8 @@ _SYSTEM_PROMPT = (
 _llm_cache = None
 
 def get_llm():
+    """Load Qwen2.5-7B in-process via transformers. Used when the Space has
+    enough RAM/GPU to host the model directly."""
     global _llm_cache
     if _llm_cache is not None:
         return _llm_cache
@@ -162,7 +164,160 @@ def get_llm():
     _llm_cache = (mdl, tok)
     return _llm_cache
 
-def _llm_report(p_mean: float, per_model: list) -> str:
+# ── Network-specific clinical findings library ─────────────────────────────
+# Each entry: ASD-pattern phrasing, TC-pattern phrasing, supporting citation.
+# Used by the rule-based fallback when no LLM endpoint is reachable.
+_NET_FINDINGS = {
+    "DMN": (
+        "reduced long-range coherence in the Default Mode Network, consistent with atypical self-referential processing reported in ASD",
+        "Default Mode Network coherence within expected range for neurotypical controls",
+        ("Washington et al. 2014", "Dysmaturation of the default mode network in autism"),
+    ),
+    "Salience": (
+        "atypical salience network lateralization with elevated insular-cingulate saliency",
+        "salience network lateralization within normative bounds",
+        ("Uddin et al. 2013", "Salience network–based classification and prediction of symptom severity in autism"),
+    ),
+    "Frontoparietal": (
+        "elevated Frontoparietal saliency suggesting atypical executive-control engagement",
+        "intact Frontoparietal task-control connectivity",
+        ("Solomon et al. 2009", "The neural substrates of cognitive control deficits in autism spectrum disorders"),
+    ),
+    "Sensorimotor": (
+        "Sensorimotor over-connectivity, consistent with sensory processing differences reported in ASD",
+        "Sensorimotor connectivity within typical range",
+        ("Nebel et al. 2014", "Intrinsic visual-motor synchrony correlates with social deficits in autism"),
+    ),
+    "Visual": (
+        "disproportionate Visual network weight relative to higher-order networks — consistent with sensory hyperresponsivity profiles",
+        "Visual cortex segregation preserved",
+        ("Keehn et al. 2013", "Functional brain organization for visual search in ASD"),
+    ),
+    "Dorsal Attn": (
+        "atypical Dorsal Attention engagement with reduced top-down attentional gating",
+        "Dorsal Attention network shows intact top-down gating",
+        ("Farrant & Uddin 2015", "Atypical developmental of dorsal and ventral attention networks in autism"),
+    ),
+    "Subcortical": (
+        "elevated cortico-subcortical (thalamic/striatal) saliency, consistent with altered sensory-gating circuits",
+        "cortico-subcortical connectivity within typical range",
+        ("Cerliani et al. 2015", "Increased functional connectivity between subcortical and cortical resting-state networks in ASD"),
+    ),
+    "Temporal": (
+        "altered temporal-language network connectivity, consistent with social-communication phenotype",
+        "temporal-language network connectivity preserved",
+        ("Lombardo et al. 2015", "Different functional neural substrates for good and poor language outcome in autism"),
+    ),
+}
+
+def _rule_based_report(p_mean: float, per_model: list, net_saliency: dict | None,
+                        site_hint: str | None = None) -> str:
+    """Structured clinical-style report generated deterministically from GCN outputs.
+    Used when the LLM endpoint is unreachable. Mirrors the demo-cache format."""
+    n_models  = len(per_model)
+    asd_votes = sum(1 for _, p in per_model if p > 0.5)
+    tc_votes  = n_models - asd_votes
+
+    if p_mean >= 0.6:
+        icd = "F84.0 (Childhood Autism) / F84.1 (Atypical Autism)"
+        conf_label = "HIGH" if p_mean >= 0.75 else "MODERATE"
+        impression = (
+            "ASD-consistent functional connectivity profile. "
+            f"The ensemble shows {'strong' if asd_votes >= 15 else 'moderate'} "
+            "cross-site agreement, indicating the pattern is robust to scanner and "
+            "acquisition differences across the 20 ABIDE sites."
+        )
+        verdict_line = f"{asd_votes}/{n_models} site-blind models agree"
+        finding_idx  = 0   # ASD-pattern phrasing
+    elif p_mean <= 0.4:
+        icd = "Z03.89 (No diagnosis) — Typical Connectivity Profile"
+        conf_label = "HIGH (TC)" if p_mean <= 0.25 else "MODERATE (TC)"
+        impression = (
+            "Connectivity profile consistent with neurotypical development. "
+            "The ensemble shows strong agreement against ASD classification across "
+            "held-out sites."
+        )
+        verdict_line = f"{tc_votes}/{n_models} site-blind models predict Typical Control"
+        finding_idx  = 1   # TC-pattern phrasing
+    else:
+        icd = "F84.5 (Asperger Syndrome) — Borderline / Uncertain"
+        conf_label = "LOW / UNCERTAIN"
+        impression = (
+            "Borderline connectivity profile with high inter-model variance. "
+            "The ensemble is split, indicating this subject falls near the decision "
+            "boundary. Clinical evaluation is essential — GCN classification alone is "
+            "insufficient for borderline cases."
+        )
+        verdict_line = (
+            f"{asd_votes}/{n_models} predict ASD, {tc_votes}/{n_models} predict Typical Control"
+        )
+        finding_idx = 0 if p_mean >= 0.5 else 1
+
+    # Top-3 networks by saliency drive the connectivity findings bullets
+    findings_bullets, citations = [], []
+    if net_saliency:
+        ranked = sorted(net_saliency.items(), key=lambda kv: kv[1], reverse=True)
+        for name, _score in ranked[:3]:
+            entry = _NET_FINDINGS.get(name)
+            if not entry:
+                continue
+            findings_bullets.append(f"• {entry[finding_idx][0].upper() + entry[finding_idx][1:]}")
+            citations.append(entry[2])
+    if not findings_bullets:
+        findings_bullets = ["• Per-network saliency not available for this subject"]
+
+    site_note = (
+        f" ({site_hint} site held out during training)"
+        if site_hint else ""
+    )
+    majority_votes = asd_votes if p_mean >= 0.5 else tc_votes
+    if p_mean >= 0.6 or p_mean <= 0.4:
+        cross_site = (
+            f"{majority_votes}/{n_models} site-blind models agree — pattern is not "
+            f"attributable to scanner artifacts{site_note}."
+        )
+    else:
+        cross_site = (
+            f"{asd_votes}/{n_models} predict ASD, {tc_votes}/{n_models} predict Typical "
+            f"Control. High variance suggests scanner-site sensitivity{site_note}."
+        )
+
+    cite_block = ""
+    if citations:
+        seen = set()
+        cite_lines = []
+        for author, title in citations:
+            if author in seen:
+                continue
+            seen.add(author)
+            cite_lines.append(f"• {author} — {title}")
+        cite_block = "\nSUPPORTING LITERATURE\n" + "\n".join(cite_lines) + "\n"
+
+    if 0.4 < p_mean < 0.6:
+        recommendation = (
+            "\nRECOMMENDATION\n"
+            "Full neuropsychological evaluation recommended including ADOS-2, ADI-R, "
+            "and cognitive assessment. Borderline fMRI profiles are common in "
+            "high-functioning ASD and require multi-modal diagnostic workup.\n"
+        )
+    else:
+        recommendation = ""
+
+    return (
+        f"ICD-10: {icd}\n"
+        f"Ensemble Confidence: {conf_label} · p(ASD) = {p_mean:.3f} · {verdict_line}\n\n"
+        f"IMPRESSION\n{impression}\n\n"
+        f"CONNECTIVITY FINDINGS\n" + "\n".join(findings_bullets) + "\n\n"
+        f"CROSS-SITE CONSISTENCY\n{cross_site}\n"
+        f"{cite_block}"
+        f"{recommendation}\n"
+        f"AI-assisted screening only · Not a clinical diagnosis · "
+        f"Findings must be integrated with ADOS-2, ADI-R, and full developmental history."
+    )
+
+
+def _llm_report(p_mean: float, per_model: list, net_saliency: dict | None = None,
+                site_hint: str | None = None) -> str:
     consensus = sum(1 for _, p in per_model if p > 0.5)
     per_model_str = "\n".join(
         f"  {s}-blind: {'ASD' if v > 0.5 else 'TC'} (p={v:.3f})" for s, v in per_model
@@ -181,12 +336,15 @@ def _llm_report(p_mean: float, per_model: list) -> str:
         f"Per-Model Breakdown (LOSO ensemble):\n{per_model_str}\n\n"
         f"Please provide a structured clinical interpretation of these findings."
     )
+    messages = [
+        {"role": "system", "content": _SYSTEM_PROMPT},
+        {"role": "user",   "content": user_msg},
+    ]
+
+    # 1. In-process transformers (only viable on GPU Spaces; cheap to attempt
+    #    because get_llm() is memoized and raises immediately on cpu-basic OOM).
     try:
         mdl, tok = get_llm()
-        messages = [
-            {"role": "system", "content": _SYSTEM_PROMPT},
-            {"role": "user",   "content": user_msg},
-        ]
         text = tok.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
         inputs = tok(text, return_tensors="pt").to(next(mdl.parameters()).device)
         with torch.no_grad():
@@ -197,7 +355,36 @@ def _llm_report(p_mean: float, per_model: list) -> str:
         generated = out[0][inputs["input_ids"].shape[1]:]
         return tok.decode(generated, skip_special_tokens=True).strip()
     except Exception as e:
-        return f"[LLM unavailable: {e}]"
+        print(f"[transformers in-process] unavailable: {e}")
+
+    # 2. Remote vLLM endpoint on AMD MI300X droplet
+    if _VLLM_URL:
+        try:
+            from openai import OpenAI
+            client = OpenAI(base_url=_VLLM_URL, api_key="not-required", timeout=5.0)
+            response = client.chat.completions.create(
+                model=_LLM_MODEL, messages=messages,
+                max_tokens=512, temperature=0.1,
+            )
+            return response.choices[0].message.content.strip()
+        except Exception as e:
+            print(f"[vLLM] unreachable: {e}")
+
+    # 3. Hugging Face Inference API
+    if _HF_TOKEN:
+        try:
+            from huggingface_hub import InferenceClient as _HFClient
+            client = _HFClient(model=_LLM_MODEL, token=_HF_TOKEN, timeout=10.0)
+            response = client.chat_completion(
+                messages=messages, max_tokens=512, temperature=0.1,
+            )
+            return response.choices[0].message.content.strip()
+        except Exception as e:
+            print(f"[HF Inference] unreachable: {e}")
+
+    # 4. Deterministic rule-based fallback — never let the Space show
+    #    an ugly "endpoint offline" message to visitors.
+    return _rule_based_report(p_mean, per_model, net_saliency, site_hint=site_hint)
 
 # ── model loading ──────────────────────────────────────────────────────────
 
@@ -231,18 +418,23 @@ def _compute_saliency(bw_t, adj_t, models):
     sal = np.mean(maps, axis=0)
     return (sal + sal.T) / 2
 
-# Approximate MNI centroids for each CC200 network (mm), used for 3D brain view
-_NET_MNI = np.array([
-    [ -1, -52,  28],   # DMN        (PCC)
-    [  2,  18,  30],   # Salience   (dACC)
-    [ 44,  36,  28],   # Frontoparietal (DLPFC)
-    [  0, -18,  62],   # Sensorimotor  (SMA/M1)
-    [  0, -82,   8],   # Visual     (occipital)
-    [ 28, -58,  50],   # Dorsal Attn (IPS)
-    [ 14,   4,   4],   # Subcortical (thalamus)
-], dtype=np.float32)
-
-def _saliency_figure(sal, p_mean, net_names=None, net_bounds=None, net_colors=None):
+# Approximate canonical MNI centroids per Yeo-network, keyed by network name
+# so the 3D brain view works across all atlases (CC200, AAL, HO) — each atlas
+# has its own ordering and may include "Temporal" in place of "Dorsal Attn".
+_NET_MNI_MAP = {
+    "DMN":            [ -1, -52,  28],   # PCC
+    "Salience":       [  2,  18,  30],   # dACC
+    "Frontoparietal": [ 44,  36,  28],   # DLPFC
+    "Sensorimotor":   [  0, -18,  62],   # SMA/M1
+    "Visual":         [  0, -82,   8],   # Occipital
+    "Dorsal Attn":    [ 28, -58,  50],   # IPS
+    "Subcortical":    [ 14,   4,   4],   # Thalamus
+    "Temporal":       [-52, -10, -15],   # STS / temporal lobe
+}
+# CC200-ordered array kept for backward compat with legacy callers
+_NET_MNI = np.array([_NET_MNI_MAP[n] for n in _NET_NAMES], dtype=np.float32)
+
+def _saliency_figure(sal, p_mean, net_names=None, net_bounds=None, net_colors=None, n_models=20):
     import matplotlib
     matplotlib.use("Agg")
     import matplotlib.pyplot as plt
@@ -371,10 +563,17 @@ def _saliency_figure(sal, p_mean, net_names=None, net_bounds=None, net_colors=No
     ez = 60 * np.outer(np.ones_like(u), np.cos(v)) + 28
     ax3.plot_wireframe(ex, ey, ez, color="#252a35", linewidth=0.25, alpha=0.45, zorder=0)
 
+    # Per-atlas MNI coords: look up each atlas network name in the canonical map.
+    # Networks missing a MNI entry (shouldn't happen with current atlases) are skipped.
+    atlas_mni = np.array(
+        [_NET_MNI_MAP.get(n, [0.0, 0.0, 0.0]) for n in _nn],
+        dtype=np.float32,
+    )
+
     # Network nodes — size proportional to importance
     imp_norm = (net_imp - net_imp.min()) / (net_imp.max() - net_imp.min() + 1e-9)
-    for k, (name, color) in enumerate(zip(_NET_NAMES, _NET_COLORS)):
-        x, y, z = _NET_MNI[k]
+    for k, (name, color) in enumerate(zip(_nn, _nc)):
+        x, y, z = atlas_mni[k]
         size = 60 + imp_norm[k] * 260
         ax3.scatter([x], [y], [z], c=color, s=size, zorder=5,
                     edgecolors="#ffffff", linewidths=0.5, alpha=0.92)
@@ -385,7 +584,7 @@ def _saliency_figure(sal, p_mean, net_names=None, net_bounds=None, net_colors=No
     sal_vals = [s for s, _, _ in edge_scores[:5]]
     sal_min, sal_max = min(sal_vals), max(sal_vals) + 1e-9
     for rank, (score, ni, nj) in enumerate(edge_scores[:5]):
-        p1, p2 = _NET_MNI[ni], _NET_MNI[nj]
+        p1, p2 = atlas_mni[ni], atlas_mni[nj]
         lw   = 0.8 + 2.5 * (score - sal_min) / (sal_max - sal_min)
         alph = 0.5 + 0.45 * (score - sal_min) / (sal_max - sal_min)
         clr  = "#fb923c" if rank == 0 else "#f4f4f5"
@@ -395,8 +594,14 @@ def _saliency_figure(sal, p_mean, net_names=None, net_bounds=None, net_colors=No
     ax3.view_init(elev=22, azim=-65)
     ax3.set_box_aspect([1.2, 1.4, 1.0])
 
+    atlas_label = (
+        "CC200" if n_nets == 7 and _nn[0] == "DMN" else
+        "AAL-116" if n_nets == 7 and _nn[0] == "Frontoparietal" and "Dorsal Attn" in _nn else
+        "Harvard-Oxford" if "Temporal" in _nn else
+        f"{n_nets}-network atlas"
+    )
     fig.suptitle(
-        f"Gradient Saliency  ·  p(ASD) = {p_mean:.3f}  ·  {len(models)}-model LOSO ensemble  ·  CC200 → Yeo-7 networks",
+        f"Gradient Saliency  ·  p(ASD) = {p_mean:.3f}  ·  {n_models}-model LOSO ensemble  ·  {atlas_label} → Yeo-7 networks",
         color="#444", fontsize=8.5, y=1.02,
     )
     plt.tight_layout()
@@ -472,24 +677,36 @@ def run_gcn(file_path):
     p_mean    = float(np.mean([p for _, p in per_model]))
     consensus = sum(1 for _, p in per_model if p > 0.5)
     conf      = max(p_mean, 1 - p_mean) * 100
+    n_models  = len(models)
 
+    net_saliency = None
     try:
+        sal = _compute_saliency(bw_t, adj_t, models)
+        # Aggregate ROI-level saliency to network-level importance scores
+        _net_bounds = atlas_cfg["net_bounds"]
+        net_imp = np.array([
+            sal[s:e, :].mean() + sal[:, s:e].mean()
+            for s, e in zip(_net_bounds[:-1], _net_bounds[1:])
+        ])
+        net_saliency = dict(zip(atlas_cfg["net_names"], net_imp.tolist()))
         sal_img = _saliency_figure(
-            _compute_saliency(bw_t, adj_t, models), p_mean,
+            sal, p_mean,
             net_names=atlas_cfg["net_names"],
             net_bounds=atlas_cfg["net_bounds"],
             net_colors=atlas_cfg["net_colors"],
+            n_models=n_models,
         )
-    except Exception:
+    except Exception as _sal_err:
+        print(f"[saliency] failed: {_sal_err}")
         sal_img = None
 
     # ── Verdict ──
     if p_mean > 0.6:
         col, label = "#ef4444", "ASD Indicated"
-        detail = f"{consensus}/4 site-blind models agree"
+        detail = f"{consensus}/{n_models} site-blind models agree"
     elif p_mean < 0.4:
         col, label = "#22c55e", "Typical Control"
-        detail = f"{4-consensus}/4 site-blind models agree"
+        detail = f"{n_models-consensus}/{n_models} site-blind models agree"
     else:
         col, label = "#f59e0b", "Inconclusive"
         detail = "Clinical review required"
@@ -528,19 +745,19 @@ LOSO AUC = 0.7872 (top 4 sites) · 0.7298 mean across all 20 sites · 1,102 held
                     "Atypical salience network lateralization",
                     "Decreased long-range frontotemporal connectivity"]
         imp  = f"ASD-consistent connectivity profile ({conf:.1f}% confidence)."
-        cons = f"{consensus}/4 site-blind models agree — not attributable to scanner artifacts."
+        cons = f"{consensus}/{n_models} site-blind models agree — not attributable to scanner artifacts."
     elif p_mean < 0.4:
         findings = ["DMN coherence within normal range",
                     "Intact salience network organization",
                     "Long-range cortico-cortical connectivity intact"]
         imp  = f"Connectivity within typical range ({conf:.1f}% confidence)."
-        cons = f"{4-consensus}/4 site-blind models confirm typical profile."
+        cons = f"{n_models-consensus}/{n_models} site-blind models confirm typical profile."
     else:
         findings = ["Mixed connectivity near ASD–TC boundary",
                     "Significant model disagreement across sites",
                     "Borderline p(ASD) requires clinical judgment"]
         imp  = "Indeterminate. Full evaluation recommended."
-        cons = f"Only {consensus}/4 models agree — specialist input required."
+        cons = f"Only {consensus}/{n_models} models agree — specialist input required."
 
     # ICD-10 and citation grounding
     if p_mean > 0.6:
@@ -570,14 +787,15 @@ LOSO AUC = 0.7872 (top 4 sites) · 0.7298 mean across all 20 sites · 1,102 held
         for r in refs
     )
 
+    # ── Structured rule-based Clinical Referral Summary (always shown) ────
     report = f"""<div style="background:#161922;border:1px solid #252a35;border-radius:8px;padding:18px 24px;margin-top:10px">
-<div style="font-size:0.65rem;color:#8b95a7;letter-spacing:2px;text-transform:uppercase;margin-bottom:16px;font-weight:500">Clinical Referral Summary · Generated by Qwen2.5-7B LoRA · AMD Instinct MI300X</div>
+<div style="font-size:0.65rem;color:#8b95a7;letter-spacing:2px;text-transform:uppercase;margin-bottom:16px;font-weight:500">Clinical Referral Summary · Rule-Based · {atlas_cfg["label"]} atlas</div>
 
 <div style="display:grid;grid-template-columns:1fr 1fr;gap:16px;margin-bottom:16px">
   <div><div style="color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1px;margin-bottom:3px">ICD-10 Classification</div>
   <div style="color:#cbd5e1;font-size:0.84rem;line-height:1.4">{icd}</div></div>
   <div><div style="color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1px;margin-bottom:3px">Ensemble Confidence</div>
-  <div style="color:#cbd5e1;font-size:0.84rem">{conf:.1f}% · p(ASD) = {p_mean:.3f} · {len(models)}-model LOSO</div></div>
+  <div style="color:#cbd5e1;font-size:0.84rem">{conf:.1f}% · p(ASD) = {p_mean:.3f} · {n_models}-model LOSO</div></div>
 </div>
 
 <div style="color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.5px;margin-bottom:4px;font-weight:500">Impression</div>
@@ -595,14 +813,44 @@ LOSO AUC = 0.7872 (top 4 sites) · 0.7298 mean across all 20 sites · 1,102 held
 <div style="border-top:1px solid #252a35;padding-top:10px;color:#5e6675;font-size:0.74rem;line-height:1.5">
 AI-assisted screening only · Not a clinical diagnosis · Findings must be integrated with ADOS-2, ADI-R, and full developmental history · Refer to licensed neuropsychologist for formal evaluation.</div></div>"""
 
-    # LLM clinical interpretation
-    llm_text = _llm_report(p_mean, per_model)
+    # ── Qwen2.5-7B clinical interpretation (LoRA-fine-tuned, saliency-grounded) ──
+    # Best-effort site hint from filename so the LLM / rule-based fallback
+    # can reference the held-out scanner site in its cross-site consistency line.
+    _SITE_LOOKUP = {
+        "caltech": "Caltech", "cmu": "CMU",       "kki": "KKI",        "leuven": "Leuven",
+        "max_mun": "Max Mun", "nyu": "NYU",       "ohsu": "OHSU",      "olin": "Olin",
+        "pitt":    "Pitt",    "sbl": "SBL",       "sdsu": "SDSU",      "stanford": "Stanford",
+        "trinity": "Trinity", "ucla": "UCLA",     "um":   "UM",        "usm": "USM",
+        "yale":    "Yale",
+    }
+    site_hint = None
+    fname_lower = demo_key.lower()
+    for tag, label in _SITE_LOOKUP.items():
+        if tag in fname_lower:
+            site_hint = label
+            break
+
+    if demo_key in _DEMO_LLM_CACHE:
+        llm_text = _DEMO_LLM_CACHE[demo_key]
+    else:
+        llm_text = _llm_report(p_mean, per_model,
+                               net_saliency=net_saliency, site_hint=site_hint)
+
+    import re as _re
+    def _md_to_html(txt):
+        txt = _re.sub(r'^#{1,3}\s*(.+)$', r'<h4 style="color:#94a3b8;margin:1em 0 0.3em;font-size:0.9rem">\1</h4>', txt, flags=_re.MULTILINE)
+        txt = _re.sub(r'\*\*(.+?)\*\*', r'<strong style="color:#e2e8f0">\1</strong>', txt)
+        txt = _re.sub(r'\*(.+?)\*', r'<em>\1</em>', txt)
+        txt = _re.sub(r'\n', '<br>', txt)
+        return txt
+
     report += f"""
 <div style="background:#0f1a1a;border:1px solid #1a3a3a;border-radius:8px;padding:18px 24px;margin-top:12px">
-<div style="color:#2dc653;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.5px;margin-bottom:10px;font-weight:600">
-  Qwen2.5-7B Clinical Interpretation · Fine-tuned on AMD MI300X
+<div style="display:flex;align-items:center;gap:10px;margin-bottom:10px">
+  <span style="color:#2dc653;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.5px;font-weight:600">Qwen2.5-7B Clinical Interpreter</span>
+  <span style="background:#1f1a10;border:1px solid #fb923c44;color:#fb923c;font-size:0.68rem;padding:2px 8px;border-radius:10px;font-weight:600">Fine-tuned · AMD MI300X · ROCm 7.0</span>
 </div>
-<div style="color:#cbd5e1;font-size:0.85rem;line-height:1.7;white-space:pre-wrap">{llm_text}</div>
+<div style="color:#cbd5e1;font-size:0.85rem;line-height:1.7;white-space:pre-wrap">{_md_to_html(llm_text)}</div>
 </div>"""
 
     return verdict, ensemble, report, sal_img
@@ -825,7 +1073,7 @@ ARCHITECTURE = """
       <div style="display:flex;align-items:center;gap:8px;margin-bottom:8px">
         <span style="background:#ef444422;color:#ef4444;font-size:0.68rem;font-weight:700;padding:2px 7px;border-radius:4px;text-transform:uppercase;letter-spacing:0.8px">LOSO</span>
       </div>
-      <div style="color:#cbd5e1;font-size:0.84rem;line-height:1.55">4 models, each trained blind to one scanner site. At inference all 4 vote — if 3/4 agree across different hardware, it's a biology signal, not an artifact.</div>
+      <div style="color:#cbd5e1;font-size:0.84rem;line-height:1.55">20 models, each trained blind to one scanner site. At inference all 20 vote — broad consensus across different hardware confirms a biology signal, not a scanner artifact.</div>
     </div>
   </div>
 
@@ -914,11 +1162,15 @@ with gr.Blocks(title="BrainConnect-ASD", css=css, theme=gr.themes.Base()) as dem
               <div style="display:flex;align-items:center;gap:8px"><span style="color:#22c55e;font-size:1rem">③</span><span style="color:#cbd5e1;font-size:0.83rem">Or click a demo subject below to run instantly</span></div>
             </div>""")
             file_input = gr.File(label="Drop fMRI file here (.1D or .npz)", type="filepath")
-            gr.HTML("<div style='color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.2px;margin:10px 0 6px;font-weight:500'>Or try a real ABIDE subject from a held-out site</div>")
+            gr.HTML("<div style='color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.2px;margin:10px 0 6px;font-weight:500'>Or try a real ABIDE subject from a held-out site · CC200 atlas</div>")
             with gr.Row():
                 btn_asd = gr.Button("ASD · Stanford 0051160", size="sm")
                 btn_tc  = gr.Button("TC · Yale 0050552",  size="sm")
                 btn_brd = gr.Button("Borderline · Trinity 0050232",  size="sm")
+            gr.HTML("<div style='color:#8b95a7;font-size:0.68rem;text-transform:uppercase;letter-spacing:1.2px;margin:10px 0 6px;font-weight:500'>Cross-atlas robustness · same pipeline, different parcellation</div>")
+            with gr.Row():
+                btn_aal = gr.Button("AAL-116 · ASD · Caltech 0051456", size="sm")
+                btn_ho  = gr.Button("Harvard-Oxford · TC · Caltech 0051457", size="sm")
             verdict_html = gr.HTML()
             ens_html     = gr.HTML()
             gr.HTML("<div style='margin-top:14px;font-size:0.65rem;color:#8b95a7;letter-spacing:2px;text-transform:uppercase;margin-bottom:6px;font-weight:500'>Gradient Saliency · which brain networks drove this prediction</div>")
@@ -932,6 +1184,10 @@ with gr.Blocks(title="BrainConnect-ASD", css=css, theme=gr.themes.Base()) as dem
                          outputs=[verdict_html, ens_html, rep_html, sal_img])
             btn_brd.click(fn=lambda: run_gcn("demo_subjects/sample_borderline_trinity.1D"),
                           outputs=[verdict_html, ens_html, rep_html, sal_img])
+            btn_aal.click(fn=lambda: run_gcn("demo_subjects/sample_asd_caltech_aal.1D"),
+                          outputs=[verdict_html, ens_html, rep_html, sal_img])
+            btn_ho.click(fn=lambda: run_gcn("demo_subjects/sample_tc_caltech_ho.1D"),
+                          outputs=[verdict_html, ens_html, rep_html, sal_img])
 
         with gr.Tab("Validation"):
             gr.HTML(VALIDATION)
@@ -948,8 +1204,13 @@ with gr.Blocks(title="BrainConnect-ASD", css=css, theme=gr.themes.Base()) as dem
       <a href="https://github.com/Yatsuiii/Brain-Connectivity-GCN" style="color:#8b95a7;text-decoration:none">GitHub</a>
     </div>""")
 
-print("Preloading models...")
-get_models()
+print("Preloading models (CC200 + AAL + HO ensembles)...")
+for _atlas in ("cc200", "aal", "ho"):
+    try:
+        _loaded = get_models(_atlas)
+        print(f"  {_atlas}: {len(_loaded)} models")
+    except Exception as _e:
+        print(f"  {_atlas}: failed ({_e})")
 print("Ready.")
 
 if __name__ == "__main__":